From 62ce515d15fd61f15985d54120265f5976fbce4d Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 19 Oct 2026 14:27:29 -0700 Subject: [PATCH 01/18] feat(router): let a deployment set a timeout per service_tier (#47406) * feat(router): let a deployment set a timeout per service_tier Adds timeout_by_service_tier and stream_timeout_by_service_tier to a deployment's litellm_params. A request whose service_tier matches a key uses that value in place of the deployment's timeout or stream_timeout. A timeout the client sends still wins and requests without a matching tier are unchanged Co-authored-by: adc77 * style(router): sort the request_timeout_resolver import * fix(router): read the service tier maps without the untyped params get * fix(router): look the service tier timeout up through a typed get --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: adc77 --- .../request_timeout_resolver.py | 23 +++++++- litellm/passthrough/timeout_utils.py | 20 +++++-- litellm/router.py | 3 + litellm/types/litellm_params.py | 2 + litellm/types/router.py | 5 ++ .../test_pass_through_endpoints.py | 26 +++++++++ tests/unit/test_router/test_router_timeout.py | 56 +++++++++++++++++++ tests/unit/types/test_litellm_params.py | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 16 ++++++ 9 files changed, 147 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/request_timeout_resolver.py b/litellm/litellm_core_utils/request_timeout_resolver.py index a32651fedee..0e011dc06a0 100644 --- a/litellm/litellm_core_utils/request_timeout_resolver.py +++ b/litellm/litellm_core_utils/request_timeout_resolver.py @@ -12,10 +12,15 @@ tell "user asked for this" from "nobody set it". This resolver answers that: from __future__ import annotations -from typing import Final +from collections.abc import Mapping +from typing import Final, Literal, Protocol + +from pydantic import TypeAdapter from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS +_TIMEOUTS_BY_SERVICE_TIER: Final[TypeAdapter[Mapping[str, float] | None]] = TypeAdapter(Mapping[str, float] | None) + def get_configured_request_timeout() -> float | None: """Return the explicitly-configured ``litellm.request_timeout``, else ``None``.""" @@ -27,3 +32,19 @@ def get_configured_request_timeout() -> float | None: if timeout != float(DEFAULT_REQUEST_TIMEOUT_SECONDS): return timeout return None + + +class _DeploymentParams(Protocol): + def get(self, key: str, /) -> object: ... + + +def get_service_tier_timeout( + service_tier: object, + litellm_params: _DeploymentParams, + key: Literal["timeout_by_service_tier", "stream_timeout_by_service_tier"], +) -> float | None: + """Return the deployment's ``key`` timeout for the request's ``service_tier``, else ``None``.""" + if not isinstance(service_tier, str): + return None + timeouts: Final = _TIMEOUTS_BY_SERVICE_TIER.validate_python(litellm_params.get(key)) + return None if timeouts is None else timeouts.get(service_tier) diff --git a/litellm/passthrough/timeout_utils.py b/litellm/passthrough/timeout_utils.py index f600c7817f2..ce13552ff49 100644 --- a/litellm/passthrough/timeout_utils.py +++ b/litellm/passthrough/timeout_utils.py @@ -5,7 +5,10 @@ from typing import Final from pydantic import TypeAdapter -from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, + get_service_tier_timeout, +) DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS: Final = 600.0 @@ -50,21 +53,27 @@ def resolve_llm_passthrough_timeout( Anthropic /v1/messages). Non-streaming precedence: kwargs timeout/request_timeout -> litellm_params - timeout/request_timeout -> router_timeout -> litellm.request_timeout (litellm_settings.request_timeout, + timeout_by_service_tier[kwargs service_tier] -> litellm_params timeout/request_timeout -> router_timeout -> litellm.request_timeout (litellm_settings.request_timeout, when explicitly set) -> general_settings.pass_through_request_timeout -> 600s default. Streaming (``kwargs["stream"]`` truthy) resolves ``stream_timeout`` at every level before any generic timeout, matching ``Router._get_stream_timeout`` on the completion route: - kwargs stream_timeout -> litellm_params stream_timeout -> router_stream_timeout, then the - non-streaming chain above. + kwargs stream_timeout -> litellm_params stream_timeout_by_service_tier[kwargs service_tier] -> + litellm_params stream_timeout -> router_stream_timeout, then the non-streaming chain above. Only the first set value is validated as seconds, so a value in a lower-precedence field never fails the call. """ request: Final = kwargs if kwargs is not None else _NO_PARAMS deployment: Final = litellm_params if litellm_params is not None else _NO_PARAMS + service_tier: Final = request.get("service_tier") stream_candidates: Final = ( - (request.get("stream_timeout"), deployment.get("stream_timeout"), router_stream_timeout) + ( + request.get("stream_timeout"), + get_service_tier_timeout(service_tier, deployment, "stream_timeout_by_service_tier"), + deployment.get("stream_timeout"), + router_stream_timeout, + ) if request.get("stream") else () ) @@ -72,6 +81,7 @@ def resolve_llm_passthrough_timeout( *stream_candidates, request.get("timeout"), request.get("request_timeout"), + get_service_tier_timeout(service_tier, deployment, "timeout_by_service_tier"), deployment.get("timeout"), deployment.get("request_timeout"), router_timeout, diff --git a/litellm/router.py b/litellm/router.py index af79fa763c4..5bf46c68b20 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -118,6 +118,7 @@ from litellm.litellm_core_utils.ptu_pricing import ( ) from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, + get_service_tier_timeout, ) from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.sensitive_data_masker import ( @@ -4773,6 +4774,7 @@ class Router: """Helper to get stream timeout from kwargs or deployment params""" return ( kwargs.get("stream_timeout", None) # the params dynamically set by user + or get_service_tier_timeout(kwargs.get("service_tier"), data, "stream_timeout_by_service_tier") or data.get("stream_timeout", None) # timeout set on litellm_params for this deployment or self.stream_timeout # timeout set on router or self.request_timeout # litellm_settings.request_timeout (per-attempt) @@ -4786,6 +4788,7 @@ class Router: return ( kwargs.get("timeout", None) # the params dynamically set by user or kwargs.get("request_timeout", None) # the params dynamically set by user + or get_service_tier_timeout(kwargs.get("service_tier"), data, "timeout_by_service_tier") or data.get("timeout", None) # timeout set on litellm_params for this deployment or data.get("request_timeout", None) # timeout set on litellm_params for this deployment or self.request_timeout # litellm_settings.request_timeout (per-attempt) diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index b05d25462aa..978eb9c4557 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -77,6 +77,8 @@ class ProviderConnection: request_timeout: float | None = None force_timeout: float | None = None stream_timeout: float | str | None = None + timeout_by_service_tier: Mapping[str, float] | None = None + stream_timeout_by_service_tier: Mapping[str, float] | None = None max_retries: int | None = None tenant_id: str | None = None client_id: str | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index 46617e1b753..f83127806db 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -505,6 +505,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): timeout: float | str | httpx.Timeout | SDKTimeout | None = None # if str, pass in as os.environ/ # read timeout of a live stream, a faked stream takes the longer of this and timeout. If str, pass in as os.environ/ stream_timeout: float | str | None = None + # per service_tier overrides of timeout / stream_timeout, used when the request's service_tier matches a key + timeout_by_service_tier: Mapping[str, float] | None = None + stream_timeout_by_service_tier: Mapping[str, float] | None = None max_retries: int | None = None drop_params: bool | str | None = None organization: str | None = None # for openai orgs @@ -709,6 +712,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): responses_api_version: ReadOnly[str | None] timeout: float | str | httpx.Timeout | SDKTimeout | None # writable-ok: preserve timeout updates stream_timeout: float | str | None + timeout_by_service_tier: ReadOnly[Mapping[str, float] | None] + stream_timeout_by_service_tier: ReadOnly[Mapping[str, float] | None] max_retries: int | None organization: list | str | None # for openai orgs provider_affinity_header: ReadOnly[str | None] diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 5412ea56ac0..4c3c7ae0f0a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1857,6 +1857,32 @@ def test_resolve_llm_passthrough_timeout_validates_only_the_winning_value( assert resolve_llm_passthrough_timeout(kwargs=kwargs, litellm_params=litellm_params) == expected +def test_resolve_llm_passthrough_timeout_uses_the_deployment_timeout_for_the_requests_service_tier(): + litellm_params = { + "timeout": 610, + "stream_timeout": 30, + "timeout_by_service_tier": {"flex": 900}, + "stream_timeout_by_service_tier": {"flex": 950}, + } + assert resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex"}, litellm_params=litellm_params) == 900.0 + assert resolve_llm_passthrough_timeout(kwargs={"service_tier": "default"}, litellm_params=litellm_params) == 610.0 + assert ( + resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex", "timeout": 20}, litellm_params=litellm_params) + == 20.0 + ) + assert ( + resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex", "stream": True}, litellm_params=litellm_params) + == 950.0 + ) + assert ( + resolve_llm_passthrough_timeout( + kwargs={"service_tier": "flex", "stream": True}, + litellm_params={"timeout": 610, "stream_timeout": 30, "timeout_by_service_tier": {"flex": 900}}, + ) + == 30.0 + ) + + def test_resolve_llm_passthrough_timeout_rejects_a_non_numeric_winner(): with pytest.raises(ValidationError): resolve_llm_passthrough_timeout(kwargs={"timeout": httpx.Timeout(30.0)}) diff --git a/tests/unit/test_router/test_router_timeout.py b/tests/unit/test_router/test_router_timeout.py index 3265179a513..8fc32620f71 100644 --- a/tests/unit/test_router/test_router_timeout.py +++ b/tests/unit/test_router/test_router_timeout.py @@ -153,6 +153,62 @@ def test_unit_test_streaming_timeout(stream): assert stream_timeout_val == normal_timeout +@pytest.mark.parametrize( + "request_kwargs, expected_timeout", + [ + ({"service_tier": "flex"}, 900.0), + ({"service_tier": "default"}, 610.0), + ({}, 610.0), + ({"service_tier": "flex", "timeout": 20}, 20), + ({"service_tier": "flex", "stream": True}, 950.0), + ({"service_tier": "default", "stream": True}, 610.0), + ({"service_tier": "flex", "stream": True, "stream_timeout": 30}, 30), + ], +) +def test_the_deployment_timeout_for_the_requests_service_tier_wins_over_its_general_timeout( + request_kwargs: dict[str, object], expected_timeout: float +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "openai-gpt-5.1", + "litellm_params": { + "model": "openai/gpt-5.1", + "api_key": "sk-test", + "timeout": 610, + "stream_timeout": 610, + "timeout_by_service_tier": {"flex": 900}, + "stream_timeout_by_service_tier": {"flex": 950}, + }, + } + ] + ) + deployment_params: Final = router.model_list[0]["litellm_params"] + + assert router._get_timeout(kwargs=request_kwargs, data=deployment_params) == expected_timeout + + +def test_a_stream_keeps_the_deployment_stream_timeout_when_only_the_request_timeout_map_names_its_tier() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "openai-gpt-5.1", + "litellm_params": { + "model": "openai/gpt-5.1", + "api_key": "sk-test", + "timeout": 610, + "stream_timeout": 30, + "timeout_by_service_tier": {"flex": 900}, + }, + } + ] + ) + deployment_params: Final = router.model_list[0]["litellm_params"] + + assert router._get_timeout(kwargs={"service_tier": "flex", "stream": True}, data=deployment_params) == 30.0 + assert router._get_timeout(kwargs={"service_tier": "flex"}, data=deployment_params) == 900.0 + + async def _drain_responses_stream(router: Router) -> None: response: Final = await router.aresponses(model="gpt", input="hi", stream=True) _ = [event async for event in response] diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index ce763378656..e88d757519c 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -75,6 +75,8 @@ CONNECTION_NAMES: Final = ( "request_timeout", "force_timeout", "stream_timeout", + "timeout_by_service_tier", + "stream_timeout_by_service_tier", "max_retries", "tenant_id", "client_id", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index daffc3c0c41..733a093bcac 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35424,6 +35424,10 @@ export interface components { } | null; /** Stream Timeout */ stream_timeout?: number | string | null; + /** Stream Timeout By Service Tier */ + stream_timeout_by_service_tier?: { + [key: string]: number; + } | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */ @@ -35436,6 +35440,10 @@ export interface components { }[] | null; /** Timeout */ timeout?: number | string | null; + /** Timeout By Service Tier */ + timeout_by_service_tier?: { + [key: string]: number; + } | null; /** Token Exchange Audience */ token_exchange_audience?: string | null; /** Token Exchange Endpoint */ @@ -50666,6 +50674,10 @@ export interface components { } | null; /** Stream Timeout */ stream_timeout?: number | string | null; + /** Stream Timeout By Service Tier */ + stream_timeout_by_service_tier?: { + [key: string]: number; + } | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */ @@ -50678,6 +50690,10 @@ export interface components { }[] | null; /** Timeout */ timeout?: number | string | null; + /** Timeout By Service Tier */ + timeout_by_service_tier?: { + [key: string]: number; + } | null; /** Token Exchange Audience */ token_exchange_audience?: string | null; /** Token Exchange Endpoint */ From 6162184b1ce2435daf7996241b349361d4342d08 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 19 Oct 2026 14:40:15 -0700 Subject: [PATCH 02/18] fix(managed_files): send provider file ids when creating a fine-tuning job from managed files (#46552) * fix(managed_files): send provider file ids when creating a fine-tuning job from managed files Fine-tuning job creation with a managed training file sent the gateway's managed id and the provider-prefixed deployment model to OpenAI, which rejected the job. The pre-call hook now builds the file id mapping for training_file and validation_file, the deployment hook swaps both, the SDK strips the provider prefix from the model, and the create response returns the caller's own file ids * fix(managed_files): route managed fine-tuning creates only to deployments holding every file Fine-tuning creates now narrow routing the same way managed batches do, so a deployment added after the upload is never picked. Adds a two-worker chaos integration test and runs tests/unit/fine_tuning in the endpoints shard. * fix(managed_files): keep fine-tuning creates on the model group when no deployment holds the files A managed file used on a model group it was never uploaded for kept failing with the provider's 400 before the routing filter. The filter turned that into a 429 no-deployments error, so fine-tuning now falls back to the whole group in that case while batch input files keep routing only to deployments that hold them. * test(fine_tuning): cover the provider prefix strip for OpenAI and Azure creates through the proxy * test(fine_tuning): start each multi-deployment case from an empty upstream log * test(fine_tuning): boot both proxies of the late deployment cell on one scratch database --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + .../proxy/hooks/managed_files.py | 63 ++-- litellm/fine_tuning/main.py | 7 +- .../proxy/fine_tuning_endpoints/endpoints.py | 3 +- .../compatibility/test_fine_tuning_jobs.py | 99 +++++- .../test_fine_tuning_managed_files_chaos.py | 289 ++++++++++++++++++ .../security/test_fine_tuning_managed_ids.py | 8 +- .../proxy/hooks/test_managed_files.py | 162 +++++++++- tests/unit/fine_tuning/__init__.py | 0 tests/unit/fine_tuning/test_main.py | 54 ++++ .../vertex_ai/fine_tuning/test_handler.py | 5 +- .../fine_tuning_endpoints/test_endpoints.py | 33 +- 12 files changed, 668 insertions(+), 56 deletions(-) create mode 100644 tests/integration/compatibility/test_fine_tuning_managed_files_chaos.py create mode 100644 tests/unit/fine_tuning/__init__.py create mode 100644 tests/unit/fine_tuning/test_main.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index cb1b595f3f0..5e42f182345 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -410,6 +410,7 @@ jobs: tests/unit/embeddings tests/unit/endpoints tests/unit/files + tests/unit/fine_tuning tests/unit/images tests/unit/interactions tests/unit/messages diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 439770390f8..bc35c4a0aba 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -6,6 +6,7 @@ import base64 import json import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -479,6 +480,13 @@ def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str, _MANAGED_FILES_TARGET: Final = "managed_files" +_PROVIDER_FILE_ID_KWARGS: Final = MappingProxyType( + { + CallTypes.acreate_batch: ("input_file_id",), + CallTypes.acreate_fine_tuning_job: ("training_file", "validation_file"), + } +) +_PROVIDER_FILE_ID_KEYS: Final = tuple(dict.fromkeys(chain.from_iterable(_PROVIDER_FILE_ID_KWARGS.values()))) _PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0) @@ -1562,6 +1570,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if isinstance(file_id, str) and file_id ) await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict) + if fine_tuning_file_ids: + data["model_file_id_mapping"] = await self.get_model_file_id_mapping( + fine_tuning_file_ids, user_api_key_dict.parent_otel_span + ) return data @@ -1576,24 +1588,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if request_kwargs is None: return healthy_deployments - input_file_id = cast(Optional[str], request_kwargs.get("input_file_id")) - model_file_id_mapping = cast( - Optional[Dict[str, Dict[str, str]]], - request_kwargs.get("model_file_id_mapping"), + request: Final = cast(Mapping[str, object], request_kwargs) + model_file_id_mapping: Final = cast(Mapping[str, Mapping[str, str]], request.get("model_file_id_mapping") or {}) + requested_file_ids: Final = (request.get(key) for key in _PROVIDER_FILE_ID_KEYS) + model_ids_per_file: Final = tuple( + frozenset(model_file_id_mapping[file_id]) + for file_id in requested_file_ids + if isinstance(file_id, str) and file_id in model_file_id_mapping + ) + allowed_model_ids: Final[frozenset[str]] = ( + model_ids_per_file[0].intersection(*model_ids_per_file[1:]) if model_ids_per_file else frozenset() ) - allowed_model_ids = [] - if input_file_id and model_file_id_mapping: - model_id_dict = model_file_id_mapping.get(input_file_id, {}) - allowed_model_ids = list(model_id_dict.keys()) if len(allowed_model_ids) == 0: return healthy_deployments - return [ + holders: Final[list[dict[str, object]]] = [ deployment for deployment in healthy_deployments if deployment.get("model_info", {}).get("id") in allowed_model_ids ] + return holders if holders or "input_file_id" in request else healthy_deployments async def async_pre_call_deployment_hook( self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] @@ -1601,25 +1616,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ Allow modifying the request just before it's sent to the deployment. """ - accessor_key: Optional[str] = None - if call_type and call_type == CallTypes.acreate_batch: - accessor_key = "input_file_id" - elif call_type and call_type == CallTypes.acreate_fine_tuning_job: - accessor_key = "training_file" - else: + accessor_keys: Final = _PROVIDER_FILE_ID_KWARGS.get(call_type) if call_type else None + if accessor_keys is None: return kwargs - if accessor_key: + model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) + # model_info may be at top-level or nested under litellm_metadata + # (batch/file operations use litellm_metadata) + model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) + if model_id is None: + model_id = cast( + Optional[str], + kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), + ) + for accessor_key in accessor_keys: input_file_id = cast(Optional[str], kwargs.get(accessor_key)) - model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) - # model_info may be at top-level or nested under litellm_metadata - # (batch/file operations use litellm_metadata) - model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) - if model_id is None: - model_id = cast( - Optional[str], - kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), - ) mapped_file_id: Optional[str] = None if input_file_id and model_file_id_mapping and model_id: mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None) @@ -1813,7 +1824,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ), ) - async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict: + async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span) -> dict: """ Get model-specific file IDs for a list of proxy file IDs. Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index 5297a8215be..f49f069b83f 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -173,6 +173,7 @@ def create_fine_tuning_job( try: _is_async: Final = kwargs.pop("acreate_fine_tuning_job", False) is True optional_params: Final = GenericLiteLLMParams(**kwargs) + provider_model: Final = model.removeprefix(f"{custom_llm_provider}/") # handle hyperparameters hyperparameters = hyperparameters or {} # original hyperparameters @@ -218,7 +219,7 @@ def create_fine_tuning_job( ) create_fine_tuning_job_data_dict = _build_fine_tuning_job_data( - model, + provider_model, training_file, _oai_hyperparameters, suffix, @@ -263,7 +264,7 @@ def create_fine_tuning_job( extra_body = _prepare_azure_extra_body(extra_body, kwargs, azure_specific_hyperparams) create_fine_tuning_job_data_dict = _build_fine_tuning_job_data( - model, + provider_model, training_file, _oai_hyperparameters, suffix, @@ -298,7 +299,7 @@ def create_fine_tuning_job( response = vertex_fine_tuning_apis_instance.create_fine_tuning_job( _is_async=_is_async, create_fine_tuning_job_data=_build_fine_tuning_job_data( - model, + provider_model, training_file, _oai_hyperparameters, suffix, diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index d015112cf3a..030f8d7a46a 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -164,7 +164,8 @@ async def create_fine_tuning_job( ) response = cast(LiteLLMFineTuningJob, await llm_router.acreate_fine_tuning_job(**data)) - response.training_file = unified_file_id + response.training_file = training_file + response.validation_file = fine_tuning_request.validation_file set_hidden_param(response, "unified_file_id", unified_file_id) ## ELSE, Route based on custom_llm_provider elif fine_tuning_request.custom_llm_provider: diff --git a/tests/integration/compatibility/test_fine_tuning_jobs.py b/tests/integration/compatibility/test_fine_tuning_jobs.py index 9678198352f..1eff386b8ba 100644 --- a/tests/integration/compatibility/test_fine_tuning_jobs.py +++ b/tests/integration/compatibility/test_fine_tuning_jobs.py @@ -27,7 +27,6 @@ MANAGED_MODEL: Final = "gpt-4.1-openai" DEPLOYMENT_KEYS: Final = ("sk-fine-tuning-deployment-a", "sk-fine-tuning-deployment-b") PROVIDER_JOB: Final = "ftjob-provider-123" RUN: Final = uuid.uuid4().hex[:12] -MANAGED_JOB: Final = f"ftjob-managed-{RUN}" TRAINING_LINE: Final = b'{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}\n' JOB_PATH: Final = re.compile(r"^/(?:openai/)?(?:v1/)?fine_tuning/jobs/([^/]+)(/cancel)?$") JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -88,8 +87,18 @@ def _job(job_id: str, status: str, training_file: str = "file-abc123", **fields: } -def _provider_file_id(bearer: str) -> str: - return f"file-{RUN}-" + bearer.removeprefix("sk-fine-tuning-") +def _provider_file_id(bearer: str, filename: str) -> str: + return f"file-{RUN}-" + bearer.removeprefix("sk-fine-tuning-") + "-" + filename.removesuffix(".jsonl") + + +def _managed_job_id(provider_training_file: str) -> str: + return provider_training_file.replace("file-", "ftjob-managed-", 1) + + +def _uploaded_filename(request: Request) -> str: + names: Final = [part.get_filename() for part in _multipart_parts(request) if part.get_filename()] + assert len(names) == 1, request.body[:300] + return str(names[0]) def _bearer(request: Request) -> str: @@ -118,19 +127,27 @@ def _respond(request: Request) -> Reply: if (request.method, path) == ("POST", "/v1/files"): return _reply( { - "id": _provider_file_id(_bearer(request)), + "id": _provider_file_id(_bearer(request), _uploaded_filename(request)), "object": "file", "bytes": len(TRAINING_LINE), "created_at": 1700000000, - "filename": "train.jsonl", + "filename": _uploaded_filename(request), "purpose": "fine-tune", "status": "processed", } ) if (request.method, path) in {("POST", "/v1/fine_tuning/jobs"), ("POST", "/openai/fine_tuning/jobs")}: sent: Final = JSON_OBJECT.validate_json(request.body) - job_id: Final = MANAGED_JOB if _bearer(request) in DEPLOYMENT_KEYS else PROVIDER_JOB - return _reply(_job(job_id, "validating_files", str(sent["training_file"]))) + managed: Final = _bearer(request) in DEPLOYMENT_KEYS + echoed: Final = {"validation_file": sent.get("validation_file")} if managed else {} + return _reply( + _job( + _managed_job_id(str(sent["training_file"])) if managed else PROVIDER_JOB, + "validating_files", + str(sent["training_file"]), + **echoed, + ) + ) matched: Final = JOB_PATH.match(path) if matched and matched.group(1) and request.method == "GET": return _reply(_job(matched.group(1), "running")) @@ -329,6 +346,30 @@ def test_azure_create_merges_finetune_settings_and_keeps_azure_fields(rig: Rig) assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "validating_files", "file-abc"), raw.text +@pytest.mark.parametrize( + ("provider", "target", "provider_defaults"), + [ + ("openai", "/v1/fine_tuning/jobs", {}), + ("azure", f"/openai/fine_tuning/jobs?api-version={AZURE_API_VERSION}", {"trainingType": 1}), + ], +) +def test_create_sends_the_model_without_its_provider_prefix( + rig: Rig, provider: str, target: str, provider_defaults: Mapping[str, JsonValue] +) -> None: + with rig.sdk("/v1") as client: + raw: Final = client.fine_tuning.jobs.with_raw_response.create( + model=f"{provider}/gpt-4o-mini", training_file="file-abc123", extra_body={"custom_llm_provider": provider} + ) + sent: Final = _only(_drained(rig), "POST", target) + assert JSON_OBJECT.validate_json(sent.body) == { + "model": "gpt-4o-mini", + "training_file": "file-abc123", + "hyperparameters": {}, + **provider_defaults, + }, sent.body + assert raw.status_code == 200, raw.text + + def test_azure_cancel_merges_finetune_settings(rig: Rig) -> None: with rig.sdk("/v1") as client: raw: Final = client.fine_tuning.jobs.with_raw_response.cancel( @@ -400,10 +441,6 @@ def test_retrieve_reads_the_documented_custom_llm_provider_header(rig: Rig) -> N def test_managed_create_sends_the_provider_file_id_of_the_pinned_deployment(rig: Rig) -> None: - pytest.skip( - "BUG: managed create sends the base64 LiteLLM file id and the prefixed model openai/gpt-4.1 to the provider" - " and returns the decoded unified id as training_file" - ) with rig.sdk("/v1") as client: uploaded: Final = client.files.with_raw_response.create( file=("train.jsonl", TRAINING_LINE, "application/jsonl"), @@ -431,12 +468,48 @@ def test_managed_create_sends_the_provider_file_id_of_the_pinned_deployment(rig: assert deployment_key in DEPLOYMENT_KEYS, sent.headers assert JSON_OBJECT.validate_json(sent.body) == { "model": "gpt-4.1", - "training_file": _provider_file_id(deployment_key), + "training_file": _provider_file_id(deployment_key, "train.jsonl"), "hyperparameters": {}, }, sent.body assert raw.status_code == 200, raw.text job: Final = Job.model_validate_json(raw.text) assert re.fullmatch( - rf"litellm_proxy;model_id:[0-9a-f]{{64}};generic_response_id:{MANAGED_JOB}", _decoded(job.id) + rf"litellm_proxy;model_id:[0-9a-f]{{64}};generic_response_id:" + rf"{_managed_job_id(_provider_file_id(deployment_key, 'train.jsonl'))}", + _decoded(job.id), ), raw.text assert job.training_file == managed_file, raw.text + + +def _managed_upload(rig: Rig, client: OpenAI, filename: str) -> str: + uploaded: Final = client.files.with_raw_response.create( + file=(filename, TRAINING_LINE, "application/jsonl"), + purpose="fine-tune", + extra_body={"target_model_names": MANAGED_MODEL}, + ) + assert uploaded.status_code == 200, uploaded.text + assert sorted((request.method, request.target, _bearer(request)) for request in _drained(rig)) == [ + ("POST", "/v1/files", key) for key in DEPLOYMENT_KEYS + ] + return str(JSON_OBJECT.validate_json(uploaded.text)["id"]) + + +def test_managed_create_sends_the_provider_ids_of_managed_training_and_validation_files(rig: Rig) -> None: + with rig.sdk("/v1") as client: + training_file: Final = _managed_upload(rig, client, "tune.jsonl") + validation_file: Final = _managed_upload(rig, client, "validate.jsonl") + raw: Final = client.fine_tuning.jobs.with_raw_response.create( + model=MANAGED_MODEL, training_file=training_file, validation_file=validation_file + ) + sent: Final = _only(_drained(rig), "POST", "/v1/fine_tuning/jobs") + deployment_key: Final = _bearer(sent) + assert deployment_key in DEPLOYMENT_KEYS, sent.headers + assert JSON_OBJECT.validate_json(sent.body) == { + "model": "gpt-4.1", + "training_file": _provider_file_id(deployment_key, "tune.jsonl"), + "validation_file": _provider_file_id(deployment_key, "validate.jsonl"), + "hyperparameters": {}, + }, sent.body + assert raw.status_code == 200, raw.text + job: Final = Job.model_validate_json(raw.text) + assert (job.training_file, job.validation_file) == (training_file, validation_file), raw.text diff --git a/tests/integration/compatibility/test_fine_tuning_managed_files_chaos.py b/tests/integration/compatibility/test_fine_tuning_managed_files_chaos.py new file mode 100644 index 00000000000..54794fe29b9 --- /dev/null +++ b/tests/integration/compatibility/test_fine_tuning_managed_files_chaos.py @@ -0,0 +1,289 @@ +"""Managed fine-tuning creates on a two-worker proxy whose model group has two deployments: a concurrent burst from +sync and async clients, traffic pinned to each worker in turn, and a worker killed and replaced, all send every +deployment its own provider file ids and hand the caller back its own managed ids. A deployment added to the group +after the upload holds no copy of the files, so it never receives the create even when the router prefers it. A model +group the files were never uploaded for still gets the create and the provider's own 400, as before routing checked +which deployments hold the files.""" + +from __future__ import annotations + +import asyncio +import json +import uuid +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import pytest +import yaml +from integration._support.client import eventually, gateway_from_environment +from integration._support.database import scratch_database +from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from openai import AsyncOpenAI, BadRequestError, OpenAI +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(480) + +MODEL: Final = "gpt-4.1-chaos" +DEPLOYMENT_KEYS: Final = ("sk-chaos-deployment-a", "sk-chaos-deployment-b") +LATE_DEPLOYMENT_KEY: Final = "sk-chaos-deployment-late" +OTHER_MODEL: Final = "gpt-4.1-chaos-other" +OTHER_DEPLOYMENT_KEY: Final = "sk-chaos-deployment-other" +RUN: Final = uuid.uuid4().hex[:12] +TRAINING_LINE: Final = b'{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}\n' +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +BURST: Final = 20 +CREATE: Final = ("POST", "/v1/fine_tuning/jobs") +WORKER_BOOT_DISCOVERY: Final = ("GET", "/v1/models") + + +@dataclass(frozen=True, slots=True) +class Rig: + owned: OwnedProxy + wire: Wire + training_file: str + validation_file: str + + def base_url(self) -> str: + return f"{str(self.owned.gateway.client.base_url).rstrip('/')}/v1" + + def sdk(self) -> OpenAI: + return OpenAI(base_url=self.base_url(), api_key=self.owned.gateway.key, max_retries=0) + + def async_sdk(self) -> AsyncOpenAI: + return AsyncOpenAI(base_url=self.base_url(), api_key=self.owned.gateway.key, max_retries=0) + + +def _bearer(request: Request) -> str: + return request.headers.get("authorization", "").removeprefix("Bearer ") + + +def _provider_file_id(bearer: str, purpose: str) -> str: + return f"file-{RUN}-{bearer.removeprefix('sk-chaos-')}-{purpose}" + + +def _uploaded_purpose(request: Request) -> str: + return "validate" if b'filename="validate.jsonl"' in request.body else "train" + + +def _invalid_training_file(training_file: str) -> Reply: + error: Final = { + "message": f"invalid training_file: {training_file}", + "type": "invalid_request_error", + "param": "training_file", + "code": None, + } + return Reply(status=400, body=json.dumps({"error": error}).encode()) + + +def _respond(request: Request) -> Reply: + path: Final = urlsplit(request.target).path + if (request.method, path) == ("POST", "/v1/files"): + return Reply( + body=json.dumps( + { + "id": _provider_file_id(_bearer(request), _uploaded_purpose(request)), + "object": "file", + "bytes": len(TRAINING_LINE), + "created_at": 1700000000, + "filename": f"{_uploaded_purpose(request)}.jsonl", + "purpose": "fine-tune", + "status": "processed", + } + ).encode() + ) + if (request.method, path) == ("POST", "/v1/fine_tuning/jobs"): + sent: Final = JSON_OBJECT.validate_json(request.body) + if not str(sent["training_file"]).startswith(f"file-{RUN}-"): + return _invalid_training_file(str(sent["training_file"])) + return Reply( + body=json.dumps( + { + "id": f"ftjob-{uuid.uuid4().hex}", + "object": "fine_tuning.job", + "created_at": 1700000000, + "error": None, + "fine_tuned_model": None, + "finished_at": None, + "hyperparameters": {"n_epochs": 3}, + "model": sent["model"], + "organization_id": "org-chaos", + "result_files": [], + "seed": 42, + "status": "validating_files", + "trained_tokens": None, + "training_file": sent["training_file"], + "validation_file": sent.get("validation_file"), + } + ).encode() + ) + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + + +def _deployment(wire: Wire, key: str, routing: Mapping[str, JsonValue], model: str = MODEL) -> dict[str, JsonValue]: + return { + "model_name": model, + "litellm_params": {"model": "openai/gpt-4.1", "api_key": key, "api_base": f"{wire.url}/v1", **routing}, + "model_info": {"supported_endpoints": ["/chat/completions", "/fine_tuning"]}, + } + + +def _config(wire: Wire, late_preferred_deployment: bool = False) -> dict[str, JsonValue]: + late: Final = (_deployment(wire, LATE_DEPLOYMENT_KEY, {"order": 1}),) if late_preferred_deployment else () + return { + "model_list": [ + *(_deployment(wire, key, {}) for key in DEPLOYMENT_KEYS), + *late, + _deployment(wire, OTHER_DEPLOYMENT_KEY, {}, OTHER_MODEL), + ], + "router_settings": {"num_retries": 0, "disable_cooldowns": True}, + "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL"}, + } + + +def _config_file(directory: Path, wire: Wire, late_preferred_deployment: bool = False) -> Path: + directory.mkdir(parents=True, exist_ok=True) + config: Final = directory / "config.yaml" + config.write_text(yaml.safe_dump(_config(wire, late_preferred_deployment))) + return config + + +def _upload(client: OpenAI, purpose: str) -> str: + uploaded: Final = client.files.with_raw_response.create( + file=(f"{purpose}.jsonl", TRAINING_LINE, "application/jsonl"), + purpose="fine-tune", + extra_body={"target_model_names": MODEL}, + ) + assert uploaded.status_code == 200, uploaded.text + return str(JSON_OBJECT.validate_json(uploaded.text)["id"]) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("fine-tuning-chaos") + with gateway_from_environment() as shared, wire_server(_respond) as wire: + config: Final = _config_file(directory, wire) + with owned_proxy_process( + shared, directory, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned: + eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=120) + unloaded: Final = Rig(owned, wire, "", "") + with unloaded.sdk() as client: + training_file: Final = _upload(client, "train") + validation_file: Final = _upload(client, "validate") + wire.drain() + yield Rig(owned, wire, training_file, validation_file) + + +def _create(rig: Rig) -> tuple[int, Mapping[str, JsonValue]]: + with rig.sdk() as client: + raw: Final = client.fine_tuning.jobs.with_raw_response.create( + model=MODEL, training_file=rig.training_file, validation_file=rig.validation_file + ) + return raw.status_code, JSON_OBJECT.validate_json(raw.text) + + +async def _async_creates(rig: Rig, count: int) -> tuple[tuple[int, Mapping[str, JsonValue]], ...]: + async with rig.async_sdk() as client: + raws: Final = await asyncio.gather( + *( + client.fine_tuning.jobs.with_raw_response.create( + model=MODEL, training_file=rig.training_file, validation_file=rig.validation_file + ) + for _ in range(count) + ) + ) + return tuple((raw.status_code, JSON_OBJECT.validate_json(raw.text)) for raw in raws) + + +def _sent_creates(rig: Rig, expected: int) -> tuple[Request, ...]: + collected: Final[list[Request]] = [] + + def creates() -> tuple[Request, ...]: + collected.extend(rig.wire.drain()) + return tuple(request for request in collected if (request.method, request.target) == CREATE) + + sent: Final = eventually(creates, lambda found: len(found) >= expected, seconds=30) + assert len(sent) == expected, [(request.method, request.target) for request in collected] + assert {(request.method, request.target) for request in collected} <= {CREATE, WORKER_BOOT_DISCOVERY}, collected + return sent + + +def _assert_each_deployment_got_its_own_ids(creates: tuple[Request, ...]) -> None: + for request in creates: + bearer = _bearer(request) + assert bearer in DEPLOYMENT_KEYS, request.headers + assert JSON_OBJECT.validate_json(request.body) == { + "model": "gpt-4.1", + "training_file": _provider_file_id(bearer, "train"), + "validation_file": _provider_file_id(bearer, "validate"), + "hyperparameters": {}, + }, request.body + + +def _assert_caller_gets_its_own_ids(rig: Rig, outcomes: tuple[tuple[int, Mapping[str, JsonValue]], ...]) -> None: + received: Final = [(status, body.get("training_file"), body.get("validation_file")) for status, body in outcomes] + assert received == [(200, rig.training_file, rig.validation_file)] * len(outcomes), outcomes + + +def test_concurrent_sync_and_async_burst_sends_each_deployment_its_own_file_ids(rig: Rig) -> None: + rig.wire.drain() + with ThreadPoolExecutor(max_workers=BURST // 2) as pool: + synced: Final = tuple(pool.map(_create, (rig,) * (BURST // 2))) + awaited: Final = asyncio.run(_async_creates(rig, BURST // 2)) + creates: Final = _sent_creates(rig, BURST) + _assert_each_deployment_got_its_own_ids(creates) + assert {_bearer(request) for request in creates} == set(DEPLOYMENT_KEYS) + _assert_caller_gets_its_own_ids(rig, synced + awaited) + + +def test_each_worker_alone_sends_the_provider_file_ids(rig: Rig) -> None: + rig.wire.drain() + for paused in sorted(rig.owned.worker_pids()): + with rig.owned.suspended(paused): + outcomes = tuple(_create(rig) for _ in range(4)) + _assert_caller_gets_its_own_ids(rig, outcomes) + _assert_each_deployment_got_its_own_ids(_sent_creates(rig, 4)) + + +def test_a_replaced_worker_still_sends_the_provider_file_ids(rig: Rig) -> None: + rig.wire.drain() + for victim in sorted(rig.owned.worker_pids()): + rig.owned.kill_worker(victim) + outcomes: Final = tuple(_create(rig) for _ in range(6)) + _assert_caller_gets_its_own_ids(rig, outcomes) + _assert_each_deployment_got_its_own_ids(_sent_creates(rig, 6)) + + +def test_a_deployment_added_after_the_upload_never_receives_the_create(tmp_path: Path) -> None: + with gateway_from_environment() as shared, wire_server(_respond) as wire, scratch_database() as database_url: + one_database: Final = {"DATABASE_URL": database_url} + before: Final = tmp_path / "before" + with owned_proxy_process(shared, before, one_database, config=_config_file(before, wire)) as uploader: + with Rig(uploader, wire, "", "").sdk() as client: + training_file: Final = _upload(client, "train") + validation_file: Final = _upload(client, "validate") + after: Final = tmp_path / "after" + with owned_proxy_process( + shared, after, one_database, config=_config_file(after, wire, late_preferred_deployment=True), workers=2 + ) as owned: + eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=120) + wire.drain() + rig: Final = Rig(owned, wire, training_file, validation_file) + outcomes: Final = tuple(_create(rig) for _ in range(6)) + creates: Final = _sent_creates(rig, 6) + _assert_caller_gets_its_own_ids(rig, outcomes) + _assert_each_deployment_got_its_own_ids(creates) + + +def test_a_model_group_without_a_copy_of_the_files_still_gets_the_providers_400(rig: Rig) -> None: + rig.wire.drain() + with rig.sdk() as client, pytest.raises(BadRequestError) as rejected: + client.fine_tuning.jobs.create(model=OTHER_MODEL, training_file=rig.training_file) + creates: Final = _sent_creates(rig, 1) + assert rejected.value.status_code == 400, rejected.value.response.text + assert [_bearer(request) for request in creates] == [OTHER_DEPLOYMENT_KEY], creates diff --git a/tests/integration/security/test_fine_tuning_managed_ids.py b/tests/integration/security/test_fine_tuning_managed_ids.py index 8f179e5a25f..791dd1d1e75 100644 --- a/tests/integration/security/test_fine_tuning_managed_ids.py +++ b/tests/integration/security/test_fine_tuning_managed_ids.py @@ -294,13 +294,9 @@ def test_owner_create_returns_the_unified_job_id(fixture: str, request: pytest.F built: Final[Rig] = request.getfixturevalue(fixture) created: Final = Job.model_validate_json(built.created_text) expected: Final = _client_job( - _unified_job_id(built.deployment_id, built.provider_job), "validating_files", built.provider_file - ) - assert created.model_dump(exclude={"training_file"}) == expected.model_dump(exclude={"training_file"}), ( - "training_file is excluded because test_managed_create_sends_the_provider_file_id_of_the_pinned_deployment " - "documents the known managed-create bug", - built.created_text, + _unified_job_id(built.deployment_id, built.provider_job), "validating_files", built.managed_file ) + assert created == expected, built.created_text def test_another_key_cannot_create_with_or_cancel_the_owners_managed_ids(rig: Rig) -> None: diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 91f6593575e..20ed8ac0c10 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -3,6 +3,7 @@ import json import logging import time import asyncio +import uuid from collections.abc import Awaitable, Callable, Mapping from types import SimpleNamespace from typing import TYPE_CHECKING, Final, cast @@ -13,6 +14,7 @@ import pytest import respx from fastapi import HTTPException from openai import APIConnectionError +from pydantic import TypeAdapter if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch @@ -22,7 +24,7 @@ from litellm_enterprise.proxy.hooks.managed_files import ( _provider_file_retrieve_credentials, ) from litellm.caching import DualCache -from litellm.proxy._types import CallTypes, LiteLLM_ManagedFileTable +from litellm.proxy._types import CallTypes, LiteLLM_ManagedFileTable, UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( BATCH_CREATE_HIDDEN_PARAM, is_base64_encoded_unified_file_id, @@ -6050,6 +6052,164 @@ async def test_create_from_raw_provider_file_id_skips_managed_ownership_lookup(c assert not any("unified_file_id" in where for where in table.find_first_calls), table.find_first_calls +def _owned_managed_file(unified_uuid: str, model_mappings: Mapping[str, str]) -> tuple[str, LiteLLM_ManagedFileTable]: + raw: Final = f"litellm_proxy:application/jsonl;unified_id,{unified_uuid};target_model_names,gpt-4o-mini" + unified_file_id: Final = base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=") + return unified_file_id, LiteLLM_ManagedFileTable( + unified_file_id=unified_file_id, + file_object=None, + model_mappings=dict(model_mappings), + flat_model_file_ids=list(model_mappings.values()), + created_by="user_a", + team_id="team_a", + ) + + +async def _provider_kwargs( + managed_files: PROXY_LiteLLMManagedFiles, call_type: str, request: Mapping[str, str], deployment_id: str +) -> Mapping[str, object]: + prepared: Final = await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()), + cache=MagicMock(), + data=dict(request), + call_type=call_type, + ) + assert isinstance(prepared, dict), prepared + sent: Final = await managed_files.async_pre_call_deployment_hook( + kwargs={**prepared, "litellm_metadata": {"model_info": {"id": deployment_id}}}, call_type=CallTypes(call_type) + ) + assert sent is not None + return sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS) +@pytest.mark.parametrize(("deployment_id", "provider_file_id"), [("model-a", "file-a"), ("model-b", "file-b")]) +async def test_create_from_managed_file_sends_the_routed_deployments_provider_file_id( + call_type: str, file_field: str, deployment_id: str, provider_file_id: str +) -> None: + unified_file_id, row = _owned_managed_file( + "0b6f2c1e-9a4d-4c3e-8f71-2d5a6b7c8d90", {"model-a": "file-a", "model-b": "file-b"} + ) + managed_files, _ = _managed_files_with_fake_prisma(row) + + sent: Final = await _provider_kwargs( + managed_files, call_type, {"model": "gpt-4o-mini", file_field: unified_file_id}, deployment_id + ) + + assert sent[file_field] == provider_file_id, sent + + +@pytest.mark.asyncio +async def test_fine_tuning_create_swaps_managed_training_and_validation_files_independently() -> None: + training_file, training_row = _owned_managed_file("5c1d9e0a-3b7f-4e2a-9c6d-1f8e7a6b5c40", {"model-a": "file-train"}) + validation_file, validation_row = _owned_managed_file( + "8e2f4a6c-1d3b-4f5e-a7c9-0b2d4f6a8c10", {"model-a": "file-validate"} + ) + managed_files, _ = _managed_files_with_fake_prisma(training_row, validation_row) + + sent: Final = await _provider_kwargs( + managed_files, + "acreate_fine_tuning_job", + {"model": "gpt-4o-mini", "training_file": training_file, "validation_file": validation_file}, + "model-a", + ) + + assert (sent["training_file"], sent["validation_file"]) == ("file-train", "file-validate"), sent + + +@pytest.mark.asyncio +async def test_fine_tuning_create_keeps_a_raw_validation_file_next_to_a_managed_training_file() -> None: + training_file, training_row = _owned_managed_file("5c1d9e0a-3b7f-4e2a-9c6d-1f8e7a6b5c40", {"model-a": "file-train"}) + managed_files, _ = _managed_files_with_fake_prisma(training_row) + + sent: Final = await _provider_kwargs( + managed_files, + "acreate_fine_tuning_job", + {"model": "gpt-4o-mini", "training_file": training_file, "validation_file": "file-raw-validate"}, + "model-a", + ) + + assert (sent["training_file"], sent["validation_file"]) == ("file-train", "file-raw-validate"), sent + + +_REQUEST_KWARGS: Final = TypeAdapter(dict[str, object]) +_DEPLOYMENTS: Final = TypeAdapter(list[dict[str, dict[str, str]]]) + + +async def _routable_deployment_ids( + managed_files: PROXY_LiteLLMManagedFiles, call_type: str, request: Mapping[str, str] +) -> tuple[str, ...]: + prepared: Final = _REQUEST_KWARGS.validate_python( + await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()), + cache=MagicMock(), + data=dict(request), + call_type=call_type, + ) + ) + routable: Final = _DEPLOYMENTS.validate_python( + await managed_files.async_filter_deployments( + model="gpt-4o-mini", + healthy_deployments=[{"model_info": {"id": model_id}} for model_id in ("model-a", "model-b", "model-c")], + messages=None, + request_kwargs=prepared, + ) + ) + return tuple(deployment["model_info"]["id"] for deployment in routable) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "files", "expected"), + [ + pytest.param( + "acreate_fine_tuning_job", + {"training_file": {"model-a": "f-1", "model-b": "f-2"}, "validation_file": {"model-b": "f-3"}}, + ("model-b",), + id="fine-tune-routes-to-the-deployment-holding-both-files", + ), + pytest.param( + "acreate_fine_tuning_job", + {"training_file": {"model-a": "f-1", "model-b": "f-2"}, "validation_file": "file-raw-validate"}, + ("model-a", "model-b"), + id="fine-tune-raw-validation-file-does-not-narrow", + ), + pytest.param( + "acreate_fine_tuning_job", + {"training_file": {"model-a": "f-1"}, "validation_file": {"model-b": "f-3"}}, + ("model-a", "model-b", "model-c"), + id="fine-tune-files-on-disjoint-deployments-keep-every-deployment", + ), + pytest.param( + "acreate_fine_tuning_job", + {"training_file": {"model-z": "f-5"}}, + ("model-a", "model-b", "model-c"), + id="fine-tune-file-of-another-model-group-keeps-every-deployment", + ), + pytest.param("acreate_batch", {"input_file_id": {"model-c": "f-4"}}, ("model-c",), id="batch-input-file"), + pytest.param( + "acreate_batch", {"input_file_id": {"model-z": "f-6"}}, (), id="batch-input-file-of-another-model-group" + ), + ], +) +async def test_create_from_managed_files_routes_only_to_deployments_holding_every_file( + call_type: str, files: Mapping[str, Mapping[str, str] | str], expected: tuple[str, ...] +) -> None: + owned: Final = { + field: _owned_managed_file(str(uuid.uuid5(uuid.NAMESPACE_URL, field)), mappings) + for field, mappings in files.items() + if not isinstance(mappings, str) + } + managed_files, _ = _managed_files_with_fake_prisma(*(row for _, row in owned.values())) + request: Final = { + "model": "gpt-4o-mini", + **{field: owned[field][0] if field in owned else str(mappings) for field, mappings in files.items()}, + } + + assert await _routable_deployment_ids(managed_files, call_type, request) == expected + + _MANAGED_READ_SALT_KEY: Final = "managed-read-unit-test-salt-key" diff --git a/tests/unit/fine_tuning/__init__.py b/tests/unit/fine_tuning/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/fine_tuning/test_main.py b/tests/unit/fine_tuning/test_main.py new file mode 100644 index 00000000000..57c845c6e1a --- /dev/null +++ b/tests/unit/fine_tuning/test_main.py @@ -0,0 +1,54 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm + +API_BASE: Final = "https://fine-tuning.test" +BASE_MODEL: Final = "gpt-4.1-nano-2025-04-14" +PROVIDER_JOB: Final = { + "id": "ftjob-abc123", + "object": "fine_tuning.job", + "created_at": 1700000000, + "error": None, + "fine_tuned_model": None, + "finished_at": None, + "hyperparameters": {"n_epochs": 3}, + "model": BASE_MODEL, + "organization_id": "org-test", + "result_files": [], + "seed": 42, + "status": "validating_files", + "trained_tokens": None, + "training_file": "file-abc123", + "validation_file": None, +} + + +@pytest.mark.asyncio +@respx.mock +@pytest.mark.parametrize( + ("custom_llm_provider", "route"), + [("openai", f"{API_BASE}/v1/fine_tuning/jobs"), ("azure", f"{API_BASE}/openai/fine_tuning/jobs")], +) +@pytest.mark.parametrize("prefixed", [True, False], ids=["provider-prefixed", "bare"]) +async def test_create_sends_the_model_without_the_provider_prefix( + monkeypatch: pytest.MonkeyPatch, custom_llm_provider: str, route: str, prefixed: bool +) -> None: + sent: Final = respx.post(url__startswith=route).mock(return_value=httpx.Response(200, json=PROVIDER_JOB)) + + async with httpx.AsyncClient() as session: + monkeypatch.setattr(litellm, "aclient_session", session) + await litellm.acreate_fine_tuning_job( + model=f"{custom_llm_provider}/{BASE_MODEL}" if prefixed else BASE_MODEL, + training_file="file-abc123", + custom_llm_provider=custom_llm_provider, + api_key="sk-test", + api_base=f"{API_BASE}/v1" if custom_llm_provider == "openai" else API_BASE, + **({"api_version": "2024-10-21"} if custom_llm_provider == "azure" else {}), + ) + + assert json.loads(sent.calls.last.request.content)["model"] == BASE_MODEL diff --git a/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py b/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py index 5447d553f30..79c96b8970a 100644 --- a/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py +++ b/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py @@ -45,7 +45,8 @@ class TestCustomLogger(CustomLogger): @pytest.mark.asyncio() -async def test_create_vertex_fine_tune_jobs_mocked(): +@pytest.mark.parametrize("model_prefix", ["", "vertex_ai/"], ids=["bare", "provider-prefixed"]) +async def test_create_vertex_fine_tune_jobs_mocked(model_prefix: str): # Define reusable variables for the test project_id = "633608382793" location = "us-central1" @@ -90,7 +91,7 @@ async def test_create_vertex_fine_tune_jobs_mocked(): ), ): create_fine_tuning_response = await litellm.acreate_fine_tuning_job( - model=base_model, + model=f"{model_prefix}{base_model}", custom_llm_provider="vertex_ai", training_file=training_file, vertex_project=project_id, diff --git a/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py index 99a79801698..06ac56b09ad 100644 --- a/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py +++ b/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py @@ -37,9 +37,9 @@ RAW_FILE_ID = "file-victim-abc123" RAW_JOB_ID = "ftjob-victim-abc123" -def _unified_file_id() -> str: +def _unified_file_id(unified_id: str = "victim-unified-id", provider_file_id: str = RAW_FILE_ID) -> str: unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + "application/json", unified_id, "gpt-4o-mini", provider_file_id, "gpt-4o-mini-id" ) return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") @@ -54,7 +54,7 @@ def _decode_unified_id(encoded_id: str) -> str: return base64.urlsafe_b64decode(f"{encoded_id}{padding}").decode() -def _job() -> LiteLLMFineTuningJob: +def _job(validation_file: str | None = None) -> LiteLLMFineTuningJob: job = LiteLLMFineTuningJob( id=RAW_JOB_ID, created_at=1234567890, @@ -69,7 +69,7 @@ def _job() -> LiteLLMFineTuningJob: status="running", trained_tokens=None, training_file=RAW_FILE_ID, - validation_file=None, + validation_file=validation_file, ) job._hidden_params = {} return job @@ -225,6 +225,31 @@ async def test_create__unified_training_file_allowed_when_managed_files_required assert response.hidden_params["unified_file_id"] == _decode_unified_id(unified_file_id) +MANAGED_VALIDATION_FILE_ID: Final = _unified_file_id("validation-unified-id", "file-provider-validation") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("validation_file", "provider_validation_file", "expected_validation_file"), + [ + (None, None, None), + (RAW_FILE_ID, RAW_FILE_ID, RAW_FILE_ID), + (MANAGED_VALIDATION_FILE_ID, "file-provider-validation", MANAGED_VALIDATION_FILE_ID), + ], + ids=["no-validation-file", "raw-validation-file", "managed-validation-file"], +) +async def test_create__managed_job_returns_the_callers_file_ids( + seams, validation_file: str | None, provider_validation_file: str | None, expected_validation_file: str | None +): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + seams.router.acreate_fine_tuning_job.return_value = _job(validation_file=provider_validation_file) + training_file: Final = _unified_file_id() + + response: Final = await _create(training_file, validation_file=validation_file) + + assert (response.training_file, response.validation_file) == (training_file, expected_validation_file) + + @pytest.mark.asyncio async def test_create__other_teams_unified_training_file_rejected(seams): seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) From cf9ebeda073bd03e6479558f7b683904054f2301 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:04:25 +0000 Subject: [PATCH 03/18] fix(proxy): stop billing every OpenAI pass-through /responses sub-route as a Responses call (#47171) * test(proxy): cover OpenAI pass-through billing for each /responses sub-route * fix(proxy): bill only the OpenAI pass-through Responses calls OpenAI bills * fix(proxy): fail closed on unknown /responses sub-routes and label unbilled spend rows Bill every OpenAI pass-through /responses call except input_tokens, input_items, delete, and a non-background retrieve or cancel, so an encoded or new sub-route is never free. Unbilled calls keep the response's model on their spend row, and a resumed background stream takes its model from the stream's final event. * fix(proxy): give billed background Responses reads their own spend log row A billed retrieve or cancel of a finished background response reused the create's resp_ id, so the spend log insert dropped the row that carried the cost. It now takes the native path's _background_cost suffix. Adds Azure path coverage and types the test recorder. * fix(proxy): bill a background Responses read only once it carries usage An in-progress poll took the _background_cost spend log id with $0, so the finished read's row that carried the cost was dropped again. Polls with no usage are now unbilled like any other read. * fix(proxy): keep OpenAI's response id for callbacks on a billed background read * fix(openai): give a streamed background create the settlement spend row id * fix(proxy): keep a request body from giving its spend row the background settlement id * test(proxy): give the forged settlement marker test a fixed time * fix(passthrough): move the background settlement marker into the pass-through types module A test that reloads spend_tracking_utils gave it a new enum class, so the handler's marker no longer matched by identity and the settlement row lost its _background_cost id. The types module is never reloaded. --- .../openai_passthrough_logging_handler.py | 88 +++- .../spend_tracking/spend_tracking_utils.py | 17 +- .../pass_through_endpoints.py | 6 + ...test_openai_passthrough_logging_handler.py | 432 +++++++++++++++++- .../test_spend_tracking_utils.py | 27 ++ 5 files changed, 550 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 2e9d92ed309..c763b5dc26e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -6,7 +6,7 @@ Handles cost tracking and logging for OpenAI passthrough endpoints, specifically from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Final +from typing import Final, Literal, TypeAlias from urllib.parse import urlparse import httpx @@ -14,6 +14,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.hidden_params import set_hidden_param +from litellm.litellm_core_utils.internal_call_metadata import is_background_response from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, @@ -31,6 +32,8 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( ) from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, EndpointType, PassthroughStandardLoggingPayload, ) @@ -56,7 +59,7 @@ _AZURE_OPENAI_HOSTNAMES: Final = ("openai.azure.com", "cognitiveservices.azure.c _AZURE_OPENAI_PATH_MARKERS: Final = ("/openai/", "/v1/") -def _hostname_matches(hostname: str, suffixes: tuple) -> bool: +def _hostname_matches(hostname: str, suffixes: tuple[str, ...]) -> bool: """True if hostname equals one of `suffixes` or is a subdomain of it. Uses suffix matching (not a bare substring test) so look-alikes such as @@ -102,6 +105,19 @@ def is_openai_compatible_url(url_route: str | None) -> bool: _is_openai_compatible_url: Final = is_openai_compatible_url +ResponsesCallBilling: TypeAlias = Literal["unbilled", "billed", "background_settlement"] + + +def _is_finished_background_response(response_body: Mapping[str, object]) -> bool: + return is_background_response(response_body) and response_body.get("usage") is not None + + +def _responses_sub_route(path: str) -> tuple[str, ...] | None: + segments: Final = path.rstrip("/").split("/") + anchor: Final = max((index for index, segment in enumerate(segments) if segment == "responses"), default=None) + return None if anchor is None else tuple(segments[anchor + 1 :]) + + def _is_remote_high_detail_image(part: object) -> bool: if not isinstance(part, Mapping) or part.get("type") != "image_url": return False @@ -186,6 +202,30 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): "/v1/responses" in parsed_url.path or "/responses" in parsed_url.path ) + @staticmethod + def openai_responses_call_billing( + method: str, url_route: str, response_body: Mapping[str, object] + ) -> ResponsesCallBilling: + if not OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route): + return "unbilled" + match (method, _responses_sub_route(urlparse(url_route).path)): + case ("POST", ("input_tokens",)) | ("GET", (_, "input_items")) | ("DELETE", (_,)): + return "unbilled" + case ("GET", (_,)) | ("POST", (_, "cancel")): + return "background_settlement" if _is_finished_background_response(response_body) else "unbilled" + case _: + return "billed" + + @staticmethod + def _label_unbilled_responses_call( + logging_obj: LiteLLMLoggingObj, + url_route: str, + response_body: Mapping[str, object], + ) -> None: + model: Final = response_body.get("model") + if OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) and isinstance(model, str): + logging_obj.model_call_details["model"] = model # rebind-ok: spend rows read the model from here + @staticmethod def is_openai_embeddings_route(url_route: str) -> bool: """Check if the URL route is an OpenAI embeddings endpoint.""" @@ -348,9 +388,13 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): is_embeddings: Final = OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(url_route) is_image_generation: Final = OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) is_image_editing: Final = OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) - is_responses: Final = OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) + responses_billing: Final = OpenAIPassthroughLoggingHandler.openai_responses_call_billing( + httpx_response.request.method, url_route, response_body + ) + is_responses: Final = responses_billing != "unbilled" if not (is_chat_completions or is_embeddings or is_image_generation or is_image_editing or is_responses): + OpenAIPassthroughLoggingHandler._label_unbilled_responses_call(logging_obj, url_route, response_body) return { "result": None, "kwargs": kwargs, @@ -505,6 +549,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): logging_obj.model_call_details["model"] = model logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider logging_obj.model_call_details["response_cost"] = response_cost + logging_obj.model_call_details[SETTLES_BACKGROUND_RESPONSE_KWARG] = ( # rebind-ok: spend rows read it here + BackgroundResponseSettlement.SETTLES if responses_billing == "background_settlement" else None + ) endpoint_type: Final = ( "chat_completions" @@ -625,16 +672,24 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): Handle logging for collected OpenAI streaming chunks with cost tracking. """ try: - # Extract model from request body - model: Final = request_body.get("model", "gpt-4o") - is_responses: Final = OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) + terminal_response: Final = ( + OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks(all_chunks=all_chunks) + if is_responses + else None + ) + requested_model: Final[object] = request_body.get("model") + model: Final = ( + requested_model + if isinstance(requested_model, str) and requested_model + else (terminal_response.model if terminal_response is not None else None) or "gpt-4o" + ) # Build complete response from chunks using our streaming handler handler: Final = OpenAIPassthroughLoggingHandler() handler_instance: Final = handler complete_response: Final = ( - OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks(all_chunks=all_chunks) + terminal_response if is_responses else handler._build_complete_streaming_response( all_chunks=all_chunks, @@ -667,17 +722,27 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) ) + settles_background_response: Final = ( + BackgroundResponseSettlement.SETTLES + if terminal_response is not None + and is_background_response(terminal_response) + and terminal_response.usage is not None + else None + ) + # Preserve existing litellm_params to maintain metadata tags existing_litellm_params: Final = litellm_logging_obj.model_call_details.get("litellm_params", {}) or {} + litellm_params: Final = existing_litellm_params.copy() # Prepare kwargs for logging kwargs: Final = { "response_cost": response_cost, + SETTLES_BACKGROUND_RESPONSE_KWARG: settles_background_response, "model": model, "custom_llm_provider": custom_llm_provider, "call_type": litellm_logging_obj.call_type, "messages": litellm_logging_obj.model_call_details.get("messages"), - "litellm_params": existing_litellm_params.copy(), + "litellm_params": litellm_params, } # Extract user information for tracking @@ -689,9 +754,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): passthrough_logging_payload=passthrough_logging_payload, ) if user: - kwargs["litellm_params"].setdefault("proxy_server_request", {}).setdefault("body", {})["user"] = ( - user - ) + litellm_params.setdefault("proxy_server_request", {}).setdefault("body", {})["user"] = user # Attach the payload to kwargs so the success handler adopts it; # its later rebuild runs on a copy whose Responses usage was @@ -712,6 +775,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): litellm_logging_obj.model_call_details["model"] = model litellm_logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider litellm_logging_obj.model_call_details["response_cost"] = response_cost + litellm_logging_obj.model_call_details[SETTLES_BACKGROUND_RESPONSE_KWARG] = ( # rebind-ok: spend row id + settles_background_response + ) verbose_proxy_logger.debug( f"OpenAI streaming passthrough cost tracking - Model: {model}, Cost: ${response_cost:.6f}" diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 78818cde266..091fb5d1305 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -69,6 +69,10 @@ from litellm.proxy._types import ( from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.router import DeploymentTypedDict, LiteLLM_Params from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -315,10 +319,12 @@ BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" -def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object]) -> str: +def _cost_row_request_id_suffix( + call_type: str, metadata: Mapping[str, object], settles_background_response: bool +) -> str: if call_type == CallTypes.aretrieve_batch.value: return BATCH_COST_REQUEST_ID_SUFFIX - if ( + if settles_background_response or ( call_type in NON_INFERENCE_CALL_TYPES and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN ): @@ -327,7 +333,7 @@ def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object]) def get_spend_logs_id( - call_type: str, response_obj: dict, kwargs: dict, metadata: Mapping[str, object] = EMPTY_MAPPING + call_type: str, response_obj: dict, kwargs: Mapping[str, object], metadata: Mapping[str, object] = EMPTY_MAPPING ) -> str | None: standard_logging_payload = kwargs.get("standard_logging_object") candidate_ids: Final = ( @@ -340,7 +346,10 @@ def get_spend_logs_id( ) if resolved_id is None: return None - return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata)}" + settles_background_response: Final = ( + kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES + ) + return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, settles_background_response)}" _MISSING_ATTRIBUTE: Final = object() diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index e634354408f..0d99cd54a6b 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -25,6 +25,12 @@ LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: Final = "litellm_pass_thro # on a collision FastAPI dispatches the built-in handler, which does not carry this flag. LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__" +SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response" + + +class BackgroundResponseSettlement(str, Enum): + SETTLES = "settles" + class EndpointType(str, Enum): VERTEX_AI = "vertex-ai" diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 86a670bcb3e..c9064caf929 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -1,19 +1,26 @@ +from dataclasses import dataclass +import asyncio import json +import logging +from collections.abc import Mapping from datetime import datetime -from typing import Any, Dict, List +from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import BaseModel, ConfigDict import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, count_relayed_prompt_tokens, ) +from litellm.proxy._types import SpendLogsPayload from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -1837,10 +1844,7 @@ class TestOpenAIPassthroughResponsesStreamingSpendLog: def _expected_spend(self) -> float: rates = litellm.model_cost[self.MODEL_MAP_KEY] - return ( - self.INPUT_TOKENS * rates["input_cost_per_token"] - + self.OUTPUT_TOKENS * rates["output_cost_per_token"] - ) + return self.INPUT_TOKENS * rates["input_cost_per_token"] + self.OUTPUT_TOKENS * rates["output_cost_per_token"] def _responses_stream_chunks(self) -> List[str]: """Usage arrives only on the terminal `response.completed` event, nested under @@ -1953,6 +1957,28 @@ class TestOpenAIPassthroughResponsesStreamingSpendLog: row_metadata = json.loads(spend_log_row["metadata"]) assert row_metadata["model_map_information"]["model_map_key"] == self.MODEL_MAP_KEY + def test_resumed_responses_stream_spend_log_is_labeled_with_the_streamed_model(self, local_model_cost_map): + result: Final = OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( + litellm_logging_obj=self._logging_obj(), + passthrough_success_handler_obj=None, + url_route=f"https://api.openai.com/v1/responses/{self.RESPONSE_ID}?stream=true", + request_body={}, + endpoint_type=None, + start_time=self.start_time, + all_chunks=self._responses_stream_chunks(), + end_time=self.end_time, + ) + + spend_log_row: Final = get_logging_payload( + kwargs=result["kwargs"], + response_obj=result["result"], + start_time=self.start_time, + end_time=self.end_time, + ) + + assert spend_log_row["model"] == self.MODEL_MAP_KEY + assert spend_log_row["spend"] == pytest.approx(self._expected_spend()) + class TestOpenAIPassthroughEmbeddingsSpendLog: """An OpenAI-passthrough `/v1/embeddings` call must write a priced spend log row @@ -2043,6 +2069,402 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: assert spend_log_row["request_id"] == self.CALL_ID +RESPONSES_URL: Final = "https://api.openai.com/v1/responses" +STORED_RESPONSE_ID: Final = "resp_080b10ff5d740c69006ad3de732e5087" +DISPATCH_TIME: Final = datetime(2026, 10, 17, 13, 47, 3) +RESPONSES_USAGE: Final = { + "input_tokens": 11, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 5, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 16, +} +CREATE_REQUEST: Final = {"model": "gpt-5.4-mini", "input": "Say hi in one word"} +COMPACT_REQUEST: Final = {"model": "gpt-5.4-mini", "input": "Remember the number 42."} +COMPACTION: Final = { + "id": STORED_RESPONSE_ID, + "object": "response.compaction", + "created_at": 1792270025, + "output": [{"id": "cmp_082492f9", "type": "compaction", "encrypted_content": "gAAAA"}], + "usage": RESPONSES_USAGE, +} +ASSISTANT_MESSAGE: Final = { + "id": "msg_080b10ff5d740c69", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hi", "annotations": []}], +} + + +def _stored_response(status: str, background: bool, usage: Mapping[str, object] | None) -> dict[str, object]: + return { + "id": STORED_RESPONSE_ID, + "object": "response", + "created_at": 1792269939, + "status": status, + "background": background, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "model": "gpt-5.4-mini-2026-03-17", + "output": [ASSISTANT_MESSAGE] if usage else [], + "parallel_tool_calls": True, + "temperature": 1.0, + "tool_choice": "auto", + "tools": [], + "top_p": 0.98, + "usage": usage, + } + + +@dataclass(frozen=True, slots=True) +class _LoggedCall: + spend_row: SpendLogsPayload + callback_response_id: str + + +class _CallbackPayloadId(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + + +class _SpendRowRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.call: Final[asyncio.Future[_LoggedCall]] = asyncio.get_running_loop().create_future() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.call.set_result( + _LoggedCall( + spend_row=get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=start_time, end_time=end_time + ), + callback_response_id=_CallbackPayloadId.model_validate(kwargs.get("standard_logging_object")).id, + ) + ) + + +class TestOpenAIPassthroughResponsesSubRouteBilling: + def _dispatch( + self, + method: str, + sub_route: str, + request_body: dict[str, object], + response_body: dict[str, object], + base_url: str = RESPONSES_URL, + ): + url_route: Final = f"{base_url}{sub_route}" + logging_obj: Final = LiteLLMLoggingObj( + model="gpt-5.4-mini", + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=DISPATCH_TIME, + litellm_call_id="6c1f3a52-93d4-4b1e-9f0e-6f2f4f6b9a10", + function_id="1234", + ) + return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=httpx.Response(200, json=response_body, request=httpx.Request(method, url_route)), + response_body=response_body, + request_body=request_body, + logging_obj=logging_obj, + url_route=url_route, + result="", + start_time=DISPATCH_TIME, + end_time=DISPATCH_TIME, + cache_hit=False, + custom_llm_provider="openai", + litellm_params={}, + call_type="pass_through_endpoint", + ) + + @pytest.mark.parametrize( + ("method", "sub_route", "request_body", "response_body"), + [ + pytest.param( + "POST", "", CREATE_REQUEST, _stored_response("completed", False, RESPONSES_USAGE), id="create" + ), + pytest.param( + "POST", + "/", + CREATE_REQUEST, + _stored_response("completed", False, RESPONSES_USAGE), + id="create_trailing_slash", + ), + pytest.param("POST", "/compact", COMPACT_REQUEST, COMPACTION, id="compact"), + pytest.param("POST", "/%2563ompact", COMPACT_REQUEST, COMPACTION, id="percent_encoded_compact"), + pytest.param( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + _stored_response("completed", True, RESPONSES_USAGE), + id="background_retrieve_completed", + ), + pytest.param( + "POST", + f"/{STORED_RESPONSE_ID}/cancel", + {}, + _stored_response("cancelled", True, RESPONSES_USAGE), + id="background_cancel", + ), + ], + ) + def test_billed_responses_call_records_cost( + self, + local_model_cost_map, + method: str, + sub_route: str, + request_body: dict[str, object], + response_body: dict[str, object], + ): + dispatched: Final = self._dispatch(method, sub_route, request_body, response_body) + + assert dispatched["standard_logging_response_object"] is not None + assert dispatched["kwargs"]["response_cost"] > 0 + + @pytest.mark.parametrize( + ("method", "sub_route", "request_body", "response_body"), + [ + pytest.param( + "POST", + "/input_tokens", + {"model": "gpt-5.4-mini", "input": "hi"}, + {"object": "response.input_tokens", "input_tokens": 7}, + id="input_tokens", + ), + pytest.param( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + _stored_response("completed", False, RESPONSES_USAGE), + id="retrieve", + ), + pytest.param( + "GET", + f"/{STORED_RESPONSE_ID}/input_items", + {}, + { + "object": "list", + "data": [ + { + "id": "msg_074176e3", + "type": "message", + "status": "completed", + "role": "user", + "content": [{"type": "input_text", "text": "Say hi in one word"}], + } + ], + "first_id": "msg_074176e3", + "last_id": "msg_074176e3", + "has_more": False, + }, + id="input_items", + ), + pytest.param( + "DELETE", + f"/{STORED_RESPONSE_ID}", + {}, + {"id": STORED_RESPONSE_ID, "object": "response.deleted", "deleted": True}, + id="delete", + ), + pytest.param( + "POST", + f"/{STORED_RESPONSE_ID}/cancel", + {}, + _stored_response("cancelled", True, None), + id="background_cancel_before_any_output", + ), + ], + ) + def test_unbilled_responses_sub_route_records_no_cost_and_no_error( + self, + local_model_cost_map, + caplog: pytest.LogCaptureFixture, + method: str, + sub_route: str, + request_body: dict[str, object], + response_body: dict[str, object], + ): + with caplog.at_level(logging.WARNING): + dispatched: Final = self._dispatch(method, sub_route, request_body, response_body) + + assert not dispatched["kwargs"].get("response_cost") + assert [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] == [] + + async def _logged_call( + self, + method: str, + sub_route: str, + request_body: dict[str, object], + response_body: dict[str, object], + litellm_call_id: str, + ) -> _LoggedCall: + url_route: Final = f"{RESPONSES_URL}{sub_route}" + recorder: Final = _SpendRowRecorder() + logging_obj: Final = LiteLLMLoggingObj( + model="unknown", + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=DISPATCH_TIME, + litellm_call_id=litellm_call_id, + function_id="1234", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="unknown", user="unknown", optional_params={}, litellm_params={}, call_type="pass_through_endpoint" + ) + + await PassThroughEndpointLogging().pass_through_async_success_handler( + httpx_response=httpx.Response(200, json=response_body, request=httpx.Request(method, url_route)), + response_body=response_body, + logging_obj=logging_obj, + url_route=url_route, + result="", + start_time=DISPATCH_TIME, + end_time=DISPATCH_TIME, + cache_hit=False, + request_body=request_body, + passthrough_logging_payload=PassthroughStandardLoggingPayload( + url=url_route, request_body=request_body, request_method=method + ), + custom_llm_provider="openai", + ) + return await asyncio.wait_for(recorder.call, timeout=5) + + @pytest.mark.asyncio + async def test_unbilled_retrieve_spend_log_keeps_the_response_model(self, local_model_cost_map): + response_body: Final = _stored_response("completed", False, RESPONSES_USAGE) + + spend_log_row: Final = ( + await self._logged_call( + "GET", f"/{STORED_RESPONSE_ID}", {}, response_body, "0f6f3d2e-6a51-4c4e-9a43-6f1b0d1e9c21" + ) + ).spend_row + + assert spend_log_row["model"] == response_body["model"] + assert not spend_log_row["spend"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("method", "sub_route", "response_body"), + [ + pytest.param( + "GET", + f"/{STORED_RESPONSE_ID}", + _stored_response("completed", True, RESPONSES_USAGE), + id="background_retrieve_completed", + ), + pytest.param( + "POST", + f"/{STORED_RESPONSE_ID}/cancel", + _stored_response("cancelled", True, RESPONSES_USAGE), + id="background_cancel", + ), + ], + ) + async def test_billed_background_read_spend_log_survives_next_to_its_create_and_poll_rows( + self, local_model_cost_map, method: str, sub_route: str, response_body: dict[str, object] + ): + create: Final = await self._logged_call( + "POST", + "", + {**CREATE_REQUEST, "background": True}, + _stored_response("queued", True, None), + "3a4f0c1e-2b7d-4c55-8e0e-0d7a51d9f6a1", + ) + poll: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + _stored_response("in_progress", True, None), + "5d8c2f71-0e4b-4a9d-b3c6-7f1e2a9d0c44", + ) + read: Final = await self._logged_call( + method, sub_route, {}, response_body, "9b2e6d40-71c3-4f8a-a5d2-4e3c8f0b7d12" + ) + + assert read.spend_row["request_id"] not in {create.spend_row["request_id"], poll.spend_row["request_id"]} + assert read.spend_row["request_id"].startswith(STORED_RESPONSE_ID) + assert read.spend_row["spend"] > 0 + assert not poll.spend_row["spend"] + assert read.callback_response_id == STORED_RESPONSE_ID + + @pytest.mark.asyncio + async def test_streamed_background_create_and_its_finished_read_share_one_spend_row(self, local_model_cost_map): + finished: Final = _stored_response("completed", True, RESPONSES_USAGE) + stream_logging_obj: Final = LiteLLMLoggingObj( + model="gpt-5.4-mini", + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=DISPATCH_TIME, + litellm_call_id="7e3b9a14-5c2d-4f60-8d1a-2b6c0e9f4a37", + function_id="1234", + ) + stream_logging_obj.model_call_details["custom_llm_provider"] = "openai" + streamed_create: Final = OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( + litellm_logging_obj=stream_logging_obj, + passthrough_success_handler_obj=None, + url_route=RESPONSES_URL, + request_body={**CREATE_REQUEST, "background": True, "stream": True}, + endpoint_type=None, + start_time=DISPATCH_TIME, + all_chunks=[ + f"data: {json.dumps({'type': 'response.completed', 'sequence_number': 1, 'response': finished})}", + "data: [DONE]", + ], + end_time=DISPATCH_TIME, + ) + create_row: Final = get_logging_payload( + kwargs=streamed_create["kwargs"], + response_obj=streamed_create["result"], + start_time=DISPATCH_TIME, + end_time=DISPATCH_TIME, + ) + + read: Final = await self._logged_call( + "GET", f"/{STORED_RESPONSE_ID}", {}, finished, "1c8d5e2a-9f47-4b3e-a6d0-5e7f2c1b8a93" + ) + + assert create_row["spend"] > 0 + assert read.spend_row["request_id"] == create_row["request_id"] + assert read.callback_response_id == STORED_RESPONSE_ID + + @pytest.mark.parametrize( + "base_url", + [ + pytest.param("https://my-resource.openai.azure.com/openai/v1/responses", id="azure_openai_v1"), + pytest.param("https://my-resource.cognitiveservices.azure.com/openai/responses", id="azure_openai_native"), + ], + ) + def test_azure_responses_sub_routes_bill_like_openai(self, local_model_cost_map, base_url: str): + retrieve: Final = self._dispatch( + "GET", f"/{STORED_RESPONSE_ID}", {}, _stored_response("completed", False, RESPONSES_USAGE), base_url + ) + input_tokens: Final = self._dispatch( + "POST", + "/input_tokens", + {"model": "gpt-5.4-mini", "input": "hi"}, + {"object": "response.input_tokens", "input_tokens": 7}, + base_url, + ) + background_retrieve: Final = self._dispatch( + "GET", f"/{STORED_RESPONSE_ID}", {}, _stored_response("completed", True, RESPONSES_USAGE), base_url + ) + compact: Final = self._dispatch("POST", "/compact", COMPACT_REQUEST, COMPACTION, base_url) + + assert not retrieve["kwargs"].get("response_cost") + assert not input_tokens["kwargs"].get("response_cost") + assert background_retrieve["kwargs"]["response_cost"] > 0 + assert compact["kwargs"]["response_cost"] > 0 + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 34e17b3cda8..11a881527e6 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -47,6 +47,10 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( spend_log_row_with_retained_metadata, ) from litellm.proxy.utils import hash_token +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -4810,6 +4814,29 @@ def test_a_background_response_billing_event_writes_a_row_apart_from_the_read_th assert request_ids == ("resp_shared", "resp_shared_background_cost") +@pytest.mark.parametrize("forged_marker", [True, BackgroundResponseSettlement.SETTLES.value]) +def test_a_request_body_cannot_move_its_spend_row_onto_a_background_settlement_id(forged_marker): + """A failed call logs its raw request body as kwargs, so a body field named like the settlement + marker must not give the failure row the id a background response's billing row owns.""" + request_body: Final = { + "model": "gpt-5.4-mini", + "input": "hi", + SETTLES_BACKGROUND_RESPONSE_KWARG: forged_marker, + "litellm_call_id": "call-failed", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + } + failed_at: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc) + + row: Final = get_logging_payload( + kwargs=request_body, + response_obj=Exception("upstream rejected the request"), + start_time=failed_at, + end_time=failed_at, + ) + + assert row["request_id"] == "call-failed" + + def test_get_spend_logs_id_prefers_the_response_id_over_the_standard_logging_id(): """An unredacted response keeps deciding its own row key, so cache-hit ids and every other call type behave exactly as they did before.""" From b110c8cfa7422be35753eb550807e1a7bdf9375b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 19 Oct 2026 15:06:02 -0700 Subject: [PATCH 04/18] test(integration): check SSO sign-in removal without reaching the public internet (#47409) The SSO removal test configured Google sign-in, which fetches its discovery document from accounts.google.com before redirecting. CircleCI blocks that egress, so the test failed on every main run since #46709 added it. Microsoft sign-in builds its endpoints locally and goes through the same SSO environment variable handling, so the test now uses it --- tests/integration/configuration/test_effective_settings.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/integration/configuration/test_effective_settings.py b/tests/integration/configuration/test_effective_settings.py index acdaacfcfb5..07477f1f048 100644 --- a/tests/integration/configuration/test_effective_settings.py +++ b/tests/integration/configuration/test_effective_settings.py @@ -1547,8 +1547,8 @@ def test_an_sso_client_removed_from_db_environment_variables_stops_sign_in_on_ev "DATABASE_URL_READ_REPLICA", "GENERIC_CLIENT_ID", "GOOGLE_CLIENT_ID", - "GOOGLE_CLIENT_SECRET", "MICROSOFT_CLIENT_ID", + "MICROSOFT_CLIENT_SECRET", "PROXY_BASE_URL", ) with ( @@ -1567,7 +1567,7 @@ def test_an_sso_client_removed_from_db_environment_variables_stops_sign_in_on_ev response: Final = first.request( "POST", "/config/update", - {"environment_variables": {"GOOGLE_CLIENT_ID": client_id, "GOOGLE_CLIENT_SECRET": secret}}, + {"environment_variables": {"MICROSOFT_CLIENT_ID": client_id, "MICROSOFT_CLIENT_SECRET": secret}}, ) assert response.status_code == 200, response.text @@ -1584,7 +1584,7 @@ def test_an_sso_client_removed_from_db_environment_variables_stops_sign_in_on_ev moved: Final = first.request( "PATCH", "/update/sso_settings", - {"google_client_id": sso_table_client, "google_client_secret": "sso-table-secret"}, + {"microsoft_client_id": sso_table_client, "microsoft_client_secret": "sso-table-secret"}, ) assert moved.status_code == 200, moved.text eventually(sign_in_clients, lambda clients: clients == (sso_table_client,) * 2, seconds=30) From 77bd7d6de841bdc746d6b45a0ba98d5f60eb7af1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 19 Oct 2026 15:06:17 -0700 Subject: [PATCH 05/18] test(integration): expect the requested model in streamed Responses events on the arize v2 rig (#47407) #46662 made streamed /v1/responses lifecycle events report the requested model in the nested response.model, matching the non-streamed body. The arize v2 caller-stream helper still expected the upstream deployment name there, so f1 and a6 failed on every main run since bbd446b476 --- tests/integration/observability/_openinference_support.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index a57e25f88ae..b46503703c7 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -386,13 +386,14 @@ def _responses_caller_body(body: dict[str, JsonValue], model: str) -> dict[str, } -def _responses_stream_caller_body(body: dict[str, JsonValue]) -> dict[str, JsonValue]: +def _responses_stream_caller_body(body: dict[str, JsonValue], model: str) -> dict[str, JsonValue]: response_id: Final = body.get("id") assert isinstance(response_id, str), body usage: Final = _json_object_value(body["usage"]) if body.get("status") == "completed" else {} return { **body, "id": "", + "model": model, **({"usage": {**usage, "cost": 4.05e-06}} if usage else {}), } @@ -413,7 +414,7 @@ def _responses_caller_stream(reply: Reply, model: str) -> tuple[dict[str, JsonVa **event, "model": model, **( - {"response": _responses_stream_caller_body(_json_object_value(event["response"]))} + {"response": _responses_stream_caller_body(_json_object_value(event["response"]), model)} if "response" in event else {} ), From 1ef59fab8b7621730384ba88ce37f5dd05c0bc65 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 19 Oct 2026 15:06:35 -0700 Subject: [PATCH 06/18] test(ui-e2e): read the virtual key from its own pre in the Save your Key dialog (#47408) #47199 added a Gateway URL pre to the created-key dialog, so the bare pre locator matched two elements and failed strict mode in four specs --- tests/e2e/ui/tests/internal-user/internalUser.spec.ts | 2 +- tests/e2e/ui/tests/proxy-admin/keys.spec.ts | 2 +- tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts | 2 +- tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index d34bc034453..cb7d9320744 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -67,7 +67,7 @@ test.describe("Internal User", () => { await page.getByRole("button", { name: "Create Key", exact: true }).click(); await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); - const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); + const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").filter({ hasText: /^sk-/ }).innerText()).trim(); expect(apiKey).toMatch(/^sk-/); await page.keyboard.press("Escape"); diff --git a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts index deb7ae70d07..d3d67bde1cf 100644 --- a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts @@ -254,7 +254,7 @@ test.describe("Proxy Admin - Keys", () => { // verify it can call /chat/completions for the model it was scoped to. // The mock LLM server (fixtures/mock_llm_server/server.py) replies with // a fixed "This is a mock response." body. - const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); + const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").filter({ hasText: /^sk-/ }).innerText()).trim(); expect(apiKey).toMatch(/^sk-/); const response = await page.request.post("/chat/completions", { diff --git a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts index 4572f71b3db..835a26af0dc 100644 --- a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts @@ -69,7 +69,7 @@ test.describe("Second proxy admin", () => { await page.getByRole("button", { name: "Create Key", exact: true }).click(); await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); - const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); + const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").filter({ hasText: /^sk-/ }).innerText()).trim(); expect(apiKey).toMatch(/^sk-/); await page.keyboard.press("Escape"); diff --git a/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts index b7978fae7fb..95000b0ff18 100644 --- a/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts +++ b/tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts @@ -130,7 +130,7 @@ test.describe("Team Admin - Member permissions", () => { await memberPage.getByRole("button", { name: "Create Key", exact: true }).click(); const saveDialog = memberPage.getByRole("dialog", { name: "Save your Key" }); await expect(saveDialog).toBeVisible({ timeout: 15_000 }); - const apiKey = (await saveDialog.locator("pre").innerText()).trim(); + const apiKey = (await saveDialog.locator("pre").filter({ hasText: /^sk-/ }).innerText()).trim(); expect(apiKey).toMatch(/^sk-/); createdKeys.push(apiKey); await memberPage.keyboard.press("Escape"); From 09e919351437851ea0ba689f7e878a7395627f57 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 19 Oct 2026 15:14:46 -0700 Subject: [PATCH 07/18] fix(cost): bill Gemini audio/image/video prompt tokens at the resolved input rate when no modality rate exists (#47373) * fix(cost): bill Gemini audio/image/video prompt tokens at the resolved input rate when no modality rate exists Co-authored-by: shivam * style(cost): ruff format the modality cost calls --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: shivam --- .../litellm_core_utils/llm_cost_calc/utils.py | 52 +++++++++++---- .../cost_calculation/cases/vertex_ai.py | 64 +++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 60 +++++++++++++++++ 3 files changed, 164 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 2afb6ed4c58..c6141be008e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1088,6 +1088,25 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu ) +def _modality_token_cost( + model_info: ModelInfo, + modality_cost_key: str, + tokens: int, + prompt_base_cost: float, + service_tier: str | None, +) -> float: + """ + Audio, image and video prompt tokens bill at their modality rate when the model has one. + Without one they are ordinary prompt tokens, so they bill at `prompt_base_cost`, which + `_get_token_base_cost` has already resolved for the service tier and long-context tier. + """ + tiered_key: Final = get_service_tier_cost_key(modality_cost_key, service_tier) + cost_key: Final = tiered_key if model_info.get(tiered_key) is not None else modality_cost_key + if model_info.get(cost_key) is None: + return float(tokens) * prompt_base_cost + return calculate_cost_component(model_info, cost_key, tokens) + + def _calculate_input_cost( prompt_tokens_details: PromptTokensDetailsResult, model_info: ModelInfo, @@ -1118,28 +1137,37 @@ def _calculate_input_cost( if prompt_tokens_details["audio_tokens"] and not ( prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None ): - audio_cost_key: Final = get_service_tier_cost_key("input_cost_per_audio_token", service_tier) - prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"]) + prompt_cost += _modality_token_cost( + model_info, + "input_cost_per_audio_token", + prompt_tokens_details["audio_tokens"], + prompt_base_cost, + service_tier, + ) ### IMAGE TOKEN COST if prompt_tokens_details["image_tokens"] and not ( prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None ): - # For image token costs: - # First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token. - image_token_cost_key = "input_cost_per_image_token" - if model_info.get(image_token_cost_key) is None: - image_token_cost_key = "input_cost_per_token" - prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"]) + prompt_cost += _modality_token_cost( + model_info, + "input_cost_per_image_token", + prompt_tokens_details["image_tokens"], + prompt_base_cost, + service_tier, + ) ### VIDEO TOKEN COST if prompt_tokens_details["video_tokens"] and not ( prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None ): - video_token_cost_key = "input_cost_per_video_token" - if model_info.get(video_token_cost_key) is None: - video_token_cost_key = "input_cost_per_token" - prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"]) + prompt_cost += _modality_token_cost( + model_info, + "input_cost_per_video_token", + prompt_tokens_details["video_tokens"], + prompt_base_cost, + service_tier, + ) ### CACHE WRITING COST - Now uses tiered pricing if ( diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 3b825bfe7b2..305c10535b7 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -946,6 +946,69 @@ GEMINI_3_1_PRO_FALLBACK_VIDEO_TOKENS_AT_INPUT_RATE: Final = CostTrackingTestCase ), ) +GEMINI_3_1_PRO_FALLBACK_VIDEO_TOKENS_ABOVE_200K: Final = CostTrackingTestCase( + name="gemini-3.1-pro-fallback_video_tokens_above_200k", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gemini-3.1-pro", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a7c1e02d9b44 summarize the attached hour-long recording in one line", + }, + { + "type": "file", + "file": { + "file_data": "data:video/mp4;base64,AAAAGGZ0eXBpc29tAAACAGlzb21pc282AAACCG1kYXQNFBsiKTA3PkVMU1phaG92fYSLkpmgp661vMPK0djf5u30+wIJEBceJSwzOkFIT1ZdZGtyeYCHjpWco6qxuL/GzdTb4unw9/4FDBMaISgvNj1ES1JZYGdudXyDipGYn6attLvCydDX3uXs8/oBCA8WHSQrMjlAR05VXGNqcXh/ho2Um6KpsLe+xczT2uHo7/b9BAsSGSAnLjU8Q0pRWF9mbXR7gomQl56lrLO6wcjP1t3k6/L5AAcOFRwjKjE4P0ZNVFtiaXB3foWMk5qhqK+2vcTL0tng5+71/AMKERgfJi00O0JJUFdeZWxzeoGIj5adpKuyucDHztXc4+rx+P8GDRQbIikwNz5FTFNaYWhvdn2Ei5KZoKeutbzDytHY3+bt9PsCCRAXHiUsMzpBSE9WXWRrcnmAh46VnKOqsbi/xs3U2+Lp8Pf+BQwTGiEoLzY9REtSWWBnbnV8g4qRmJ+mrbS7wsnQ197l7PP6AQgPFh0kKzI5QEdOVVxjanF4f4aNlJuiqbC3vsXM09rh6O/2/QQLEhkgJy41PENKUVhfZm10e4KJkJeepayzusHIz9bd5Ovy+QAHDhUcIyoxOD9GTVRbYmlwd36FjJOaoaivtr3Ey9LZ4Ofu9fwDChEYHyYtNDtCSVBXXmVsc3qBiI+WnaSrsrnAx87V3OPq8fj/Bg==", + "format": "mp4", + }, + }, + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "candidates": [ + { + "content": {"parts": [{"text": "scripted answer a7c1e02d9b44"}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 250000, + "candidatesTokenCount": 300, + "totalTokenCount": 250300, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 10000}, + {"modality": "VIDEO", "tokenCount": 240000}, + ], + }, + "modelVersion": "gemini-3.1-pro", + }, + ), + expected=ExactExpected( + spend=1.05567, input_cost=1.05, output_cost=0.00567, prompt_tokens=250000, completion_tokens=300 + ), +) + GEMINI_3_1_PRO_STREAM: Final = CostTrackingTestCase( name="gemini-3.1-pro-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", @@ -3587,6 +3650,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_WEB_SEARCH_MEDIUM, GEMINI_3_1_PRO_GOOGLE_MAPS_GROUNDING, GEMINI_3_1_PRO_FALLBACK_VIDEO_TOKENS_AT_INPUT_RATE, + GEMINI_3_1_PRO_FALLBACK_VIDEO_TOKENS_ABOVE_200K, GEMINI_3_1_PRO_STREAM, GEMINI_3_1_PRO_STREAM_NO_USAGE, GEMINI_3_1_PRO_STREAM_NO_USAGE_TOOL_CALL, diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 67a3e2ddd5c..15712cc7d9c 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -450,6 +450,66 @@ def test_video_tokens_fallback_to_base_cost(): assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12) +def test_prompt_audio_and_video_tokens_without_modality_rates_bill_at_tiered_input_rate(): + """Prompt audio/video tokens with no modality rate are prompt tokens, so a >200k prompt bills them above-200k.""" + from unittest.mock import patch + + mock_model_info = { + "input_cost_per_token": 2e-6, + "input_cost_per_token_above_200k_tokens": 4e-6, + "output_cost_per_token": 1.2e-5, + "output_cost_per_token_above_200k_tokens": 1.8e-5, + } + + usage = Usage( + prompt_tokens=816032, + completion_tokens=4559, + total_tokens=820591, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=9033, audio_tokens=14999, video_tokens=792000), + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=4559), + ) + + with patch( + "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", + return_value=mock_model_info, + ): + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", usage=usage, custom_llm_provider="vertex_ai" + ) + + assert round(prompt_cost, 9) == round(816032 * 4e-6, 9) + assert round(completion_cost, 9) == round(4559 * 1.8e-5, 9) + + +def test_prompt_modality_rates_still_win_over_tiered_input_rate(): + """An explicit input_cost_per_audio_token / input_cost_per_video_token beats the fallback, even above 200k.""" + from unittest.mock import patch + + mock_model_info = { + "input_cost_per_token": 2e-6, + "input_cost_per_token_above_200k_tokens": 4e-6, + "input_cost_per_audio_token": 3e-6, + "input_cost_per_video_token": 1e-6, + "output_cost_per_token": 1.2e-5, + } + + usage = Usage( + prompt_tokens=300000, + completion_tokens=10, + total_tokens=300010, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100000, audio_tokens=100000, video_tokens=100000), + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10), + ) + + with patch( + "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", + return_value=mock_model_info, + ): + prompt_cost, _ = generic_cost_per_token(model="test-model", usage=usage, custom_llm_provider="vertex_ai") + + assert round(prompt_cost, 9) == round(100000 * 4e-6 + 100000 * 3e-6 + 100000 * 1e-6, 9) + + def test_generic_cost_per_token_above_200k_tokens(_local_model_cost_map): # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing model = "gemini-2.5-pro" From d7b28840eb2d7c7f33ad831badee05ba02ae6974 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 19 Oct 2026 15:19:40 -0700 Subject: [PATCH 08/18] test(integration): read stream stall spend rows from the proxy that served them (#47410) The cooldown cases in test_hosted_vllm_stream_stall_wire.py boot an owned proxy, which writes to a private database since #47213. The file was written on a base without #47213 and read its rows from the worker database through read_rows, so every cooldown case that waits for a failure row timed out on main with State did not converge: [] Every row read now goes through Gateway.rows on the gateway that served the request --- .../test_hosted_vllm_stream_stall_wire.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py index fb58885e011..00e5388b474 100644 --- a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py +++ b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py @@ -16,7 +16,6 @@ from integration._support.client import ( object_value, string_value, ) -from integration._support.database import read_rows from integration._support.openai_wire import ( answering_model_discovery, chat_stream, @@ -159,14 +158,14 @@ def _error_class(message: JsonValue) -> str: return string_value(message).split(":", 1)[0] -def _failure_row(query: str, groups: tuple[str, ...]) -> dict[str, JsonValue]: - rows: Final = eventually(lambda: read_rows(query, groups), lambda found: len(found) >= 1, seconds=_ROW_SECONDS) +def _failure_row(gateway: Gateway, query: str, groups: tuple[str, ...]) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: gateway.rows(query, groups), lambda found: len(found) >= 1, seconds=_ROW_SECONDS) assert [row["status"] for row in rows] == ["failure"], rows return object_value(object_value(rows[0]["metadata"])["error_information"]) -def _logged_failure(model: str) -> tuple[JsonValue, JsonValue]: - error: Final = _failure_row(_ROWS, (model,)) +def _logged_failure(gateway: Gateway, model: str) -> tuple[JsonValue, JsonValue]: + error: Final = _failure_row(gateway, _ROWS, (model,)) return error["error_code"], error["error_class"] @@ -188,7 +187,7 @@ def test_a_chat_stream_that_stalls_past_its_stream_timeout_ends_with_a_408_timeo events: Final = parse_sse(response.text) assert _chat_errors(events[:-1]) == () and "error" in events[-1].data, response.text error: Final = object_value(events[-1].data["error"]) - observed: Final = (error["code"], _error_class(error["message"]), *_logged_failure(model)) + observed: Final = (error["code"], _error_class(error["message"]), *_logged_failure(gateway, model)) assert observed == ("408", "litellm.Timeout", "408", "Timeout"), (error, observed) assert _upstream_targets(wire) == (target,) @@ -205,7 +204,7 @@ def test_a_native_responses_stream_that_stalls_past_its_stream_timeout_fails_wit assert "response.output_text.delta" in types, response.text assert types[-1] == "response.failed" and "response.completed" not in types, response.text failed: Final = object_value(object_value(events[-1].data["response"])["error"]) - observed: Final = (failed["code"], *_logged_failure(model)) + observed: Final = (failed["code"], *_logged_failure(gateway, model)) assert observed == ("request_timeout", "408", "Timeout"), (failed, observed) assert _upstream_targets(wire) == (_RESPONSES_TARGET,) @@ -222,7 +221,7 @@ def test_a_messages_stream_that_stalls_past_its_stream_timeout_ends_with_a_timeo assert "content_block_delta" in types and "message_stop" not in types, response.text assert types[-1] == "error", response.text error: Final = object_value(events[-1].data["error"]) - observed: Final = (_error_class(error["message"]), *_logged_failure(model)) + observed: Final = (_error_class(error["message"]), *_logged_failure(gateway, model)) assert observed == ("litellm.Timeout", "408", "Timeout"), (error, observed) assert _upstream_targets(wire) == (_CHAT_TARGET,) @@ -238,7 +237,7 @@ def test_a_chat_stream_that_never_pauses_is_answered_in_full_under_the_same_stre assert _chat_errors(events) == (), response.text assert _finish_reasons(events) == ("stop",), response.text assert _upstream_targets(wire) == (_CHAT_TARGET,) - rows: Final = eventually(lambda: read_rows(_ROWS, (model,)), lambda found: len(found) >= 1, _ROW_SECONDS) + rows: Final = eventually(lambda: gateway.rows(_ROWS, (model,)), lambda found: len(found) >= 1, _ROW_SECONDS) assert [row["status"] for row in rows] == ["success"], rows @@ -302,7 +301,7 @@ def test_a_deployment_whose_stream_stalls_stays_out_of_cooldown( with _peer(shape) as wire, cooldown_rig.gateway.scenario() as scenario: group, identity = _deployment(scenario, wire, backend) response: Final = _streamed(cooldown_rig.gateway, endpoint, group) - error: Final = _failure_row(_ROWS, (group,)) + error: Final = _failure_row(cooldown_rig.gateway, _ROWS, (group,)) assert not cooled(cooldown_rig.cache, identity), (error, response.text) @@ -324,7 +323,7 @@ def _fallback_cooled(cooldown_rig: _CooldownRig, fallback_shape: Shape) -> tuple primary, _ = _deployment(scenario, failing) fallback, identity = _deployment(scenario, fallback_wire) response: Final = _streamed(cooldown_rig.gateway, "chat", primary, fallbacks=[fallback]) - error: Final = _failure_row(_EITHER_GROUP_ROWS, (primary, fallback)) + error: Final = _failure_row(cooldown_rig.gateway, _EITHER_GROUP_ROWS, (primary, fallback)) assert _upstream_targets(fallback_wire) == (_CHAT_TARGET,), response.text return cooled(cooldown_rig.cache, identity), error, response.text From 1525e4bce50cbe80eb811541129529819e0fd4cd Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 19 Oct 2026 15:33:29 -0700 Subject: [PATCH 09/18] feat(cli): keep the claude.ai login with `lite claude --subscription` and save pkce in config (#47385) `lite claude --subscription` sends the LiteLLM key in the x-litellm-api-key header through ANTHROPIC_CUSTOM_HEADERS and leaves ANTHROPIC_AUTH_TOKEN and ANTHROPIC_API_KEY unset, so Claude Code stays on its claude.ai login and its connectors keep working. The pre-launch key check sends the key the same way. `claude_subscription` and `pkce` are now `lite config` keys that default the new `--subscription/--no-subscription` flag on `lite claude` and the `--pkce/--no-pkce` flag on `lite login`. The flags still win for a single run. Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/agents.py | 237 ++++++++- litellm/proxy/client/cli/commands/auth.py | 12 +- litellm/proxy/client/cli/commands/config.py | 38 ++ litellm/proxy/client/cli/commands/up.py | 3 +- tests/unit/proxy/client/cli/test_agents.py | 496 +++++++++++++++++- .../proxy/client/cli/test_auth_commands.py | 37 +- .../proxy/client/cli/test_config_commands.py | 54 +- .../unit/proxy/client/cli/test_up_commands.py | 26 +- 8 files changed, 873 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 03540a7395e..52e36964bc1 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -7,13 +7,14 @@ import sys import tempfile from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from itertools import chain from pathlib import Path from types import MappingProxyType -from typing import Final, TypeAlias +from typing import Final, TypeAlias, TypeVar import click import requests -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm.proxy.common_utils.codex_model_catalog import ( CODEX_BASE_INSTRUCTIONS_PATH, @@ -25,8 +26,15 @@ from litellm.proxy.common_utils.codex_model_catalog import ( from litellm.types.llms.base_model import LiteLLMBaseModel from .auth import cli_context_obj, context_secret_vault, get_stored_api_key, login -from .claude_settings import ClaudeSettingsError, install_statusline_script +from .claude_settings import ( + ENV_KEY, + ClaudeSettingsError, + claude_settings_path, + install_statusline_script, + load_json_or_empty, +) from .cmd_quoting import quote_for_cmd +from .config import claude_subscription_default from .pi import ( LITELLM_PROXY_API_KEY_ENV, PI_PROVIDER_NAME, @@ -41,6 +49,8 @@ from .pi import ( ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_API_KEY_ENV: Final = "ANTHROPIC_API_KEY" +ANTHROPIC_CUSTOM_HEADERS_ENV: Final = "ANTHROPIC_CUSTOM_HEADERS" +LITELLM_KEY_HEADER: Final = "x-litellm-api-key" ENABLE_TOOL_SEARCH_ENV: Final = "ENABLE_TOOL_SEARCH" ENABLE_TOOL_SEARCH_VALUE: Final = "true" ENABLE_GATEWAY_MODEL_DISCOVERY_ENV: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" @@ -73,6 +83,7 @@ _INSTALL_DOCS: Final[dict[str, str]] = { } _HIDDEN_AGENTS: Final = frozenset({"pi"}) +_SUBSCRIPTION_AGENTS: Final = frozenset({"claude"}) CODEX_PROXY_PROVIDER: Final = "litellm" CODEX_HOME_ENV: Final = "CODEX_HOME" @@ -96,11 +107,24 @@ def agent_profile(command: str) -> tuple[str, frozenset[str]]: return base, frozenset({PROFILE_ANTHROPIC, PROFILE_OPENAI}) +def _custom_headers_with_litellm_key(existing: str, api_key: str) -> str: + kept: Final = tuple( + line for line in existing.splitlines() if line.strip() and not _is_header_line(line, LITELLM_KEY_HEADER) + ) + return "\n".join((*kept, f"{LITELLM_KEY_HEADER}: Bearer {api_key}")) + + +def _is_header_line(line: str, name: str) -> bool: + return line.partition(":")[0].strip().lower() == name + + def build_agent_env( base_env: Mapping[str, str], base_url: str, api_key: str, profiles: frozenset[str], + *, + subscription: bool = False, ) -> dict[str, str]: """Return a copy of base_env wired to route the agent through the proxy. @@ -115,12 +139,22 @@ def build_agent_env( proxy's /v1/models; likewise left alone when already set. pi ignores both base URL variables and instead resolves $LITELLM_PROXY_API_KEY from its synced models.json provider entry. + With subscription=True the Anthropic profile keeps Claude Code on its claude.ai login: the key + travels in the x-litellm-api-key header through ANTHROPIC_CUSTOM_HEADERS, and no + ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY is left in the environment, since either one makes + Claude Code stop using the subscription (and its connectors). """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") if PROFILE_ANTHROPIC in profiles: env[ANTHROPIC_BASE_URL_ENV] = root - env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key + if subscription: + env[ANTHROPIC_CUSTOM_HEADERS_ENV] = _custom_headers_with_litellm_key( + env.get(ANTHROPIC_CUSTOM_HEADERS_ENV, ""), api_key + ) + env.pop(ANTHROPIC_AUTH_TOKEN_ENV, None) + else: + env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key env.pop(ANTHROPIC_API_KEY_ENV, None) if ENABLE_TOOL_SEARCH_ENV not in env: env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE @@ -146,6 +180,130 @@ def codex_proxy_provider(base_url: str) -> Mapping[str, str | bool]: ) +_SETTINGS_FLAG: Final = "--settings" +_SUBSCRIPTION_SETTINGS_ENV: Final = MappingProxyType({ANTHROPIC_AUTH_TOKEN_ENV: "", ANTHROPIC_API_KEY_ENV: ""}) +_SETTINGS_OVERRIDE_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) + + +def _settings_flag_spans(args: Sequence[str]) -> tuple[tuple[int, int, str], ...]: + """(start, end, value) of every `--settings X` or `--settings=X` in args, in order.""" + return tuple( + (i, i + 1, arg.partition("=")[2]) if arg != _SETTINGS_FLAG else (i, i + 2, "".join(args[i + 1 : i + 2])) + for i, arg in enumerate(args) + if arg == _SETTINGS_FLAG or arg.startswith(_SETTINGS_FLAG + "=") + ) + + +def _split_settings_flags(args: Sequence[str]) -> tuple[str | None, tuple[str, ...]]: + """The value of the last `--settings` the user passed (the one Claude Code would keep), and args without any.""" + spans: Final = _settings_flag_spans(args) + covered: Final = frozenset(chain.from_iterable(range(start, end) for start, end, _ in spans)) + rest: Final = tuple(arg for i, arg in enumerate(args) if i not in covered) + return (spans[-1][2] if spans else None), rest + + +def _load_settings_override(value: str) -> Mapping[str, JsonValue]: + """The user's `--settings` value: inline JSON when it starts with `{`, else a settings file that must exist.""" + if value.lstrip().startswith("{"): + try: + return _SETTINGS_OVERRIDE_ADAPTER.validate_json(value) + except ValidationError as e: + raise AgentRunError(f"`{_SETTINGS_FLAG}` is not a JSON object: {value}") from e + path: Final = Path(value).expanduser() + try: + return load_json_or_empty(path) if path.is_file() else _missing_settings_file(path) + except ClaudeSettingsError as e: + raise AgentRunError(str(e)) from e + except OSError as e: + raise AgentRunError(f"`{_SETTINGS_FLAG}` must be a settings file or a JSON object, got: {value}") from e + + +def _missing_settings_file(path: Path) -> Mapping[str, JsonValue]: + raise AgentRunError(f"`{_SETTINGS_FLAG}` file not found: {path}") + + +def _settings_env(settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + env: Final = settings.get(ENV_KEY) + return env if isinstance(env, dict) else {} + + +def _settings_json(base: Mapping[str, JsonValue], api_key: str, inherited_headers: str) -> str: + base_env: Final = _settings_env(base) + own_headers: Final = base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV) + headers: Final = _custom_headers_with_litellm_key( + own_headers if isinstance(own_headers, str) else inherited_headers, api_key + ) + overlay: Final = {**base_env, **_SUBSCRIPTION_SETTINGS_ENV, ANTHROPIC_CUSTOM_HEADERS_ENV: headers} + return json.dumps({**base, ENV_KEY: overlay}, separators=(",", ":")) + + +def _inherited_custom_headers(settings_files: Sequence[Mapping[str, JsonValue]], environ: Mapping[str, str]) -> str: + """The ANTHROPIC_CUSTOM_HEADERS Claude Code would apply without the overlay: the first settings file in + precedence order that sets it, else the process environment.""" + from_files: Final = (_settings_env(settings).get(ANTHROPIC_CUSTOM_HEADERS_ENV) for settings in settings_files) + return next( + (value for value in from_files if isinstance(value, str)), environ.get(ANTHROPIC_CUSTOM_HEADERS_ENV, "") + ) + + +def claude_settings_files( + read_settings: Callable[[Path], Mapping[str, JsonValue]], environ: Mapping[str, str], project_dir: Path +) -> tuple[Mapping[str, JsonValue], ...]: + """The settings files Claude Code merges below `--settings`, highest precedence first: project local, project, + then user. A missing file reads as empty, a malformed one is a clean CLI error.""" + paths: Final = ( + project_dir / ".claude" / "settings.local.json", + project_dir / ".claude" / "settings.json", + claude_settings_path(environ), + ) + try: + return tuple(read_settings(path) for path in paths) + except ClaudeSettingsError as e: + raise AgentRunError(str(e)) from e + + +OVERLAY_FILE_NAME: Final = "litellm-subscription-overlay.json" + + +def write_settings_overlay(text: str, environ: Mapping[str, str]) -> Path: + """Write the overlay next to Claude Code's own settings.json (so under CLAUDE_CONFIG_DIR when set) and return + its path. One file per user, rewritten on every launch, mode 0600 whether or not it already existed. It carries + the LiteLLM key, which is why it is a private file and not an inline `--settings` argument that every local + user could read from the process list, and why it does not pile up in the temp dir.""" + path: Final = claude_settings_path(environ).parent / OVERLAY_FILE_NAME + path.parent.mkdir(parents=True, exist_ok=True) + fd: Final = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + if os.name == "posix": + os.fchmod(fd, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(text) + return path + + +def subscription_launch_args( + command: Sequence[str], + api_key: str, + environ: Mapping[str, str], + settings_files: Sequence[Mapping[str, JsonValue]], + write_overlay: Callable[[str, Mapping[str, str]], Path], +) -> tuple[tuple[str, ...], tuple[str, ...]]: + """The `--settings` that keeps Claude Code on its claude.ai login, plus the user's remaining args. + + Claude Code applies `env` from its settings files over the process environment, so a key that `lite up` or + `lite configure claude` left in ANTHROPIC_AUTH_TOKEN would land back in `Authorization`, and an + ANTHROPIC_CUSTOM_HEADERS there would drop the x-litellm-api-key header. The overlay blanks both credential + keys (read as unset for this session only) and carries the custom headers with the LiteLLM key appended, + taking the headers from the user's `--settings`, else the settings files in Claude Code's precedence order, + else the process environment. Any `--settings` the user passed is stripped and the last one (inline JSON or a + file, the one Claude Code would keep) is merged in, since Claude Code keeps only the last `--settings`. + """ + inherited: Final = _inherited_custom_headers(settings_files, environ) + user_value, rest = _split_settings_flags(tuple(command[1:])) + base: Final[Mapping[str, JsonValue]] = _load_settings_override(user_value) if user_value is not None else {} + overlay: Final = write_overlay(_settings_json(base, api_key, inherited), environ) + return (_SETTINGS_FLAG, str(overlay)), rest + + def _codex_proxy_args(base_url: str) -> list[str]: """Codex `-c` overrides that point it at the proxy. @@ -570,6 +728,7 @@ def verify_proxy_key( api_key: str, *, get: Callable[..., requests.Response] = requests.get, + key_header: str = "Authorization", ) -> None: """Probe the proxy with the key so bad creds fail here, not inside the agent. @@ -578,7 +737,7 @@ def verify_proxy_key( """ url: Final = base_url.rstrip("/") + "/v1/models" try: - resp: Final = get(url, headers={"Authorization": f"Bearer {api_key}"}, timeout=10) + resp: Final = get(url, headers={key_header: f"Bearer {api_key}"}, timeout=10) except requests.RequestException as e: raise AgentRunError( f"Could not reach the LiteLLM proxy at {base_url.rstrip('/')}: {e}. " @@ -591,6 +750,15 @@ def verify_proxy_key( ) +def verify_subscription_key( + base_url: str, + api_key: str, + *, + get: Callable[..., requests.Response] = requests.get, +) -> None: + verify_proxy_key(base_url, api_key, get=get, key_header=LITELLM_KEY_HEADER) + + _WINDOWS_SHIM_SUFFIXES: Final[frozenset[str]] = frozenset({".cmd", ".bat"}) _CMD_LINE_BREAKS: Final = ("\r", "\n") @@ -686,14 +854,19 @@ def run_agent( command: Sequence[str], *, skip_verify: bool = False, + subscription: bool = False, base_env: Mapping[str, str] | None = None, which: Callable[[str], str | None] = shutil.which, verify: Callable[[str, str], None] = verify_proxy_key, + verify_subscription: Callable[[str, str], None] = verify_subscription_key, sync_models: Callable[[str, Mapping[str, str], str, str, bool], ModelSyncResult] = agent_model_sync_env, warn: Callable[[str], None] = _warn, launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off, reattach_terminal: Callable[[], None] | None = None, preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS), + read_settings: Callable[[Path], Mapping[str, JsonValue]] = load_json_or_empty, + write_overlay: Callable[[str, Mapping[str, str]], Path] = write_settings_overlay, + project_dir: Path | None = None, ) -> None: """Validate, wire the environment, and hand off to the agent. @@ -713,7 +886,7 @@ def run_agent( raise AgentRunError(f"Could not find `{command[0]}` on your PATH.{hint}") if not skip_verify: - verify(base_url, api_key) + (verify_subscription if subscription else verify)(base_url, api_key) env_before_sync: Final = base_env if base_env is not None else os.environ synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify) @@ -725,15 +898,28 @@ def run_agent( env: Final = MappingProxyType( { - **build_agent_env(env_before_sync, base_url, api_key, profiles), + **build_agent_env(env_before_sync, base_url, api_key, profiles, subscription=subscription), **(synced if isinstance(synced, Mapping) else _NO_EXTRA_ENV), } ) synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else () extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args) + settings_args, forwarded = ( + subscription_launch_args( + command, + api_key, + env_before_sync, + claude_settings_files( + read_settings, env_before_sync, project_dir if project_dir is not None else Path.cwd() + ), + write_overlay, + ) + if subscription and os.path.basename(command[0]) in _SUBSCRIPTION_AGENTS + else ((), tuple(command[1:])) + ) if reattach_terminal is not None: reattach_terminal() - launcher(binary, [command[0], *extra_args, *command[1:]], env) + launcher(binary, [command[0], *extra_args, *settings_args, *forwarded], env) def _is_interactive() -> bool: @@ -762,16 +948,27 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." +_SUBSCRIPTION_HELP: Final = ( + "Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the " + "x-litellm-api-key header instead of replacing Claude Code's own credential. " + "Make it the default with `lite config set claude_subscription true`." +) -def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: +def _launch( + ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool, subscription: bool | None = None +) -> None: ctx_obj: Final = cli_context_obj(ctx) base_url: Final = ctx_obj["base_url"] started_interactive: Final = _is_interactive() api_key: Final = resolve_api_key(ctx) display_name, _profiles = agent_profile(binary) - click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") + use_subscription: Final = ( + subscription if subscription is not None else binary in _SUBSCRIPTION_AGENTS and claude_subscription_default() + ) + login_note: Final = " (keeping your claude.ai login)" if use_subscription else "" + click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}{login_note}") try: run_agent( @@ -779,12 +976,22 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify api_key, [binary, *args], skip_verify=skip_verify, + subscription=use_subscription, reattach_terminal=(_restore_controlling_terminal if started_interactive else None), ) except AgentRunError as e: raise click.ClickException(str(e)) +_Decorated: Final = TypeVar("_Decorated", bound=Callable[..., object] | click.Command) + + +def _subscription_option(binary: str) -> Callable[[_Decorated], _Decorated]: + if binary not in _SUBSCRIPTION_AGENTS: + return lambda func: func + return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP) + + def _make_agent_command(binary: str, display_name: str) -> click.Command: @click.command( name=binary, @@ -793,10 +1000,11 @@ def _make_agent_command(binary: str, display_name: str) -> click.Command: hidden=binary in _HIDDEN_AGENTS, ) @click.option("--skip-verify", is_flag=True, default=False, help=_SKIP_VERIFY_HELP) + @_subscription_option(binary) @click.argument("args", nargs=-1, type=click.UNPROCESSED) @click.pass_context - def _command(ctx: click.Context, skip_verify: bool, args: Sequence[str]) -> None: - _launch(ctx, binary, list(args), skip_verify=skip_verify) + def _command(ctx: click.Context, skip_verify: bool, args: Sequence[str], subscription: bool | None = None) -> None: + _launch(ctx, binary, list(args), skip_verify=skip_verify, subscription=subscription) _command.help = ( f"Run {display_name} routed through your LiteLLM proxy.\n\n" @@ -813,6 +1021,7 @@ def agent_commands() -> tuple[click.Command, ...]: __all__ = [ + "OVERLAY_FILE_NAME", "AgentRunError", "ListedModel", "ModelSyncSkipped", @@ -821,10 +1030,14 @@ __all__ = [ "agent_model_sync_env", "agent_profile", "build_agent_env", + "claude_settings_files", "opencode_model_sync_env", "opencode_provider_config", "prepare_pi", "resolve_api_key", "run_agent", + "subscription_launch_args", "verify_proxy_key", + "verify_subscription_key", + "write_settings_overlay", ] diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index ad54cab6460..b5302d1f561 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -52,6 +52,7 @@ from .claude_settings import ( refuse_while_owned, settings_file_owners, ) +from .config import pkce_default from .pkce_login import ( Http, PkceFailure, @@ -884,17 +885,16 @@ def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None: ), ) @click.option( - "--pkce", - is_flag=True, - default=False, + "--pkce/--no-pkce", + default=None, help=( "Sign in with OAuth authorization code + PKCE through your system browser (loopback redirect), " "with a refresh token that renews the key automatically. Requires a proxy that serves " - "/.well-known/litellm-cli-auth." + "/.well-known/litellm-cli-auth. Make it the default with `lite config set pkce true`." ), ) @click.pass_context -def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: +def login(ctx: click.Context, config_claude: bool, pkce: bool | None) -> None: """Login to LiteLLM proxy using SSO authentication""" from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER @@ -908,7 +908,7 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: raise click.ClickException(f"Cannot configure Claude Code, so not logging in: {e}") try: - if pkce: + if pkce if pkce is not None else pkce_default(): _pkce_login(base_url, config_claude, context_secret_vault(ctx)) return cli_sso_flow: Final = _start_cli_sso_flow(base_url=base_url) diff --git a/litellm/proxy/client/cli/commands/config.py b/litellm/proxy/client/cli/commands/config.py index de22251ea8c..539d88dce4a 100644 --- a/litellm/proxy/client/cli/commands/config.py +++ b/litellm/proxy/client/cli/commands/config.py @@ -13,6 +13,8 @@ from pydantic import TypeAdapter from litellm.litellm_core_utils.private_json import ensure_private_dir, write_private_json HIDDEN_COMMANDS_KEY: Final = "hidden_commands" +CLAUDE_SUBSCRIPTION_KEY: Final = "claude_subscription" +PKCE_KEY: Final = "pkce" _config_adapter: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str]) @@ -88,10 +90,46 @@ def _normalize_hidden_commands(value: str) -> str: return ",".join(sorted(names)) +_BOOLEAN_WORDS: Final[Mapping[str, str]] = MappingProxyType( + { + "true": "true", + "on": "true", + "1": "true", + "yes": "true", + "false": "false", + "off": "false", + "0": "false", + "no": "false", + } +) + + +def _normalize_boolean(key: str) -> Callable[[str], str]: + def normalize(value: str) -> str: + normalized: Final = _BOOLEAN_WORDS.get(value.strip().lower()) + if normalized is None: + raise click.UsageError(f"{key} must be true or false, e.g. `lite config set {key} true`") + return normalized + + return normalize + + +def claude_subscription_default() -> bool: + """Whether `lite claude` keeps the claude.ai login unless told otherwise.""" + return get_config_value(CLAUDE_SUBSCRIPTION_KEY) == "true" + + +def pkce_default() -> bool: + """Whether `lite login` signs in with browser PKCE unless told otherwise.""" + return get_config_value(PKCE_KEY) == "true" + + _NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType( { "base_url": normalize_base_url, HIDDEN_COMMANDS_KEY: _normalize_hidden_commands, + CLAUDE_SUBSCRIPTION_KEY: _normalize_boolean(CLAUDE_SUBSCRIPTION_KEY), + PKCE_KEY: _normalize_boolean(PKCE_KEY), } ) diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index d1e966d6eae..4043761f255 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -30,6 +30,7 @@ from .claude_settings import ( merge_claude_settings, write_claude_settings, ) +from .config import pkce_default class UpError(ClaudeSettingsError): @@ -131,7 +132,7 @@ def ensure_fresh_login(ctx: click.Context) -> None: if _usable_login(_key_resolved_on_the_way_in(ctx_obj, base_url, vault), vault): return - pkce: Final = _stored_login_is_pkce(vault) + pkce: Final = _stored_login_is_pkce(vault) or pkce_default() login_command: Final = "lite login --pkce" if pkce else "lite login" if not sys.stdin.isatty(): raise UpError(f"No fresh LiteLLM login found for this proxy. Run `{login_command}` first.") diff --git a/tests/unit/proxy/client/cli/test_agents.py b/tests/unit/proxy/client/cli/test_agents.py index 7be98521f8e..3481346076f 100644 --- a/tests/unit/proxy/client/cli/test_agents.py +++ b/tests/unit/proxy/client/cli/test_agents.py @@ -1,6 +1,7 @@ import inspect import json import os +import stat import subprocess import sys from pathlib import Path @@ -27,7 +28,9 @@ from litellm.proxy.client.cli.commands.agents import ( opencode_model_sync_env, run_agent, verify_proxy_key, + verify_subscription_key, ) +from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError from litellm.proxy.client.cli.main import cli as lite AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" @@ -173,6 +176,19 @@ class _FakeJsonResponse: return self._payload +def _overlay_writer(tmp_path): + def write(text, environ): + target = tmp_path / "overlay.json" + target.write_text(text) + return target + + return write + + +def _overlay(args): + return json.loads(Path(args[2]).read_text()) + + class TestAgentProfile: def test_claude_is_anthropic(self): name, profiles = agent_profile("claude") @@ -263,6 +279,55 @@ class TestBuildAgentEnv: assert env["PATH"] == "/usr/bin" assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} + def test_subscription_sends_the_key_in_a_header_and_sets_no_credential(self): + env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"}), subscription=True) + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["ANTHROPIC_CUSTOM_HEADERS"] == "x-litellm-api-key: Bearer sk-key" + assert "ANTHROPIC_AUTH_TOKEN" not in env + assert "ANTHROPIC_API_KEY" not in env + + def test_subscription_clears_inherited_credentials(self): + base = {"ANTHROPIC_AUTH_TOKEN": "stale-token", "ANTHROPIC_API_KEY": "real-key"} + env = build_agent_env(base, "http://localhost:4000", "sk-key", frozenset({"anthropic"}), subscription=True) + assert "ANTHROPIC_AUTH_TOKEN" not in env + assert "ANTHROPIC_API_KEY" not in env + assert base == {"ANTHROPIC_AUTH_TOKEN": "stale-token", "ANTHROPIC_API_KEY": "real-key"} + + def test_subscription_keeps_the_users_other_custom_headers(self): + env = build_agent_env( + {"ANTHROPIC_CUSTOM_HEADERS": "x-team: blue\n\nx-trace: 1\n"}, + "http://localhost:4000", + "sk-key", + frozenset({"anthropic"}), + subscription=True, + ) + assert env["ANTHROPIC_CUSTOM_HEADERS"].splitlines() == [ + "x-team: blue", + "x-trace: 1", + "x-litellm-api-key: Bearer sk-key", + ] + + def test_subscription_replaces_a_stale_litellm_key_header_whatever_its_case(self): + env = build_agent_env( + {"ANTHROPIC_CUSTOM_HEADERS": "X-LiteLLM-Api-Key: Bearer old\nx-team: blue"}, + "http://localhost:4000", + "sk-new", + frozenset({"anthropic"}), + subscription=True, + ) + assert env["ANTHROPIC_CUSTOM_HEADERS"].splitlines() == ["x-team: blue", "x-litellm-api-key: Bearer sk-new"] + + def test_subscription_leaves_openai_agents_on_the_bearer_key(self): + env = build_agent_env({}, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True) + assert env["OPENAI_API_KEY"] == "sk-key" + assert "ANTHROPIC_CUSTOM_HEADERS" not in env + + def test_default_mode_does_not_touch_custom_headers(self): + env = build_agent_env( + {"ANTHROPIC_CUSTOM_HEADERS": "x-team: blue"}, "http://localhost:4000", "sk-key", frozenset({"anthropic"}) + ) + assert env["ANTHROPIC_CUSTOM_HEADERS"] == "x-team: blue" + class TestAgentLaunchArgs: def test_claude_and_opencode_get_no_extra_args(self): @@ -330,6 +395,32 @@ class TestVerifyProxyKey: get=lambda *a, **k: _FakeResponse(500), ) + def test_key_header_moves_the_bearer_off_authorization(self): + captured = {} + + def fake_get(url, headers, timeout): + captured["headers"] = headers + return _FakeResponse(200) + + verify_proxy_key("http://localhost:4000", "sk-key", get=fake_get, key_header="x-litellm-api-key") + + assert captured["headers"] == {"x-litellm-api-key": "Bearer sk-key"} + + def test_subscription_probe_sends_the_key_only_in_the_litellm_header(self): + captured = {} + + def fake_get(url, headers, timeout): + captured["headers"] = headers + return _FakeResponse(200) + + verify_subscription_key("http://localhost:4000", "sk-key", get=fake_get) + + assert captured["headers"] == {"x-litellm-api-key": "Bearer sk-key"} + + def test_subscription_probe_reports_a_rejected_key(self): + with pytest.raises(AgentRunError, match="rejected your key"): + verify_subscription_key("http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(401)) + class TestOpencodeModelSync: @staticmethod @@ -1195,6 +1286,327 @@ class TestRunAgent: with pytest.raises(AgentRunError): run_agent("http://localhost:4000", "sk-key", []) + def test_subscription_launch_probes_with_the_header_and_exports_it(self): + probed = [] + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + subscription=True, + base_env={"ANTHROPIC_AUTH_TOKEN": "stale"}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: probed.append(("bearer", a)), + verify_subscription=lambda *a: probed.append(("header", a)), + launcher=lambda p, a, e: launched.update(env=dict(e)), + read_settings=lambda path: {}, + ) + assert probed == [("header", ("http://localhost:4000", "sk-key"))] + assert launched["env"]["ANTHROPIC_CUSTOM_HEADERS"] == "x-litellm-api-key: Bearer sk-key" + assert "ANTHROPIC_AUTH_TOKEN" not in launched["env"] + + def test_subscription_launch_blanks_the_credentials_claude_code_would_load_from_settings_json(self, tmp_path): + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--resume"], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda p, a, e: launched.update(args=tuple(a)), + read_settings=lambda path: {}, + write_overlay=_overlay_writer(tmp_path), + ) + assert launched["args"][:2] == ("claude", "--settings") + assert _overlay(launched["args"]) == { + "env": { + "ANTHROPIC_AUTH_TOKEN": "", + "ANTHROPIC_API_KEY": "", + "ANTHROPIC_CUSTOM_HEADERS": "x-litellm-api-key: Bearer sk-key", + } + } + assert launched["args"][3:] == ("--resume",) + + def test_subscription_launch_keeps_the_key_out_of_argv_and_in_one_private_file_next_to_settings_json( + self, tmp_path + ): + config_dir = tmp_path / "cfg" + config_dir.mkdir() + stale = config_dir / "litellm-subscription-overlay.json" + stale.write_text('{"env": {"ANTHROPIC_CUSTOM_HEADERS": "x-litellm-api-key: Bearer old-key"}}') + stale.chmod(0o644) + launched = {} + for _ in range(2): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + subscription=True, + base_env={"CLAUDE_CONFIG_DIR": str(config_dir)}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda p, a, e: launched.update(args=list(a)), + read_settings=lambda path: {}, + ) + assert "sk-key" not in " ".join(launched["args"]) + assert Path(launched["args"][2]) == stale + assert sorted(path.name for path in config_dir.iterdir()) == ["litellm-subscription-overlay.json"] + assert json.loads(stale.read_text())["env"]["ANTHROPIC_CUSTOM_HEADERS"] == "x-litellm-api-key: Bearer sk-key" + if os.name == "posix": + assert stat.S_IMODE(stale.stat().st_mode) == 0o600 + + def test_subscription_launch_accepts_inline_settings_json_longer_than_a_file_name(self, tmp_path): + long_value = "a" * 400 + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--settings", json.dumps({"env": {"FOO": long_value}})], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda p, a, e: launched.update(args=list(a)), + read_settings=lambda path: {}, + write_overlay=_overlay_writer(tmp_path), + ) + assert _overlay(launched["args"])["env"]["FOO"] == long_value + + def test_subscription_launch_reports_a_missing_settings_file_as_a_clean_cli_error(self, tmp_path): + with pytest.raises(AgentRunError, match="file not found"): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--settings", str(tmp_path / "nope.json")], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda *a: None, + read_settings=lambda path: {}, + write_overlay=_overlay_writer(tmp_path), + ) + + def test_subscription_launch_reads_the_settings_files_claude_code_merges_below_the_flag(self, tmp_path): + seen = [] + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + subscription=True, + base_env={"CLAUDE_CONFIG_DIR": str(tmp_path / "cfg")}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda *a: None, + read_settings=lambda path: seen.append(path) or {}, + write_overlay=_overlay_writer(tmp_path), + project_dir=tmp_path / "repo", + ) + assert seen == [ + tmp_path / "repo" / ".claude" / "settings.local.json", + tmp_path / "repo" / ".claude" / "settings.json", + tmp_path / "cfg" / "settings.json", + ] + + @pytest.mark.parametrize("separate_flag", [True, False]) + def test_subscription_launch_merges_the_blanks_into_the_users_settings_flag(self, tmp_path, separate_flag): + settings_file = tmp_path / "mine.json" + settings_file.write_text('{"model": "opus", "env": {"ANTHROPIC_AUTH_TOKEN": "mine", "FOO": "bar"}}') + user_flag = ["--settings", str(settings_file)] if separate_flag else [f"--settings={settings_file}"] + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--resume", *user_flag, "-p", "hi"], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda p, a, e: launched.update(args=list(a)), + read_settings=lambda path: {}, + write_overlay=_overlay_writer(tmp_path), + ) + assert launched["args"][:2] == ["claude", "--settings"] + assert _overlay(launched["args"]) == { + "model": "opus", + "env": { + "ANTHROPIC_AUTH_TOKEN": "", + "FOO": "bar", + "ANTHROPIC_API_KEY": "", + "ANTHROPIC_CUSTOM_HEADERS": "x-litellm-api-key: Bearer sk-key", + }, + } + assert launched["args"][3:] == ["--resume", "-p", "hi"] + + def test_subscription_launch_merges_the_blanks_into_inline_settings_json(self, tmp_path): + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--settings", '{"env": {"ANTHROPIC_API_KEY": "mine"}, "theme": "dark"}'], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda p, a, e: launched.update(args=list(a)), + read_settings=lambda path: {}, + write_overlay=_overlay_writer(tmp_path), + ) + assert launched["args"][:2] == ["claude", "--settings"] + assert _overlay(launched["args"]) == { + "theme": "dark", + "env": { + "ANTHROPIC_API_KEY": "", + "ANTHROPIC_AUTH_TOKEN": "", + "ANTHROPIC_CUSTOM_HEADERS": "x-litellm-api-key: Bearer sk-key", + }, + } + assert launched["args"][3:] == [] + + @pytest.mark.parametrize( + ("flag_headers", "local_headers", "project_headers", "user_headers", "expected"), + [ + ("X-Flag: f", "X-Local: l", "X-Proj: p", "X-User: u", "X-Flag: f\nx-litellm-api-key: Bearer sk-key"), + (None, "X-Local: l", "X-Proj: p", "X-User: u", "X-Local: l\nx-litellm-api-key: Bearer sk-key"), + ( + None, + None, + "X-Proj: p\nx-litellm-api-key: Bearer stale", + "X-User: u", + "X-Proj: p\nx-litellm-api-key: Bearer sk-key", + ), + (None, None, None, "X-User: u", "X-User: u\nx-litellm-api-key: Bearer sk-key"), + (None, None, None, None, "X-Proc: h\nx-litellm-api-key: Bearer sk-key"), + ], + ) + def test_subscription_overlay_keeps_the_custom_headers_claude_code_would_apply_and_adds_the_key( + self, tmp_path, flag_headers, local_headers, project_headers, user_headers, expected + ): + """Claude Code applies settings `env` over the process env, in precedence order, so the overlay must too.""" + user_flag = ( + ["--settings", json.dumps({"env": {"ANTHROPIC_CUSTOM_HEADERS": flag_headers}})] if flag_headers else [] + ) + files = { + tmp_path / "repo" / ".claude" / "settings.local.json": local_headers, + tmp_path / "repo" / ".claude" / "settings.json": project_headers, + tmp_path / "cfg" / "settings.json": user_headers, + } + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", *user_flag], + subscription=True, + base_env={"ANTHROPIC_CUSTOM_HEADERS": "X-Proc: h", "CLAUDE_CONFIG_DIR": str(tmp_path / "cfg")}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda p, a, e: launched.update(args=list(a), env=dict(e)), + read_settings=lambda path: {"env": {"ANTHROPIC_CUSTOM_HEADERS": files[path]}} if files[path] else {}, + write_overlay=_overlay_writer(tmp_path), + project_dir=tmp_path / "repo", + ) + assert _overlay(launched["args"])["env"]["ANTHROPIC_CUSTOM_HEADERS"] == expected + assert launched["env"]["ANTHROPIC_CUSTOM_HEADERS"] == "X-Proc: h\nx-litellm-api-key: Bearer sk-key" + + def test_subscription_launch_strips_every_settings_flag_and_merges_the_last_one_claude_code_would_keep( + self, tmp_path + ): + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + [ + "claude", + "--settings", + '{"env": {"ANTHROPIC_AUTH_TOKEN": "first"}, "theme": "light"}', + "--resume", + '--settings={"theme": "dark"}', + ], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda p, a, e: launched.update(args=list(a)), + read_settings=lambda path: {}, + write_overlay=_overlay_writer(tmp_path), + ) + assert launched["args"][:2] == ["claude", "--settings"] + assert _overlay(launched["args"]) == { + "theme": "dark", + "env": { + "ANTHROPIC_AUTH_TOKEN": "", + "ANTHROPIC_API_KEY": "", + "ANTHROPIC_CUSTOM_HEADERS": "x-litellm-api-key: Bearer sk-key", + }, + } + assert launched["args"][3:] == ["--resume"] + + def test_subscription_launch_reports_a_malformed_settings_file_as_a_clean_cli_error(self): + def broken(path): + raise ClaudeSettingsError(f"{path} contains invalid JSON") + + with pytest.raises(AgentRunError, match="contains invalid JSON"): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda *a: None, + read_settings=broken, + ) + + def test_subscription_launch_rejects_a_settings_flag_that_is_not_a_json_object(self): + with pytest.raises(AgentRunError, match="is not a JSON object"): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--settings", '{"env": 1, '], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=lambda *a: None, + launcher=lambda *a: None, + read_settings=lambda path: {}, + ) + + def test_default_launch_still_probes_with_the_bearer(self): + probed = [] + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: probed.append(("bearer", a)), + verify_subscription=lambda *a: probed.append(("header", a)), + launcher=lambda *a: None, + ) + assert probed == [("bearer", ("http://localhost:4000", "sk-key"))] + + def test_subscription_rejection_aborts_before_launch(self): + launched = [] + + def boom(*a): + raise AgentRunError("rejected") + + with pytest.raises(AgentRunError): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify_subscription=boom, + launcher=lambda *a: launched.append(a), + ) + assert launched == [] + def test_reattach_terminal_runs_just_before_launch(self): order = [] run_agent( @@ -1552,6 +1964,71 @@ class TestAgentCommands: assert captured["command"] == ["opencode"] assert "routing OpenCode through proxy at http://localhost:4000" in result.output + def test_claude_subscription_flag_reaches_run_agent_and_is_not_forwarded(self): + captured = {} + + def fake_run_agent(base_url, api_key, command, **kwargs): + captured["command"] = list(command) + captured["subscription"] = kwargs.get("subscription") + + with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent): + result = self.runner.invoke( + _agent_command("claude"), + ["--subscription", "--resume"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + + assert result.exit_code == 0, result.output + assert captured["subscription"] is True + assert captured["command"] == ["claude", "--resume"] + assert "keeping your claude.ai login" in result.output + + @pytest.mark.parametrize( + ("configured", "args", "expected"), + [ + (False, [], False), + (True, [], True), + (True, ["--no-subscription"], False), + (False, ["--subscription"], True), + ], + ) + def test_subscription_comes_from_the_flag_else_the_saved_setting(self, configured, args, expected): + captured = {} + with ( + patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=configured), + patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)), + ), + ): + result = self.runner.invoke( + _agent_command("claude"), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} + ) + assert result.exit_code == 0, result.output + assert captured["sub"] is expected + assert captured["command"] == ["claude"] + assert ("keeping your claude.ai login" in result.output) is expected + + @pytest.mark.parametrize("agent", ["codex", "opencode"]) + def test_the_saved_claude_setting_does_not_change_other_agents(self, agent): + captured = {} + with ( + patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=True), + patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")), + ), + ): + result = self.runner.invoke( + _agent_command(agent), [], obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} + ) + assert result.exit_code == 0, result.output + assert captured["sub"] is False + + @pytest.mark.parametrize("agent", ["codex", "opencode", "pi"]) + def test_only_claude_offers_the_subscription_flag(self, agent): + assert "--subscription" not in {opt for param in _agent_command(agent).params for opt in param.opts} + def test_skip_verify_is_consumed_not_forwarded(self): captured = {} @@ -1676,7 +2153,9 @@ class TestPrepareCodex: def test_registers_the_installed_script_as_a_session_scoped_stop_hook(self): from litellm.proxy.client.cli.commands.agents import prepare_codex - args = prepare_codex("http://localhost:4000", "sk-key", {}, install=lambda: "/py /home/me/.litellm/statusline.py") + args = prepare_codex( + "http://localhost:4000", "sk-key", {}, install=lambda: "/py /home/me/.litellm/statusline.py" + ) assert args == ( "-c", 'hooks.Stop=[{hooks=[{type="command",command="/py /home/me/.litellm/statusline.py"}]}]', @@ -1697,11 +2176,20 @@ class TestPrepareCodex: warnings = [] env = {"CODEX_HOME": str(tmp_path)} - for body in ('[[hooks.Stop]]\nhooks = [{ type = "command", command = "mine" }]\n', 'hooks.Stop = []\n', "[hooks]\n"): + for body in ( + '[[hooks.Stop]]\nhooks = [{ type = "command", command = "mine" }]\n', + "hooks.Stop = []\n", + "[hooks]\n", + ): (tmp_path / "config.toml").write_text(body) - assert prepare_codex("http://localhost:4000", "sk", env, install=lambda: "/py /s.py", warn=warnings.append) == () + assert ( + prepare_codex("http://localhost:4000", "sk", env, install=lambda: "/py /s.py", warn=warnings.append) + == () + ) (tmp_path / "config.toml").write_text('model = "gpt-5.6-sol"\n[projects."/x"]\ntrust_level = "trusted"\n') - assert prepare_codex("http://localhost:4000", "sk", env, install=lambda: "/py /s.py", warn=warnings.append) != () + assert ( + prepare_codex("http://localhost:4000", "sk", env, install=lambda: "/py /s.py", warn=warnings.append) != () + ) assert len(warnings) == 3 and "already declares hooks" in warnings[0] def test_a_config_that_cannot_be_read_or_decoded_still_lets_codex_launch(self, tmp_path): diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index 27edf9f3eb3..cb1fcba9df7 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -5,7 +5,6 @@ import time from pathlib import Path from unittest.mock import Mock, patch - import pytest import responses from click.testing import CliRunner @@ -26,7 +25,6 @@ from litellm.litellm_core_utils.cli_token_utils import ( ) from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module -from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner from litellm.proxy.client.cli.commands.auth import ( get_stored_api_key, login, @@ -34,6 +32,7 @@ from litellm.proxy.client.cli.commands.auth import ( print_token, whoami, ) +from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner @pytest.fixture @@ -1718,6 +1717,35 @@ class TestPkceLoginCommand: save.assert_called_once() assert _FakeSession.instances[0].posts == [] + @pytest.mark.parametrize( + ("saved", "args", "uses_pkce"), + [ + ("true", [], True), + ("false", [], False), + (None, [], False), + ("true", ["--no-pkce"], False), + ("false", ["--pkce"], True), + ], + ) + def test_login_uses_pkce_when_the_flag_or_the_saved_setting_says_so(self, saved, args, uses_pkce): + from litellm.proxy.client.cli.commands.config import save_config + + if saved is not None: + save_config({"pkce": saved}) + with ( + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()), + patch( + "litellm.proxy.client.cli.commands.auth._start_cli_sso_flow", + side_effect=RuntimeError("classic sso flow"), + ), + patch("litellm.proxy.client.cli.commands.auth.save_token", return_value=SecretStored()), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + result = self.runner.invoke(login, args, obj={"base_url": PKCE_BASE_URL}) + + assert ("Login successful!" in result.output) is uses_pkce + assert ("Authentication failed: classic sso flow" in result.output) is (not uses_pkce) + def test_pkce_login_saves_the_refreshable_record_and_skips_the_sso_poll(self): with ( patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()) as run, @@ -2050,7 +2078,10 @@ class TestGetStoredApiKeyRefresh: @pytest.mark.parametrize( "body, expected_detail", [ - ('{"detail": "Too many CLI login attempts.", "retry": {"after": [30, null]}}', ": Too many CLI login attempts."), + ( + '{"detail": "Too many CLI login attempts.", "retry": {"after": [30, null]}}', + ": Too many CLI login attempts.", + ), ('{"detail": 5}', ""), ('{"detail": ""}', ""), ('{"detail": {"message": "nested"}}', ""), diff --git a/tests/unit/proxy/client/cli/test_config_commands.py b/tests/unit/proxy/client/cli/test_config_commands.py index 611307635e0..2cd7a76ab95 100644 --- a/tests/unit/proxy/client/cli/test_config_commands.py +++ b/tests/unit/proxy/client/cli/test_config_commands.py @@ -6,16 +6,16 @@ from unittest.mock import patch import pytest from click.testing import CliRunner - - +from litellm.litellm_core_utils.private_json import write_private_json from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands.config import ( + claude_subscription_default, get_config_file_path, get_config_value, load_config, + pkce_default, save_config, ) -from litellm.litellm_core_utils.private_json import write_private_json from litellm.proxy.client.cli.interface import show_commands @@ -360,3 +360,51 @@ class TestWritePrivateJson: assert not target.exists() assert list(tmp_path.glob(".tmp-*")) == [] + + +class TestClaudeSubscriptionSetting: + @pytest.mark.parametrize("word", ["true", "TRUE", "on", "1", "yes"]) + def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, word): + result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word]) + + assert result.exit_code == 0, result.output + assert load_config()["claude_subscription"] == "true" + assert claude_subscription_default() is True + + @pytest.mark.parametrize("word", ["false", "Off", "0", "no"]) + def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, word): + cli_runner.invoke(cli, ["config", "set", "claude_subscription", "true"]) + result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word]) + + assert result.exit_code == 0, result.output + assert load_config()["claude_subscription"] == "false" + assert claude_subscription_default() is False + + def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home): + result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", "maybe"]) + + assert result.exit_code != 0 + assert "true or false" in result.output + assert "claude_subscription" not in load_config() + + def test_unset_means_off(self, isolated_home): + assert claude_subscription_default() is False + + +class TestPkceSetting: + @pytest.mark.parametrize(("word", "expected"), [("true", True), ("on", True), ("false", False), ("off", False)]) + def test_the_saved_value_drives_the_default(self, cli_runner, isolated_home, word, expected): + result = cli_runner.invoke(cli, ["config", "set", "pkce", word]) + + assert result.exit_code == 0, result.output + assert pkce_default() is expected + + def test_a_non_boolean_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home): + result = cli_runner.invoke(cli, ["config", "set", "pkce", "sometimes"]) + + assert result.exit_code != 0 + assert "true or false" in result.output + assert "pkce" not in load_config() + + def test_unset_means_off(self, isolated_home): + assert pkce_default() is False diff --git a/tests/unit/proxy/client/cli/test_up_commands.py b/tests/unit/proxy/client/cli/test_up_commands.py index ddd54dd1374..5c2e726315b 100644 --- a/tests/unit/proxy/client/cli/test_up_commands.py +++ b/tests/unit/proxy/client/cli/test_up_commands.py @@ -15,8 +15,8 @@ from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsErro from litellm.proxy.client.cli.commands.up import ( BackupRecord, UpError, - ensure_fresh_login, down, + ensure_fresh_login, load_json_or_empty, merge_claude_settings, read_backup, @@ -372,6 +372,30 @@ class TestEnsureFreshLogin: assert login_calls == [("http://proxy-a:4000", True)] + def test_logs_in_with_pkce_when_the_saved_pkce_default_is_on(self, monkeypatch): + """`lite config set pkce true` is honoured by the login `lite up` starts, not only by `lite login`.""" + monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: True) + store = _FakeTokenStore(monkeypatch, {"key": "sk-a", "base_url": "http://proxy-a:4000"}, {}) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + monkeypatch.setattr(up_module, "pkce_default", lambda: True) + login_calls = _capture_login( + monkeypatch, + on_login=lambda: store.log_in(_pkce_record("http://proxy-b:4000", seconds_left=86_400), "sk-pkce-fresh"), + ) + + ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + + assert login_calls == [("http://proxy-b:4000", True)] + + def test_names_the_pkce_login_in_the_non_interactive_error_when_the_saved_pkce_default_is_on(self, monkeypatch): + monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: False) + _FakeTokenStore(monkeypatch, {"key": "sk-a", "base_url": "http://proxy-a:4000"}, {}) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + monkeypatch.setattr(up_module, "pkce_default", lambda: True) + + with pytest.raises(UpError, match="Run `lite login --pkce` first"): + ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + def test_names_the_pkce_login_in_the_non_interactive_error_for_a_pkce_credential(self, monkeypatch): monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: False) _FakeTokenStore(monkeypatch, _pkce_record("http://proxy-a:4000", seconds_left=-10), {}) From b67003162a6f6ffb97576da42cb434f947b43f29 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 19 Oct 2026 15:35:29 -0700 Subject: [PATCH 10/18] test(integration): wait for every worker to report the websocket session cap instead of sleeping (#47412) The override test slept a fixed 10 s after writing responses_websocket_session_limit_seconds and then opened sockets on both workers. A worker whose config resync ran later than that kept the old 7200 s cap, so its socket never closed and the test failed after 75 s A test-only /health/integration-general-setting route returns the value a worker holds along with its pid. The test polls it over fresh connections until every worker reports the new cap, and after the delete until every worker reports it gone --- tests/integration/_support/runtime.py | 9 +++++ .../test_responses_websocket_session_limit.py | 34 ++++++++++++++----- 2 files changed, 34 insertions(+), 9 deletions(-) diff --git a/tests/integration/_support/runtime.py b/tests/integration/_support/runtime.py index a657ce4dff0..c53c23d42b3 100644 --- a/tests/integration/_support/runtime.py +++ b/tests/integration/_support/runtime.py @@ -33,8 +33,17 @@ async def worker_deployment_ids(model_name: str) -> JSONResponse: return JSONResponse(content={"model_ids": model_ids}, headers={"x-integration-worker-pid": str(os.getpid())}) +async def worker_general_setting(name: str) -> JSONResponse: + from litellm.proxy.proxy_server import proxy_config + + return JSONResponse( + content={"value": proxy_config.settings.get(name)}, headers={"x-integration-worker-pid": str(os.getpid())} + ) + + def configure_readiness() -> None: from litellm.proxy.proxy_server import app app.add_api_route("/health/integration-readiness", worker_readiness, include_in_schema=False) app.add_api_route("/health/integration-deployment-ids", worker_deployment_ids, include_in_schema=False) + app.add_api_route("/health/integration-general-setting", worker_general_setting, include_in_schema=False) diff --git a/tests/integration/providers/test_responses_websocket_session_limit.py b/tests/integration/providers/test_responses_websocket_session_limit.py index f2b3a48c6d1..4369a84e1cc 100644 --- a/tests/integration/providers/test_responses_websocket_session_limit.py +++ b/tests/integration/providers/test_responses_websocket_session_limit.py @@ -9,6 +9,7 @@ import threading import time import uuid from collections.abc import Generator, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path @@ -53,7 +54,7 @@ JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) SESSION_LIMIT_FIELD: Final = "responses_websocket_session_limit_seconds" LIMIT_CLOSE_REASON: Final = "Session duration limit reached" RELOAD_INTERVAL_SECONDS: Final = 3 -WORKER_SYNC_SECONDS: Final = RELOAD_INTERVAL_SECONDS + 7 +WORKER_SYNC_DEADLINE_SECONDS: Final = 60 CAP_SECONDS: Final = 60 CAP_CLOSE_SLACK_SECONDS: Final = 15 DEAF_CLOSE_SLACK_SECONDS: Final = 3 @@ -839,11 +840,28 @@ def _holding_worker(workers: frozenset[int], client_port: int) -> int | None: return next((pid for pid in sorted(workers) if holds(pid)), None) -def _settled_on_every_worker(written_at: float) -> None: +def _session_limit_on_a_worker(candidate: Gateway) -> tuple[int, JsonValue]: + response: Final = candidate.request( + "GET", + "/health/integration-general-setting", + params={"name": SESSION_LIMIT_FIELD}, + headers={"connection": "close"}, + ) + assert response.status_code == 200, f"Session limit on a worker: {response.status_code} {response.text}" + return int(response.headers["x-integration-worker-pid"]), JSON_OBJECT.validate_json(response.content)["value"] + + +def _session_limits_by_worker(candidate: Gateway) -> tuple[tuple[int, JsonValue], ...]: + with ThreadPoolExecutor(max_workers=8) as pool: + probes: Final = tuple(pool.submit(_session_limit_on_a_worker, candidate) for _ in range(16)) + return tuple(probe.result() for probe in probes) + + +def _settled_on_every_worker(candidate: Gateway, workers: frozenset[int], limit: JsonValue) -> None: eventually( - lambda: time.monotonic() - written_at, - lambda elapsed: elapsed >= WORKER_SYNC_SECONDS, - seconds=WORKER_SYNC_SECONDS + 5, + lambda: _session_limits_by_worker(candidate), + lambda answers: frozenset(pid for pid, _ in answers) >= workers and all(value == limit for _, value in answers), + seconds=WORKER_SYNC_DEADLINE_SECONDS, ) @@ -939,16 +957,14 @@ async def _override_workload(owned: OwnedProxy, key: str, model: str, database_u ) as earlier: accepted_at: Final = time.monotonic() updates: Final = await _update_sequence(candidate, (7200, CAP_SECONDS, CAP_SECONDS)) - written_at: Final = time.monotonic() stored_after_updates: Final = await asyncio.to_thread(_stored_limit, database_url) - await asyncio.to_thread(_settled_on_every_worker, written_at) + await asyncio.to_thread(_settled_on_every_worker, candidate, workers, CAP_SECONDS) opened: Final = await _sockets_on_every_worker(proxy, key, workers) capped: Final = await asyncio.gather(*tuple(_capped_close(socket) for socket in opened)) earlier_result: Final = await _held_turn(earlier, model, accepted_at) delete: Final = await asyncio.to_thread(_delete_limit, candidate) - deleted_at: Final = time.monotonic() stored_after_delete: Final = await asyncio.to_thread(_general_settings_row, database_url) - await asyncio.to_thread(_settled_on_every_worker, deleted_at) + await asyncio.to_thread(_settled_on_every_worker, candidate, workers, None) restored: Final = await _idle_turn( proxy, key, From ca0ec627b60b2b3db3ad9644c45a1f86141ac2c3 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 19 Oct 2026 15:45:56 -0700 Subject: [PATCH 11/18] fix(exa): return highlights, highlight scores, score and the response envelope from Exa search (#47423) Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../websearch_interception/transformation.py | 10 +- litellm/llms/exa_ai/search/transformation.py | 14 +- .../test_websearch_interception_wire.py | 78 ++++++++++ .../integration/providers/test_search_wire.py | 133 +++++++++++++++++- .../sdk/test_exa_search_envelope_sdk.py | 82 +++++++++++ .../test_websearch_native_blocks.py | 58 ++++++++ .../llms/exa_ai/search/test_transformation.py | 77 +++++++++- 7 files changed, 442 insertions(+), 10 deletions(-) create mode 100644 tests/integration/sdk/test_exa_search_envelope_sdk.py diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index bbb1b4d494b..58ae8bd6e48 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -5,7 +5,7 @@ Transforms between Anthropic/OpenAI tool_use format and LiteLLM search format. """ import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final from pydantic import TypeAdapter @@ -539,7 +539,8 @@ class WebSearchTransformation: result: SearchResponse from litellm.asearch() Returns: - Formatted text with Title, URL, Snippet for each result + Formatted text with Title, URL, Snippet for each result, or the response + without its provider_metadata (the caller's billing envelope) when there are no results """ # Convert SearchResponse to string if hasattr(result, "results") and result.results: @@ -548,6 +549,9 @@ class WebSearchTransformation: [f"Title: {r.title}\nURL: {r.url}\nSnippet: {r.snippet}" for r in result.results] ) else: - search_result_text = str(result) + fields: Final[Mapping[str, object]] = dict(result) + search_result_text = " ".join( + f"{key}={value!r}" for key, value in fields.items() if key != "provider_metadata" + ) return search_result_text diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index fbde1377180..89073409599 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -21,6 +21,7 @@ from litellm.secret_managers.main import get_secret_str _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) _STRS: Final = TypeAdapter(Iterable[str], config=ConfigDict(hide_input_in_errors=True)) +_RESULT_EXTRAS: Final = (("highlights", "highlights"), ("highlightScores", "highlight_scores"), ("score", "score")) class _ExaAISearchRequestRequired(TypedDict): @@ -173,6 +174,9 @@ class ExaAISearchConfig(BaseSearchConfig): - results[].summary → SearchResult.snippet (fallback when "text" and "highlights" are absent) - results[].publishedDate → SearchResult.date - No last_updated field in Exa AI response (set to None) + - results[].highlights / highlightScores / score → SearchResult.highlights / highlight_scores / score, + each set only when Exa returned it + - every top-level key except results (requestId, costDollars, ...) → SearchResponse.provider_metadata Args: raw_response: Raw httpx response from Exa AI API @@ -197,11 +201,15 @@ class ExaAISearchConfig(BaseSearchConfig): or "", "date": result.get("publishedDate"), # ISO 8601 datetime string "last_updated": None, # Exa AI doesn't provide last_updated in response + **{field: result[key] for key, field in _RESULT_EXTRAS if key in result}, } ) results.append(search_result) - return SearchResponse( - results=results, - object="search", + return SearchResponse.model_validate( + { + "results": results, + "object": "search", + "provider_metadata": {key: value for key, value in response_json.items() if key != "results"}, + } ) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py index 4e5ef521c0e..5ec860910e5 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py @@ -409,3 +409,81 @@ def test_database_created_search_tool_backend_receives_the_intercepted_query_ove assert targets[-3:] == (("POST", "/v1/messages"), ("GET", "/database/search"), ("POST", "/v1/messages")), ( targets ) + + +def test_empty_exa_search_reaches_the_model_without_the_exa_response_envelope(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "websearch-exa-" + uuid.uuid4().hex + tool_name: Final = "integration-exa-" + uuid.uuid4().hex + envelope: Final = { + "requestId": "exa-request-" + identity, + "resolvedSearchType": "neural", + "searchTime": 98.5, + "costDollars": {"total": 0.005, "search": {"neural": 0.005}}, + } + + def respond(request: Request) -> Reply: + parts: Final = urlsplit(request.target) + if parts.path == "/exa/search": + assert (request.method, json.loads(request.body)["query"]) == ("POST", _QUERY), request.body + return Reply(body=json.dumps({**envelope, "results": []}).encode()) + assert request.method == "POST" and parts.path == "/v1/messages", request.target + body: Final = json.loads(request.body) + assert [tool["name"] for tool in body["tools"]] == ["litellm_web_search"], body["tools"] + results: Final = [ + block + for message in body["messages"] + if isinstance(message["content"], list) + for block in message["content"] + if block["type"] == "tool_result" + ] + if not results: + return _anthropic_reply(identity, [_TEXT_BLOCK, _search_tool_use(identity)], "tool_use") + assert results == [{"type": "tool_result", "tool_use_id": identity, "content": "results=[] object='search'"}], ( + results + ) + return _anthropic_reply(identity, [_FINAL_BLOCK], "end_turn") + + with wire_server(respond) as wire: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["search_tools"] = [ + { + "search_tool_name": tool_name, + "litellm_params": { + "search_provider": "exa_ai", + "api_key": "synthetic-exa-key", + "api_base": wire.url + "/exa", + }, + } + ] + config["litellm_settings"].update( + { + "callbacks": ["websearch_interception"], + "websearch_interception_params": { + "enabled": True, + "enabled_providers": ["anthropic"], + "search_tool_name": tool_name, + }, + } + ) + path: Final = tmp_path / "websearch-exa.yaml" + path.write_text(yaml.safe_dump(config)) + environment: Final = {"ANTHROPIC_API_BASE": wire.url} + with owned_proxy(gateway, tmp_path, environment, config=path) as candidate, candidate.scenario() as models: + model: Final = models.model( + model="anthropic/claude-sonnet-4-5-20250929", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + response: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": identity}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + ) + assert response.status_code == 200, response.text + assert response.json()["content"][-1] == _FINAL_BLOCK, response.text + assert envelope["requestId"] not in response.text, response.text + targets: Final = tuple((request.method, urlsplit(request.target).path) for request in wire.drain()) + assert targets[-3:] == (("POST", "/v1/messages"), ("POST", "/exa/search"), ("POST", "/v1/messages")), targets diff --git a/tests/integration/providers/test_search_wire.py b/tests/integration/providers/test_search_wire.py index 5ee80f7f5db..005539cc331 100644 --- a/tests/integration/providers/test_search_wire.py +++ b/tests/integration/providers/test_search_wire.py @@ -1,11 +1,12 @@ import json import uuid +from concurrent.futures import ThreadPoolExecutor from typing import Final import pytest from integration._support.client import Gateway, Scenario from integration._support.wire import Reply, Request, wire_server -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue _EXA_KEY: Final = "synthetic-exa-key" _QUERY: Final = "latest AI developments 2024" @@ -43,6 +44,72 @@ class _SearchResponse(BaseModel): results: tuple[_SearchResult, ...] +_EXA_HIGHLIGHTS: Final = {"highlights": {"numSentences": 2}} +_EXA_ENVELOPE: Final = { + "requestId": "exa-request-5d1f0c", + "resolvedSearchType": "neural", + "searchTime": 412.5, + "costDollars": {"total": 0.007, "search": {"neural": 0.005}, "contents": {"highlights": 0.002}}, +} +_EXA_SCORED_RESULTS: Final = ( + { + "id": "https://result.invalid/margins", + "title": "Margin outlook", + "url": "https://result.invalid/margins", + "publishedDate": "2026-01-15T00:00:00.000Z", + "score": 0.91, + "highlights": ["Margins hold steady.", "Deposit costs ease."], + "highlightScores": [0.82, 0.64], + }, + { + "id": "https://result.invalid/deposits", + "title": "Deposit pricing", + "url": "https://result.invalid/deposits", + "score": 0.47, + "highlights": ["Deposit betas fall."], + "highlightScores": [0.55], + }, +) + + +class _ScoredSearchResult(_SearchResult): + highlights: tuple[str, ...] | None = None + highlight_scores: tuple[float, ...] | None = None + score: float | None = None + + +class _ScoredSearchResponse(BaseModel): + object: str + results: tuple[_ScoredSearchResult, ...] + provider_metadata: JsonValue = None + + +_EXPECTED_SCORED_RESPONSE: Final = _ScoredSearchResponse( + object="search", + results=( + _ScoredSearchResult( + title="Margin outlook", + url="https://result.invalid/margins", + snippet="Margins hold steady.\n\nDeposit costs ease.", + date="2026-01-15T00:00:00.000Z", + last_updated=None, + highlights=("Margins hold steady.", "Deposit costs ease."), + highlight_scores=(0.82, 0.64), + score=0.91, + ), + _ScoredSearchResult( + title="Deposit pricing", + url="https://result.invalid/deposits", + snippet="Deposit betas fall.", + date=None, + last_updated=None, + highlights=("Deposit betas fall.",), + highlight_scores=(0.55,), + score=0.47, + ), + ), + provider_metadata=_EXA_ENVELOPE, +) _EXPECTED_RESPONSE: Final = _SearchResponse( object="search", results=( @@ -162,6 +229,70 @@ def test_list_query_reaches_exa_as_one_space_joined_query(gateway: Gateway) -> N assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/exa/search")] +def test_exa_highlights_scores_and_response_envelope_reach_the_caller_on_every_route_spelling( + gateway: Gateway, +) -> None: + def respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/exa/search"), request.target + assert json.loads(request.body) == {"query": _QUERY, "numResults": 2, "contents": _EXA_HIGHLIGHTS}, request.body + return Reply(body=json.dumps({**_EXA_ENVELOPE, "results": list(_EXA_SCORED_RESULTS)}).encode()) + + tool: Final = f"exa-tool-{uuid.uuid4().hex}" + body: Final = {"query": _QUERY, "max_results": 2, "contents": _EXA_HIGHLIGHTS} + with wire_server(respond) as wire, gateway.scenario() as scenario: + _create_exa_tool(gateway, scenario, tool, f"{wire.url}/exa") + spellings: Final = ( + (f"/v1/search/{tool}", body), + (f"/search/{tool}", body), + ("/v1/search", {"search_tool_name": tool, **body}), + ("/search", {"model": tool, **body}), + ) + for path, payload in spellings: + response = gateway.request("POST", path, payload) + assert response.status_code == 200, f"{path}: {response.text}" + assert _ScoredSearchResponse.model_validate_json(response.content) == _EXPECTED_SCORED_RESPONSE, ( + f"{path}: {response.text}" + ) + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/exa/search")] * len( + spellings + ) + + +def test_concurrent_exa_searches_each_get_their_own_response_envelope(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + query: Final = json.loads(request.body)["query"] + results: Final = [] if query.endswith("-empty") else [_EXA_SCORED_RESULTS[1]] + return Reply( + body=json.dumps({"requestId": f"exa-{query}", "costDollars": {"total": 0.005}, "results": results}).encode() + ) + + tool: Final = f"exa-tool-{uuid.uuid4().hex}" + queries: Final = tuple(f"burst-{uuid.uuid4().hex}{'-empty' if index % 2 else ''}" for index in range(16)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + _create_exa_tool(gateway, scenario, tool, f"{wire.url}/exa") + with ThreadPoolExecutor(max_workers=len(queries)) as pool: + responses: Final = tuple( + pool.map(lambda query: gateway.request("POST", f"/v1/search/{tool}", {"query": query}), queries) + ) + observed: Final = tuple( + ( + response.status_code, + _ScoredSearchResponse.model_validate_json(response.content).provider_metadata, + tuple(result.score for result in _ScoredSearchResponse.model_validate_json(response.content).results), + ) + for response in responses + ) + assert observed == tuple( + ( + 200, + {"requestId": f"exa-{query}", "costDollars": {"total": 0.005}}, + () if query.endswith("-empty") else (0.47,), + ) + for query in queries + ), [response.text for response in responses] + assert len(wire.drain()) == len(queries) + + def _not_found(name: str) -> str: return f"Search tool '{name}' not found in router.search_tools" diff --git a/tests/integration/sdk/test_exa_search_envelope_sdk.py b/tests/integration/sdk/test_exa_search_envelope_sdk.py new file mode 100644 index 00000000000..73c3538fa45 --- /dev/null +++ b/tests/integration/sdk/test_exa_search_envelope_sdk.py @@ -0,0 +1,82 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse + +pytestmark: Final = pytest.mark.timeout(60) + +_ENVELOPE: Final[dict[str, JsonValue]] = { + "requestId": "exa-request-sdk", + "resolvedSearchType": "neural", + "searchTime": 211.0, + "costDollars": {"total": 0.005, "search": {"neural": 0.005}}, +} +_RESULT: Final[dict[str, JsonValue]] = { + "id": "https://result.invalid/sdk", + "title": "SDK result", + "url": "https://result.invalid/sdk", + "publishedDate": "2026-02-01T00:00:00.000Z", + "score": 0.73, + "highlights": ["First line.", "Second line."], + "highlightScores": [0.61, 0.42], +} +_EXPECTED: Final[dict[str, JsonValue]] = { + "object": "search", + "results": [ + { + "title": "SDK result", + "url": "https://result.invalid/sdk", + "snippet": "First line.\n\nSecond line.", + "date": "2026-02-01T00:00:00.000Z", + "last_updated": None, + "highlights": ["First line.", "Second line."], + "highlight_scores": [0.61, 0.42], + "score": 0.73, + } + ], + "provider_metadata": _ENVELOPE, +} + + +def _respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/exa/search"), request.target + return Reply(body=json.dumps({**_ENVELOPE, "results": [_RESULT]}).encode()) + + +def test_sync_exa_search_returns_highlights_scores_and_the_response_envelope() -> None: + with wire_server(_respond) as wire: + response: Final = litellm.search( + query=f"sdk sync {uuid.uuid4().hex}", + search_provider="exa_ai", + api_key="synthetic-exa-key", + api_base=f"{wire.url}/exa", + max_results=1, + contents={"highlights": {"numSentences": 2}}, + ) + assert len(wire.drain()) == 1 + + assert isinstance(response, SearchResponse), type(response) + assert response.model_dump(mode="json") == _EXPECTED + + +@pytest.mark.asyncio +async def test_async_exa_search_returns_highlights_scores_and_the_response_envelope() -> None: + with wire_server(_respond) as wire: + response: Final = await litellm.asearch( + query=f"sdk async {uuid.uuid4().hex}", + search_provider="exa_ai", + api_key="synthetic-exa-key", + api_base=f"{wire.url}/exa", + max_results=1, + contents={"highlights": {"numSentences": 2}}, + ) + assert len(wire.drain()) == 1 + + assert isinstance(response, SearchResponse), type(response) + assert response.model_dump(mode="json") == _EXPECTED diff --git a/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py index 92feab86dc0..c43a51cfd52 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py @@ -670,3 +670,61 @@ class TestLegacyPathMatchesNewPath: "text", ] assert out["content"][1]["tool_use_id"] == out["content"][0]["id"] + + +@pytest.mark.parametrize( + ("response", "expected_text"), + [ + (SearchResponse(results=[], object="search"), "results=[] object='search'"), + ( + SearchResponse.model_validate( + { + "results": [], + "object": "search", + "provider_metadata": {"requestId": "exa-request", "costDollars": {"total": 0.005}}, + } + ), + "results=[] object='search'", + ), + ( + SearchResponse.model_validate( + { + "results": [], + "object": "search", + "search_id": "search-1", + "provider_metadata": {"requestId": "exa-request"}, + } + ), + "results=[] object='search' search_id='search-1'", + ), + ( + SearchResponse.model_validate({"results": [], "object": "search", "query": "q", "cls": 1}), + "results=[] object='search' query='q' cls=1", + ), + ( + SearchResponse.model_validate( + { + "results": [ + SearchResult.model_validate( + { + "title": "T", + "url": "https://example.com", + "snippet": "S", + "highlights": ["S"], + "highlight_scores": [0.5], + "score": 0.9, + } + ) + ], + "object": "search", + "provider_metadata": {"requestId": "exa-request"}, + } + ), + "Title: T\nURL: https://example.com\nSnippet: S", + ), + ], +) +def test_format_search_response_keeps_the_provider_metadata_out_of_the_model_text( + response: SearchResponse, expected_text: str +): + assert WebSearchTransformation.format_search_response(response) == expected_text diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 456ed26cc86..a5a5d5e22dc 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -123,12 +123,83 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order(): assert response.object == "search" assert response.results == [ - SearchResult( - title="First", url="https://example.com/first", snippet="First text", date="2024-01-05T00:00:00.000Z" + SearchResult.model_validate( + { + "title": "First", + "url": "https://example.com/first", + "snippet": "First text", + "date": "2024-01-05T00:00:00.000Z", + "highlights": ["ignored when text is present"], + "highlight_scores": [0.5], + "score": 0.42, + } + ), + SearchResult.model_validate( + { + "title": "Second", + "url": "https://example.com/second", + "snippet": "one\n\ntwo", + "highlights": ["one", "two"], + } ), - SearchResult(title="Second", url="https://example.com/second", snippet="one\n\ntwo"), SearchResult(title="Third", url="https://example.com/third", snippet="Summary"), ] + assert response.model_extra == { + "provider_metadata": { + "requestId": "b5947044c4b78efa9552a7c89b306d95", + "resolvedSearchType": "neural", + "costDollars": {"total": 0.005}, + } + } + + +@pytest.mark.parametrize( + ("exa_fields", "expected_extras"), + [ + ({"text": "full text"}, {}), + ({"score": 0.3}, {"score": 0.3}), + ({"highlights": ["a highlight"]}, {"highlights": ["a highlight"]}), + ({"highlightScores": [0.1]}, {"highlight_scores": [0.1]}), + ({"score": None}, {"score": None}), + ({"highlights": [], "highlightScores": [], "score": 0}, {"highlights": [], "highlight_scores": [], "score": 0}), + ({"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, {}), + ], +) +def test_transform_search_response_attaches_only_the_exa_result_fields_it_returned( + exa_fields: dict[str, object], expected_extras: dict[str, object] +): + response: Final = _transform( + { + "results": [ + {"title": "Plain", "url": "https://example.com/plain", "text": "plain"}, + {"title": "Title", "url": "https://example.com", **exa_fields}, + ] + } + ) + + assert [result.model_extra for result in response.results] == [{}, expected_extras] + + +@pytest.mark.parametrize( + ("payload", "expected_metadata"), + [ + ({}, {}), + ({"results": []}, {}), + ({"requestId": "abc"}, {"requestId": "abc"}), + ({"context": "joined result text", "results": []}, {"context": "joined result text"}), + ( + {"requestId": "abc", "searchTime": 412.5, "costDollars": {"total": 0.007}, "results": [{"text": "t"}]}, + {"requestId": "abc", "searchTime": 412.5, "costDollars": {"total": 0.007}}, + ), + ], +) +def test_transform_search_response_keeps_the_exa_envelope_as_provider_metadata( + payload: dict[str, object], expected_metadata: dict[str, object] +): + response: Final = _transform(payload) + + assert response.model_extra == {"provider_metadata": expected_metadata} + assert response.model_dump(mode="json")["provider_metadata"] == expected_metadata def test_transform_search_response_result_without_optional_fields_gets_empty_defaults(): From 9b4f44135dfe25c7de077d51280a60fc761f2af0 Mon Sep 17 00:00:00 2001 From: hamzashah-dev Date: Tue, 20 Oct 2026 03:50:33 +0500 Subject: [PATCH 12/18] fix(proxy): refresh cached team on /team/block and /team/unblock (#35643) * fix(proxy): refresh cached team on /team/block and /team/unblock Both endpoints wrote `blocked` to Postgres and returned without touching the management-object cache. Auth resolves teams through `team_id:{team_id}` before reading the DB and `common_checks` enforces `blocked` off that cached object, so a team whose traffic had already warmed the cache kept serving requests after being blocked, and kept getting rejected after being unblocked, until the entry expired. `_refresh_cached_team` already exists for exactly this and its docstring says every endpoint mutating `litellm_teamtable` must call it; these two were the only mutating endpoints that didn't. The update now also asks for `object_permission` so the refreshed cache entry doesn't drop the relation, matching `team_model_add` and `team_model_delete`. * fix(proxy): keep a team block from being undone by in-flight spend writes update_cache read the cached team when a request finished and wrote that whole copy back in a background task, so a team blocked or invalidated while requests were in flight came back unblocked on every busy worker until its 60 s cache entry expired. The team spend now goes onto the team held in this worker's memory at write time, and a missing entry is left missing for the next auth read to load. Block and unblock responses keep litellm_model_table and object_permission null, as before this PR, while the cache still gets both relations. * test(proxy): type the team spend cache test parameters * test(proxy): hold a team block under in-flight team spend and keep team aliases routing after block and unblock --------- Co-authored-by: Hamza Shah Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../proxy/hooks/proxy_track_cost_callback.py | 1 - .../management_endpoints/team_endpoints.py | 96 +++++------ litellm/proxy/proxy_server.py | 41 ++--- .../integration/management/test_team_block.py | 150 +++++++++++++++++- tests/unit/proxy/auth/test_jwt.py | 29 ---- .../test_team_endpoints.py | 104 +++++++++++- .../proxy/proxy_server/test_spend_counters.py | 73 ++++++++- 7 files changed, 381 insertions(+), 113 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3195394b9aa..4434fe57ab6 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -457,7 +457,6 @@ class _ProxyDBLogger(CustomLogger): update_cache_read_keys=update_cache_read_keys( user_id=user_id, end_user_id=end_user_id, - team_id=team_id, tags=tags, response_cost=response_cost, ), diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0f40662ad83..0595d23bd9a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5343,6 +5343,52 @@ async def team_member_me( ) +async def _set_team_blocked( + team_id: str, + blocked: bool, + user_api_key_dict: UserAPIKeyAuth, +) -> "prisma_models.LiteLLM_TeamTable": + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise Exception("No DB Connected.") + + existing_team: Final = await _team_db(prisma_client).find_unique(where={"team_id": team_id}) + if existing_team is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={team_id}"}, + ) + + if not await get_team_access().allows( + user_api_key_dict, LiteLLM_TeamTable.model_validate(existing_team.model_dump()), TEAM_OR_ORG_ADMIN + ): + team_access_denied() + + record: Final = await _team_db(prisma_client).update( + where={"team_id": team_id}, + data={"blocked": blocked}, + include=_TEAM_UPDATE_INCLUDE, + ) + if record is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={team_id}"}, + ) + + await refresh_cached_team( + team_row=record, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + return record.model_copy(update=dict.fromkeys(_TEAM_UPDATE_INCLUDE)) + + @router.post("/team/block", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def block_team( @@ -5372,30 +5418,7 @@ async def block_team( """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise Exception("No DB Connected.") - - existing_team: Final = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) - if existing_team is None: - raise HTTPException( - status_code=404, - detail={"error": f"Team not found, passed team_id={data.team_id}"}, - ) - - # Verify caller has access to manage this team - if not await get_team_access().allows( - user_api_key_dict, LiteLLM_TeamTable.model_validate(existing_team.model_dump()), TEAM_OR_ORG_ADMIN - ): - team_access_denied() - - record: Final = await _team_db(prisma_client).update( - where={"team_id": data.team_id}, - data={"blocked": True}, - ) - - return record + return await _set_team_blocked(team_id=data.team_id, blocked=True, user_api_key_dict=user_api_key_dict) @router.post("/team/unblock", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -5421,30 +5444,7 @@ async def unblock_team( }' ``` """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise Exception("No DB Connected.") - - existing_team: Final = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) - if existing_team is None: - raise HTTPException( - status_code=404, - detail={"error": f"Team not found, passed team_id={data.team_id}"}, - ) - - # Verify caller has access to manage this team - if not await get_team_access().allows( - user_api_key_dict, LiteLLM_TeamTable.model_validate(existing_team.model_dump()), TEAM_OR_ORG_ADMIN - ): - team_access_denied() - - record: Final = await _team_db(prisma_client).update( - where={"team_id": data.team_id}, - data={"blocked": False}, - ) - - return record + return await _set_team_blocked(team_id=data.team_id, blocked=False, user_api_key_dict=user_api_key_dict) @router.get( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c1102cee8d1..4ea3100f869 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4194,7 +4194,6 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) - def update_cache_read_keys( user_id: str | None, end_user_id: str | None, - team_id: str | None, tags: Sequence[object] | None, response_cost: float | None, ) -> tuple[str, ...]: @@ -4202,9 +4201,8 @@ def update_cache_read_keys( return () user_keys: tuple[str, ...] = (user_id, GLOBAL_PROXY_SPEND_CACHE_KEY) if user_id is not None else () end_user_keys: tuple[str, ...] = (end_user_cache_key(end_user_id),) if end_user_id is not None else () - team_keys: tuple[str, ...] = (f"team_id:{team_id}",) if team_id is not None else () tag_keys: tuple[str, ...] = tuple(tag_cache_key(tag) for tag in tags or () if isinstance(tag, str) and tag) - return user_keys + end_user_keys + team_keys + tag_keys + return user_keys + end_user_keys + tag_keys _UPDATE_CACHE_PREFETCH_SLOT: Final = "update_cache_read" @@ -4281,7 +4279,7 @@ async def update_cache( with service_target(AUTH_OBJECTS_TARGET): cached_values: Final = await _read_update_cache_values( keys=update_cache_read_keys( - user_id=user_id, end_user_id=end_user_id, team_id=team_id, tags=tags, response_cost=response_cost + user_id=user_id, end_user_id=end_user_id, tags=tags, response_cost=response_cost ), parent_otel_span=parent_otel_span, ) @@ -4431,35 +4429,28 @@ async def update_cache( ) ### UPDATE TEAM SPEND ### - async def _update_team_cache(): + def _update_team_cache(): if team_id is None or response_cost is None: return _id: Final = f"team_id:{team_id}" try: - cached_team: Final = cached_values.get(_id) - if cached_team is None: - # do nothing if team not in api key cache - return - existing_spend_obj: Final[LiteLLM_TeamTableCachedObj | None] = CacheCodec.deserialize( - cached_team, LiteLLM_TeamTableCachedObj + current_team: Final = CacheCodec.deserialize( + user_api_key_cache.get_cache(key=_id, local_only=True), LiteLLM_TeamTableCachedObj ) - if existing_spend_obj is None: + if current_team is None: return verbose_proxy_logger.debug( - "_update_team_db: existing spend: %s; response_cost: %s", existing_spend_obj, response_cost + "_update_team_db: existing spend: %s; response_cost: %s", current_team, response_cost ) - - existing_spend: Final[float] = existing_spend_obj.spend or 0.0 - # Calculate the new cost by adding the existing cost and response_cost - new_spend: Final = existing_spend + response_cost - - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append( - ( - _id, - CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj), - ) + user_api_key_cache.set_cache( + key=_id, + value=CacheCodec.serialize( + current_team.model_copy(update={"spend": (current_team.spend or 0.0) + response_cost}), + model_type=LiteLLM_TeamTableCachedObj, + ), + local_only=True, + ttl=get_management_object_ttl(user_api_key_cache), ) except Exception as e: verbose_proxy_logger.warning( @@ -4536,7 +4527,7 @@ async def update_cache( await _update_end_user_cache() if team_id is not None: - await _update_team_cache() + _update_team_cache() if tags is not None: await _update_tag_cache() diff --git a/tests/integration/management/test_team_block.py b/tests/integration/management/test_team_block.py index 3d1f9e430ea..9a340944569 100644 --- a/tests/integration/management/test_team_block.py +++ b/tests/integration/management/test_team_block.py @@ -1,14 +1,21 @@ +import itertools +import threading import uuid from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor from typing import Final import httpx import pytest -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue -from tests.integration._support.client import Gateway, string_value +from tests.integration._support.client import Gateway, eventually, held, string_value from tests.integration._support.database import read_rows +IN_FLIGHT_CALLERS_PER_PROXY: Final = 8 +BLOCK_CYCLES: Final = 4 +PROBES_PER_PROXY: Final = 6 + class _ObservedRequest(BaseModel): body: dict[str, JsonValue] @@ -35,6 +42,16 @@ class _ErrorResponse(BaseModel): error: _ProxyError +class _Message(BaseModel): + model_config = ConfigDict(frozen=True) + content: str + + +class _ChatBody(BaseModel): + model_config = ConfigDict(frozen=True) + messages: tuple[_Message, ...] + + @pytest.fixture def upstream(gateway: Gateway) -> Iterator[httpx.Client]: with httpx.Client(base_url=gateway.upstream_url, timeout=15, trust_env=False) as client: @@ -58,6 +75,87 @@ def _models(client: httpx.Client) -> tuple[str, ...]: return tuple(string_value(request.body["model"]) for request in observations.requests) +def _texts(client: httpx.Client) -> tuple[str, ...]: + response: Final = client.get("/__observations") + assert response.status_code == 200, response.text + observations: Final = _Observations.model_validate_json(response.text) + return tuple( + message.content + for message in itertools.chain.from_iterable( + _ChatBody.model_validate(request.body).messages for request in observations.requests + ) + ) + + +def _keep_calling(proxy: Gateway, model: str, key: str, stop: threading.Event) -> tuple[int, ...]: + return tuple( + _chat(proxy, model, key, f"in flight {uuid.uuid4().hex}").status_code + for _ in itertools.takewhile(lambda _: not stop.is_set(), itertools.count()) + ) + + +def _statuses_on_fresh_connections(proxies: tuple[Gateway, ...], model: str, key: str, marker: str) -> tuple[int, ...]: + return tuple( + proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"{marker} {uuid.uuid4().hex}"}]}, + key=key, + headers={"connection": "close"}, + ).status_code + for proxy in proxies * PROBES_PER_PROXY + ) + + +def _block_while_team_traffic_is_in_flight( + gateway: Gateway, peer: Gateway, upstream: httpx.Client, model: str, key: str, team: str +) -> tuple[httpx.Response, tuple[int, ...]]: + stop: Final = threading.Event() + with ThreadPoolExecutor(max_workers=2 * IN_FLIGHT_CALLERS_PER_PROXY) as pool: + callers: Final = tuple( + pool.submit(_keep_calling, proxy, model, key, stop) + for proxy in (gateway, peer) * IN_FLIGHT_CALLERS_PER_PROXY + ) + try: + eventually( + lambda: len(_models(upstream)), lambda served: served >= 4 * IN_FLIGHT_CALLERS_PER_PROXY, seconds=30 + ) + blocked: Final = gateway.request("POST", "/team/block", {"team_id": team}) + finally: + stop.set() + statuses: Final = tuple(itertools.chain.from_iterable(caller.result() for caller in callers)) + return blocked, statuses + + +def _block_and_unblock_under_team_traffic( + gateway: Gateway, peer: Gateway, upstream: httpx.Client, model: str, key: str, team: str, cycle: int +) -> None: + upstream.delete("/__observations").raise_for_status() + blocked, in_flight = _block_while_team_traffic_is_in_flight(gateway, peer, upstream, model, key, team) + assert blocked.status_code == 200, blocked.text + assert in_flight.count(200) >= 4 * IN_FLIGHT_CALLERS_PER_PROXY, repr(in_flight) + assert set(in_flight) <= {200, 401}, repr(in_flight) + + marker: Final = f"after block {cycle} {uuid.uuid4().hex}" + refused: Final = held( + lambda: _statuses_on_fresh_connections((gateway, peer), model, key, marker), + lambda statuses: set(statuses) == {401}, + holding=5, + seconds=15, + ) + assert refused == (401,) * (2 * PROBES_PER_PROXY), repr(refused) + leaked: Final = tuple(text for text in _texts(upstream) if text.startswith(marker)) + assert leaked == (), repr(leaked) + + unblocked: Final = gateway.request("POST", "/team/unblock", {"team_id": team}) + assert unblocked.status_code == 200, unblocked.text + eventually( + lambda: _statuses_on_fresh_connections((gateway, peer), model, key, f"after unblock {cycle}"), + lambda statuses: set(statuses) == {200}, + seconds=15, + ) + + def test_team_block_and_unblock_update_endpoint_info_and_db_state( gateway: Gateway, peer: Gateway, upstream: httpx.Client ) -> None: @@ -126,8 +224,6 @@ def test_team_block_and_unblock_update_endpoint_info_and_db_state( def test_unblocking_a_team_restores_warmed_keys_on_both_proxies( gateway: Gateway, peer: Gateway, upstream: httpx.Client ) -> None: - pytest.skip("BUG: SLP-605 /team/unblock does not restore warmed team-key access after blocked auth is cached") - with gateway.scenario() as scenario: provider_model: Final = f"team-block-{uuid.uuid4().hex}" model: Final = scenario.model(model=f"openai/{provider_model}") @@ -166,8 +262,6 @@ def test_unblocking_a_team_restores_warmed_keys_on_both_proxies( def test_warmed_team_keys_are_refused_on_both_proxies_after_block( gateway: Gateway, peer: Gateway, upstream: httpx.Client ) -> None: - pytest.skip("BUG: SLP-605 blocking a team does not refresh warmed team-key auth state") - with gateway.scenario() as scenario: provider_model: Final = f"team-block-{uuid.uuid4().hex}" model: Final = scenario.model(model=f"openai/{provider_model}") @@ -202,8 +296,6 @@ def test_warmed_team_keys_are_refused_on_both_proxies_after_block( def test_a_new_team_key_is_refused_on_a_peer_that_cached_the_team_before_block( gateway: Gateway, peer: Gateway, upstream: httpx.Client ) -> None: - pytest.skip("BUG: SLP-605 a peer proxy that cached the team before /team/block admits a never-seen team key") - with gateway.scenario() as scenario: provider_model: Final = f"team-block-{uuid.uuid4().hex}" model: Final = scenario.model(model=f"openai/{provider_model}") @@ -229,3 +321,45 @@ def test_a_new_team_key_is_refused_on_a_peer_that_cached_the_team_before_block( assert peer_denied.status_code == 401, peer_denied.text assert _ErrorResponse.model_validate_json(peer_denied.text).error.type == "auth_error", peer_denied.text assert after_block == (), repr(after_block) + + +def test_a_team_model_alias_still_routes_on_both_proxies_after_block_and_unblock( + gateway: Gateway, peer: Gateway, upstream: httpx.Client +) -> None: + with gateway.scenario() as scenario: + provider_model: Final = f"team-block-{uuid.uuid4().hex}" + model: Final = scenario.model(model=f"openai/{provider_model}") + alias: Final = f"alias-{uuid.uuid4().hex}" + team: Final = scenario.team(models=[model], model_aliases={alias: model}) + key: Final = scenario.key(team_id=team, models=[model]) + + warmed: Final = tuple( + _chat(proxy, alias, key, f"team alias warm {uuid.uuid4().hex}").status_code for proxy in (gateway, peer) + ) + assert warmed == (200, 200), repr(warmed) + assert _models(upstream) == (provider_model, provider_model) + upstream.delete("/__observations").raise_for_status() + + blocked: Final = gateway.request("POST", "/team/block", {"team_id": team}) + assert blocked.status_code == 200, blocked.text + unblocked: Final = gateway.request("POST", "/team/unblock", {"team_id": team}) + assert unblocked.status_code == 200, unblocked.text + + after_unblock: Final = _statuses_on_fresh_connections((gateway, peer), alias, key, "team alias after unblock") + assert after_unblock == (200,) * (2 * PROBES_PER_PROXY), repr(after_unblock) + assert _models(upstream) == (provider_model,) * (2 * PROBES_PER_PROXY) + + +def test_a_block_issued_while_team_traffic_is_in_flight_holds_on_every_worker( + gateway: Gateway, peer: Gateway, upstream: httpx.Client +) -> None: + with gateway.scenario() as scenario: + provider_model: Final = f"team-block-{uuid.uuid4().hex}" + model: Final = scenario.model( + model=f"openai/{provider_model}", input_cost_per_token=0.000001, output_cost_per_token=0.000002 + ) + team: Final = scenario.team(models=[model]) + key: Final = scenario.key(team_id=team) + + for cycle in range(BLOCK_CYCLES): + _block_and_unblock_under_team_traffic(gateway, peer, upstream, model, key, team, cycle) diff --git a/tests/unit/proxy/auth/test_jwt.py b/tests/unit/proxy/auth/test_jwt.py index a7aa379ad37..fb512539126 100644 --- a/tests/unit/proxy/auth/test_jwt.py +++ b/tests/unit/proxy/auth/test_jwt.py @@ -860,35 +860,6 @@ async def test_allowed_routes_admin( import pytest -@pytest.mark.asyncio -async def test_team_cache_update_called(): - import litellm - from litellm.proxy.proxy_server import user_api_key_cache - - # Use setattr to replace the method on the user_api_key_cache object - cache = DualCache() - - setattr( - litellm.proxy.proxy_server, - "user_api_key_cache", - cache, - ) - - with patch.object(cache, "async_batch_get_cache", new=AsyncMock(return_value=[None])) as mock_call_cache: - # Call the function under test - await litellm.proxy.proxy_server.update_cache( - token=None, - user_id=None, - end_user_id=None, - team_id="1234", - response_cost=20, - parent_otel_span=None, - ) # type: ignore - - await asyncio.sleep(3) - mock_call_cache.assert_awaited_once_with(keys=["team_id:1234"], parent_otel_span=None, throttle_redis=False) - - @pytest.fixture def public_jwt_key(): import json diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index a379b424aeb..5063ae0b357 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -2636,10 +2636,99 @@ async def test_team_model_add_delete_keep_model_aliases_in_team_cache(endpoint_n assert team_model_aliases(cached_team) == {"fast": "gpt-4o"} +@pytest.mark.asyncio +@pytest.mark.parametrize(("endpoint_name", "blocked"), [("block_team", True), ("unblock_team", False)]) +async def test_team_block_unblock_reach_warmed_auth_on_every_worker( + endpoint_name: str, blocked: bool, monkeypatch: pytest.MonkeyPatch +): + import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module + from litellm.proxy._types import BlockTeamRequest + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.auth.team_grants import team_model_aliases + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints import team_endpoints + + team_id: Final = "team-kill-switch" + columns: Final = {"team_id": team_id, "team_alias": "kill-switch", "blocked": not blocked} + alias_table: Final = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"} + + async def update( + where: Mapping[str, str], data: Mapping[str, bool], include: Mapping[str, bool] | None = None + ) -> LiteLLM_TeamTable: + relations: Final = {"litellm_model_table": alias_table} if (include or {}).get("litellm_model_table") else {} + return LiteLLM_TeamTable.model_validate({**columns, **data, **relations}) + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(**columns)) + prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=update) + cache: Final = UserApiKeyCache() + redis: Final = _BroadcastRecordingRedis() + monkeypatch.setattr(pubsub_module, "coordination_redis_cache", lambda: redis) + monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + + warmed: Final = await get_team_object(team_id=team_id, prisma_client=prisma_client, user_api_key_cache=cache) + assert warmed.blocked is not blocked + + response: Final = await getattr(team_endpoints, endpoint_name)( + data=BlockTeamRequest(team_id=team_id), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + await asyncio.gather(*pubsub_module._pending_publishes) # pyright: ignore[reportPrivateUsage] # drain module-level tasks + + seen_by_auth: Final = await get_team_object(team_id=team_id, prisma_client=prisma_client, user_api_key_cache=cache) + assert seen_by_auth.blocked is blocked + assert team_model_aliases(seen_by_auth) == {"fast": "gpt-4o"} + assert redis.broadcast_keys == [f"team_id:{team_id}", "team_alias:kill-switch"] + assert (response.blocked, response.litellm_model_table, response.object_permission) == (blocked, None, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint_name", ["block_team", "unblock_team"]) +@pytest.mark.parametrize( + ("caller", "stored_team", "status_code"), + [ + ( + UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice"), + LiteLLM_TeamTable(team_id="team-1234", members_with_roles=[Member(user_id="alice", role="user")]), + 403, + ), + (UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), None, 404), + ], +) +async def test_team_block_unblock_refuse_before_writing( + endpoint_name: str, + caller: UserAPIKeyAuth, + stored_team: LiteLLM_TeamTable | None, + status_code: int, + monkeypatch: pytest.MonkeyPatch, +): + from litellm.proxy._types import BlockTeamRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints import team_endpoints + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=stored_team) + prisma_client.db.litellm_teamtable.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + + with pytest.raises(HTTPException) as exc_info: + await getattr(team_endpoints, endpoint_name)( + data=BlockTeamRequest(team_id="team-1234"), http_request=MagicMock(), user_api_key_dict=caller + ) + + assert exc_info.value.status_code == status_code + prisma_client.db.litellm_teamtable.update.assert_not_awaited() + + @pytest.mark.asyncio @pytest.mark.parametrize( "endpoint_name", - ["team_model_add", "team_model_delete", "update_team_member_permissions"], + ["team_model_add", "team_model_delete", "update_team_member_permissions", "block_team", "unblock_team"], ) async def test_team_write_404s_when_row_vanishes_before_update(endpoint_name): """A team deleted between the read and the write must 404. @@ -2662,9 +2751,12 @@ async def test_team_write_404s_when_row_vanishes_before_update(endpoint_name): TeamModelDeleteRequest, UserAPIKeyAuth, ) + from litellm.proxy._types import BlockTeamRequest from litellm.proxy.management_endpoints.team_endpoints import ( + block_team, team_model_add, team_model_delete, + unblock_team, update_team_member_permissions, ) @@ -2701,6 +2793,16 @@ async def test_team_write_404s_when_row_vanishes_before_update(endpoint_name): http_request=mock_request, user_api_key_dict=mock_user_api_key_dict, ), + "block_team": lambda: block_team( + data=BlockTeamRequest(team_id="team-1234"), + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ), + "unblock_team": lambda: unblock_team( + data=BlockTeamRequest(team_id="team-1234"), + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ), }[endpoint_name] with ( diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 01400384842..3ff4293ba93 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -22,6 +22,7 @@ Pins covered: from __future__ import annotations import asyncio +from collections.abc import Awaitable, Callable from datetime import datetime from typing import Final, cast from unittest.mock import AsyncMock, MagicMock @@ -29,6 +30,10 @@ from unittest.mock import AsyncMock, MagicMock import pytest import litellm.proxy.proxy_server as ps +from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team from .conftest import normalize @@ -1575,7 +1580,6 @@ async def test_update_cache_no_cached_entities_schedules_pipeline_flush(monkeypa "u1", f"{ps.litellm_proxy_admin_name}:spend", "end_user_id:eu1", - "team_id:t1", "tag:x", ] @@ -1600,3 +1604,70 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey ) assert result is None + + +async def _drain_spend_cache_writes() -> None: + for _ in range(3): + await asyncio.sleep(0) + + +async def _update_team_spend(team_id: str, response_cost: float) -> None: + await ps.update_cache( + token=None, + user_id=None, + end_user_id=None, + team_id=team_id, + response_cost=response_cost, + parent_otel_span=None, + ) + + +@pytest.mark.asyncio +async def test_update_cache_adds_spend_to_the_cached_team_only(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = UserApiKeyCache() + cache.set_cache(key="team_id:warm", value=LiteLLM_TeamTableCachedObj(team_id="warm", spend=5.0, blocked=True)) + monkeypatch.setattr(ps, "user_api_key_cache", cache) + + await _update_team_spend(team_id="warm", response_cost=2.5) + await _update_team_spend(team_id="cold", response_cost=2.5) + await _drain_spend_cache_writes() + + warm: Final = cache.get_cache(key="team_id:warm", model_type=LiteLLM_TeamTableCachedObj) + assert warm is not None + assert (warm.spend, warm.blocked) == (7.5, True) + assert cache.get_cache(key="team_id:cold") is None + + +async def _block_on_this_worker(cache: UserApiKeyCache, team_id: str) -> None: + await refresh_cached_team( + team_row=LiteLLM_TeamTable(team_id=team_id, blocked=True), + user_api_key_cache=cache, + proxy_logging_obj=ps.proxy_logging_obj, + ) + + +async def _invalidated_by_a_peer(cache: UserApiKeyCache, team_id: str) -> None: + evict_local(cache_keys=(f"team_id:{team_id}",), user_api_key_cache=cache) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("management_write", "expected_blocked"), + [(_block_on_this_worker, True), (_invalidated_by_a_peer, None)], +) +async def test_update_cache_never_restores_a_team_a_management_write_replaced( + management_write: Callable[[UserApiKeyCache, str], Awaitable[None]], + expected_blocked: bool | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + team_id: Final = "kill-switch" + cache: Final = UserApiKeyCache() + cache.set_cache(key=f"team_id:{team_id}", value=LiteLLM_TeamTableCachedObj(team_id=team_id, blocked=False)) + monkeypatch.setattr(ps, "user_api_key_cache", cache) + + await _update_team_spend(team_id=team_id, response_cost=1.0) + await management_write(cache, team_id) + await _drain_spend_cache_writes() + + seen_by_auth: Final = cache.get_cache(key=f"team_id:{team_id}", model_type=LiteLLM_TeamTableCachedObj) + assert (seen_by_auth.blocked if seen_by_auth is not None else None) == expected_blocked From 48763ccb505a98b04933d72f50cfdd89ae56f5da Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 19 Oct 2026 15:58:29 -0700 Subject: [PATCH 13/18] fix(scim)!: store PATCH enterprise extension attributes in metadata.scim_enterprise (#46738) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/scim/scim_v2.py | 259 +++++++- .../scim/test_scim_patch_user.py | 556 ++++++++++++++++++ 2 files changed, 805 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 97a6a29ef66..ec905a9f187 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -5,13 +5,13 @@ This is an enterprise feature and requires a premium license. """ import re -from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass -from functools import partial +from functools import partial, reduce from itertools import chain, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypeAlias, overload +from typing import TYPE_CHECKING, Final, NamedTuple, NoReturn, Protocol, TypeAlias, overload from fastapi import ( APIRouter, @@ -24,6 +24,7 @@ from fastapi import ( Response, ) from pydantic import TypeAdapter, ValidationError +from pydantic.fields import FieldInfo from typing_extensions import ReadOnly, TypedDict, assert_never import litellm @@ -2271,6 +2272,217 @@ def _handle_multi_valued_attribute_update(path: str, op_type: str, value: object metadata[metadata_key] = dumped +def _enterprise_extension_sub_path(path: str) -> str | None: + schema: Final = SCIM_ENTERPRISE_USER_SCHEMA + if path.casefold() == schema.casefold(): + return "" + prefix: Final = f"{schema}:" + if path.casefold().startswith(prefix.casefold()): + return path[len(prefix) :] + return None + + +def _scim_model_field_names(model_fields: Mapping[str, FieldInfo]) -> Mapping[str, str]: + return { + **{field_name.casefold(): field_name for field_name in model_fields}, + **{ + field.alias.casefold(): field_name + for field_name, field in model_fields.items() + if isinstance(field.alias, str) + }, + } + + +def _canonicalize_scim_model_data( + value: Mapping[str, object], + model_fields: Mapping[str, FieldInfo], +) -> Mapping[str, object]: + field_names: Final = _scim_model_field_names(model_fields) + return {field_names.get(key.casefold(), key): field_value for key, field_value in value.items()} + + +def _canonicalize_enterprise_data(value: Mapping[str, object]) -> Mapping[str, object]: + enterprise_data: Final = _canonicalize_scim_model_data(value, SCIMEnterpriseUser.model_fields) + manager: Final = enterprise_data.get("manager") + if isinstance(manager, str): + return {**enterprise_data, "manager": {"value": manager}} + if not isinstance(manager, dict): + return enterprise_data + return { + **enterprise_data, + "manager": _canonicalize_scim_model_data( + _JSON_OBJECT_ADAPTER.validate_python(manager), + SCIMUserManager.model_fields, + ), + } + + +@dataclass(frozen=True, slots=True) +class _InvalidEnterpriseValue: + detail: str + + +def _merge_enterprise_data( + existing: Mapping[str, object], + incoming: Mapping[str, object], + incoming_manager_is_object: bool, +) -> Mapping[str, object]: + merged: Final[Mapping[str, object]] = {**existing, **incoming} + existing_manager: Final = existing.get("manager") + incoming_manager: Final = incoming.get("manager") + if incoming_manager_is_object and isinstance(existing_manager, dict) and isinstance(incoming_manager, dict): + existing_manager_data: Final = _canonicalize_scim_model_data( + _JSON_OBJECT_ADAPTER.validate_python(existing_manager), + SCIMUserManager.model_fields, + ) + incoming_manager_data: Final = _JSON_OBJECT_ADAPTER.validate_python(incoming_manager) + return {**merged, "manager": {**existing_manager_data, **incoming_manager_data}} + return merged + + +def _validate_enterprise_update( + value: Mapping[str, object], +) -> Mapping[str, object] | None | _InvalidEnterpriseValue: + try: + enterprise: Final = SCIMEnterpriseUser.model_validate(value) + except ValidationError: + return _InvalidEnterpriseValue(detail="Invalid value for the enterprise extension") + dumped: Final = enterprise.model_dump(by_alias=True, exclude_none=True) + return dumped if dumped else None + + +def _raise_invalid_enterprise_value(error: _InvalidEnterpriseValue) -> NoReturn: + raise HTTPException(status_code=400, detail={"error": error.detail}) + + +def _handle_enterprise_update( + sub_path: str, + op_type: str, + value: object, + existing_enterprise: Mapping[str, object], +) -> Mapping[str, object] | None | _InvalidEnterpriseValue: + if op_type == "remove" and not sub_path: + return None + + if not sub_path: + if not isinstance(value, dict): + return _InvalidEnterpriseValue(detail="The enterprise extension value must be an object") + incoming_value: Final = _JSON_OBJECT_ADAPTER.validate_python(value) + incoming: Final = _canonicalize_enterprise_data(incoming_value) + incoming_manager_is_object: Final = isinstance( + _canonicalize_scim_model_data(incoming_value, SCIMEnterpriseUser.model_fields).get("manager"), + dict, + ) + return _validate_enterprise_update( + _merge_enterprise_data(existing_enterprise, incoming, incoming_manager_is_object) + ) + + path_parts: Final = sub_path.split(".") + if len(path_parts) == 1: + field_name: Final = _scim_model_field_names(SCIMEnterpriseUser.model_fields).get(sub_path.casefold()) + if field_name is None: + return _validate_enterprise_update(existing_enterprise) + if op_type == "remove": + return _validate_enterprise_update( + {key: attr for key, attr in existing_enterprise.items() if key != field_name} + ) + if value is None: + return _InvalidEnterpriseValue(detail=f"The {op_type} operation requires a 'value' member") + field_incoming: Final = _canonicalize_enterprise_data({field_name: value}) + field_incoming_manager_is_object: Final = field_name == "manager" and isinstance(value, dict) + return _validate_enterprise_update( + _merge_enterprise_data( + existing_enterprise, + field_incoming, + field_incoming_manager_is_object, + ) + ) + + if len(path_parts) == 2 and path_parts[0].casefold() == "manager": + manager_field: Final = _scim_model_field_names(SCIMUserManager.model_fields).get(path_parts[1].casefold()) + if manager_field is None: + return _validate_enterprise_update(existing_enterprise) + existing_manager: Final = existing_enterprise.get("manager") + manager_data: Final[Mapping[str, object]] = ( + _canonicalize_scim_model_data( + _JSON_OBJECT_ADAPTER.validate_python(existing_manager), + SCIMUserManager.model_fields, + ) + if isinstance(existing_manager, dict) + else {} + ) + if op_type == "remove": + removed_manager: Final[Mapping[str, object]] = { + key: attr for key, attr in manager_data.items() if key != manager_field + } + updated_enterprise: Final[Mapping[str, object]] = ( + {**existing_enterprise, "manager": removed_manager} + if removed_manager + else {key: attr for key, attr in existing_enterprise.items() if key != "manager"} + ) + return _validate_enterprise_update(updated_enterprise) + if value is None: + return _InvalidEnterpriseValue(detail=f"The {op_type} operation requires a 'value' member") + updated_manager: Final[Mapping[str, object]] = {**manager_data, manager_field: value} + return _validate_enterprise_update({**existing_enterprise, "manager": updated_manager}) + + return _validate_enterprise_update(existing_enterprise) + + +def _enterprise_update_for_pathless_key( + key_and_value: tuple[str, object], + *, + op_type: str, +) -> Iterator[tuple[str, str, object]]: + key, value = key_and_value + sub_path: Final = _enterprise_extension_sub_path(key) + if sub_path is not None: + yield sub_path, op_type, value + + +def _enterprise_updates_for_operation( + op: SCIMPatchOperation, +) -> Iterator[tuple[str, str, object]]: + path: Final = op.path.lower() if isinstance(op.path, str) else "" + value: Final[object] = TypeAdapter(object).validate_python(op.value) + if not path and isinstance(value, dict): + pathless_values: Final[Mapping[str, object]] = _JSON_OBJECT_ADAPTER.validate_python(value) + yield from chain.from_iterable( + map(partial(_enterprise_update_for_pathless_key, op_type=op.op), pathless_values.items()) + ) + return + sub_path: Final = _enterprise_extension_sub_path(path) + if sub_path is not None: + yield sub_path, op.op, value + + +def _enterprise_updates(operations: Sequence[SCIMPatchOperation]) -> Iterator[tuple[str, str, object]]: + yield from chain.from_iterable(map(_enterprise_updates_for_operation, operations)) + + +def _apply_enterprise_update( + state: Mapping[str, object] | None | _InvalidEnterpriseValue, + update: tuple[str, str, object], +) -> Mapping[str, object] | None | _InvalidEnterpriseValue: + if isinstance(state, _InvalidEnterpriseValue): + return state + sub_path, op_type, value = update + existing_enterprise: Final[Mapping[str, object]] = {} if state is None else state + return _handle_enterprise_update(sub_path, op_type, value, existing_enterprise) + + +def _apply_enterprise_updates( + updates: Sequence[tuple[str, str, object]], + stored_enterprise: object, +) -> Mapping[str, object] | None | _InvalidEnterpriseValue: + initial_enterprise: Final[Mapping[str, object] | None] = ( + _canonicalize_enterprise_data(_JSON_OBJECT_ADAPTER.validate_python(stored_enterprise)) + if isinstance(stored_enterprise, dict) + else None + ) + return reduce(_apply_enterprise_update, updates, initial_enterprise) + + def _handle_generic_metadata(path: str, op_type: str, value: object, metadata: dict[str, object]) -> None: """Handle generic metadata operations for unknown paths.""" if op_type == "remove": @@ -2285,20 +2497,37 @@ def _apply_patch_ops( ) -> tuple[dict[str, object], set[str]]: """Apply patch operations and return update data and final team set.""" update_data: Final[dict[str, object]] = {} - metadata: Final = existing_user.metadata or {} - scim_metadata: Final = metadata.get("scim_metadata", {}) + stored_metadata: Final[object] = TypeAdapter(object).validate_python(existing_user.model_dump().get("metadata")) + metadata: Final = _JSON_OBJECT_ADAPTER.validate_python(stored_metadata) if isinstance(stored_metadata, dict) else {} + stored_scim_metadata: Final = metadata.get("scim_metadata") + scim_metadata: Final = ( + _JSON_OBJECT_ADAPTER.validate_python(stored_scim_metadata) if isinstance(stored_scim_metadata, dict) else {} + ) + enterprise_updates: Final = tuple(_enterprise_updates(patch_ops.Operations)) + enterprise: Final[Mapping[str, object] | None | _InvalidEnterpriseValue] = ( + _apply_enterprise_updates( + enterprise_updates, + metadata.get(SCIM_ENTERPRISE_METADATA_KEY), + ) + if enterprise_updates + else None + ) + if isinstance(enterprise, _InvalidEnterpriseValue): + _raise_invalid_enterprise_value(enterprise) teams_set: Final[set[str]] = set(existing_user.teams or []) replace_team_set: set[str] | None = None for op in patch_ops.Operations: - path = (op.path or "").lower() + path = op.path.lower() if isinstance(op.path, str) else "" value = op.value op_type = op.op # Handle SCIM operations without path where value contains the fields if not path and isinstance(value, dict): - for key, val in value.items(): + for key, val in _JSON_OBJECT_ADAPTER.validate_python(value).items(): + if _enterprise_extension_sub_path(key) is not None: + continue key_lower = key.lower() if key_lower == "active": _handle_active_update(op_type, val, metadata) @@ -2309,7 +2538,7 @@ def _apply_patch_ops( elif key_lower in SCIM_MULTI_VALUED_ATTRIBUTE_METADATA_KEYS: _handle_multi_valued_attribute_update(key_lower, op_type, val, metadata) elif key_lower == "name" and isinstance(val, dict): - for name_key, name_val in val.items(): + for name_key, name_val in _JSON_OBJECT_ADAPTER.validate_python(val).items(): name_key_lower = name_key.lower() if name_key_lower in ("givenname", "familyname"): _handle_name_update( @@ -2334,12 +2563,22 @@ def _apply_patch_ops( new_replace_set = _handle_group_operations(op_type, value, teams_set, op.path) if new_replace_set is not None: replace_team_set = new_replace_set + elif _enterprise_extension_sub_path(path) is not None: + continue else: _handle_generic_metadata(path, op_type, value, metadata) final_team_set: Final = replace_team_set if replace_team_set is not None else teams_set - metadata["scim_metadata"] = scim_metadata - update_data["metadata"] = metadata + final_metadata: Final = ( + { + **{key: metadata_value for key, metadata_value in metadata.items() if key != SCIM_ENTERPRISE_METADATA_KEY}, + "scim_metadata": scim_metadata, + **({SCIM_ENTERPRISE_METADATA_KEY: enterprise} if enterprise is not None else {}), + } + if enterprise_updates + else {**metadata, "scim_metadata": scim_metadata} + ) + update_data["metadata"] = final_metadata return update_data, final_team_set diff --git a/tests/unit/proxy/management_endpoints/scim/test_scim_patch_user.py b/tests/unit/proxy/management_endpoints/scim/test_scim_patch_user.py index edcce16ab41..31c5a54a753 100644 --- a/tests/unit/proxy/management_endpoints/scim/test_scim_patch_user.py +++ b/tests/unit/proxy/management_endpoints/scim/test_scim_patch_user.py @@ -1,11 +1,17 @@ +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from pydantic import TypeAdapter from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.management_endpoints.scim.scim_v2 import _apply_patch_ops, patch_user from litellm.types.proxy.management_endpoints.scim_v2 import ( + SCIM_ENTERPRISE_METADATA_KEY, + SCIM_ENTERPRISE_USER_SCHEMA, + SCIMEnterpriseUser, SCIMPatchOp, SCIMPatchOperation, SCIMUser, @@ -536,3 +542,553 @@ def test_apply_patch_ops_replace_groups_empty_value_does_not_use_path_filter(): _, final_team_set = _apply_patch_ops(existing_user=user, patch_ops=patch_ops) assert final_team_set == set() + + +@pytest.mark.asyncio +async def test_patch_user_pathless_enterprise_extension_merges_existing_metadata(): + mock_user: Final = _user_with_metadata({SCIM_ENTERPRISE_METADATA_KEY: {"employeeNumber": "42"}}) + mock_client: Final = MagicMock() + mock_db: Final = MagicMock() + mock_client.db = mock_db + mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) + + async def mock_update(*, where: Mapping[str, object], data: Mapping[str, object]) -> LiteLLM_UserTable: + assert where == {"user_id": "user-mva"} + metadata_json: Final = data["metadata"] + assert isinstance(metadata_json, str) + metadata: Final = TypeAdapter(dict[str, object]).validate_json(metadata_json) + return _user_with_metadata(metadata) + + mock_db.litellm_usertable.update = AsyncMock(side_effect=mock_update) + mock_db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="replace", + value={ + SCIM_ENTERPRISE_USER_SCHEMA: { + "department": "Eng", + "costCenter": "CC-10", + } + }, + ) + ] + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_client), + ): + result: Final = await patch_user(user_id="user-mva", patch_ops=patch_ops) + + update_data: Final = mock_db.litellm_usertable.update.call_args.kwargs["data"] + metadata_json: Final = update_data["metadata"] + assert isinstance(metadata_json, str) + metadata: Final = TypeAdapter(dict[str, object]).validate_json(metadata_json) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "employeeNumber": "42", + "department": "Eng", + "costCenter": "CC-10", + } + assert not any(isinstance(key, str) and key.casefold().startswith("urn:") for key in metadata) + assert result.enterprise_user == SCIMEnterpriseUser( + employeeNumber="42", + department="Eng", + costCenter="CC-10", + ) + + +def test_apply_patch_ops_replace_enterprise_urn_path_from_idp(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="replace", + path="urn:ietf:params:scim:schemas:extension:enterprise:2.0:User:department", + value="Eng", + ) + ] + ) + + update_data, _ = _apply_patch_ops(existing_user=_user_with_metadata({}), patch_ops=patch_ops) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == {"department": "Eng"} + assert not any(isinstance(key, str) and key.casefold().startswith("urn:") for key in metadata) + + +def test_apply_patch_ops_pathless_enterprise_manager_sub_attribute_merges(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="replace", + value={ + f"{SCIM_ENTERPRISE_USER_SCHEMA.upper()}:MANAGER.VALUE": "manager-1", + }, + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + {SCIM_ENTERPRISE_METADATA_KEY: {"manager": {"displayName": "Engineering Manager"}}} + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "manager": {"value": "manager-1", "displayName": "Engineering Manager"} + } + + +def test_apply_patch_ops_replace_enterprise_manager_normalizes_case_insensitive_keys(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="replace", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:manager", + value={"VALUE": "mgr-2"}, + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + { + SCIM_ENTERPRISE_METADATA_KEY: { + "employeeNumber": "42", + "manager": {"value": "mgr-1", "displayName": "Boss"}, + } + } + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "employeeNumber": "42", + "manager": {"value": "mgr-2", "displayName": "Boss"}, + } + + +@pytest.mark.parametrize( + ("path", "value"), + ( + pytest.param( + f"{SCIM_ENTERPRISE_USER_SCHEMA}:manager", + {"displayName": "Q"}, + id="pathed", + ), + pytest.param( + None, + {SCIM_ENTERPRISE_USER_SCHEMA: {"manager": {"displayName": "Q"}}}, + id="pathless", + ), + pytest.param( + None, + {SCIM_ENTERPRISE_USER_SCHEMA: {"MANAGER": {"displayName": "Q"}}}, + id="pathless-uppercase-key", + ), + ), +) +def test_apply_patch_ops_replace_manager_object_preserves_unspecified_sub_attributes( + path: str | None, + value: object, +) -> None: + patch_ops: Final = SCIMPatchOp( + Operations=[SCIMPatchOperation(op="replace", path=path, value=value)] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + { + SCIM_ENTERPRISE_METADATA_KEY: { + "department": "A", + "manager": {"value": "m", "displayName": "B"}, + } + } + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "department": "A", + "manager": {"value": "m", "displayName": "Q"}, + } + + +@pytest.mark.parametrize( + ("manager_path", "manager_value"), + ( + pytest.param( + f"{SCIM_ENTERPRISE_USER_SCHEMA}:manager", + {"$ref": "new"}, + id="pathed", + ), + pytest.param( + None, + {SCIM_ENTERPRISE_USER_SCHEMA: {"manager": {"$ref": "new"}}}, + id="pathless", + ), + ), +) +def test_apply_patch_ops_batched_manager_ref_update_replaces_previous_link( + manager_path: str | None, + manager_value: object, +) -> None: + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="replace", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:department", + value="B", + ), + SCIMPatchOperation( + op="replace", + path=manager_path, + value=manager_value, + ), + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + { + SCIM_ENTERPRISE_METADATA_KEY: { + "department": "A", + "manager": {"value": "m", "$ref": "old"}, + } + } + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "department": "B", + "manager": {"value": "m", "$ref": "new"}, + } + + +def test_apply_patch_ops_add_enterprise_manager_merges_existing_sub_attributes(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="add", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:manager", + value={"displayName": "New Boss"}, + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + { + SCIM_ENTERPRISE_METADATA_KEY: { + "employeeNumber": "42", + "manager": { + "value": "mgr-1", + "$ref": "https://idp.example/Users/mgr-1", + }, + } + } + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "employeeNumber": "42", + "manager": { + "value": "mgr-1", + "$ref": "https://idp.example/Users/mgr-1", + "displayName": "New Boss", + }, + } + + +def test_apply_patch_ops_pathless_add_enterprise_manager_merges_existing_sub_attributes(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="add", + value={ + SCIM_ENTERPRISE_USER_SCHEMA: { + "manager": {"displayName": "New Boss"}, + } + }, + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + { + SCIM_ENTERPRISE_METADATA_KEY: { + "employeeNumber": "42", + "manager": { + "value": "mgr-1", + "$ref": "https://idp.example/Users/mgr-1", + }, + } + } + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "employeeNumber": "42", + "manager": { + "value": "mgr-1", + "$ref": "https://idp.example/Users/mgr-1", + "displayName": "New Boss", + }, + } + + +def test_apply_patch_ops_remove_enterprise_manager_sub_attribute_preserves_other_fields(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="remove", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:manager.value", + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + { + SCIM_ENTERPRISE_METADATA_KEY: { + "employeeNumber": "42", + "manager": {"value": "mgr-1", "displayName": "Boss"}, + } + } + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "employeeNumber": "42", + "manager": {"displayName": "Boss"}, + } + + +def test_apply_patch_ops_remove_enterprise_attribute_preserves_other_attributes(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="remove", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:costCenter", + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata({SCIM_ENTERPRISE_METADATA_KEY: {"costCenter": "CC-10", "department": "Eng"}}), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == {"department": "Eng"} + + +def test_apply_patch_ops_remove_last_enterprise_attribute_drops_metadata_key(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="remove", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:employeeNumber", + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata({SCIM_ENTERPRISE_METADATA_KEY: {"employeeNumber": "42"}}), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert SCIM_ENTERPRISE_METADATA_KEY not in metadata + + +def test_apply_patch_ops_unknown_enterprise_attribute_preserves_existing_and_applies_other_ops() -> None: + existing_enterprise: Final = { + "employeeNumber": "42", + "manager": {"value": "mgr-1", "displayName": "Boss"}, + } + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="replace", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:shoeSize", + value="10", + ), + SCIMPatchOperation(op="replace", path="displayName", value="Updated User"), + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata({SCIM_ENTERPRISE_METADATA_KEY: existing_enterprise}), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == existing_enterprise + assert not any(isinstance(key, str) and key.casefold().startswith("urn:") for key in metadata) + assert update_data["user_alias"] == "Updated User" + + +def test_apply_patch_ops_display_name_only_preserves_enterprise_value_unchanged() -> None: + existing_enterprise: Final = {"Department": "Eng", "manager": "mgr-1"} + patch_ops: Final = SCIMPatchOp( + Operations=[SCIMPatchOperation(op="replace", path="displayName", value="Ada")] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata({SCIM_ENTERPRISE_METADATA_KEY: existing_enterprise}), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == {"Department": "Eng", "manager": "mgr-1"} + + +@pytest.mark.parametrize("path", ("manager.foo", "manager.a.b", "department.x")) +def test_apply_patch_ops_unknown_enterprise_nested_attribute_preserves_existing(path: str) -> None: + existing_enterprise: Final = { + "department": "Engineering", + "manager": {"value": "mgr-1", "displayName": "Boss"}, + } + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="replace", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:{path}", + value="ignored", + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata({SCIM_ENTERPRISE_METADATA_KEY: existing_enterprise}), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == existing_enterprise + + +@pytest.mark.parametrize("manager_id", ("mgr-id", "")) +def test_apply_patch_ops_add_enterprise_manager_string_replaces_existing_manager(manager_id: str) -> None: + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="add", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:manager", + value=manager_id, + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + { + SCIM_ENTERPRISE_METADATA_KEY: { + "employeeNumber": "42", + "manager": {"value": "old", "displayName": "Old Boss"}, + } + } + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "employeeNumber": "42", + "manager": {"value": manager_id}, + } + + +def test_apply_patch_ops_pathless_add_enterprise_manager_string_replaces_existing_manager() -> None: + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="add", + value={SCIM_ENTERPRISE_USER_SCHEMA: {"manager": "mgr-id"}}, + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + { + SCIM_ENTERPRISE_METADATA_KEY: { + "employeeNumber": "42", + "manager": {"value": "old", "displayName": "Old Boss"}, + } + } + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert metadata[SCIM_ENTERPRISE_METADATA_KEY] == { + "employeeNumber": "42", + "manager": {"value": "mgr-id"}, + } + + +def test_apply_patch_ops_remove_enterprise_manager_drops_manager() -> None: + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="remove", + path=f"{SCIM_ENTERPRISE_USER_SCHEMA}:manager", + ) + ] + ) + + update_data, _ = _apply_patch_ops( + existing_user=_user_with_metadata( + {SCIM_ENTERPRISE_METADATA_KEY: {"manager": {"value": "mgr-1", "displayName": "Boss"}}} + ), + patch_ops=patch_ops, + ) + + metadata: Final = update_data["metadata"] + assert isinstance(metadata, dict) + assert SCIM_ENTERPRISE_METADATA_KEY not in metadata + + +def test_apply_patch_ops_pathless_enterprise_value_must_be_a_dict(): + patch_ops: Final = SCIMPatchOp( + Operations=[ + SCIMPatchOperation( + op="replace", + value={SCIM_ENTERPRISE_USER_SCHEMA: "not-a-dict"}, + ) + ] + ) + + with pytest.raises(HTTPException) as exc_info: + _apply_patch_ops(existing_user=_user_with_metadata({}), patch_ops=patch_ops) + + assert exc_info.value.status_code == 400 From 24d92f77d5816ddfdf4cbd1e76794207c47e9360 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:58:35 +0000 Subject: [PATCH 14/18] test(perf): stop flagging load-balanced worker reads as restarts in database upgrade test (#47425) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/perf/test_database_upgrade.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py index f165b9b738f..8079aa0274d 100644 --- a/tests/e2e/perf/test_database_upgrade.py +++ b/tests/e2e/perf/test_database_upgrade.py @@ -55,9 +55,9 @@ def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: Resource monitor.assert_complete(state) assert_finished(connection, state) rss_after: Final = read_rss_everywhere(proxy) - assert {r.worker for r in rss_before.readings} == {r.worker for r in rss_after.readings}, ( - "Gateway workers restarted" - ) + assert {(r.replica, r.hostname) for r in rss_before.readings} == { + (r.replica, r.hostname) for r in rss_after.readings + }, "Gateway pods replaced during the upgrade" observed: Final = Observation( records=outcome.records, mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)), From ad1f3c5e91b8c7fcb2bcb067e05bafb3dd4bd1be Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 19 Oct 2026 16:01:23 -0700 Subject: [PATCH 15/18] fix(mcp): retry upstream 429 on tools/list and error when every server fails (#46649) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 7 + .../proxy/_experimental/mcp_server/catalog.py | 28 +- .../mcp_server/faults/list_outcomes.py | 121 +++++++- .../mcp_server/mcp_server_manager.py | 13 +- .../_experimental/mcp_server/operations.py | 22 +- .../_experimental/mcp_server/subscriptions.py | 12 +- .../mcp_server/faults/test_list_outcomes.py | 285 ++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 150 ++++++++- .../test_mcp_server_tool_calls_and_headers.py | 201 +++++++++++- .../mcp_server/test_operations.py | 39 +++ .../mcp_server/test_subscriptions.py | 50 ++- 11 files changed, 903 insertions(+), 25 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c7c1576bc20..2d54a873b89 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -252,6 +252,13 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0) MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0) +MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES: Final = get_env_int("LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES", 3) +MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS: Final = get_env_float( + "LITELLM_MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS", 0.5 +) +MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS: Final = get_env_float( + "LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS", 5.0 +) MCP_METADATA_TIMEOUT: Final = get_env_float("LITELLM_MCP_METADATA_TIMEOUT", 10.0) MCP_HEALTH_CHECK_TIMEOUT: Final = get_env_float("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", 10.0) AGENT_HEALTH_CHECK_TIMEOUT_SECONDS: Final = get_env_float("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", 5.0) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index d0383be61a8..71eb2112bd5 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -17,6 +17,7 @@ from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeAlias, TypeVar, cast +from mcp.shared.exceptions import MCPError from mcp.types import CacheableResult from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter @@ -1055,12 +1056,20 @@ async def aggregate_gateway_tools( ) +class ToolListingFailedOutright(MCPError): + pass + + async def list_gateway_tools( context: OperationContext, params: PaginatedRequestParams, *, log_list_tools_to_spendlogs: bool = True ) -> ListToolsResult: - from mcp.types import ListToolsResult + from mcp.types import INTERNAL_ERROR, ListToolsResult - from litellm.proxy._experimental.mcp_server.faults.list_outcomes import SERVER_OUTCOMES_META_KEY, outcome_wire_value + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + listing_failed_outright, + outcome_wire_value, + ) from litellm.proxy._experimental.mcp_server.operations import _list_mcp_tools caller, auth, servers, server_headers, oauth_headers, headers, client_ip = context.legacy_auth() @@ -1078,16 +1087,21 @@ async def list_gateway_tools( record_listing=True, list_tools_log_source="mcp_protocol", ) + wire_outcomes: Final = ( + {key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items()} if listing.outcomes else None + ) + if params.cursor is None and listing_failed_outright(listing.outcomes): + raise ToolListingFailedOutright( + code=INTERNAL_ERROR, + message="Every MCP server in scope failed to list tools", + data={SERVER_OUTCOMES_META_KEY: wire_outcomes}, + ) return ListToolsResult( tools=listing.tools, next_cursor=listing.next_cursor, ttl_ms=listing.ttl_ms, cache_scope="private", - _meta={ - SERVER_OUTCOMES_META_KEY: {key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items()} - } - if listing.outcomes - else None, + _meta={SERVER_OUTCOMES_META_KEY: wire_outcomes} if wire_outcomes else None, ) diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index b94ed1ee251..84673c8c738 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -10,17 +10,27 @@ becomes an outcome, never a second failure. from __future__ import annotations +import email.utils +import random import re -from collections.abc import Iterator -from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias +from collections.abc import Awaitable, Callable, Iterator, Mapping +from datetime import datetime, timezone +from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias, TypeVar from urllib.request import parse_http_list, parse_keqv_list +import anyio import httpx import httpx2 from mcp.types import Tool as MCPTool from pydantic import ConfigDict, TypeAdapter from typing_extensions import assert_never +from litellm._logging import verbose_logger +from litellm.constants import ( + MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS, + MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS, + MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES, +) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, MCPUpstreamAuthError, @@ -176,6 +186,113 @@ def classify_list_exception(exc: BaseException) -> ServerListFault: return ServerListFault(tag="internal") +class UpstreamRateLimit(NamedTuple): + retry_after_seconds: float | None + + +def _parse_retry_after(value: str | None) -> float | None: + if value is None: + return None + stripped: Final = value.strip() + if re.fullmatch(r"\d+", stripped): + return float(stripped) + try: + parsed: Final = email.utils.parsedate_to_datetime(stripped) + except (TypeError, ValueError): + return None + aware: Final = parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc) + return max(0.0, (aware - datetime.now(timezone.utc)).total_seconds()) + + +def upstream_rate_limit(exc: BaseException) -> UpstreamRateLimit | None: + """The upstream 429 signal for the tools/list retry loop: an auth challenge wins over rate + limiting so a 401/403 anywhere in the tree keeps its own channel, then the first 429 response + in deliberate order carries its Retry-After hint. Anything else is not retryable here.""" + if upstream_auth_challenge(exc) is not None: + return None + for response in _iter_upstream_responses(exc): + if response.status_code == 429: + return UpstreamRateLimit(_parse_retry_after(response.headers.get("retry-after"))) + return None + + +def rate_limit_retry_delay( + attempt: int, + retry_after_seconds: float | None, + *, + base_delay: float, + max_delay: float, + jitter: float, +) -> float: + if retry_after_seconds is not None: + return min(retry_after_seconds, max_delay) + return min(max_delay, base_delay * 2**attempt * (1 + jitter)) + + +T: Final = TypeVar("T") + + +async def retry_upstream_rate_limit( + fetch: Callable[[], Awaitable[T]], + server_name: str, + *, + max_retries: int = MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES, + base_delay: float = MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS, + max_delay: float = MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS, + sleep: Callable[[float], Awaitable[None]] = anyio.sleep, + jitter: Callable[[], float] = random.random, + deadline_seconds: float | None = None, + clock: Callable[[], float] = anyio.current_time, +) -> T: + attempts: Final = max(0, max_retries) + deadline: Final = None if deadline_seconds is None else clock() + deadline_seconds + last_rate_limit: BaseException | None = None # rebind-ok: each 429 replaces the cause a retry deadline re-raises + for attempt in range(attempts + 1): + try: + if attempt > 0 and deadline is not None: + try: + with anyio.fail_after(max(0.0, deadline - clock())): + return await fetch() + except TimeoutError as timeout: + if last_rate_limit is None: + raise + raise last_rate_limit from timeout + return await fetch() + except Exception as exc: + limit = upstream_rate_limit(exc) + if limit is not None: + last_rate_limit = exc + if limit is None or attempt == attempts: + raise + delay = rate_limit_retry_delay( + attempt, + limit.retry_after_seconds, + base_delay=base_delay, + max_delay=max_delay, + jitter=jitter(), + ) + if deadline is not None and clock() + delay >= deadline: + raise + verbose_logger.warning( + "MCP server %s answered 429 to tools/list; retry %d/%d in %.2fs", + server_name, + attempt + 1, + max_retries, + delay, + ) + await sleep(delay) + raise AssertionError("unreachable") + + +def listing_failed_outright(outcomes: Mapping[str, ServerOutcome]) -> bool: + """True only when every server in scope contributed a non-auth fault: an auth outcome means a + caller can fix the listing by authenticating, so the aggregate stays a 200 with outcomes.""" + return bool(outcomes) and all( + isinstance(outcome, ServerListFault) and outcome.tag not in ("auth_required", "forbidden") + for outcome in outcomes.values() + ) + + def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: """The client-visible form of one outcome, for the tools/list result ``_meta`` and the REST response: category plus status code only, never upstream prose or URLs.""" diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 96cf6ee11aa..aba5161d56d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -104,6 +104,7 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( ServerListFault, raise_classified_list_failure, + retry_upstream_rate_limit, upstream_auth_challenge, upstream_insufficient_scope, ) @@ -4531,7 +4532,11 @@ class MCPServerManager: ) else: page: Final = ( - await self._fetch_tools_page_with_timeout(client, server, params) + await retry_upstream_rate_limit( + lambda: self._fetch_tools_page_with_timeout(client, server, params), + server.name, + deadline_seconds=MCP_TOOL_LISTING_TIMEOUT, + ) if params is not None else ListToolsResult(tools=await self._fetch_tools_with_timeout(client, server.name)) ) @@ -5600,7 +5605,11 @@ class MCPServerManager: """ try: with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): - tools: Final = await client.list_tools(raise_on_error=True) + tools: Final = await retry_upstream_rate_limit( + lambda: client.list_tools(raise_on_error=True), + server_name, + deadline_seconds=MCP_TOOL_LISTING_TIMEOUT, + ) verbose_logger.debug("Tools from %s: %s", server_name, tools) return tools except TimeoutError as e: diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 90f01368127..c45faf38f51 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -3041,6 +3041,19 @@ async def _execute_handle_list_tools( return ListToolsResult(tools=[]) +async def _discover_tools_listing( + context: OperationContext, host_progress_callback: ProgressCallback | None +) -> ListToolsResult: + from litellm.proxy._experimental.mcp_server.catalog import ToolListingFailedOutright + + try: + return await _execute_handle_list_tools( + context, PaginatedRequestParams(), host_progress_callback, log_list_tools_to_spendlogs=False + ) + except ToolListingFailedOutright: + return ListToolsResult(tools=[]) + + async def _execute_mcp_server_tool_call( context: OperationContext, params: CallToolRequestParams, host_progress_callback: ProgressCallback | None = None ) -> CallToolResult | InputRequiredResult: @@ -3517,14 +3530,7 @@ class GatewayOperations: ) memo_token: Final = _mcp_server_admission_memo.set({}) tasks: Final = ( - asyncio.create_task( - _execute_handle_list_tools( - context, - PaginatedRequestParams(), - self._host_progress_callback, - log_list_tools_to_spendlogs=False, - ) - ), + asyncio.create_task(_discover_tools_listing(context, self._host_progress_callback)), *(asyncio.create_task(self.execute(listing, context)) for listing in listings), ) try: diff --git a/litellm/proxy/_experimental/mcp_server/subscriptions.py b/litellm/proxy/_experimental/mcp_server/subscriptions.py index db622ea3cbf..e9de5692d83 100644 --- a/litellm/proxy/_experimental/mcp_server/subscriptions.py +++ b/litellm/proxy/_experimental/mcp_server/subscriptions.py @@ -50,6 +50,9 @@ from mcp.types import ( from litellm.constants import MCP_TOOL_LISTING_MAX_PAGES, MCP_TOOL_LISTING_TIMEOUT from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server import operations +from litellm.proxy._experimental.mcp_server.catalog import ( + ToolListingFailedOutright, +) from litellm.proxy._experimental.mcp_server.catalog import ( _configuration_identity as configuration_identity, # pyright: ignore[reportPrivateUsage] # reuse catalog-owned configuration identity without treating discovered metadata as configuration ) @@ -283,9 +286,12 @@ async def _listing(event: ServerEvent, context: OperationContext) -> str | None: cursor: str | None = None # rebind-ok: advance the signed gateway cursor for each catalog page with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): for _ in range(MCP_TOOL_LISTING_MAX_PAGES): - result = await gateway.execute( - operation.model_copy(update={"params": PaginatedRequestParams(cursor=cursor)}), context - ) + try: + result = await gateway.execute( + operation.model_copy(update={"params": PaginatedRequestParams(cursor=cursor)}), context + ) + except ToolListingFailedOutright: + return digest.hexdigest() items = ( result.tools if isinstance(result, ListToolsResult) diff --git a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index ee33c41afb6..f6902c778e5 100644 --- a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -326,3 +326,288 @@ def test_managed_scope_step_up_only_accepts_unambiguous_bearer_challenges(status wrapped = RuntimeError("transport failure") wrapped.__cause__ = cause assert upstream_insufficient_scope(wrapped) == expected + + +def _rate_limited_error(retry_after: str | None = None) -> httpx.HTTPStatusError: + request: Final = httpx.Request("POST", "https://mcp.example.com/mcp") + headers: Final = {"retry-after": retry_after} if retry_after is not None else {} + return httpx.HTTPStatusError( + "rate limited", + request=request, + response=httpx.Response(429, headers=headers, request=request), + ) + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_backs_off_then_reraises() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + calls: Final[list[httpx.HTTPStatusError]] = [] + failure: Final = _rate_limited_error() + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def always_limited() -> None: + calls.append(failure) + raise failure + + expected_error: Final = pytest.raises(httpx.HTTPStatusError) + with expected_error as exc_info: + await retry_upstream_rate_limit( + always_limited, + "srv", + max_retries=3, + base_delay=0.5, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + ) + + assert exc_info.value is failure + assert sleeps == [0.5, 1.0, 2.0] + assert len(calls) == 4 + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_honors_retry_after_up_to_cap() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def always_limited() -> None: + raise _rate_limited_error("120") + + with pytest.raises(httpx.HTTPStatusError): + await retry_upstream_rate_limit( + always_limited, + "srv", + max_retries=3, + base_delay=0.5, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + ) + + assert sleeps == [5.0, 5.0, 5.0] + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_parses_http_date_retry_after() -> None: + import email.utils + import time + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + attempts: Final[list[None]] = [] + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def limited() -> str: + attempts.append(None) + if len(attempts) == 1: + raise _rate_limited_error(email.utils.formatdate(time.time() + 3, usegmt=True)) + raise _rate_limited_error("abc") + + with pytest.raises(httpx.HTTPStatusError): + await retry_upstream_rate_limit( + limited, + "srv", + max_retries=3, + base_delay=0.5, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + ) + + assert 0 < sleeps[0] <= 3.5 + assert sleeps[1:] == [1.0, 2.0] + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_passes_through_non_429() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + calls = 0 + failure: Final = ValueError("broken") + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def broken() -> None: + nonlocal calls + calls += 1 + raise failure + + expected_error: Final = pytest.raises(ValueError, match="broken") + with expected_error as exc_info: + await retry_upstream_rate_limit( + broken, "srv", max_retries=3, base_delay=0.5, max_delay=5.0, sleep=record, jitter=lambda: 0.0 + ) + + assert exc_info.value is failure + assert calls == 1 + assert sleeps == [] + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_negative_max_retries_is_a_single_attempt() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + calls = 0 + failure: Final = _rate_limited_error("0") + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def always_limited() -> None: + nonlocal calls + calls += 1 + raise failure + + expected_error: Final = pytest.raises(httpx.HTTPStatusError) + with expected_error as exc_info: + await retry_upstream_rate_limit( + always_limited, + "srv", + max_retries=-1, + base_delay=0.5, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + ) + + assert exc_info.value is failure + assert calls == 1 + assert sleeps == [] + + +@pytest.mark.asyncio +async def test_retry_deadline_skips_sleep_that_would_pass_deadline() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + now: Final[list[float]] = [100.0] + sleeps: Final[list[float]] = [] + calls = 0 + failure: Final = _rate_limited_error() + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def always_limited() -> None: + nonlocal calls + calls += 1 + raise failure + + expected_error: Final = pytest.raises(httpx.HTTPStatusError) + with expected_error as exc_info: + await retry_upstream_rate_limit( + always_limited, + "srv", + max_retries=3, + base_delay=2.0, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + deadline_seconds=1.0, + clock=lambda: now[0], + ) + + assert exc_info.value is failure + assert calls == 1 + assert sleeps == [] + + +@pytest.mark.asyncio +async def test_retry_attempt_is_interrupted_at_deadline() -> None: + import anyio + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + calls = 0 + failure: Final = _rate_limited_error("0") + + async def fetch() -> None: + nonlocal calls + calls += 1 + if calls == 1: + raise failure + await anyio.sleep(5) + + expected_error: Final = pytest.raises(httpx.HTTPStatusError) + with anyio.fail_after(5): + with expected_error as exc_info: + await retry_upstream_rate_limit( + fetch, "srv", max_retries=3, base_delay=0.0, max_delay=5.0, deadline_seconds=0.2 + ) + + assert exc_info.value is failure + + +@pytest.mark.asyncio +async def test_retry_deadline_does_not_bound_first_attempt() -> None: + import anyio + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + async def slow() -> str: + await anyio.sleep(0.15) + return "tools" + + result: Final = await retry_upstream_rate_limit( + slow, "srv", max_retries=3, base_delay=0.5, max_delay=5.0, deadline_seconds=0.05 + ) + + assert result == "tools" + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_never_retries_cancellation() -> None: + import asyncio + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + calls = 0 + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def cancelled() -> None: + nonlocal calls + calls += 1 + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await retry_upstream_rate_limit( + cancelled, "srv", max_retries=3, base_delay=0.5, max_delay=5.0, sleep=record, jitter=lambda: 0.0 + ) + + assert calls == 1 + assert sleeps == [] + + +def test_listing_failed_outright_only_when_every_server_faulted_non_auth() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + ServerListOk, + listing_failed_outright, + ) + + assert not listing_failed_outright({}) + assert listing_failed_outright({"a": ServerListFault(tag="upstream_error", status_code=429)}) + assert listing_failed_outright( + {"a": ServerListFault(tag="timeout"), "b": ServerListFault(tag="rate_limited", status_code=429)} + ) + assert not listing_failed_outright( + {"a": ServerListFault(tag="upstream_error", status_code=500), "b": ServerListOk(tool_count=2)} + ) + assert not listing_failed_outright({"a": ServerListFault(tag="auth_required", status_code=401)}) + assert not listing_failed_outright({"a": ServerListFault(tag="forbidden", status_code=403)}) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d47b2f73a48..d314291ef74 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -7,6 +7,7 @@ import logging import os import sys import time +import anyio from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from datetime import datetime from pathlib import Path @@ -37,6 +38,8 @@ from mcp.types import ( CallToolRequestParams as MCPCallToolRequestParams, CallToolResult, GetPromptResult, + ListToolsResult, + PaginatedRequestParams, Prompt, ResourceTemplate, TextResourceContents, @@ -44,7 +47,7 @@ from mcp.types import ( from mcp.types import Tool as MCPTool from pydantic import AnyUrl, TypeAdapter -from litellm.constants import MCP_METADATA_TIMEOUT +from litellm.constants import MCP_METADATA_TIMEOUT, MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES from litellm.proxy._experimental.mcp_server import discoverable_endpoints from litellm.proxy._experimental.mcp_server.legacy_callbacks import ElicitationCallback from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var @@ -20427,3 +20430,148 @@ async def test_legacy_upstream_elicitation_rejects_modern_downstream_without_con finally: active_mcp_request_ctx_var.reset(request_token) legacy_server.active_mcp_session_var.reset(session_token) + + +def _upstream_http_status_error(status_code: int, headers: dict[str, str] | None = None) -> httpx.HTTPStatusError: + request: Final = httpx.Request("POST", "https://upstream.example/mcp") + return httpx.HTTPStatusError( + "upstream rejected", + request=request, + response=httpx.Response(status_code, headers=headers or {}, request=request), + ) + + +class TestMCPUpstreamRateLimitRetry: + """An upstream 429 on tools/list is retried inside the listing deadline instead of folding + into a silent empty outcome on the first attempt.""" + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_retries_429_then_returns_tools(self) -> None: + manager: Final = MCPServerManager() + tool: Final = MCPTool(name="ping", inputSchema={"type": "object"}) + client: Final = MagicMock() + client.list_tools = AsyncMock( + side_effect=[ + _upstream_http_status_error(429, {"Retry-After": "0"}), + _upstream_http_status_error(429, {"Retry-After": "0"}), + [tool], + ] + ) + + tools: Final = await manager._fetch_tools_with_timeout(client, "flaky-server") + + assert tools == [tool] + assert client.list_tools.await_count == 3 + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_gives_up_after_max_retries(self) -> None: + manager: Final = MCPServerManager() + client: Final = MagicMock() + client.list_tools = AsyncMock(side_effect=_upstream_http_status_error(429, {"Retry-After": "0"})) + + expected_error: Final = pytest.raises(MCPServerListError) + with expected_error as exc_info: + await manager._fetch_tools_with_timeout(client, "saturated-server") + + assert exc_info.value.fault == ServerListFault(tag="upstream_error", status_code=429) + assert client.list_tools.await_count == 1 + MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_does_not_retry_auth_error(self) -> None: + manager: Final = MCPServerManager() + client: Final = MagicMock() + client.list_tools = AsyncMock(side_effect=_upstream_http_status_error(401)) + + with pytest.raises(MCPUpstreamAuthError): + await manager._fetch_tools_with_timeout(client, "srv") + + assert client.list_tools.await_count == 1 + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_does_not_retry_server_error(self) -> None: + manager: Final = MCPServerManager() + client: Final = MagicMock() + client.list_tools = AsyncMock(side_effect=_upstream_http_status_error(500)) + + expected_error: Final = pytest.raises(MCPServerListError) + with expected_error as exc_info: + await manager._fetch_tools_with_timeout(client, "srv") + + assert exc_info.value.fault == ServerListFault(tag="upstream_error", status_code=500) + assert client.list_tools.await_count == 1 + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_reports_429_when_backoff_would_pass_listing_timeout( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCP_TOOL_LISTING_TIMEOUT", 0.5 + ) + manager: Final = MCPServerManager() + client: Final = MagicMock() + client.list_tools = AsyncMock(side_effect=_upstream_http_status_error(429, {"Retry-After": "5"})) + + with anyio.fail_after(3): + expected_error: Final = pytest.raises(MCPServerListError) + with expected_error as exc_info: + await manager._fetch_tools_with_timeout(client, "srv") + + assert exc_info.value.fault == ServerListFault(tag="upstream_error", status_code=429) + assert client.list_tools.await_count == 1 + + @pytest.mark.asyncio + async def test_get_tools_page_retries_429_then_returns_page(self) -> None: + tool: Final = MCPTool(name="list_notes", inputSchema={"type": "object"}) + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="notes", name="notes", transport=MCPTransport.http) + manager.registry = {server.server_id: server} + manager.create_mcp_client = AsyncMock( + return_value=SimpleNamespace( + _last_initialize_instructions=None, + list_tools_page=AsyncMock( + side_effect=[ + _upstream_http_status_error(429, {"Retry-After": "0"}), + ListToolsResult(tools=[tool]), + ] + ), + ) + ) + + result: Final = await manager.get_tools_page(server, params=PaginatedRequestParams()) + + assert [t.name for t in result.tools] == ["notes-list_notes"] + client: Final = manager.create_mcp_client.return_value + assert client.list_tools_page.await_count == 2 + + @pytest.mark.asyncio + async def test_get_tools_page_retry_attempt_is_bounded_by_listing_timeout( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCP_TOOL_LISTING_TIMEOUT", 0.05 + ) + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="notes", name="notes", transport=MCPTransport.http) + manager.registry = {server.server_id: server} + first_seen: Final = anyio.Event() + never: Final = anyio.Event() + error: Final = _upstream_http_status_error(429, {"Retry-After": "0"}) + + async def _list_tools_page(params: PaginatedRequestParams) -> ListToolsResult: + if not first_seen.is_set(): + first_seen.set() + raise error + await never.wait() + raise AssertionError("unreachable") + + list_tools_page: Final = AsyncMock(side_effect=_list_tools_page) + client: Final = SimpleNamespace(_last_initialize_instructions=None, list_tools_page=list_tools_page) + manager.create_mcp_client = AsyncMock(return_value=client) + + with anyio.fail_after(5): + expected_error: Final = pytest.raises(MCPServerListError) + with expected_error as exc_info: + await manager.get_tools_page(server, params=PaginatedRequestParams()) + + assert exc_info.value.fault == ServerListFault(tag="upstream_error", status_code=429) + assert list_tools_page.await_count == 2 diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 48106b1de5f..66d66271f71 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -6,19 +6,22 @@ import json import os from datetime import datetime, timedelta from types import SimpleNamespace -from typing import Final, Literal +from typing import Callable, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from fastapi import HTTPException from mcp import ReadResourceResult, Resource +from mcp.server.context import ServerRequestContext from mcp.server.models import InitializationOptions from mcp.types import ( INVALID_REQUEST, METHOD_NOT_FOUND, BlobResourceContents, CallToolResult, + ListToolsResult, + PaginatedRequestParams, Prompt, ResourceTemplate, TextContent, @@ -32,6 +35,7 @@ from starlette.types import Message, Receive, Scope, Send from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var from litellm.proxy._experimental.mcp_server.mcp_server_manager import ListedToolsCaller +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.proxy._types import ( LiteLLM_MCPServerTable, MCPTransport, @@ -5876,6 +5880,76 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab assert spend_meta["per_server_list_outcomes"] == {"server_a": {"status": "ok", "tool_count": 1}} +@pytest.mark.asyncio +async def test_get_tools_from_mcp_servers_all_failed_listing_never_runs_proxy_failure_hooks() -> None: + from litellm.proxy._types import UserAPIKeyAuth + + servers: Final = ( + MCPServer( + server_id="server_a", + name="server_a", + alias="server_a", + server_name="server_a", + transport=MCPTransport.http, + ), + MCPServer( + server_id="server_b", + name="server_b", + alias="server_b", + server_name="server_b", + transport=MCPTransport.http, + ), + ) + user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}} + logging_obj.async_success_handler = AsyncMock() + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock() + manager: Final = MagicMock() + manager.get_tools_from_server = AsyncMock(side_effect=RuntimeError("server unavailable")) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", + new=AsyncMock(return_value=servers), + ), + patch( + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", + return_value=(None, None), + ), + patch("litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", manager), + patch("litellm.proxy._experimental.mcp_server.operations.filter_tools_by_allowed_tools", return_value=[]), + patch( + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_key_team_permissions", + new=AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.operations.function_setup", + return_value=(logging_obj, None), + ), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj), + ): + listing: Final = await mcp_operations._get_tools_from_mcp_servers( + user_api_key_auth=user_auth, + mcp_auth_header=None, + mcp_servers=["server_a", "server_b"], + mcp_server_auth_headers=None, + params=None, + log_list_tools_to_spendlogs=True, + enforce_rate_limits=False, + ) + + assert listing.tools == [] + logging_obj.async_success_handler.assert_awaited_once() + spend_meta: Final = logging_obj.model_call_details["metadata"]["spend_logs_metadata"] + assert spend_meta["per_server_list_outcomes"] == { + "server_a": {"status": "internal"}, + "server_b": {"status": "internal"}, + } + proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + + @pytest.mark.asyncio async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fails(): """ @@ -10071,6 +10145,131 @@ async def test_handle_list_tools_attaches_outcome_meta(_mcp_request_ctx): assert outcomes_meta["broken"] == {"status": "unreachable"} +async def _list_tools_with_listing( + _mcp_request_ctx: Callable[..., ServerRequestContext], + listing: AggregateToolListing, + params: PaginatedRequestParams, +) -> ListToolsResult: + try: + from litellm.proxy._experimental.mcp_server.server import handle_list_tools + except ImportError: + pytest.skip("MCP server not available") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + ), + patch( + "litellm.proxy._experimental.mcp_server.operations._list_mcp_tools", + new=AsyncMock(return_value=listing), + ), + ): + return await handle_list_tools(_mcp_request_ctx(), params) + + +@pytest.mark.asyncio +async def test_handle_list_tools_errors_when_every_server_failed( + _mcp_request_ctx: Callable[..., ServerRequestContext], +) -> None: + from mcp.shared.exceptions import MCPError + from mcp.types import INTERNAL_ERROR + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListFault, + ) + + listing: Final = AggregateToolListing( + tools=[], + outcomes={ + "a": ServerListFault(tag="upstream_error", status_code=429), + "b": ServerListFault(tag="timeout"), + }, + ) + + expected_error: Final = pytest.raises(MCPError) + with expected_error as exc_info: + await _list_tools_with_listing(_mcp_request_ctx, listing, _paged_params()) + + assert exc_info.value.error.code == INTERNAL_ERROR + assert exc_info.value.error.data[SERVER_OUTCOMES_META_KEY] == { + "a": {"status": "upstream_error", "http_status": 429}, + "b": {"status": "timeout"}, + } + + +@pytest.mark.asyncio +async def test_handle_list_tools_auth_fault_keeps_degraded_result( + _mcp_request_ctx: Callable[..., ServerRequestContext], +) -> None: + from mcp.types import ListToolsResult + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListFault, + ) + + listing: Final = AggregateToolListing( + tools=[], + outcomes={"a": ServerListFault(tag="auth_required", status_code=401), "b": ServerListFault(tag="timeout")}, + ) + + result: Final = await _list_tools_with_listing(_mcp_request_ctx, listing, _paged_params()) + + assert isinstance(result, ListToolsResult) + assert result.tools == [] + wire: Final = result.model_dump(by_alias=True) + assert wire["_meta"][SERVER_OUTCOMES_META_KEY]["a"] == {"status": "auth_required", "http_status": 401} + + +@pytest.mark.asyncio +async def test_handle_list_tools_all_faults_on_continuation_stays_degraded( + _mcp_request_ctx: Callable[..., ServerRequestContext], +) -> None: + from mcp.types import ListToolsResult, PaginatedRequestParams + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListFault, + ) + + listing: Final = AggregateToolListing( + tools=[], + outcomes={"a": ServerListFault(tag="upstream_error", status_code=429)}, + ) + + result: Final = await _list_tools_with_listing( + _mcp_request_ctx, listing, PaginatedRequestParams(cursor="cursor-1") + ) + + assert isinstance(result, ListToolsResult) + assert result.tools == [] + wire: Final = result.model_dump(by_alias=True) + assert wire["_meta"][SERVER_OUTCOMES_META_KEY]["a"] == {"status": "upstream_error", "http_status": 429} + + +@pytest.mark.asyncio +async def test_handle_list_tools_empty_outcomes_is_a_clean_empty_result( + _mcp_request_ctx: Callable[..., ServerRequestContext], +) -> None: + from mcp.types import ListToolsResult + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing + + listing: Final = AggregateToolListing(tools=[], outcomes={}) + + result: Final = await _list_tools_with_listing(_mcp_request_ctx, listing, _paged_params()) + + assert isinstance(result, ListToolsResult) + assert result.tools == [] + wire: Final = result.model_dump(by_alias=True) + assert wire.get("_meta") is None + + def _make_oauth2_server( alias: str, *, diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index 3957f486c4c..6f7437c90dc 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -1263,6 +1263,45 @@ async def test_discovery_preserves_caller_scope_and_proxy_restrictions(): assert listed.await_args.args[0].mcp_servers == ("only-this",) +@pytest.mark.asyncio +async def test_discovery_outright_tools_failure_still_advertises_other_capabilities() -> None: + from mcp.types import DiscoverRequest, ListPromptsResult, ListToolsResult, Prompt + + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.catalog import ToolListingFailedOutright + + tools: Final = AsyncMock(side_effect=ToolListingFailedOutright(code=-32603, message="every server failed")) + prompts: Final = AsyncMock(return_value=ListPromptsResult(prompts=[Prompt(name="p")])) + context: Final = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["broken"]) + with ( + patch.object(operations, "_execute_handle_list_tools", tools), + patch.object(operations, "_execute_list_prompts", prompts), + patch.object(operations, "_execute_list_resources", AsyncMock(return_value=ListToolsResult(tools=[]))), + patch.object(operations, "_execute_list_resource_templates", AsyncMock(return_value=ListToolsResult(tools=[]))), + ): + result: Final = await GatewayOperations().execute(DiscoverRequest(), context) + assert result.capabilities.tools is None + assert result.capabilities.prompts is not None + + +@pytest.mark.asyncio +async def test_discovery_propagates_non_outright_mcp_errors_from_tools_listing() -> None: + from mcp import MCPError + from mcp.types import INVALID_PARAMS, DiscoverRequest + + from litellm.proxy._experimental.mcp_server import operations + + tools: Final = AsyncMock(side_effect=MCPError(code=INVALID_PARAMS, message="bad cursor")) + context: Final = prepare_context( + UserAPIKeyAuth(user_id="scoped"), mcp_servers=["broken"], mcp_proxy_mode=True, protocol_version="2025-06-18" + ) + with patch.object(operations, "_execute_handle_list_tools", tools): + expected_error: Final = pytest.raises(MCPError) + with expected_error as exc_info: + await GatewayOperations().execute(DiscoverRequest(), context) + assert exc_info.value.error.code == INVALID_PARAMS + + @pytest.mark.asyncio async def test_discovery_denial_cannot_advertise_tools(): from mcp.types import DiscoverRequest diff --git a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py index a97ab04c506..fadb709e00f 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py @@ -10,7 +10,7 @@ from mcp.server.context import ServerRequestContext from mcp.server.session import ServerSession from mcp.server.subscriptions import ListenHandler from mcp.shared.subscriptions import PromptsListChanged, ToolsListChanged -from mcp.types import SubscriptionFilter, SubscriptionsListenRequestParams +from mcp.types import ListToolsRequest, ListToolsResult, SubscriptionFilter, SubscriptionsListenRequestParams from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server import subscriptions @@ -45,6 +45,18 @@ def context(session: ServerSession, revision: str = "2026-07-28") -> ServerReque ) +class _ListingGateway: + def __init__(self, *, result: ListToolsResult | None = None, error: MCPError | None = None) -> None: + self._result: ListToolsResult | None = result + self._error: MCPError | None = error + + async def execute(self, _operation: ListToolsRequest, _context: OperationContext) -> ListToolsResult: + if self._error is not None: + raise self._error + assert self._result is not None + return self._result + + @pytest.mark.asyncio async def test_queued_notification_rechecks_authority_before_delivery() -> None: bus: Final = subscriptions._ListenerBus() @@ -486,3 +498,39 @@ async def test_legacy_listener_catalog_refreshes_do_not_consume_server_rate_limi finally: listener.cancel() await asyncio.gather(listener, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_all_failed_tools_listing_keeps_empty_catalog_digest(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._experimental.mcp_server.catalog import ToolListingFailedOutright + + listener_context: Final = OperationContext(_caller=UserAPIKeyAuth(user_id="listener")) + empty_gateway: Final = _ListingGateway(result=ListToolsResult(tools=[])) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.operations.GatewayOperations", lambda: empty_gateway + ) + empty_digest: Final = await subscriptions._listing(ToolsListChanged(), listener_context) + + failed_gateway: Final = _ListingGateway( + error=ToolListingFailedOutright(code=-32603, message="Every MCP server in scope failed") + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.operations.GatewayOperations", lambda: failed_gateway + ) + failed_digest: Final = await subscriptions._listing(ToolsListChanged(), listener_context) + + assert failed_digest == empty_digest + + +@pytest.mark.asyncio +async def test_tools_listing_propagates_other_mcp_errors(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = MCPError(code=-32602, message="Bad cursor") + gateway: Final = _ListingGateway(error=failure) + monkeypatch.setattr("litellm.proxy._experimental.mcp_server.operations.GatewayOperations", lambda: gateway) + listener_context: Final = OperationContext(_caller=UserAPIKeyAuth(user_id="listener")) + + expected_error: Final = pytest.raises(MCPError) + with expected_error as exc_info: + await subscriptions._listing(ToolsListChanged(), listener_context) + + assert exc_info.value is failure From 6f66b7a2014e673f60d886a32d5915e99d349e26 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 19 Oct 2026 16:01:40 -0700 Subject: [PATCH 16/18] feat(proxy)!: personal keys inherit MCP server, MCP tool and agent access from the owner's teams (#46293) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 648 +++++++-- .../_experimental/mcp_server/contracts.py | 1 + .../mcp_server/mcp_server_manager.py | 71 +- .../_experimental/mcp_server/operations.py | 375 ++++- .../mcp_server/rest_endpoints.py | 249 +++- .../_experimental/mcp_server/tool_search.py | 31 +- .../proxy/agent_endpoints/a2a_endpoints.py | 14 +- .../auth/agent_permission_handler.py | 255 +++- litellm/proxy/auth/auth_checks.py | 259 +++- litellm/proxy/auth/team_grants.py | 86 ++ litellm/proxy/common_request_processing.py | 13 +- .../hooks/parallel_request_limiter_v3.py | 23 +- litellm/proxy/litellm_pre_call_utils.py | 4 +- litellm/proxy/proxy_server.py | 6 +- .../auth/test_user_api_key_auth_mcp.py | 1020 +++++++++++++- .../mcp_server/test_mcp_server_manager.py | 1067 +++++++++++--- .../test_mcp_server_tool_calls_and_headers.py | 217 ++- .../mcp_server/test_rest_endpoints.py | 274 +++- .../auth/test_agent_permission_handler.py | 1245 ++++++++++++++--- .../agent_endpoints/test_a2a_endpoints.py | 118 ++ ...st_auth_checks_object_access_and_lookup.py | 264 ++++ .../hooks/test_parallel_request_limiter_v3.py | 120 ++ 22 files changed, 5600 insertions(+), 760 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 4090666e16f..53b56c7b31d 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,3 +1,4 @@ +import asyncio import functools import re import secrets @@ -5,8 +6,9 @@ from collections.abc import Awaitable, Callable, 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, cast +from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException from pydantic import SecretStr, TypeAdapter @@ -41,11 +43,15 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( is_session_bearer_shaped, ) -from litellm.proxy._experimental.mcp_server.utils import asgi_scope_fields +from litellm.proxy._experimental.mcp_server.utils import ( + asgi_scope_fields, + strip_known_server_prefix, +) from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_AccessGroupTable, LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, LiteLLM_TeamTable, ProxyException, SpecialHeaders, @@ -466,6 +472,24 @@ class MCPServerAccess: scope: Literal["unscoped", "scoped", "unresolved"] = "unscoped" +def _mcp_server_access_after_error( + error: Exception, + key_object_permission: LiteLLM_ObjectPermissionTable | None, + *, + require_readable_team_policy: bool, +) -> MCPServerAccess: + if require_readable_team_policy: + raise error + if isinstance(error, UnloadableEntitlementError): + verbose_logger.warning("Denying MCP access, entitlement unreadable: %s", error) + else: + verbose_logger.warning("Failed to get allowed MCP servers: %s", error) + return MCPServerAccess( + server_ids=(), + scope="scoped" if getattr(key_object_permission, "mcp_servers", None) is not None else "unresolved", + ) + + @dataclass(frozen=True, slots=True) class DcrBridgeTarget: """The single DCR-bridge server a request targets, paired with the exact name the caller @@ -1804,11 +1828,28 @@ class MCPRequestHandler: ) return list(access.server_ids) + @staticmethod + async def _apply_agent_caller_ceiling_to_multi_source_server_access( + access: MCPServerAccess, + user_api_key_auth: UserAPIKeyAuth | None, + ) -> MCPServerAccess: + if user_api_key_auth is None or not MCPRequestHandler.is_personal_key_inheriting_team_grants(user_api_key_auth): + return access + caller_capped, caller_restricts = await MCPRequestHandler.apply_agent_caller_ceiling( + access.server_ids, + user_api_key_auth, + ) + return MCPServerAccess( + server_ids=caller_capped, + scope="scoped" if caller_restricts else access.scope, + ) + @staticmethod async def get_mcp_server_access( user_api_key_auth: UserAPIKeyAuth | None = None, *, keyless_source: bool = False, + require_readable_team_policy: bool = False, ) -> MCPServerAccess: """ Get list of allowed MCP servers for the given user/key based on permissions. @@ -1844,9 +1885,11 @@ class MCPRequestHandler: # matters: the no_mcp_servers opt-out below reads the caller's own object_permission, so above # this branch a user's own opt-out would wrongly zero their TEAMS' grants too (each source is # independent; an opt-out silences only its own source, inside the recursive call). - if is_mcp_admitted_user_subject(user_api_key_auth) and user_api_key_auth is not None: - return MCPServerAccess( - server_ids=tuple(await MCPRequestHandler.resolve_admitted_subject_servers(user_api_key_auth)), + multi_source_access: Final = await MCPRequestHandler._multi_source_server_access(user_api_key_auth) + if multi_source_access is not None: + return await MCPRequestHandler._apply_agent_caller_ceiling_to_multi_source_server_access( + multi_source_access, + user_api_key_auth, ) # Get allowed servers from key and team @@ -1861,7 +1904,10 @@ class MCPRequestHandler: ): return MCPServerAccess(server_ids=(), scope="scoped") - allowed_mcp_servers_for_team = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_api_key_auth) + allowed_mcp_servers_for_team = await MCPRequestHandler._get_allowed_mcp_servers_for_policy( + user_api_key_auth, + require_readable_policy=require_readable_team_policy, + ) key_access_group_grants = await MCPRequestHandler._get_key_access_group_mcp_server_extras(user_api_key_auth) @@ -1993,15 +2039,10 @@ class MCPRequestHandler: ), ) except Exception as e: - if isinstance(e, UnloadableEntitlementError): - # A ceiling we KNOW exists and cannot read. Denying is the only answer that does not - # widen this caller past what an operator configured, for both caller shapes. - verbose_logger.warning("Denying MCP access, entitlement unreadable: %s", e) - else: - verbose_logger.warning("Failed to get allowed MCP servers: %s", e) - return MCPServerAccess( - server_ids=(), - scope="scoped" if getattr(key_object_permission, "mcp_servers", None) is not None else "unresolved", + return _mcp_server_access_after_error( + e, + key_object_permission, + require_readable_team_policy=require_readable_team_policy, ) @staticmethod @@ -2144,8 +2185,6 @@ class MCPRequestHandler: # whole union to deny-all over one momentarily-unreadable row. verbose_logger.warning("MCP admitted-subject source team %r unresolvable, skipping: %s", team_id, e) return None - if team_obj is None: - return None member_user_ids: Final = {getattr(m, "user_id", None) for m in (team_obj.members_with_roles or [])} - {None} if auth.user_id not in member_user_ids: return None @@ -2240,35 +2279,156 @@ class MCPRequestHandler: return list(reachable) @staticmethod - async def billing_auth_for_tool_call(auth: UserAPIKeyAuth, tool_name: str) -> UserAPIKeyAuth: + async def billing_auth_for_tool_call( + auth: UserAPIKeyAuth, + tool_name: str, + *, + mcp_servers: Sequence[str] | None = None, + scoped_server: MCPServer | None = None, + native_tool_name: str | None = None, + allowed_mcp_servers: Sequence[MCPServer] | None = None, + ) -> UserAPIKeyAuth: """The auth object a tool call's SPEND should be recorded against. - ``auth`` unchanged for any non-admitted caller (key/JWT billing byte-identical). For an admitted - subject whose call is reached through a team's grant, a copy carrying that team's ``team_id`` and - owning ``org_id`` so the team's budget accumulates and the right org is charged. Falls back to - user-level attribution (rather than guessing a team) when the tool name does not resolve to a - server, reusing the manager's own tool-name lookup.""" - if not is_mcp_admitted_user_subject(auth): + ``auth`` unchanged for callers without inherited team grants. For an admitted subject or personal + key whose call is reached through a team's grant, a copy carrying that team's ``team_id`` and + owning ``org_id`` so the team's budget accumulates and the right org is charged.""" + from litellm.proxy.auth.auth_checks import personal_key_inherits_team_grants + + if not is_mcp_admitted_user_subject(auth) and not personal_key_inherits_team_grants(auth): return auth try: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.operations import ( + _get_allowed_mcp_servers, + _resolve_display_name_to_original, + resolve_scoped_mcp_server, ) - server: Final = global_mcp_server_manager.get_mcp_server_from_tool_name(tool_name) + servers_for_name_resolution: Final = ( + allowed_mcp_servers + if allowed_mcp_servers is not None + else await _get_allowed_mcp_servers(auth, mcp_servers) + ) + resolved_tool_name: Final = _resolve_display_name_to_original(tool_name, servers_for_name_resolution) + server: Final = ( + scoped_server + or (await resolve_scoped_mcp_server(auth, mcp_servers) if mcp_servers is not None else None) + or global_mcp_server_manager.get_mcp_server_from_tool_name(resolved_tool_name) + or global_mcp_server_manager.server_owning_tool_name_prefix(resolved_tool_name) + ) if server is None: return auth - source: Final = await MCPRequestHandler.attributing_source_for_server(auth, server.server_id) + source: Final = await ( + MCPRequestHandler.attributing_source_for_server(auth, server.server_id) + if is_mcp_admitted_user_subject(auth) + else MCPRequestHandler.personal_key_attributing_source( + auth, + server.server_id, + native_tool_name + if native_tool_name is not None + else strip_known_server_prefix(resolved_tool_name, server), + ) + ) if source is None or not source.team_id: return auth - billed: Final = auth.model_copy() - billed.team_id = source.team_id - billed.org_id = source.org_id - return billed + if is_mcp_admitted_user_subject(auth): + return auth.model_copy( + update=MappingProxyType( + { + "team_id": source.team_id, + "org_id": source.org_id, + "team_metadata": source.team_metadata, + } + ) + ) + source_auth_update: Final = MappingProxyType( + { + "team_id": source.team_id, + "org_id": source.org_id, + "team_metadata": source.team_metadata, + } + ) + try: + team_auth_context: Final = await MCPRequestHandler._personal_key_team_auth_context(auth, source) + except Exception as e: # noqa: BLE001 # keep the selected team's attribution on follow-up read failures + verbose_logger.warning("MCP billing context lookup failed for team %r: %s", source.team_id, e) + return auth.model_copy(update=source_auth_update) + if team_auth_context is None: + return auth.model_copy(update=source_auth_update) + team_object, team_membership = team_auth_context + from litellm.proxy.auth.team_grants import team_grant_auth_update, team_grant_billing_auth_update + + try: + billed_update: Final = await team_grant_billing_auth_update(auth, team_object, team_membership) + except Exception as e: # noqa: BLE001 # retain team controls when organization attribution fails + verbose_logger.warning("MCP billing controls lookup failed for team %r: %s", team_object.team_id, e) + return auth.model_copy( + update=team_grant_auth_update( + team_object, + team_membership, + auth.user_id, + key_org_id=auth.org_id, + ) + ) + return auth.model_copy(update=billed_update) except Exception as e: # noqa: BLE001 # attribution must never fail an authorized call verbose_logger.warning("MCP billing attribution failed for %r, billing the user: %s", tool_name, e) return auth + @staticmethod + async def guardrail_auth_for_call( + guardrail_auth: UserAPIKeyAuth | None, + user_api_key_auth: UserAPIKeyAuth | None, + server: MCPServer | None, + tool_name: str, + allowed_mcp_servers: Sequence[MCPServer] | None = None, + ) -> UserAPIKeyAuth | None: + if guardrail_auth is not None or user_api_key_auth is None or server is None: + return guardrail_auth if guardrail_auth is not None else user_api_key_auth + if not MCPRequestHandler.is_personal_key_inheriting_team_grants(user_api_key_auth): + return user_api_key_auth + return await MCPRequestHandler.billing_auth_for_tool_call( + user_api_key_auth, + tool_name, + scoped_server=server, + native_tool_name=tool_name, + allowed_mcp_servers=allowed_mcp_servers, + ) + + @staticmethod + async def _personal_key_team_auth_context( + auth: UserAPIKeyAuth, + source: UserAPIKeyAuth, + ) -> tuple[LiteLLM_TeamTable, LiteLLM_TeamMembership | None] | None: + from litellm.proxy.auth.auth_checks import get_team_membership, get_team_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None or source.team_id is None: + return None + team: Final = await get_team_object( + team_id=source.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=auth.requires_fresh_policy, + ) + membership: Final = ( + await get_team_membership( + user_id=auth.user_id, + team_id=source.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=auth.requires_fresh_policy, + ) + if auth.user_id is not None + else None + ) + return team, membership + @staticmethod async def attributing_source_for_server( auth: UserAPIKeyAuth, @@ -2328,6 +2488,201 @@ class MCPRequestHandler: allowed.update(tools) return await MCPRequestHandler._narrow_tools_to_toolset(sorted(allowed), server_id, auth) + @staticmethod + async def _multi_source_server_access(user_api_key_auth: UserAPIKeyAuth | None) -> MCPServerAccess | None: + """Server access for the two caller shapes pinned to no single team, each resolved per grant + source: a keyless admitted subject (its own grants plus its teams') and, under + ``personal_key_model_access_from_teams``, a personal key (its owner's teams). None for every + other caller, which the single-source rules then resolve.""" + from litellm.proxy.auth.auth_checks import personal_key_inherits_team_grants + + if user_api_key_auth is None: + return None + if is_mcp_admitted_user_subject(user_api_key_auth): + return MCPServerAccess( + server_ids=tuple(await MCPRequestHandler.resolve_admitted_subject_servers(user_api_key_auth)), + ) + if personal_key_inherits_team_grants(user_api_key_auth): + return await MCPRequestHandler.resolve_personal_key_team_servers(user_api_key_auth) + return None + + @staticmethod + def is_personal_key_inheriting_team_grants(auth: UserAPIKeyAuth) -> bool: + from litellm.proxy.auth.auth_checks import personal_key_inherits_team_grants + + return personal_key_inherits_team_grants(auth) + + @staticmethod + async def personal_key_team_sources(auth: UserAPIKeyAuth) -> tuple[UserAPIKeyAuth, ...]: + """One source per current team of a personal key's owner, each resolving as the same key pinned + to that team, so key/team inheritance and intersection hold exactly as for a team key. The + pinned ``team_id`` keeps a source out of this branch. A team of another organization grants + nothing to a key issued under its own org, so neither org ceiling is dropped. No team, or an owner that cannot be + loaded, yields no source: the key reaches nothing through its owner.""" + from litellm.proxy.auth.auth_checks import personal_key_owner_teams + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + teams: Final = await personal_key_owner_teams( + valid_token=auth, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(auth.requires_fresh_policy), + ) + return tuple( + auth.model_copy( + update=MappingProxyType( + { + "team_id": team.team_id, + "org_id": auth.org_id or team.organization_id, + "team_metadata": TypeAdapter(Mapping[str, object]).validate_python( + team.model_dump().get("metadata") or {} + ), + } + ) + ) + for team in teams + if MCPRequestHandler._personal_key_team_permission_is_readable(team) + and (auth.org_id is None or team.organization_id in (None, auth.org_id)) + ) + + @staticmethod + def _personal_key_team_permission_is_readable(team: LiteLLM_TeamTable) -> bool: + if team.object_permission_id is None or team.object_permission is not None: + return True + verbose_logger.warning( + "Personal key owner team %s has an unreadable object_permission, granting nothing", + team.team_id, + ) + return False + + @staticmethod + async def personal_key_team_server_grants( + auth: UserAPIKeyAuth, + ) -> tuple[tuple[UserAPIKeyAuth, MCPServerAccess], ...]: + sources: Final = await MCPRequestHandler.personal_key_team_sources(auth) + + async def resolve_source(source: UserAPIKeyAuth) -> tuple[UserAPIKeyAuth, MCPServerAccess]: + try: + return source, await MCPRequestHandler.get_mcp_server_access( + source, + require_readable_team_policy=True, + ) + except Exception as e: # noqa: BLE001 # an unreadable team grant must not widen inherited access + verbose_logger.warning("MCP access lookup failed for owner team %r: %s", source.team_id, e) + return source, MCPServerAccess(server_ids=(), scope="scoped") + + return tuple(await asyncio.gather(*(resolve_source(source) for source in sources))) + + @staticmethod + async def personal_key_attributing_source( + auth: UserAPIKeyAuth, + server_id: str, + tool_name: str, + grants: tuple[tuple[UserAPIKeyAuth, MCPServerAccess], ...] | None = None, + ) -> UserAPIKeyAuth | None: + source_grants: Final = grants or await MCPRequestHandler.personal_key_team_server_grants(auth) + + async def _source_if_tool_allowed(source: UserAPIKeyAuth, access: MCPServerAccess) -> UserAPIKeyAuth | None: + if server_id not in access.server_ids: + return None + try: + allowed_tools: Final = await MCPRequestHandler._single_credential_tools(server_id, source) + except Exception as e: # noqa: BLE001 # attribution metadata must not fail an authorized call + verbose_logger.warning("MCP tool attribution failed for team %r: %s", source.team_id, e) + return None + return source if allowed_tools is None or tool_name in allowed_tools else None + + eligible_sources: Final = tuple( + source + for source in await asyncio.gather( + *(_source_if_tool_allowed(source, access) for source, access in source_grants) + ) + if source is not None + ) + return min(eligible_sources, key=lambda source: source.team_id or "") if eligible_sources else None + + @staticmethod + async def personal_key_team_rpm_limits( + auth: UserAPIKeyAuth, server: MCPServer, tool_name: str + ) -> Mapping[str, Mapping[str, int]] | None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + try: + source: Final = await MCPRequestHandler.personal_key_attributing_source(auth, server.server_id, tool_name) + if source is None or source.team_id is None or prisma_client is None: + return None + team: Final = await get_team_object( + team_id=source.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(auth.requires_fresh_policy), + ) + team_metadata: Final[Mapping[str, object]] = TypeAdapter(Mapping[str, object]).validate_python( + team.model_dump().get("metadata") or {} + ) + team_limit: Final = team_metadata.get("mcp_rpm_limit") + if not isinstance(team_limit, dict) or not team_limit: + return None + validated_limit: Final = TypeAdapter(dict[str, int]).validate_python(team_limit) + applicable: Final = { + name: rpm + for name, rpm in validated_limit.items() + if server.server_id in global_mcp_server_manager.expand_permission_list([name]) + } + except Exception as e: # noqa: BLE001 # throttling metadata must not fail an allowed request + verbose_logger.warning("Failed to resolve per-team MCP rpm limits for personal key: %s", e) + return None + else: + return {source.team_id: applicable} if applicable else None + + @staticmethod + async def resolve_personal_key_team_servers(auth: UserAPIKeyAuth) -> MCPServerAccess: + """Union of what the owner's teams grant the personal key, each answered by the canonical + resolver. Unscoped only when some team leaves the key unscoped, as a team key of it would be.""" + grants: Final = await MCPRequestHandler.personal_key_team_server_grants(auth) + server_ids: Final = tuple(dict.fromkeys(chain.from_iterable(access.server_ids for _source, access in grants))) + unscoped: Final = any(access.scope == "unscoped" for _source, access in grants) + return MCPServerAccess(server_ids=server_ids, scope="unscoped" if unscoped else "scoped") + + @staticmethod + async def resolve_personal_key_team_tools(server_id: str, auth: UserAPIKeyAuth) -> tuple[str, ...] | None: + """Effective tool allowlist on ``server_id`` for a personal key: the union over the owner's teams + that grant the server, where a granting team without a tool rule allows every tool. A server no + team grants is reachable only through an open channel (allow_all_keys, the owner's BYOM, the + admin-view role) and + then carries the key's OWN tool rules, never a team's; otherwise nothing is allowed. Each team + resolves through the raising ``_single_credential_tools`` so one team's fault reaches the + caller's fail-closed handler instead of becoming an allow-all ``None`` in the union.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + grants: Final = await MCPRequestHandler.personal_key_team_server_grants(auth) + granting: Final = tuple(source for source, access in grants if server_id in access.server_ids) + if not granting: + reachable: Final = server_id in await global_mcp_server_manager.operator_open_server_ids( + auth + ) or await MCPRequestHandler.admin_view_unscoped(auth) + if not reachable: + return () + own_tools: Final = await MCPRequestHandler._single_credential_tools(server_id, auth) + return None if own_tools is None else tuple(own_tools) + tool_lists: Final = await asyncio.gather( + *(MCPRequestHandler._single_credential_tools(server_id, source) for source in granting) + ) + if any(tools is None for tools in tool_lists): + return None + return tuple(sorted(frozenset(chain.from_iterable(tools for tools in tool_lists if tools is not None)))) + @staticmethod def _get_key_object_permission( user_api_key_auth: UserAPIKeyAuth | None = None, @@ -2693,6 +3048,8 @@ class MCPRequestHandler: return await managed_agent_tools(server_id, user_api_key_auth) + from litellm.proxy.auth.auth_checks import personal_key_inherits_team_grants + try: # FIRST statement, mirroring get_allowed_mcp_servers: a keyless admitted subject resolves per # source and shares nothing with the single-credential prelude below. Ordering is the invariant: @@ -2700,97 +3057,11 @@ class MCPRequestHandler: if is_mcp_admitted_user_subject(user_api_key_auth): return await MCPRequestHandler.resolve_admitted_subject_tools(server_id, user_api_key_auth) - # Get key and team object permissions (already loaded in main auth flow) - key_obj_perm: Final = MCPRequestHandler._get_key_object_permission(user_api_key_auth) - team_obj_perm: Final = await MCPRequestHandler._get_team_object_permission(user_api_key_auth) + if personal_key_inherits_team_grants(user_api_key_auth): + return _as_list(await MCPRequestHandler.resolve_personal_key_team_tools(server_id, user_api_key_auth)) - # Extract tool permissions for this server. Dict keys may be - # server_ids OR names/aliases; normalize to server_id-keyed form - # before lookup so a name-based key does not silently drop its - # tool restrictions when server_id is the resolved uuid. - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - - key_direct_tools: Final = ( - global_mcp_server_manager.expand_tool_permissions(key_obj_perm.mcp_tool_permissions).get(server_id) - if key_obj_perm - else None - ) - - # Tools granted through the key's toolsets restrict this server exactly - # as direct tool permissions do; union with any direct grants so the - # tool-level check sees the key's full effective tool scope - key_toolset_ids: Final = (key_obj_perm.mcp_toolsets or []) if key_obj_perm else [] - key_toolset_tools: Final = ( - ( - await global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=key_toolset_ids, - requires_fresh_policy=user_api_key_auth.requires_fresh_policy, - ) - ).get(server_id) - if key_toolset_ids - else None - ) - - key_group_tools: Final = await MCPRequestHandler._key_access_group_tools_for_server( - user_api_key_auth, server_id - ) - key_direct_servers: Final = await MCPRequestHandler._key_direct_servers( - key_obj_perm, requires_fresh_policy=bool(user_api_key_auth.requires_fresh_policy) - ) - key_tools: Final = _as_list( - MCPRequestHandler._key_tool_scope( - key_direct_tools, key_toolset_tools, key_group_tools, server_id in key_direct_servers - ) - ) - team_direct_tools: Final = ( - global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id) - if team_obj_perm - else None - ) - - # Tools granted through the team's toolsets restrict this server exactly - # as the team's direct tool permissions do, mirroring the key path above - team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( - team_obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy - ) - team_group_tools: Final = await MCPRequestHandler._team_access_group_tools_for_server( - user_api_key_auth, server_id - ) - team_tools: Final = MCPRequestHandler._union_tool_grants( - team_direct_tools, team_toolset_tools, team_group_tools - ) - - # Apply same inheritance logic as get_allowed_mcp_servers. ``is not None`` so an explicit - # empty allowlist stays a deny instead of reading as "no restriction". - if team_tools is not None: - if key_tools is not None: - # Both have restrictions → intersection - allowed_tools = list(set(team_tools) & set(key_tools)) - else: - # Only team has restrictions → inherit from team - allowed_tools = team_tools - else: - # No team restrictions → use key restrictions - allowed_tools = cast(list[str], key_tools) - - allowed_tools = _as_list( - await MCPRequestHandler._apply_end_user_tool_ceiling(allowed_tools, server_id, user_api_key_auth) - ) - - allowed_tools = _as_list( - await MCPRequestHandler._apply_user_tool_ceiling( - allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source - ) - ) - - allowed_tools = _as_list( - await MCPRequestHandler.apply_agent_caller_tool_ceiling(allowed_tools, server_id, user_api_key_auth) - ) - - return await MCPRequestHandler._apply_agent_and_org_tool_ceilings( - allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source + return await MCPRequestHandler._single_credential_tools( + server_id, user_api_key_auth, keyless_source=keyless_source ) except Exception as e: @@ -2806,10 +3077,108 @@ class MCPRequestHandler: # keyless_source AND the marker are needed: each source resolves through an UNMARKED auth, so # without keyless_source a fault under a source returns None and wins the union as allow-all. deny_all: Final = ( - unreadable_entitlement or keyless_source or is_mcp_admitted_user_subject(user_api_key_auth) + unreadable_entitlement + or keyless_source + or is_mcp_admitted_user_subject(user_api_key_auth) + or personal_key_inherits_team_grants(user_api_key_auth) ) return [] if deny_all else None + @staticmethod + async def _single_credential_tools( + server_id: str, + user_api_key_auth: UserAPIKeyAuth, + *, + keyless_source: bool = False, + ) -> list[str] | None: # mutable-ok: resolver returns a list + """Tool allowlist on ``server_id`` for one credential pinned to at most one team: the key's own + and its team's tool rules under the end-user, user, agent-caller, agent and org ceilings.""" + # Get key and team object permissions (already loaded in main auth flow) + key_obj_perm: Final = MCPRequestHandler._get_key_object_permission(user_api_key_auth) + team_obj_perm: Final = await MCPRequestHandler._get_team_object_permission(user_api_key_auth) + + # Extract tool permissions for this server. Dict keys may be + # server_ids OR names/aliases; normalize to server_id-keyed form + # before lookup so a name-based key does not silently drop its + # tool restrictions when server_id is the resolved uuid. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + key_direct_tools: Final = ( + global_mcp_server_manager.expand_tool_permissions(key_obj_perm.mcp_tool_permissions).get(server_id) + if key_obj_perm + else None + ) + + # Tools granted through the key's toolsets restrict this server exactly + # as direct tool permissions do; union with any direct grants so the + # tool-level check sees the key's full effective tool scope + key_toolset_ids: Final = (key_obj_perm.mcp_toolsets or []) if key_obj_perm else [] + key_toolset_tools: Final = ( + ( + await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=key_toolset_ids, + requires_fresh_policy=user_api_key_auth.requires_fresh_policy, + ) + ).get(server_id) + if key_toolset_ids + else None + ) + + key_group_tools: Final = await MCPRequestHandler._key_access_group_tools_for_server( + user_api_key_auth, server_id + ) + key_direct_servers: Final = await MCPRequestHandler._key_direct_servers( + key_obj_perm, requires_fresh_policy=bool(user_api_key_auth.requires_fresh_policy) + ) + key_tools: Final = _as_list( + MCPRequestHandler._key_tool_scope( + key_direct_tools, key_toolset_tools, key_group_tools, server_id in key_direct_servers + ) + ) + team_direct_tools: Final = ( + global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id) + if team_obj_perm + else None + ) + + # Tools granted through the team's toolsets restrict this server exactly + # as the team's direct tool permissions do, mirroring the key path above + team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( + team_obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) + team_group_tools: Final = await MCPRequestHandler._team_access_group_tools_for_server( + user_api_key_auth, server_id + ) + team_tools: Final = MCPRequestHandler._union_tool_grants( + team_direct_tools, team_toolset_tools, team_group_tools + ) + + # Apply same inheritance logic as get_allowed_mcp_servers. ``is not None`` so an explicit + # empty allowlist stays a deny instead of reading as "no restriction". + initial_allowed_tools: Final = ( + list(set(team_tools) & set(key_tools)) + if team_tools is not None and key_tools is not None + else team_tools + if team_tools is not None + else key_tools + ) + end_user_allowed_tools: Final = _as_list( + await MCPRequestHandler._apply_end_user_tool_ceiling(initial_allowed_tools, server_id, user_api_key_auth) + ) + user_allowed_tools: Final = _as_list( + await MCPRequestHandler._apply_user_tool_ceiling( + end_user_allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source + ) + ) + caller_allowed_tools: Final = _as_list( + await MCPRequestHandler.apply_agent_caller_tool_ceiling(user_allowed_tools, server_id, user_api_key_auth) + ) + return await MCPRequestHandler._apply_agent_and_org_tool_ceilings( + caller_allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source + ) + @staticmethod async def _apply_agent_and_org_tool_ceilings( allowed_tools: list[str] | None, @@ -3069,6 +3438,8 @@ class MCPRequestHandler: @staticmethod async def _get_allowed_mcp_servers_for_team( user_api_key_auth: UserAPIKeyAuth | None = None, + *, + require_readable_policy: bool = False, ) -> list[str]: """Get allowed MCP servers a caller inherits from the team it is pinned to. @@ -3080,7 +3451,24 @@ class MCPRequestHandler: team_ids: Final = await MCPRequestHandler._team_ids_for_mcp_grant(user_api_key_auth) if not team_ids: return [] - return await MCPRequestHandler._allowed_mcp_servers_for_single_team(team_ids[0], user_api_key_auth) + return await MCPRequestHandler._allowed_mcp_servers_for_single_team( + team_ids[0], + user_api_key_auth, + require_readable_policy=require_readable_policy, + ) + + @staticmethod + async def _get_allowed_mcp_servers_for_policy( + user_api_key_auth: UserAPIKeyAuth | None, + *, + require_readable_policy: bool, + ) -> list[str]: + if require_readable_policy: + return await MCPRequestHandler._get_allowed_mcp_servers_for_team( + user_api_key_auth, + require_readable_policy=True, + ) + return await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_api_key_auth) @staticmethod async def _team_ids_for_mcp_grant(user_api_key_auth: UserAPIKeyAuth | None) -> list[str]: @@ -3134,6 +3522,7 @@ class MCPRequestHandler: team_access_group_servers: list[str], *, requires_fresh_policy: bool = False, + require_readable_policy: bool = False, ) -> AbstractSet[str]: """The raw MCP-server set a team grants (before any org ceiling): its object_permission (direct ``mcp_servers``, the ``all_proxy_servers`` sentinel → the full registry, legacy access groups, @@ -3150,10 +3539,11 @@ class MCPRequestHandler: return set(global_mcp_server_manager.get_registry().keys()) legacy_access_group_servers: Final = await MCPRequestHandler.get_mcp_servers_from_access_groups( object_permissions.mcp_access_groups or [], - requires_fresh_policy=requires_fresh_policy, + requires_fresh_policy=requires_fresh_policy or require_readable_policy, ) toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( - object_permissions, requires_fresh_policy=requires_fresh_policy + object_permissions, + requires_fresh_policy=requires_fresh_policy or require_readable_policy, ) servers: Final = ( set(global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or [])) @@ -3168,6 +3558,8 @@ class MCPRequestHandler: async def _allowed_mcp_servers_for_single_team( team_id: str, user_api_key_auth: UserAPIKeyAuth | None, + *, + require_readable_policy: bool = False, ) -> list[str]: """Allowed MCP servers granted by ONE team (its raw grant, then capped by the team's own org for a keyless admitted subject). @@ -3216,17 +3608,19 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), + strict_errors=require_readable_policy, ) servers: Final = await MCPRequestHandler._team_granted_servers( team_obj, team_access_group_servers, requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), + require_readable_policy=require_readable_policy, ) return list(servers) except Exception as e: if isinstance(e, UnloadableEntitlementError) or ( - user_api_key_auth is not None and user_api_key_auth.requires_fresh_policy + require_readable_policy or (user_api_key_auth is not None and user_api_key_auth.requires_fresh_policy) ): raise verbose_logger.warning("Failed to get allowed MCP servers for team: %s", e) diff --git a/litellm/proxy/_experimental/mcp_server/contracts.py b/litellm/proxy/_experimental/mcp_server/contracts.py index 5c3fbde67a6..84cd2751444 100644 --- a/litellm/proxy/_experimental/mcp_server/contracts.py +++ b/litellm/proxy/_experimental/mcp_server/contracts.py @@ -138,5 +138,6 @@ class AuthorizedToolCall: host_progress_callback: ProgressCallback | None guardrail_context: Mapping[str, object] | None logging_data: Mapping[str, object] + guardrail_auth: UserAPIKeyAuth | None = None input_responses: InputResponses | None = None request_state: str | None = None diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index aba5161d56d..f65310a77e3 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3580,7 +3580,7 @@ class MCPServerManager: *, allow_all_server_ids: list[str] | None = None, submitted_server_ids: list[str] | None = None, - ) -> set: + ) -> frozenset[str]: """Servers reachable through OPEN channels rather than a grant: operator-opened ``allow_all_keys`` servers, plus the caller's own active BYOM submissions when the caller carries no explicit ``mcp_servers`` scope. @@ -3601,10 +3601,10 @@ class MCPServerManager: if user_api_key_auth is not None and ( user_api_key_auth.mcp_toolset_id is not None or user_api_key_auth.mcp_explicit_grants_only ): - return set() - if allow_all_server_ids is None: - allow_all_server_ids = self.get_allow_all_keys_server_ids() - open_ids: Final = set(allow_all_server_ids) + return frozenset() + open_ids: Final = frozenset( + allow_all_server_ids if allow_all_server_ids is not None else self.get_allow_all_keys_server_ids() + ) key_object_permission: Final = user_api_key_auth.object_permission if user_api_key_auth else None # "Explicitly scoped, so do not widen with BYOM" is a rule about a CREDENTIAL that carries # its own mcp_servers list. It does not describe a keyless admitted subject: its @@ -3616,11 +3616,14 @@ class MCPServerManager: and key_object_permission is not None and (key_object_permission.mcp_servers is not None) ) - if not has_explicit_object_permission: - if submitted_server_ids is None: - submitted_server_ids = await self._get_active_submitted_mcp_server_ids_for_user(user_api_key_auth) - open_ids.update(submitted_server_ids) - return open_ids + if has_explicit_object_permission: + return open_ids + submitted_ids: Final = ( + submitted_server_ids + if submitted_server_ids is not None + else await self._get_active_submitted_mcp_server_ids_for_user(user_api_key_auth) + ) + return open_ids | frozenset(submitted_ids) @staticmethod def _admitted_session_resource_scope(user_api_key_auth: UserAPIKeyAuth | None) -> frozenset[str] | None: @@ -6051,6 +6054,7 @@ class MCPServerManager: guardrail_context: Mapping[str, object] | None = None, tool: MCPTool | None = None, incoming_bearer_token: SecretStr | None | EllipsisType = ..., + guardrail_auth: UserAPIKeyAuth | None = None, ) -> dict[str, Any]: """ Run pre-call checks and guardrail hooks for an MCP tool call. @@ -6107,6 +6111,27 @@ class MCPServerManager: if proxy_logging_obj is None: return hook_result + base_guardrail_auth: Final = await MCPRequestHandler.guardrail_auth_for_call( + guardrail_auth=guardrail_auth, + user_api_key_auth=user_api_key_auth, + server=server, + tool_name=name, + ) + checked_auth: Final = ( + base_guardrail_auth.model_copy( + update={ + "mcp_source_team_rpm_limits": await MCPRequestHandler.personal_key_team_rpm_limits( + user_api_key_auth, server, name + ) + } + ) + if base_guardrail_auth is not None + and user_api_key_auth is not None + and MCPRequestHandler.is_personal_key_inheriting_team_grants(user_api_key_auth) + and base_guardrail_auth.mcp_source_team_rpm_limits is None + else base_guardrail_auth + ) + caller_bearer: Final = ( MCPRequestHandler.get_incoming_bearer_token(raw_headers or {}) if incoming_bearer_token is ... @@ -6119,13 +6144,11 @@ class MCPServerManager: "arguments": arguments, "server_name": server_name, "mcp_rate_limit_server_name": server.alias or server.server_name or server.name, - "user_api_key_auth": user_api_key_auth, - "user_api_key_user_id": (getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None), - "user_api_key_team_id": (getattr(user_api_key_auth, "team_id", None) if user_api_key_auth else None), - "user_api_key_end_user_id": ( - getattr(user_api_key_auth, "end_user_id", None) if user_api_key_auth else None - ), - "user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None), + "user_api_key_auth": checked_auth, + "user_api_key_user_id": (getattr(checked_auth, "user_id", None) if checked_auth else None), + "user_api_key_team_id": (getattr(checked_auth, "team_id", None) if checked_auth else None), + "user_api_key_end_user_id": (getattr(checked_auth, "end_user_id", None) if checked_auth else None), + "user_api_key_hash": (getattr(checked_auth, "api_key_hash", None) if checked_auth else None), "incoming_bearer_token": caller_bearer.get_secret_value() if caller_bearer is not None else None, "headers": logging_safe_mcp_headers(raw_headers), "tool_description": tool.description if tool is not None else None, @@ -6146,11 +6169,11 @@ class MCPServerManager: try: # Use standard pre_call_hook modified_data: Final = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_auth, + user_api_key_dict=checked_auth, data=synthetic_llm_data, call_type=CallTypes.call_mcp_tool.value, ) - await proxy_logging_obj.enforce_mcp_server_rate_limits(user_api_key_auth, server) + await proxy_logging_obj.enforce_mcp_server_rate_limits(checked_auth, server) if modified_data: # Convert response back to MCP format and apply modifications modified_kwargs: Final = proxy_logging_obj.convert_mcp_hook_response_to_kwargs( @@ -6911,6 +6934,7 @@ class MCPServerManager: input_responses: InputResponses | None = None, request_state: str | None = None, *, + guardrail_auth: UserAPIKeyAuth | None = None, catalog_auth_header: str | None | EllipsisType = ..., listed_tool: MCPTool | None | EllipsisType = ..., incoming_bearer_token: SecretStr | None | EllipsisType = ..., @@ -6938,6 +6962,12 @@ class MCPServerManager: """ start_time: Final = datetime.datetime.now() mcp_server: Final = self._resolve_mcp_server_for_tool_call(server_name, name) + effective_guardrail_auth: Final = await MCPRequestHandler.guardrail_auth_for_call( + guardrail_auth=guardrail_auth, + user_api_key_auth=user_api_key_auth, + server=mcp_server, + tool_name=name, + ) client_auth_header: Final = _catalog_auth_header(mcp_auth_header, catalog_auth_header) # Resolved before any hook runs so a missing BYOK credential (401) never @@ -6969,6 +6999,7 @@ class MCPServerManager: guardrail_context=guardrail_context, tool=self.get_listed_tool(mcp_server, name, listed_caller) if listed_tool is ... else listed_tool, incoming_bearer_token=incoming_bearer_token, + guardrail_auth=effective_guardrail_auth, ) if "arguments" in hook_result: arguments = hook_result["arguments"] @@ -6980,7 +7011,7 @@ class MCPServerManager: name=name, arguments=arguments, server_name_from_prefix=server_name, - user_api_key_auth=user_api_key_auth, + user_api_key_auth=effective_guardrail_auth, proxy_logging_obj=proxy_logging_obj, start_time=start_time, litellm_logging_obj=litellm_logging_obj, diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index c45faf38f51..39e882fcab0 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -5,11 +5,11 @@ import time import traceback import types import uuid -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextvars import ContextVar from dataclasses import dataclass, replace from datetime import datetime -from functools import partial +from functools import partial, reduce from types import EllipsisType from typing import Any, Final, NoReturn, TypeAlias, overload @@ -47,7 +47,7 @@ from mcp.types import ( ToolsCapability, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl, ConfigDict, Field, SecretStr, TypeAdapter +from pydantic import AnyUrl, ConfigDict, Field, JsonValue, SecretStr, TypeAdapter from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_logger @@ -317,6 +317,7 @@ async def _build_virtual_call_logging_obj( name: str, arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth, + logging_auth: UserAPIKeyAuth | None = None, raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, ) -> LiteLLMLoggingObj | None: @@ -336,7 +337,7 @@ async def _build_virtual_call_logging_obj( raw_headers=raw_headers, client_ip=client_ip, ) - _, virtual_logging_obj = await ProxyBaseLLMRequestProcessing( + pre_call_result: Final = await ProxyBaseLLMRequestProcessing( data={"name": name, "arguments": arguments} ).common_processing_pre_call_logic( request=request, @@ -345,10 +346,67 @@ async def _build_virtual_call_logging_obj( route_type=CallTypes.call_mcp_tool.value, proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, + logging_auth=logging_auth, ) + _, logging_candidate = TypeAdapter(tuple[dict[str, object], object]).validate_python(pre_call_result) + if not isinstance(logging_candidate, LiteLLMLoggingObj): + verbose_logger.warning("MCP request processing did not return a logging object") + return None + virtual_logging_obj: Final = logging_candidate return virtual_logging_obj +async def _log_virtual_mcp_tool_failure( + *, + logging_obj: LiteLLMLoggingObj | None, + exc: Exception, + call_start: datetime, + name: str, + arguments: dict[str, object] | None, + user_api_key_auth: UserAPIKeyAuth, + raw_headers: Mapping[str, str] | None, + client_ip: str | None, +) -> None: + try: + failure_logging_obj: Final = ( + logging_obj + if logging_obj is not None + else await _build_virtual_call_logging_obj( + name=name, + arguments=arguments or {}, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + ) + if failure_logging_obj is None: + return + from litellm.proxy.proxy_server import proxy_logging_obj as request_logging_obj + + failure_end: Final = datetime.now() # noqa: DTZ005 # matches the logging pipeline start time + failure_traceback: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + failure_logging_obj.failure_handler(exc, failure_traceback, call_start, failure_end) + await failure_logging_obj.async_failure_handler(exc, failure_traceback, call_start, failure_end) + if isinstance(exc, MCPUpstreamAuthError): + return + failure_hook: Final = TypeAdapter(Callable[..., Awaitable[HTTPException | None]]).validate_python( + request_logging_obj.post_call_failure_hook + ) + await failure_hook( + request_data={ + "name": name, + "arguments": arguments, + "litellm_logging_obj": failure_logging_obj, + }, + original_exception=exc, + user_api_key_dict=user_api_key_auth, + route="/mcp/call_tool", + traceback_str=failure_traceback, + ) + except Exception: # noqa: BLE001 # failure logging must not mask the tool call error + verbose_logger.exception("Error logging failed MCP proxy tool call") + + async def _dispatch_virtual_mcp_tool( name: str, arguments: dict[str, object] | None, @@ -393,11 +451,51 @@ async def _dispatch_virtual_mcp_tool( if mcp_proxy_mode and name in MCP_PROXY_TOOL_NAMES: assert user_api_key_auth is not None proxy_call_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes + proxy_arguments: Final = arguments or {} + proxy_arguments_json: Final = TypeAdapter(dict[str, JsonValue]).validate_python(proxy_arguments) + proxy_tool_id: Final = proxy_arguments.get("tool_id") + try: + proxy_target_tool: Final = ( + await find_mcp_proxy_tool( + tool_id=proxy_tool_id, + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=client_ip, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + if name == MCP_PROXY_CALL_TOOL_NAME and isinstance(proxy_tool_id, str) + else None + ) + proxy_guardrail_auth: Final = ( + await MCPRequestHandler.billing_auth_for_tool_call( + user_api_key_auth, + tool_name=proxy_target_tool.name, + mcp_servers=mcp_servers, + ) + if proxy_target_tool is not None + else user_api_key_auth + ) + except Exception as exc: + await _log_virtual_mcp_tool_failure( + logging_obj=None, + exc=exc, + call_start=proxy_call_start, + name=name, + arguments=arguments, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + raise proxy_logging_obj: Final = ( await _build_virtual_call_logging_obj( name=name, - arguments=arguments or {}, + arguments=proxy_arguments, user_api_key_auth=user_api_key_auth, + logging_auth=proxy_guardrail_auth, raw_headers=raw_headers, client_ip=client_ip, ) @@ -407,7 +505,7 @@ async def _dispatch_virtual_mcp_tool( try: proxy_result: Final = await handle_mcp_proxy_tool( name=name, - arguments=arguments or {}, + arguments=proxy_arguments_json, user_api_key_dict=user_api_key_auth, client_ip=client_ip, mcp_servers=mcp_servers, @@ -417,30 +515,21 @@ async def _dispatch_virtual_mcp_tool( raw_headers=raw_headers, litellm_logging_obj=proxy_logging_obj, incoming_bearer_token=incoming_bearer_token, + guardrail_auth=proxy_guardrail_auth, + resolved_tool=proxy_target_tool if name == MCP_PROXY_CALL_TOOL_NAME else ..., ) except Exception as exc: if proxy_logging_obj is not None: - from litellm.proxy.proxy_server import proxy_logging_obj as request_logging_obj - - failure_end: Final = datetime.now() # noqa: DTZ005 # matches the logging pipeline start time - failure_traceback: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) - try: - proxy_logging_obj.failure_handler(exc, failure_traceback, proxy_call_start, failure_end) - await proxy_logging_obj.async_failure_handler(exc, failure_traceback, proxy_call_start, failure_end) - if not isinstance(exc, MCPUpstreamAuthError): - await request_logging_obj.post_call_failure_hook( - request_data={ - "name": name, - "arguments": arguments, - "litellm_logging_obj": proxy_logging_obj, - }, - original_exception=exc, - user_api_key_dict=user_api_key_auth, - route="/mcp/call_tool", - traceback_str=failure_traceback, - ) - except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error - verbose_logger.exception("Error logging failed MCP proxy tool call") + await _log_virtual_mcp_tool_failure( + logging_obj=proxy_logging_obj, + exc=exc, + call_start=proxy_call_start, + name=name, + arguments=arguments, + user_api_key_auth=proxy_guardrail_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) raise if proxy_logging_obj is not None: return await _fire_mcp_tool_call_logging( @@ -448,8 +537,8 @@ async def _dispatch_virtual_mcp_tool( result=proxy_result, start_time=proxy_call_start, end_time=datetime.now(), # noqa: DTZ005 # matches the logging pipeline start time - user_api_key_auth=user_api_key_auth, - request_data=types.MappingProxyType({"name": name, "arguments": arguments}), + user_api_key_auth=proxy_guardrail_auth, + request_data=types.MappingProxyType({"name": name, "arguments": proxy_arguments}), ) return proxy_result @@ -498,16 +587,22 @@ async def _dispatch_virtual_mcp_tool( top_k=coerce_top_k(args.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K), user_api_key_dict=user_api_key_auth, ) + tool_request: Final = TypeAdapter(CallToolRequestParams).validate_python( + {"name": args.get("tool_name", ""), "arguments": args.get("arguments") or {}} + ) + tool_guardrail_auth: Final = await MCPRequestHandler.billing_auth_for_tool_call( + user_api_key_auth, + tool_name=tool_request.name, + mcp_servers=mcp_servers, + ) virtual_logging_obj: Final = await _build_virtual_call_logging_obj( - name=name, - arguments=args, + name=tool_request.name, + arguments=tool_request.arguments or {}, user_api_key_auth=user_api_key_auth, + logging_auth=tool_guardrail_auth, raw_headers=raw_headers, client_ip=client_ip, ) - tool_request: Final = CallToolRequestParams.model_validate( - types.MappingProxyType({"name": args.get("tool_name", ""), "arguments": args.get("arguments") or {}}) - ) return await handle_mcp_tool_call( tool_name=tool_request.name, arguments=tool_request.arguments or {}, @@ -520,6 +615,7 @@ async def _dispatch_virtual_mcp_tool( raw_headers=raw_headers, litellm_logging_obj=virtual_logging_obj, incoming_bearer_token=incoming_bearer_token, + guardrail_auth=tool_guardrail_auth, ) @@ -1775,7 +1871,7 @@ async def _list_mcp_resource_templates( def _resolve_display_name_to_original( name: str, - allowed_mcp_servers: list[MCPServer], + allowed_mcp_servers: Sequence[MCPServer], ) -> str: """Translate a display-name override back to the original prefixed tool name. @@ -1792,6 +1888,15 @@ def _resolve_display_name_to_original( return name +def resolve_native_tool_name_for_server( + name: str, + server: MCPServer, + allowed_mcp_servers: Sequence[MCPServer], +) -> str: + display_name_resolved: Final = _resolve_display_name_to_original(name, allowed_mcp_servers) + return strip_known_server_prefix(display_name_resolved, server) + + async def _get_byok_credential( mcp_server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, @@ -1988,6 +2093,7 @@ async def execute_mcp_tool( raw_headers: dict[str, str] | None = None, host_progress_callback: ProgressCallback | None = None, guardrail_context: Mapping[str, object] | None = None, + guardrail_auth: UserAPIKeyAuth | None = None, client_ip: str | None = None, wire_compat: WireCompat = WireCompat.LEGACY, input_responses: InputResponses | None = None, @@ -2012,6 +2118,7 @@ async def execute_mcp_tool( start_time=start_time, host_progress_callback=host_progress_callback, guardrail_context=guardrail_context, + guardrail_auth=guardrail_auth, logging_data=types.MappingProxyType(kwargs), input_responses=input_responses, request_state=request_state, @@ -2031,6 +2138,7 @@ async def _execute_mcp_tool( raw_headers: dict[str, str] | None = None, host_progress_callback: ProgressCallback | None = None, guardrail_context: Mapping[str, object] | None = None, + guardrail_auth: UserAPIKeyAuth | None = None, client_ip: str | None = None, wire_compat: WireCompat = WireCompat.LEGACY, input_responses: InputResponses | None = None, @@ -2138,7 +2246,7 @@ async def _execute_mcp_tool( break if mcp_server is not None: server_name = mcp_server.name - original_tool_name = strip_known_server_prefix(name, mcp_server) + original_tool_name = resolve_native_tool_name_for_server(name, mcp_server, allowed_mcp_servers) if requested_server is not None: if mcp_server is not None and mcp_server.server_id != requested_server.server_id: @@ -2156,7 +2264,7 @@ async def _execute_mcp_tool( if mcp_server is None: mcp_server = requested_server server_name = requested_server.name - original_tool_name = strip_known_server_prefix(name, requested_server) + original_tool_name = resolve_native_tool_name_for_server(name, requested_server, allowed_mcp_servers) # Only enforce server-level permissions when we can resolve a server if server_name: @@ -2187,6 +2295,14 @@ async def _execute_mcp_tool( name ) + effective_guardrail_auth: Final = await MCPRequestHandler.guardrail_auth_for_call( + guardrail_auth=guardrail_auth, + user_api_key_auth=user_api_key_auth, + server=mcp_server, + tool_name=original_tool_name, + allowed_mcp_servers=allowed_mcp_servers, + ) + client_auth_header: Final = mcp_auth_header if mcp_server: standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get("mcp_server_cost_info") @@ -2274,6 +2390,7 @@ async def _execute_mcp_tool( ), ), incoming_bearer_token=incoming_bearer_token, + guardrail_auth=effective_guardrail_auth, ) # `pre_call_tool_check` may return guardrail-modified # arguments; honor them on the local path too. @@ -2322,6 +2439,7 @@ async def _execute_mcp_tool( name=original_tool_name, arguments=arguments, user_api_key_auth=user_api_key_auth, + guardrail_auth=effective_guardrail_auth, mcp_auth_header=mcp_auth_header, catalog_auth_header=client_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, @@ -2331,7 +2449,11 @@ async def _execute_mcp_tool( litellm_logging_obj=litellm_logging_obj, guardrail_context=guardrail_context, host_progress_callback=_guarded_progress_callback( - host_progress_callback, litellm_logging_obj, user_api_key_auth, kwargs + callback=host_progress_callback, + logging_obj=litellm_logging_obj, + original_auth=user_api_key_auth, + guardrail_auth=effective_guardrail_auth, + request_data=kwargs, ), wire_compat=wire_compat, input_responses=input_responses, @@ -2410,6 +2532,7 @@ async def _execute_mcp_tool( ), ), incoming_bearer_token=incoming_bearer_token, + guardrail_auth=effective_guardrail_auth, ) if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args @@ -2422,7 +2545,8 @@ async def _execute_mcp_tool( return await _run_post_mcp_call_guardrails( result=converted, litellm_logging_obj=litellm_logging_obj, - user_api_key_auth=user_api_key_auth, + original_auth=user_api_key_auth, + guardrail_auth=effective_guardrail_auth, request_data=kwargs, ) @@ -2430,7 +2554,8 @@ async def _execute_mcp_tool( def _guarded_progress_callback( callback: ProgressCallback | None, logging_obj: LiteLLMLoggingObj | None, - caller: UserAPIKeyAuth | None, + original_auth: UserAPIKeyAuth | None, + guardrail_auth: UserAPIKeyAuth | None, request_data: Mapping[str, object], ) -> ProgressCallback | None: if callback is None: @@ -2442,7 +2567,11 @@ def _guarded_progress_callback( return try: checked: Final = await _run_post_mcp_call_guardrails( - CallToolResult(content=[TextContent(type="text", text=message)]), logging_obj, caller, request_data + CallToolResult(content=[TextContent(type="text", text=message)]), + litellm_logging_obj=logging_obj, + original_auth=original_auth, + guardrail_auth=guardrail_auth, + request_data=request_data, ) filtered: Final = TextContent.model_validate(checked.content[0]).text except Exception: # noqa: BLE001 # any guardrail failure must withhold unverified progress text @@ -2457,7 +2586,8 @@ def _guarded_progress_callback( async def _run_post_mcp_call_guardrails( result: CallToolResult, litellm_logging_obj: LiteLLMLoggingObj | None, - user_api_key_auth: UserAPIKeyAuth | None, + original_auth: UserAPIKeyAuth | None, + guardrail_auth: UserAPIKeyAuth | None, request_data: Mapping[str, object], ) -> CallToolResult: """Run ``post_mcp_call`` guardrails over an executed tool result. @@ -2472,15 +2602,139 @@ async def _run_post_mcp_call_guardrails( if proxy_logging_obj is None: return result + post_auth: Final = guardrail_auth if guardrail_auth is not None else original_auth + if post_auth is original_auth: + request_data_for_hook: Final = ( + litellm_logging_obj.model_call_details if litellm_logging_obj is not None else dict(request_data) + ) + return await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data=request_data_for_hook, + user_api_key_dict=post_auth, + ) + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + + base_request_data: Final[Mapping[str, object]] = ( + litellm_logging_obj.model_call_details if litellm_logging_obj is not None else request_data + ) + request_metadata: Final[object | None] = base_request_data.get("metadata") + metadata: Final[Mapping[str, object]] = ( + TypeAdapter(Mapping[str, object]).validate_python(request_metadata) + if isinstance(request_metadata, Mapping) + else {} + ) + post_auth_metadata: Final[Mapping[str, object]] = ( + TypeAdapter(Mapping[str, object]).validate_python(post_auth.metadata) + if post_auth is not None and post_auth.metadata is not None + else {} + ) + post_auth_team_metadata: Final[Mapping[str, object]] = ( + TypeAdapter(Mapping[str, object]).validate_python(post_auth.team_metadata) + if post_auth is not None and post_auth.team_metadata is not None + else {} + ) + post_auth_project_metadata: Final[Mapping[str, object]] = ( + TypeAdapter(Mapping[str, object]).validate_python(post_auth.project_metadata) + if post_auth is not None and post_auth.project_metadata is not None + else {} + ) + auth_guardrail_data: Final[dict[str, object]] = {"metadata": {}} # mutable-ok: resolver writes metadata + if post_auth is not None: + auth_for_guardrail_resolution: Final = post_auth.model_copy( + update={ + "metadata": _guardrail_metadata_for_legacy_resolver(post_auth_metadata), + "team_metadata": _guardrail_metadata_for_legacy_resolver(post_auth_team_metadata), + "project_metadata": _guardrail_metadata_for_legacy_resolver(post_auth_project_metadata), + } + ) + add_guardrails_from_auth_metadata( + user_api_key_dict=auth_for_guardrail_resolution, + data=auth_guardrail_data, + metadata_variable_name="metadata", + ) + auth_metadata: Final = auth_guardrail_data.get("metadata") + resolved_auth_metadata: Final[Mapping[str, object]] = ( + TypeAdapter(Mapping[str, object]).validate_python(auth_metadata) if isinstance(auth_metadata, Mapping) else {} + ) + request_guardrails: Final = _guardrails_from_metadata(metadata) + request_guardrail_names: Final = frozenset[str]().union( + *(_guardrail_names_from_entry(entry) for entry in request_guardrails) + ) + auth_guardrails: Final = _deduplicate_auth_guardrails( + ( + *_auth_guardrails_from_metadata(post_auth_metadata), + *_auth_guardrails_from_metadata(post_auth_team_metadata), + *_auth_guardrails_from_metadata(post_auth_project_metadata), + *_auth_guardrails_from_metadata(resolved_auth_metadata), + ), + request_guardrail_names, + ) + merged_guardrails: Final = (*request_guardrails, *auth_guardrails) + guardrail_request_data: Final = { + **base_request_data, + "metadata": { + **metadata, + **resolved_auth_metadata, + "guardrails": list(merged_guardrails), + }, + } return await proxy_logging_obj.post_mcp_call_hook( response=result, - request_data=( - litellm_logging_obj.model_call_details if litellm_logging_obj is not None else dict(request_data) - ), - user_api_key_dict=user_api_key_auth, + request_data=guardrail_request_data, + user_api_key_dict=post_auth, ) +def _guardrails_from_metadata(metadata: Mapping[str, object] | None) -> tuple[object, ...]: + if metadata is None: + return () + guardrails_value: Final = metadata.get("guardrails") + if not isinstance(guardrails_value, list): + return () + return TypeAdapter(tuple[object, ...]).validate_python(guardrails_value) + + +def _guardrail_metadata_for_legacy_resolver(metadata: Mapping[str, object]) -> dict[str, object]: + return { + **metadata, + "guardrails": [entry for entry in _guardrails_from_metadata(metadata) if isinstance(entry, str)], + } + + +def _auth_guardrails_from_metadata(metadata: Mapping[str, object]) -> tuple[object, ...]: + return tuple( + TypeAdapter(Mapping[object, object]).validate_python(guardrail) if isinstance(guardrail, Mapping) else guardrail + for guardrail in _guardrails_from_metadata(metadata) + if isinstance(guardrail, (str, Mapping)) + ) + + +def _guardrail_names_from_entry(entry: object) -> frozenset[str]: + if isinstance(entry, str): + return frozenset({entry}) + if not isinstance(entry, Mapping): + return frozenset() + typed_entry: Final[Mapping[object, object]] = TypeAdapter(Mapping[object, object]).validate_python(entry) + return frozenset(name for name in typed_entry if isinstance(name, str)) + + +def _deduplicate_auth_guardrails( + guardrails: tuple[object, ...], + request_guardrail_names: frozenset[str], +) -> tuple[object, ...]: + def keep_distinct( + accumulated: tuple[tuple[object, ...], frozenset[str]], + entry: object, + ) -> tuple[tuple[object, ...], frozenset[str]]: + kept_entries, kept_names = accumulated + entry_names: Final = _guardrail_names_from_entry(entry) + if entry_names & (request_guardrail_names | kept_names): + return accumulated + return (*kept_entries, entry), kept_names | entry_names + + return reduce(keep_distinct, guardrails, ((), frozenset()))[0] + + def suppress_completed_success_logging(logging_obj: LiteLLMLoggingObj) -> None: """An interim ``InputRequiredResult`` is not a completed call, so the ``@client`` wrapper on ``call_mcp_tool`` must not run the success handlers for it when the coroutine returns.""" @@ -2607,6 +2861,14 @@ def _single_scoped_server( return allowed_mcp_servers[0] +async def resolve_scoped_mcp_server( + user_api_key_auth: UserAPIKeyAuth, + mcp_servers: Sequence[str], +) -> MCPServer | None: + allowed_mcp_servers: Final = await _get_allowed_mcp_servers(user_api_key_auth, mcp_servers) + return _single_scoped_server(mcp_servers, allowed_mcp_servers) + + @client @catalog_operation(lambda: global_mcp_server_manager) async def call_mcp_tool( @@ -2623,6 +2885,7 @@ async def call_mcp_tool( input_responses: InputResponses | None = None, request_state: str | None = None, incoming_bearer_token: SecretStr | None | EllipsisType = ..., + guardrail_auth: UserAPIKeyAuth | None = None, **kwargs: Any, ) -> CallToolResult | InputRequiredResult: """ @@ -2686,6 +2949,7 @@ async def call_mcp_tool( request_state=request_state, incoming_bearer_token=incoming_bearer_token, requested_server_id=scoped_server.server_id if scoped_server is not None else None, + guardrail_auth=guardrail_auth, **kwargs, ) except Exception as e: @@ -2874,6 +3138,7 @@ async def _handle_managed_mcp_tool( litellm_logging_obj: LiteLLMLoggingObj | None = None, host_progress_callback: ProgressCallback | None = None, guardrail_context: Mapping[str, object] | None = None, + guardrail_auth: UserAPIKeyAuth | None = None, client_ip: str | None = None, wire_compat: WireCompat = WireCompat.LEGACY, input_responses: InputResponses | None = None, @@ -2903,6 +3168,7 @@ async def _handle_managed_mcp_tool( host_progress_callback=host_progress_callback, litellm_logging_obj=litellm_logging_obj, guardrail_context=guardrail_context, + guardrail_auth=guardrail_auth, wire_compat=wire_compat, input_responses=input_responses, request_state=request_state, @@ -3112,7 +3378,16 @@ async def _execute_mcp_server_tool_call( raw_headers=raw_headers, client_ip=_client_ip, ) - if user_api_key_auth is not None: + guardrail_auth: Final = ( + await MCPRequestHandler.billing_auth_for_tool_call( + user_api_key_auth, + tool_name=params.name, + mcp_servers=mcp_servers, + ) + if user_api_key_auth is not None + else None + ) + if guardrail_auth is not None: data = await add_litellm_data_to_request( data=body_data, request=request, @@ -3124,9 +3399,7 @@ async def _execute_mcp_server_tool_call( # half (an already-over-budget team stops granting) lives in the source gate. # Authorization is unaffected: it ran before this, and the union is resolved # from the untouched auth object passed to call_mcp_tool below. - user_api_key_dict=await MCPRequestHandler.billing_auth_for_tool_call( - user_api_key_auth, tool_name=params.name - ), + user_api_key_dict=guardrail_auth, proxy_config=proxy_config, ) else: @@ -3145,6 +3418,7 @@ async def _execute_mcp_server_tool_call( input_responses=params.input_responses, request_state=params.request_state, incoming_bearer_token=context.incoming_bearer_token, + guardrail_auth=guardrail_auth, **data, # for logging ) except MCPMissingUserEnvVarsError as e: @@ -3599,6 +3873,7 @@ class GatewayOperations: client_ip=_client_ip, host_progress_callback=operation.host_progress_callback, guardrail_context=operation.guardrail_context, + guardrail_auth=operation.guardrail_auth, wire_compat=context.wire_compat, input_responses=operation.input_responses, request_state=operation.request_state, diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index b96b6745951..ee063b81ae4 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,5 +1,6 @@ import asyncio import importlib +import sys from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime @@ -12,7 +13,7 @@ import anyio import httpx import httpx2 from fastapi import APIRouter, Depends, HTTPException, Query, Request, status -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError from starlette.datastructures import Headers from typing_extensions import ReadOnly @@ -48,6 +49,7 @@ from litellm.proxy._experimental.mcp_server.ui_session_utils import ( ) from litellm.proxy._experimental.mcp_server.utils import ( MCPMissingUserEnvVarsError, + add_server_prefix_to_name, get_server_prefix, merge_mcp_headers, ) @@ -62,7 +64,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.responses.mcp.request_context import MCPRequestContext if TYPE_CHECKING: - from mcp.types import CallToolResult + from mcp.types import CallToolResult, InputRequiredResult from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload @@ -226,6 +228,7 @@ if MCP_AVAILABLE: from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.litellm_proxy.skills.skill_search import ( DEFAULT_SKILL_SEARCH_TOP_K, ) @@ -249,6 +252,7 @@ if MCP_AVAILABLE: filter_tools_by_allowed_tools, filter_tools_by_key_team_permissions, fire_mcp_tool_call_failure_logging, + resolve_native_tool_name_for_server, ) from litellm.proxy._experimental.mcp_server.server import ( # noqa: F401 # legacy module exports _apply_toolset_scope, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -1215,6 +1219,102 @@ if MCP_AVAILABLE: "message": f"An unexpected error occurred: {e}", } + async def _execute_rest_mcp_tool( + request: Request, + tool_name: str, + tool_arguments: Mapping[str, object], + allowed_mcp_servers: Sequence[MCPServer], + canonical_server_id: str, + target_server: MCPServer | None, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + logging_obj: "LiteLLMLoggingObj | None", + guardrail_auth: UserAPIKeyAuth, + ) -> "tuple[CallToolResult | InputRequiredResult, datetime]": + ( + mcp_auth_header, + mcp_server_auth_headers, + raw_headers_from_request, + caller_oauth2_headers_from_request, + ) = _extract_mcp_headers_from_request(request, user_api_key_dict) + typed_data: Final = TypeAdapter(dict[str, object]).validate_python(data) + request_metadata: Final = TypeAdapter(dict[str, object]).validate_python(typed_data.get("metadata") or {}) + user_api_key_auth_from_metadata: Final = request_metadata.get("user_api_key_auth") + request_auth: Final = ( + {"user_api_key_auth": TypeAdapter(UserAPIKeyAuth).validate_python(user_api_key_auth_from_metadata)} + if user_api_key_auth_from_metadata is not None + else {} + ) + request_headers: Final = { + **({"mcp_auth_header": mcp_auth_header} if mcp_auth_header else {}), + **({"mcp_server_auth_headers": mcp_server_auth_headers} if mcp_server_auth_headers else {}), + "raw_headers": raw_headers_from_request, + } + execution_data: Final = MappingProxyType({**typed_data, **request_auth, **request_headers}) + _stamp_mcp_tool_metadata(logging_obj, canonical_server_id, tool_name) + user_oauth_extra_headers: Final = ( + await _get_user_oauth_extra_headers(target_server, user_api_key_dict) if target_server is not None else None + ) + caller_oauth2_headers: Final = ( + caller_oauth2_headers_from_request + if target_server is not None and target_server.auth_type in _CLIENT_FORWARDED_TOKEN_AUTH_TYPES + else None + ) + tool_start_time: Final = datetime.now() + executed: Final = await execute_mcp_tool( + name=tool_name, + arguments=TypeAdapter(dict[str, object]).validate_python(tool_arguments), + allowed_mcp_servers=list(allowed_mcp_servers), + start_time=tool_start_time, + user_api_key_auth=TypeAdapter(UserAPIKeyAuth | None).validate_python( + execution_data.get("user_api_key_auth") + ), + mcp_auth_header=TypeAdapter(str | None).validate_python(execution_data.get("mcp_auth_header")), + mcp_server_auth_headers=TypeAdapter(dict[str, dict[str, str]] | None).validate_python( + execution_data.get("mcp_server_auth_headers") + ), + oauth2_headers=user_oauth_extra_headers or caller_oauth2_headers, + raw_headers=TypeAdapter(dict[str, str] | None).validate_python(execution_data.get("raw_headers")), + client_ip=IPAddressUtils.get_mcp_client_ip(request), + litellm_logging_obj=logging_obj, + guardrail_context=MCPRequestContext.resolve_guardrail_context(execution_data), + requested_server_id=canonical_server_id, + incoming_bearer_token=MCPRequestHandler.get_guardrail_bearer_token(request.headers, user_api_key_dict), + guardrail_auth=guardrail_auth, + ) + return executed, tool_start_time + + async def _rest_tool_call_server_and_auth( + user_api_key_auth: UserAPIKeyAuth, + allowed_mcp_servers: Sequence[MCPServer], + canonical_server_id: str, + tool_name: str, + ) -> tuple[MCPServer | None, UserAPIKeyAuth]: + + target_server: Final = next( + (server for server in allowed_mcp_servers if server.server_id == canonical_server_id), + None, + ) + if target_server is None: + return None, user_api_key_auth + native_tool_name: Final = resolve_native_tool_name_for_server( + tool_name, + target_server, + allowed_mcp_servers, + ) + billing_tool_name: Final = add_server_prefix_to_name( + native_tool_name, + get_server_prefix(target_server), + ) + guardrail_auth: Final = await MCPRequestHandler.billing_auth_for_tool_call( + user_api_key_auth, + tool_name=billing_tool_name, + scoped_server=target_server, + native_tool_name=native_tool_name, + allowed_mcp_servers=allowed_mcp_servers, + ) + return target_server, guardrail_auth + @router.post("/tools/call", dependencies=[Depends(user_api_key_auth)]) @catalog_operation(global_manager) async def call_tool_rest_api( @@ -1279,9 +1379,48 @@ if MCP_AVAILABLE: _request_start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the tool start time below try: ( - data, - logging_obj, - ) = await proxy_base_llm_response_processor.common_processing_pre_call_logic( + allowed_mcp_servers, + canonical_server_id, + ) = await _resolve_allowed_mcp_servers_with_ip_filter(request, user_api_key_dict, server_id) + except Exception as access_error: + try: + failure_pre_call_result: Final = ( + await proxy_base_llm_response_processor.common_processing_pre_call_logic( + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + route_type=CallTypes.call_mcp_tool.value, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + skip_guardrails=True, + logging_auth=user_api_key_dict, + ) + ) + failure_data, failure_logging_candidate = TypeAdapter( + tuple[dict[str, object], object] + ).validate_python(failure_pre_call_result) + if isinstance(failure_logging_candidate, LiteLLMLoggingObj): + _stamp_mcp_tool_metadata(failure_logging_candidate, server_id, tool_name) + await _safe_fire_mcp_tool_call_failure_logging( + failure_logging_candidate, + access_error, + _request_start_time, + user_api_key_dict, + failure_data, + ) + else: + verbose_logger.warning("MCP request processing did not return a logging object") + except Exception as logging_error: # noqa: BLE001 # logging failures must not mask the access error + verbose_logger.warning("MCP tool call access-failure logging failed: %s", logging_error) + raise + target_server, guardrail_auth = await _rest_tool_call_server_and_auth( + user_api_key_dict, + allowed_mcp_servers, + canonical_server_id, + tool_name, + ) + try: + pre_call_result: Final = await proxy_base_llm_response_processor.common_processing_pre_call_logic( request=request, user_api_key_dict=user_api_key_dict, proxy_config=proxy_config, @@ -1289,81 +1428,53 @@ if MCP_AVAILABLE: proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, skip_guardrails=True, + logging_auth=guardrail_auth, ) - - # Extract MCP auth headers from request and add to data dict - ( - mcp_auth_header, - mcp_server_auth_headers, - raw_headers_from_request, - caller_oauth2_headers_from_request, - ) = _extract_mcp_headers_from_request(request, user_api_key_dict) - if mcp_auth_header: - data["mcp_auth_header"] = mcp_auth_header - if mcp_server_auth_headers: - data["mcp_server_auth_headers"] = mcp_server_auth_headers - data["raw_headers"] = raw_headers_from_request - - # Extract user_api_key_auth from metadata and add to top level - # call_mcp_tool expects user_api_key_auth as a top-level parameter - if "metadata" in data and "user_api_key_auth" in data["metadata"]: - data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"] - - _stamp_mcp_tool_metadata(logging_obj, server_id, tool_name) - - # Resolve allowed MCP servers with IP filtering - ( - allowed_mcp_servers, - canonical_server_id, - ) = await _resolve_allowed_mcp_servers_with_ip_filter(request, user_api_key_dict, server_id) - - # Look up per-user OAuth headers for this server (mirrors list_tool_rest_api). - user_oauth_extra_headers: dict[str, str] | None = None - target_server: Final = next( - (s for s in allowed_mcp_servers if s.server_id == canonical_server_id), - None, - ) - if target_server is not None: - user_oauth_extra_headers = await _get_user_oauth_extra_headers(target_server, user_api_key_dict) - caller_oauth2_headers: Final = ( - caller_oauth2_headers_from_request - if target_server is not None and target_server.auth_type in _CLIENT_FORWARDED_TOKEN_AUTH_TYPES - else None + data, logging_candidate = TypeAdapter(tuple[dict[str, object], object]).validate_python(pre_call_result) + logging_obj: Final[LiteLLMLoggingObj | None] = ( + logging_candidate if isinstance(logging_candidate, LiteLLMLoggingObj) else None ) + if logging_obj is None: + verbose_logger.warning("MCP request processing did not return a logging object") - # Call execute_mcp_tool directly (permission checks already done) - _tool_start_time: Final = datetime.now() - executed: Final = await execute_mcp_tool( - name=tool_name, - arguments=tool_arguments, + executed, _tool_start_time = await _execute_rest_mcp_tool( + request=request, + tool_name=tool_name, + tool_arguments=tool_arguments, allowed_mcp_servers=allowed_mcp_servers, - start_time=_tool_start_time, - user_api_key_auth=data.get("user_api_key_auth"), - mcp_auth_header=data.get("mcp_auth_header"), - mcp_server_auth_headers=data.get("mcp_server_auth_headers"), - oauth2_headers=user_oauth_extra_headers or caller_oauth2_headers, - raw_headers=data.get("raw_headers"), - client_ip=IPAddressUtils.get_mcp_client_ip(request), - litellm_logging_obj=data.get("litellm_logging_obj"), - guardrail_context=MCPRequestContext.resolve_guardrail_context(data), - requested_server_id=canonical_server_id, - incoming_bearer_token=MCPRequestHandler.get_guardrail_bearer_token( - request.headers, user_api_key_dict - ), + canonical_server_id=canonical_server_id, + target_server=target_server, + user_api_key_dict=user_api_key_dict, + data=data, + logging_obj=logging_obj, + guardrail_auth=guardrail_auth, ) result: Final = complete_call_tool_result(executed, WireCompat.LEGACY) - except Exception as e: - request_data: Final = proxy_base_llm_response_processor.data - await _safe_fire_mcp_tool_call_failure_logging( - request_data.get("litellm_logging_obj"), e, _request_start_time, user_api_key_dict, request_data - ) - raise + finally: + failure: Final = sys.exc_info()[1] + if isinstance(failure, Exception): + request_data: Final[Mapping[str, object]] = TypeAdapter(Mapping[str, object]).validate_python( + proxy_base_llm_response_processor.data + ) + request_logging_obj_candidate: Final = request_data.get("litellm_logging_obj") + request_logging_obj: Final[LiteLLMLoggingObj | None] = ( + request_logging_obj_candidate + if isinstance(request_logging_obj_candidate, LiteLLMLoggingObj) + else None + ) + await _safe_fire_mcp_tool_call_failure_logging( + request_logging_obj, + failure, + _request_start_time, + guardrail_auth, + request_data, + ) return await _safe_fire_mcp_tool_call_logging( logging_obj, result, _tool_start_time, datetime.now(), - user_api_key_auth=user_api_key_dict, + user_api_key_auth=guardrail_auth, request_data=data, ) except MCPMissingUserEnvVarsError as e: diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 18642a9162c..371a1ff6fc9 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -600,6 +600,8 @@ async def handle_mcp_proxy_tool( raw_headers: dict[str, str] | None = None, # mutable-ok: preserve request headers litellm_logging_obj: LiteLLMLoggingObj | None = None, incoming_bearer_token: SecretStr | None | EllipsisType = ..., + guardrail_auth: UserAPIKeyAuth | None = None, + resolved_tool: Tool | None | EllipsisType = ..., ) -> CallToolResult: from fastapi import HTTPException @@ -646,18 +648,22 @@ async def handle_mcp_proxy_tool( tool_id: Final = arguments.get("tool_id") tool: Final = ( - await find_mcp_proxy_tool( - tool_id, - user_api_key_auth=user_api_key_dict, - mcp_servers=mcp_servers, - client_ip=client_ip, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, + resolved_tool + if resolved_tool is not ... + else ( + await find_mcp_proxy_tool( + tool_id, + user_api_key_auth=user_api_key_dict, + mcp_servers=mcp_servers, + client_ip=client_ip, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + if isinstance(tool_id, str) + else None ) - if isinstance(tool_id, str) - else None ) if tool is None: return _text_tool_result("Unknown or unauthorized tool_id", is_error=True) @@ -687,6 +693,7 @@ async def handle_mcp_proxy_tool( raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, incoming_bearer_token=incoming_bearer_token, + guardrail_auth=guardrail_auth, ) @@ -704,6 +711,7 @@ async def handle_mcp_tool_call( requested_server_id: str | None = None, guardrail_context: Mapping[str, object] | None = None, incoming_bearer_token: SecretStr | None | EllipsisType = ..., + guardrail_auth: UserAPIKeyAuth | None = None, ) -> CallToolResult: from litellm.proxy._experimental.mcp_server.operations import ( _get_allowed_mcp_servers, @@ -745,6 +753,7 @@ async def handle_mcp_tool_call( litellm_logging_obj=litellm_logging_obj, requested_server_id=requested_server_id, guardrail_context=guardrail_context, + guardrail_auth=guardrail_auth, incoming_bearer_token=incoming_bearer_token, ) return complete_call_tool_result(result, WireCompat.LEGACY) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 5c8849e50ca..c11d296837e 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -731,6 +731,10 @@ async def invoke_agent_a2a( ) user_api_key_dict.invoked_agent_id = agent.agent_id + strict_policy: Final = agent.identity_managed or bool(user_api_key_dict.requires_fresh_policy) + user_api_key_auth_for_call: Final = await AgentRequestHandler.personal_key_agent_billing_auth( + user_api_key_dict, agent.agent_id, strict=strict_policy + ) _enforce_inbound_trace_id(agent, request) @@ -790,7 +794,7 @@ async def invoke_agent_a2a( data, logging_obj = await processor.common_processing_pre_call_logic( request=request, general_settings=general_settings, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=user_api_key_auth_for_call, proxy_logging_obj=proxy_logging_obj, proxy_config=proxy_config, route_type="asend_message", @@ -889,7 +893,7 @@ async def invoke_agent_a2a( try: response = await proxy_logging_obj.post_call_success_hook( - user_api_key_dict=user_api_key_dict, + user_api_key_dict=user_api_key_auth_for_call, data=data, response=response, ) @@ -924,7 +928,7 @@ async def invoke_agent_a2a( metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), agent_extra_headers=agent_extra_headers, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=user_api_key_auth_for_call, request_data=data, proxy_logging_obj=proxy_logging_obj, served_version=served_version, @@ -990,7 +994,7 @@ async def invoke_agent_a2a( response = LiteLLMSendMessageResponse.from_dict(result, request_id=request_id) response = await proxy_logging_obj.post_call_success_hook( - user_api_key_dict=user_api_key_dict, + user_api_key_dict=user_api_key_auth_for_call, data=data, response=response, ) @@ -1017,7 +1021,7 @@ async def invoke_agent_a2a( request_id=request_id, extra_headers=agent_extra_headers, proxy_logging_obj=proxy_logging_obj, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=user_api_key_auth_for_call, request_data=data, served_version=served_version, ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 48213625c4f..dd002a2763e 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -8,6 +8,7 @@ Follows the same pattern as MCP permission handling. import asyncio from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass +from itertools import chain from types import MappingProxyType from typing import Final, Protocol, TypeAlias @@ -45,6 +46,11 @@ class RestrictedAgentAccess: agent_ids: frozenset[str] +@dataclass(frozen=True, slots=True) +class AgentPolicyUnavailable: + team_id: str + + AgentAccess: TypeAlias = UnrestrictedAgentAccess | RestrictedAgentAccess @@ -83,6 +89,16 @@ def _restricted_ids(access: AgentAccess) -> frozenset[str] | None: return _to_stable_ids(access.agent_ids) +def _union_agent_access(accesses: Sequence[AgentAccess]) -> AgentAccess: + if accesses and all(isinstance(access, UnrestrictedAgentAccess) for access in accesses): + return UnrestrictedAgentAccess() + return RestrictedAgentAccess( + frozenset( + chain.from_iterable(access.agent_ids for access in accesses if isinstance(access, RestrictedAgentAccess)) + ) + ) + + def _intersect_agent_access(key_access: AgentAccess, team_access: AgentAccess) -> AgentAccess: key_ids: Final = _restricted_ids(key_access) team_ids: Final = _restricted_ids(team_access) @@ -251,7 +267,11 @@ class AgentRequestHandler: ) return target.agent_id in explicit - match await AgentRequestHandler.resolve_agent_access(user_api_key_auth, resolve_ceiling): + match await AgentRequestHandler.resolve_agent_access( + user_api_key_auth, + resolve_ceiling, + strict=bool(user_api_key_auth is not None and user_api_key_auth.requires_fresh_policy), + ): case UnrestrictedAgentAccess(): return True case RestrictedAgentAccess(allowed_agent_ids): @@ -370,6 +390,200 @@ class AgentRequestHandler: verbose_logger.warning("Failed to get allowed agents for key: %s", e) return UnrestrictedAgentAccess() + @staticmethod + async def _personal_key_owner_team_agent_access(user_api_key_auth: UserAPIKeyAuth, *, strict: bool) -> AgentAccess: + """Union owner-team access; teams without a rule add nothing once any team declares one.""" + from litellm.proxy.auth.auth_checks import personal_key_owner_teams_with_status + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + try: + owner_teams: Final = await personal_key_owner_teams_with_status( + valid_token=user_api_key_auth, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=strict, + ) + except Exception as e: # noqa: BLE001 # fail closed: an owner whose teams cannot be read grants nothing + if strict: + raise HTTPException(503, "Agent invocation policy is unavailable") from e + verbose_logger.warning("Failed to load owner teams for personal key agent access: %s", e) + return RestrictedAgentAccess(frozenset()) + team_results: Final = await asyncio.gather( + *(AgentRequestHandler._owner_team_agent_grants(team, strict=strict) for team in owner_teams.teams) + ) + withheld_results: Final = await asyncio.gather( + *(AgentRequestHandler._owner_team_agent_grants(team, strict=strict) for team in owner_teams.withheld) + ) + if any(isinstance(access, AgentPolicyUnavailable) for access in team_results): + raise HTTPException(503, "Agent invocation policy is unavailable") + team_accesses: Final = tuple( + access for access in team_results if not isinstance(access, AgentPolicyUnavailable) + ) + withheld_restrictions: Final = tuple( + RestrictedAgentAccess(frozenset()) + for access in withheld_results + if isinstance(access, (AgentPolicyUnavailable, RestrictedAgentAccess)) + ) + accesses: Final = ( + (*team_accesses, *withheld_restrictions, RestrictedAgentAccess(frozenset())) + if owner_teams.unreadable + else (*team_accesses, *withheld_restrictions) + ) + return _union_agent_access(accesses) + + @staticmethod + async def personal_key_agent_billing_auth( + auth: UserAPIKeyAuth, + agent_id: str, + *, + strict: bool = False, + ) -> UserAPIKeyAuth: + try: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.auth.auth_checks import ( + personal_key_inherits_team_grants, + personal_key_owner_teams, + ) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if not personal_key_inherits_team_grants(auth): + return auth + teams: Final = await personal_key_owner_teams( + valid_token=auth, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=strict, + ) + accesses: Final = await asyncio.gather( + *(AgentRequestHandler._owner_team_agent_grants(team, strict=strict) for team in teams) + ) + stable_agent_id: Final = global_agent_registry.stable_agent_id(agent_id) + + restricted_sources: Final = tuple( + (team, access) + for team, access in zip(teams, accesses, strict=True) + if isinstance(access, RestrictedAgentAccess) and stable_agent_id in _to_stable_ids(access.agent_ids) + ) + unrestricted_sources: Final = ( + tuple(zip(teams, accesses, strict=True)) + if not restricted_sources + and teams + and all(isinstance(access, UnrestrictedAgentAccess) for access in accesses) + else () + ) + sources: Final = restricted_sources or unrestricted_sources + if not sources: + verbose_logger.warning("No owner team grants agent %s; billing against the personal key", agent_id) + return auth + team: Final = min((source[0] for source in sources), key=lambda source: source.team_id) + except Exception as e: # noqa: BLE001 # attribution must never fail an authorized call + verbose_logger.warning("Failed to attribute agent %s spend to an owner team: %s", agent_id, e) + return auth + return await AgentRequestHandler._personal_key_agent_billing_auth_for_team( + auth, + agent_id, + team, + strict=strict, + ) + + @staticmethod + async def _personal_key_agent_billing_auth_for_team( + auth: UserAPIKeyAuth, + agent_id: str, + team: LiteLLM_TeamTable, + *, + strict: bool, + ) -> UserAPIKeyAuth: + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.auth.team_grants import team_grant_auth_update, team_grant_billing_auth_update + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + try: + membership: Final = ( + await get_team_membership( + user_id=auth.user_id, + team_id=team.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=strict, + ) + if auth.user_id and prisma_client is not None + else None + ) + team_auth_update: Final = await team_grant_billing_auth_update(auth, team, membership) + except Exception as e: # noqa: BLE001 # a chosen team's controls must survive billing lookup failures + verbose_logger.warning( + "Failed to load billing controls for agent %s from owner team %s: %s", + agent_id, + team.team_id, + e, + ) + return auth.model_copy(update=team_grant_auth_update(team, None, auth.user_id, key_org_id=auth.org_id)) + return auth.model_copy(update=team_auth_update) + + @staticmethod + async def _owner_team_agent_grants( + team: LiteLLM_TeamTable, + *, + strict: bool, + ) -> AgentAccess | AgentPolicyUnavailable: + if team.object_permission_id is not None and team.object_permission is None: + if strict: + return AgentPolicyUnavailable(team_id=team.team_id) + verbose_logger.warning("Owner team %s has an unreadable object_permission, granting nothing", team.team_id) + return RestrictedAgentAccess(frozenset()) + try: + return await AgentRequestHandler._team_agent_grants(team, strict=strict) + except Exception as e: # noqa: BLE001 # per-team isolation: a team whose grants cannot be read grants nothing + if strict: + return AgentPolicyUnavailable(team_id=team.team_id) + verbose_logger.warning("Failed to get allowed agents for owner team %s: %s", team.team_id, e) + return RestrictedAgentAccess(frozenset()) + + @staticmethod + async def _team_agent_grants(team: LiteLLM_TeamTable, *, strict: bool) -> AgentAccess: + """The agents one loaded team grants: its object_permission agents plus its access groups. + Declared-but-empty grants stay restricted; a team with no grant at all does not restrict.""" + # 1. Get agents from object_permission (native permissions) + object_permissions: Final = team.object_permission + direct_agents: Final = tuple(object_permissions.agents or () if object_permissions is not None else ()) + declared_access_groups: Final = tuple( + object_permissions.agent_access_groups or () if object_permissions is not None else () + ) + # 2. Also include agents from team's access_group_ids (unified access groups) + team_access_group_ids: Final = tuple(team.access_group_ids or ()) + + if not direct_agents and not declared_access_groups and not team_access_group_ids: + return UnrestrictedAgentAccess() + + access_group_agents: Final = ( + tuple( + await AgentRequestHandler._get_agents_from_access_groups(declared_access_groups, check_db_only=strict) + ) + if declared_access_groups + else () + ) + unified_agents: Final = ( + tuple( + await AgentRequestHandler._get_unified_access_group_agents(team_access_group_ids, check_db_only=strict) + ) + if team_access_group_ids + else () + ) + + return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) + @staticmethod async def _get_allowed_agents_for_team( user_api_key_auth: UserAPIKeyAuth | None = None, @@ -389,7 +603,11 @@ class AgentRequestHandler: return UnrestrictedAgentAccess() if user_api_key_auth.team_id is None: - return UnrestrictedAgentAccess() + from litellm.proxy.auth.auth_checks import personal_key_inherits_team_grants + + if not personal_key_inherits_team_grants(user_api_key_auth): + return UnrestrictedAgentAccess() + return await AgentRequestHandler._personal_key_owner_team_agent_access(user_api_key_auth, strict=strict) try: from litellm.proxy.auth.auth_checks import get_team_object @@ -415,38 +633,7 @@ class AgentRequestHandler: if team_obj is None: return RestrictedAgentAccess(frozenset()) if strict else UnrestrictedAgentAccess() - # 1. Get agents from object_permission (native permissions) - object_permissions: Final = team_obj.object_permission - direct_agents: Final = tuple(object_permissions.agents or () if object_permissions is not None else ()) - declared_access_groups: Final = tuple( - object_permissions.agent_access_groups or () if object_permissions is not None else () - ) - # 2. Also include agents from team's access_group_ids (unified access groups) - team_access_group_ids: Final = tuple(team_obj.access_group_ids or ()) - - if not direct_agents and not declared_access_groups and not team_access_group_ids: - return UnrestrictedAgentAccess() - - access_group_agents: Final = ( - tuple( - await AgentRequestHandler._get_agents_from_access_groups( - declared_access_groups, check_db_only=strict - ) - ) - if declared_access_groups - else () - ) - unified_agents: Final = ( - tuple( - await AgentRequestHandler._get_unified_access_group_agents( - team_access_group_ids, check_db_only=strict - ) - ) - if team_access_group_ids - else () - ) - - return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) + return await AgentRequestHandler._team_agent_grants(team_obj, strict=strict) except Exception as e: if strict: raise HTTPException(503, "Agent invocation policy is unavailable") from e diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a04fba2d905..0f94d2952ec 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -18,7 +18,7 @@ from dataclasses import dataclass from functools import partial from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, Protocol, TypeAlias +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, Protocol, TypeAlias from fastapi import HTTPException, Request, status from pydantic import BaseModel, TypeAdapter, ValidationError @@ -2585,6 +2585,10 @@ def _membership_from_cached_payload( return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS +def _raise_no_db_connected() -> NoReturn: + raise Exception("No db connected") + + @log_db_metrics @with_service_target(AUTH_OBJECTS_TARGET) async def _fetch_team_membership_from_db( @@ -2594,9 +2598,13 @@ async def _fetch_team_membership_from_db( user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, + *, + use_writer: bool = False, ) -> LiteLLM_TeamMembership | None: _ = parent_otel_span, proxy_logging_obj - response: Final = await _dictable_table(TeamMembershipRepository(prisma_client), "team_membership").find_unique( + response: Final = await _dictable_table( + TeamMembershipRepository(prisma_client, use_writer=use_writer), "team_membership" + ).find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) @@ -2650,6 +2658,8 @@ async def get_team_membership( user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, + *, + check_db_only: bool = False, ) -> Optional["LiteLLM_TeamMembership"]: """ Returns team membership object if user is member of team. @@ -2661,6 +2671,19 @@ async def get_team_membership( _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) + if check_db_only: + if prisma_client is None: + _raise_no_db_connected() + return await _fetch_team_membership_from_db( + user_id=user_id, + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + use_writer=True, + ) + l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True) l1_membership: Final = _membership_from_cached_payload(l1_cached) if not isinstance(l1_membership, _TeamMembershipCacheMiss): @@ -2671,7 +2694,7 @@ async def get_team_membership( return _membership_from_shared_load(await asyncio.shield(inflight)) if prisma_client is None: - raise Exception("No db connected") + _raise_no_db_connected() task: Final = asyncio.ensure_future( _load_team_membership_on_cache_miss( @@ -4674,6 +4697,8 @@ async def get_mcp_server_ids_from_access_groups( user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, check_db_only: bool = False, + *, + strict_errors: bool = False, ) -> list[str]: """ Collect MCP server IDs from unified access groups. @@ -4686,6 +4711,7 @@ async def get_mcp_server_ids_from_access_groups( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, check_db_only=check_db_only, + strict_errors=strict_errors, ) @@ -5470,6 +5496,8 @@ async def _personal_key_owner( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + *, + check_db_only: bool = False, ) -> LiteLLM_UserTable | None: if valid_token.user_id is None: return None @@ -5481,6 +5509,7 @@ async def _personal_key_owner( user_id_upsert=False, parent_otel_span=valid_token.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) except Exception as e: # noqa: BLE001 # fail closed: an owner that cannot be loaded grants nothing verbose_proxy_logger.warning( @@ -5870,13 +5899,32 @@ def personal_key_team_model_access_applies(valid_token: UserAPIKeyAuth | None) - ) -async def _load_personal_key_team( +def personal_key_inherits_team_grants(valid_token: UserAPIKeyAuth | None) -> bool: + """MCP server, MCP tool and agent inheritance is for a virtual key only. A keyless auth with no + team (one grant source of an admitted subject, a JWT user) resolves as the plain caller it is.""" + return ( + personal_key_team_model_access_applies(valid_token) + and valid_token is not None + and valid_token.api_key is not None + ) + + +@dataclass(frozen=True, slots=True) +class _PersonalKeyTeamLoad: + team: LiteLLM_TeamTableCachedObj | None + unreadable: bool + withheld_team: LiteLLM_TeamTableCachedObj | None + + +async def _load_personal_key_team_with_status( team_id: str, valid_token: UserAPIKeyAuth, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, -) -> LiteLLM_TeamTableCachedObj | None: + *, + check_db_only: bool = False, +) -> _PersonalKeyTeamLoad: try: team_object: Final = await get_team_object( team_id=team_id, @@ -5884,13 +5932,36 @@ async def _load_personal_key_team( user_api_key_cache=user_api_key_cache, parent_otel_span=valid_token.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) except Exception as e: # noqa: BLE001 # fail closed: a team that cannot be loaded grants nothing verbose_proxy_logger.warning( "Personal key team model access: team=%s lookup failed, granting nothing: %s", team_id, e ) - return None - return None if team_object.blocked is True else team_object + return _PersonalKeyTeamLoad(team=None, unreadable=True, withheld_team=None) + if team_object.blocked is True: + return _PersonalKeyTeamLoad(team=None, unreadable=False, withheld_team=team_object) + return _PersonalKeyTeamLoad(team=team_object, unreadable=False, withheld_team=None) + + +async def _load_personal_key_team( + team_id: str, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + *, + check_db_only: bool = False, +) -> LiteLLM_TeamTableCachedObj | None: + loaded: Final = await _load_personal_key_team_with_status( + team_id=team_id, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + return loaded.team async def _load_personal_key_teams( @@ -5899,6 +5970,8 @@ async def _load_personal_key_teams( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, + *, + check_db_only: bool = False, ) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: return tuple( await asyncio.gather( @@ -5909,6 +5982,7 @@ async def _load_personal_key_teams( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) for team_id in dict.fromkeys(user_object.teams) ) @@ -5988,6 +6062,8 @@ async def _personal_key_listing_teams( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache | None, proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, ) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: if ( user_api_key_dict.user_id is None @@ -6001,6 +6077,7 @@ async def _personal_key_listing_teams( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) if owner is None: return () @@ -6010,9 +6087,177 @@ async def _personal_key_listing_teams( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) +async def _personal_key_team_within_budget( + team: LiteLLM_TeamTableCachedObj, + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + *, + check_db_only: bool = False, +) -> bool: + source: Final = valid_token.model_copy( + update={"team_id": team.team_id, "org_id": valid_token.org_id or team.organization_id} + ) + try: + await team_max_budget_check( + team_object=team, + valid_token=source, + proxy_logging_obj=proxy_logging_obj, + ) + await _team_multi_budget_check(team_object=team) + if check_db_only: + team_member_user_id: Final = valid_token.user_id + if team_member_user_id is None: + return False + team_membership: Final = await get_team_membership( + user_id=team_member_user_id, + team_id=team.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + await _check_team_member_budget( + team_object=team, + user_object=None, + valid_token=source, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=True, + ) + else: + await _check_team_member_budget( + team_object=team, + user_object=None, + valid_token=source, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await organization_max_budget_check( + valid_token=source, + team_object=team, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except litellm.BudgetExceededError: + return False + except Exception as e: # noqa: BLE001 # fail closed: a team whose budget cannot be checked grants nothing + verbose_proxy_logger.warning("Failed to check budget for personal key team %s, skipping: %s", team.team_id, e) + return False + return True + + +@dataclass(frozen=True, slots=True) +class PersonalKeyOwnerTeams: + teams: tuple[LiteLLM_TeamTableCachedObj, ...] + unreadable: bool + withheld: tuple[LiteLLM_TeamTableCachedObj, ...] + + +async def personal_key_owner_teams_with_status( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> PersonalKeyOwnerTeams: + """The owner's current, loadable, unblocked, within-budget teams: the grant sources a personal key inherits + MCP server, MCP tool and agent access from under ``personal_key_model_access_from_teams``.""" + if valid_token.user_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None: + return PersonalKeyOwnerTeams(teams=(), unreadable=False, withheld=()) + owner: Final = await _personal_key_owner( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + if owner is None: + return PersonalKeyOwnerTeams(teams=(), unreadable=False, withheld=()) + team_ids: Final = tuple(team_id for team_id in dict.fromkeys(owner.teams) if team_id and team_id != UI_TEAM_ID) + team_loads: Final = await asyncio.gather( + *( + _load_personal_key_team_with_status( + team_id=team_id, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + for team_id in team_ids + ) + ) + blocked_teams: Final = tuple( + load.withheld_team + for load in team_loads + if load.withheld_team is not None + and (valid_token.org_id is None or load.withheld_team.organization_id in (None, valid_token.org_id)) + ) + loaded_teams: Final = tuple(load.team for load in team_loads if load.team is not None) + eligible: Final = tuple( + team + for team in loaded_teams + if valid_token.org_id is None or team.organization_id in (None, valid_token.org_id) + ) + if not eligible: + return PersonalKeyOwnerTeams( + teams=(), + unreadable=any(load.unreadable for load in team_loads), + withheld=blocked_teams, + ) + within_budget: Final = await asyncio.gather( + *( + _personal_key_team_within_budget( + team=team, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + for team in eligible + ) + ) + return PersonalKeyOwnerTeams( + teams=tuple(team for team, allowed in zip(eligible, within_budget, strict=True) if allowed), + unreadable=any(load.unreadable for load in team_loads), + withheld=( + *blocked_teams, + *(team for team, allowed in zip(eligible, within_budget, strict=True) if not allowed), + ), + ) + + +async def personal_key_owner_teams( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> tuple[LiteLLM_TeamTableCachedObj, ...]: + return ( + await personal_key_owner_teams_with_status( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + ).teams + + async def personal_key_team_visible_models( models: Sequence[str], user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py index 450a3ce456e..a48bb12cfb7 100644 --- a/litellm/proxy/auth/team_grants.py +++ b/litellm/proxy/auth/team_grants.py @@ -19,6 +19,7 @@ from litellm.proxy._types import ( LiteLLM_TeamMembership, LiteLLM_TeamTable, Member, + UserAPIKeyAuth, ) from litellm.types.llms.base import LiteLLMBaseModel @@ -74,6 +75,91 @@ class TeamGrants(TypedDict, total=False): team_member_rpm_limit: ReadOnly[int | None] +_TEAM_GRANT_AUTH_FIELDS: Final[frozenset[str]] = frozenset( + { + "team_alias", + "team_tpm_limit", + "team_rpm_limit", + "team_tpd_limit", + "team_max_budget", + "team_soft_budget", + "team_spend", + "team_metadata", + "team_member", + "team_member_spend", + "team_member_tpm_limit", + "team_member_rpm_limit", + } +) + + +def team_grant_auth_update( + team_object: LiteLLM_TeamTable, + team_membership: LiteLLM_TeamMembership | None, + user_id: str | None, + *, + key_org_id: str | None, +) -> Mapping[str, object]: + grant_values: Final = team_grants(team_object, team_membership, user_id) + return MappingProxyType( + { + "team_id": team_object.team_id, + "org_id": key_org_id or team_object.organization_id, + **{field: value for field, value in grant_values.items() if field in _TEAM_GRANT_AUTH_FIELDS}, + } + ) + + +async def team_grant_billing_auth_update( + auth: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable, + team_membership: LiteLLM_TeamMembership | None, +) -> Mapping[str, object]: + team_auth_update: Final = team_grant_auth_update( + team_object, + team_membership, + auth.user_id, + key_org_id=auth.org_id, + ) + if auth.org_id is not None or team_object.organization_id is None: + return team_auth_update + + organization_auth_fields: Final[Mapping[str, object]] = MappingProxyType( + { + "organization_alias": None, + "organization_metadata": None, + "organization_max_budget": None, + "organization_tpm_limit": None, + "organization_rpm_limit": None, + } + ) + from litellm.proxy.auth.auth_checks import get_org_object_for_request + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + return MappingProxyType({**team_auth_update, **organization_auth_fields}) + org_object: Final = await get_org_object_for_request( + org_id=team_object.organization_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if org_object is None: + return MappingProxyType({**team_auth_update, **organization_auth_fields}) + budget: Final = org_object.litellm_budget_table + return MappingProxyType( + { + **team_auth_update, + "organization_alias": org_object.organization_alias, + "organization_metadata": org_object.metadata, + "organization_max_budget": budget.max_budget if budget is not None else None, + "organization_tpm_limit": budget.tpm_limit if budget is not None else None, + "organization_rpm_limit": budget.rpm_limit if budget is not None else None, + } + ) + + def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns: try: return _TeamJsonColumns.model_validate(team_object.model_dump(include=_JSON_COLUMNS)) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index b443711752a..24a7b46da4b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -253,7 +253,9 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( AnthropicStreamModelRestamper, ) from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, add_litellm_data_to_request, + get_metadata_variable_name, refresh_proxy_server_request_body_snapshot, reject_url_valued_destination, sanitize_for_log, @@ -2244,6 +2246,7 @@ class ProxyBaseLLMRequestProcessing: llm_router: Router | None = None, rate_limited_model: str | None = None, skip_guardrails: bool = False, + logging_auth: UserAPIKeyAuth | None = None, ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks @@ -2263,6 +2266,14 @@ class ProxyBaseLLMRequestProcessing: version=version, proxy_config=proxy_config, ) + if logging_auth is not None: + metadata_variable_name: Final = get_metadata_variable_name(request) + request_metadata: Final = self.data.get(metadata_variable_name) + if isinstance(request_metadata, dict): + self.data[metadata_variable_name] = { + **TypeAdapter(dict[str, object]).validate_python(request_metadata), + **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(logging_auth), + } if not general_settings.get("expose_fallback_errors_to_caller"): self.data.pop("include_fallback_errors", None) if route_type in {"aresponses", "_aresponses_websocket"}: @@ -2285,8 +2296,6 @@ class ProxyBaseLLMRequestProcessing: # Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved if queue_time_seconds is not None: - from litellm.proxy.litellm_pre_call_utils import get_metadata_variable_name - _metadata_variable_name: Final = get_metadata_variable_name(request) if _metadata_variable_name not in self.data: self.data[_metadata_variable_name] = {} diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 81e77283ed8..d858eff3348 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -105,6 +105,7 @@ else: _REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object]) +_MCP_KEY_RPM_LIMITS: Final[TypeAdapter[dict[str, int] | None]] = TypeAdapter(dict[str, int] | None) @dataclass(frozen=True, slots=True) @@ -3000,12 +3001,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): MCP tool calls have no token usage, so only requests_per_unit is set; tokens_per_unit stays None so the TPM reservation path is never engaged. """ - from litellm.proxy.auth.auth_utils import get_key_mcp_rpm_limit - if not mcp_server_name or not user_api_key_dict.api_key: return - mcp_rpm_limit: Final = get_key_mcp_rpm_limit(user_api_key_dict) + metadata: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(user_api_key_dict.metadata or {}) + mcp_rpm_limit: Final = _MCP_KEY_RPM_LIMITS.validate_python(metadata.get("mcp_rpm_limit")) if not mcp_rpm_limit: return @@ -3046,11 +3046,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # team's mcp_rpm_limit. Every applicable team is charged rather than one being picked: the # limiter enforces all descriptors, so each team's own ceiling binds on a call made through # its grant, and there is no arbitrary attribution when several teams grant the same server. - team_limits: Final[list[tuple[str | None, dict[str, int] | None]]] = [] - if user_api_key_dict.team_id: - team_limits.append((user_api_key_dict.team_id, get_team_mcp_rpm_limit(user_api_key_dict))) - for source_team_id, source_limit in (user_api_key_dict.mcp_source_team_rpm_limits or {}).items(): - team_limits.append((source_team_id, source_limit)) + source_team_limits: Final = user_api_key_dict.mcp_source_team_rpm_limits or {} + primary_team_ids: Final = (user_api_key_dict.team_id,) if user_api_key_dict.team_id else () + team_ids: Final = tuple(dict.fromkeys((*primary_team_ids, *source_team_limits))) + team_limits: Final = tuple( + ( + team_id, + get_team_mcp_rpm_limit(user_api_key_dict) or source_team_limits.get(team_id) + if team_id == user_api_key_dict.team_id + else source_team_limits[team_id], + ) + for team_id in team_ids + ) for team_id, mcp_rpm_limit in team_limits: if not team_id or not mcp_rpm_limit: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 06033f092dd..dbeda4e27e3 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3718,7 +3718,7 @@ def _add_guardrails_from_policies_in_metadata( def add_guardrails_from_auth_metadata( user_api_key_dict: UserAPIKeyAuth, - data: dict, # mutable-ok: writes guardrails into the live request dict, same contract as the helpers it wraps + data: dict[str, object], # mutable-ok: writes guardrails into the live request dict, same contract as helpers metadata_variable_name: str, ) -> None: """Resolve key, team, and project guardrails, direct and via policies, onto the request metadata.""" @@ -3776,7 +3776,7 @@ async def move_guardrails_to_metadata( add_guardrails_from_auth_metadata( user_api_key_dict=user_api_key_dict, - data=data, + data=data, # pyright: ignore[reportUnknownArgumentType] # request data retains a legacy dynamic shape metadata_variable_name=_metadata_variable_name, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4ea3100f869..5030b1512a6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -20023,9 +20023,9 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie "personal_key_model_access_from_teams": { "type": "Boolean", "description": ( - "Caps each personal key (a key with no team) to the models its owner's current teams can call, " - "re-evaluated on every request. A user in no team gets no model access. A user in several teams " - "can call any model one of those teams can call. Team keys are unaffected." + "Caps each personal key (a key with no team) to the models, MCP servers, MCP tools and agents its " + "owner's current teams can reach, re-evaluated on every request. A user in no team gets none of them. " + "A user in several teams gets the union of what those teams grant. Team keys are unaffected." ), }, "max_ui_session_budget": { diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c4414003ca5..f73a8547158 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -23,6 +23,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( is_mcp_admitted_user_subject, ) from litellm.proxy._types import ( + UI_TEAM_ID, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable, @@ -8649,7 +8650,11 @@ class TestUserSubjectTeamUnion: "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.get_mcp_server_from_tool_name", MagicMock(return_value=server), ): - billed = await MCPRequestHandler.billing_auth_for_tool_call(auth, tool_name="t-grant/tool_a") + billed = await MCPRequestHandler.billing_auth_for_tool_call( + auth, + tool_name="t-grant/tool_a", + allowed_mcp_servers=(), + ) assert (billed.team_id, billed.org_id) == ("t-grant", "org-team") assert (auth.team_id, auth.org_id) == (None, "org-user-primary"), "authz object must be untouched" @@ -11666,3 +11671,1016 @@ async def test_token_minted_credential_lists_tools_for_every_scope_shape(monkeyp assert database.writer_db.litellm_usertable.find_unique.await_count >= 1 database.writer_db.litellm_verificationtoken.find_unique.assert_not_awaited() assert credential.team_id == "team-eng" and credential.allowed_routes == list(granted.allowed_routes or ()) + + +_PERSONAL_KEY_TEAMS: Final = { + "team-a": ("srv-a", {"srv-a": ["add"]}), + "team-b": ("srv-b", None), +} + + +def _personal_key_team(team_id: str) -> LiteLLM_TeamTable: + server_id, tool_permissions = _PERSONAL_KEY_TEAMS[team_id] + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=f"org-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + mcp_servers=[server_id], + mcp_tool_permissions=tool_permissions, + ), + ) + + +def _personal_key_team_granting_tools(team_id: str, tools: tuple[str, ...]) -> LiteLLM_TeamTable: + return _personal_key_team(team_id).model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + mcp_servers=["srv-b"], + mcp_tool_permissions={"srv-b": list(tools)}, + ) + } + ) + + +@pytest.fixture +def personal_key_owner_teams(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest): + """A personal key whose owner is on ``request.param`` teams, resolved against a registry of + srv-a (team-a, tool ``add`` only), srv-b (team-b, every tool), srv-c (granted to no team) and + srv-open (``allow_all_keys``, granted to no team).""" + import litellm + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.auth import auth_checks + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return _personal_key_team(team_id) + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(return_value=LiteLLM_UserTable(user_id="alice", teams=list(request.param))), + ) + for server_id in ("srv-a", "srv-b", "srv-c", "srv-open"): + global_mcp_server_manager.registry[server_id] = MCPServer( + server_id=server_id, + name=server_id, + server_name=server_id, + url=f"https://{server_id}.example.com", + transport=MCPTransport.http, + allow_all_keys=server_id == "srv-open", + ) + yield UserAPIKeyAuth(api_key="sk-personal", user_id="alice", team_id=None) + for server_id in ("srv-a", "srv-b", "srv-c", "srv-open"): + global_mcp_server_manager.registry.pop(server_id, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "personal_key_owner_teams, expected_servers", + [(("team-a", "team-b"), {"srv-a", "srv-b"}), (("team-b",), {"srv-b"}), ((), set())], + indirect=["personal_key_owner_teams"], +) +async def test_personal_key_mcp_servers_are_the_union_of_its_owner_team_grants( + personal_key_owner_teams: UserAPIKeyAuth, expected_servers: set[str] +) -> None: + assert set(await MCPRequestHandler.get_allowed_mcp_servers(personal_key_owner_teams)) == expected_servers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_mcp_listing_applies_agent_caller_ceiling( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPServerAccess + from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers + + auth: Final = personal_key_owner_teams.model_copy(update={"agent_id": "agent-1"}) + caller: Final = agent_caller_from_headers( + {"x-litellm-team-id": "team-b"}, + auth, + ) + auth_with_caller: Final = auth.model_copy(update={"agent_caller": caller}) + + capped: Final = await MCPRequestHandler.get_mcp_server_access(auth_with_caller) + uncapped: Final = await MCPRequestHandler.get_mcp_server_access(auth) + + assert capped == MCPServerAccess(server_ids=("srv-b",), scope="scoped") + assert uncapped == MCPServerAccess(server_ids=("srv-a", "srv-b"), scope="scoped") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", UI_TEAM_ID)], indirect=True) +async def test_personal_key_ui_team_does_not_expand_mcp_access( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + if team_id == UI_TEAM_ID: + return LiteLLM_TeamTable(team_id=team_id) + return _personal_key_team(team_id) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + assert await MCPRequestHandler.get_allowed_mcp_servers(personal_key_owner_teams) == ["srv-a"] + assert (await MCPRequestHandler.get_mcp_server_access(personal_key_owner_teams)).scope == "scoped" + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-c", personal_key_owner_teams) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_tool_call_bills_granting_team( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.team_membership import LiteLLM_TeamMembership + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR, get_server_prefix + from litellm.proxy.auth import auth_checks + + team_membership: Final = LiteLLM_TeamMembership( + user_id=personal_key_owner_teams.user_id, + team_id="team-a", + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=3), + ) + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + team: Final = _personal_key_team_granting_tools(team_id, ("divide",) if team_id == "team-a" else ("multiply",)) + return team.model_copy( + update={ + "metadata": {"guardrails": ["team-guardrail"]} if team_id == "team-b" else team.metadata, + "rpm_limit": 7 if team_id == "team-a" else team.rpm_limit, + } + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=team_membership)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + server_b: Final = global_mcp_server_manager.registry["srv-b"] + server_open: Final = global_mcp_server_manager.registry["srv-open"] + tool_name: Final = f"{get_server_prefix(server_b)}{MCP_TOOL_PREFIX_SEPARATOR}multiply" + with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server_b): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call(personal_key_owner_teams, tool_name) + assert billed is not personal_key_owner_teams + assert (billed.team_id, billed.org_id) == ("team-b", "org-team-b") + assert billed.team_metadata == {"guardrails": ["team-guardrail"]} + assert personal_key_owner_teams.team_id is None + + team_a_tool_name: Final = f"{get_server_prefix(server_b)}{MCP_TOOL_PREFIX_SEPARATOR}divide" + with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server_b): + team_a_billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, team_a_tool_name + ) + assert team_a_billed.team_id == "team-a" + assert team_a_billed.team_rpm_limit == 7 + assert team_a_billed.team_member_rpm_limit == 3 + + open_tool_name: Final = f"{get_server_prefix(server_open)}{MCP_TOOL_PREFIX_SEPARATOR}tool" + with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server_open): + unbilled: Final = await MCPRequestHandler.billing_auth_for_tool_call(personal_key_owner_teams, open_tool_name) + assert unbilled is personal_key_owner_teams + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_org_id,expected_org_id", + [("org-1", None, "org-1"), (None, "org-2", "org-2")], +) +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_mcp_billing_preserves_key_organization( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, + key_org_id: str | None, + team_org_id: str | None, + expected_org_id: str, +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR, get_server_prefix + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=team_org_id, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + mcp_servers=["srv-a"], + mcp_tool_permissions={"srv-a": ["add"]}, + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + auth: Final = personal_key_owner_teams.model_copy(update={"org_id": key_org_id}) + server: Final = global_mcp_server_manager.registry["srv-a"] + tool_name: Final = f"{get_server_prefix(server)}{MCP_TOOL_PREFIX_SEPARATOR}add" + + with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call(auth, tool_name) + + assert billed.team_id == "team-a" + assert billed.org_id == expected_org_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_display_name_billing_uses_original_tool_grant( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + + server: Final = global_mcp_server_manager.registry["srv-a"].model_copy( + update={ + "name": "math_a", + "server_name": "math_a", + "tool_name_to_display_name": {"add": "Add Numbers"}, + } + ) + monkeypatch.setitem(global_mcp_server_manager.registry, "srv-a", server) + monkeypatch.setitem(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, "math_a-add", "math_a") + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + tool_name="Add Numbers", + ) + + assert billed.team_id == "team-a" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_display_name_billing_uses_allowed_servers( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + hidden_server: Final = MCPServer( + server_id="srv_hidden", + name="srv_hidden", + server_name="srv_hidden", + transport=MCPTransport.http, + tool_name_to_display_name={"foo": "Shared"}, + ) + allowed_server: Final = MCPServer( + server_id="srv_a", + name="srv_a", + server_name="srv_a", + transport=MCPTransport.http, + tool_name_to_display_name={"bar": "Shared"}, + ) + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return _personal_key_team(team_id).model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + mcp_servers=["srv_a"], + mcp_tool_permissions={"srv_a": ["bar"]}, + ) + } + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setitem(global_mcp_server_manager.registry, hidden_server.server_id, hidden_server) + monkeypatch.setitem(global_mcp_server_manager.registry, allowed_server.server_id, allowed_server) + + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + tool_name="Shared", + ) + + assert billed.team_id == "team-a" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_tool_call_bills_lowest_team_when_both_allow( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR, get_server_prefix + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return _personal_key_team_granting_tools(team_id, ("multiply",)) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + server_b: Final = global_mcp_server_manager.registry["srv-b"] + tool_name: Final = f"{get_server_prefix(server_b)}{MCP_TOOL_PREFIX_SEPARATOR}multiply" + with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server_b): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call(personal_key_owner_teams, tool_name) + assert (billed.team_id, billed.org_id) == ("team-a", "org-team-a") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_attribution_preserves_hyphenated_native_tool_names( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return _personal_key_team_granting_tools( + team_id, ("math-divide",) if team_id == "team-a" else ("math-multiply",) + ).model_copy(update={"metadata": {"mcp_rpm_limit": {"srv-b": 3 if team_id == "team-b" else 9}}}) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + server: Final = global_mcp_server_manager.registry["srv-b"].model_copy(update={"alias": "math"}) + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server), + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server), + ): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, "math-math-multiply" + ) + assert (billed.team_id, billed.org_id) == ("team-b", "org-team-b") + assert personal_key_owner_teams.team_id is None + assert await MCPRequestHandler.personal_key_team_rpm_limits(personal_key_owner_teams, server, "math-multiply") == { + "team-b": {"srv-b": 3} + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_guardrail_auth_preserves_native_tool_name_matching_server_alias( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + + 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",) + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + server: Final = global_mcp_server_manager.registry["srv-b"].model_copy(update={"alias": "math"}) + + billed: Final = await MCPRequestHandler.guardrail_auth_for_call( + None, + personal_key_owner_teams, + server, + "math-multiply", + allowed_mcp_servers=[server], + ) + + assert billed is not None + assert billed.team_id == "team-b" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_mcp_billing_inherits_adopted_organization_limits( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + from litellm.proxy._types import LiteLLM_OrganizationTable + + organization: Final = LiteLLM_OrganizationTable( + organization_id="org-1", + organization_alias="org-one", + budget_id="budget-1", + created_by="test", + updated_by="test", + metadata={"k": "v"}, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100, rpm_limit=5, tpm_limit=50), + ) + + 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"} + ) + + 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) + server: Final = global_mcp_server_manager.registry["srv-b"] + tool_name: Final = "srv-b-divide" + + with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call(personal_key_owner_teams, tool_name) + + assert ( + billed.org_id, + billed.organization_rpm_limit, + billed.organization_tpm_limit, + billed.organization_metadata, + ) == ("org-1", 5, 50, {"k": "v"}) + assert billed.organization_max_budget == 100 + assert ( + personal_key_owner_teams.org_id, + personal_key_owner_teams.organization_rpm_limit, + personal_key_owner_teams.organization_tpm_limit, + personal_key_owner_teams.organization_metadata, + ) == (None, None, None, None) + org_loader.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_mcp_billing_keeps_key_organization_limits( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + from litellm.proxy._types import LiteLLM_OrganizationTable + + organization: Final = LiteLLM_OrganizationTable( + organization_id="org-1", + budget_id="budget-1", + created_by="test", + updated_by="test", + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=5), + ) + + 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"} + ) + + 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} + ) + 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): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call(auth, "srv-b-divide") + + assert (billed.org_id, billed.organization_rpm_limit) == ("org-k", 9) + org_loader.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_mcp_billing_keeps_team_controls_when_org_lookup_fails( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + 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", + "rpm_limit": 5, + "max_budget": 40, + "metadata": {"team_private": True}, + } + ) + + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + org_loader: Final = AsyncMock(side_effect=RuntimeError("org lookup failed")) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object_for_request", org_loader) + 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): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + "srv-b-divide", + scoped_server=server, + native_tool_name="divide", + allowed_mcp_servers=[server], + ) + + assert billed is not personal_key_owner_teams + assert billed.team_id == "team-a" + assert billed.team_rpm_limit == 5 + assert billed.team_metadata == {"team_private": True} + assert billed.team_max_budget == 40 + assert billed.org_id == "org-1" + assert personal_key_owner_teams.team_id is None + assert personal_key_owner_teams.team_rpm_limit is None + org_loader.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_mcp_billing_keeps_team_id_when_membership_lookup_fails( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return _personal_key_team_granting_tools(team_id, ("divide",)) + + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + membership_loader: Final = AsyncMock(side_effect=RuntimeError("membership lookup failed")) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr( + auth_checks, + "get_team_membership", + membership_loader, + ) + server: Final = global_mcp_server_manager.registry["srv-b"] + + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + "srv-b-divide", + scoped_server=server, + native_tool_name="divide", + allowed_mcp_servers=[server], + ) + + membership_loader.assert_awaited_once() + assert billed is not personal_key_owner_teams + assert billed.team_id == "team-a" + assert billed.org_id == "org-team-a" + assert personal_key_owner_teams.team_id is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requires_fresh_policy, expected_team_rpm, expected_member_rpm, expected_metadata", + [ + pytest.param(False, 2, 3, {"policy_source": "cache"}, id="default-uses-cache"), + pytest.param(True, 7, 8, {"policy_source": "database"}, id="strict-uses-database"), + ], +) +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_mcp_billing_uses_fresh_team_and_membership_for_strict_policy( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, + requires_fresh_policy: bool, + expected_team_rpm: int, + expected_member_rpm: int, + expected_metadata: dict[str, str], +) -> None: + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.team_membership import LiteLLM_TeamMembership + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + async def _get_team_object(team_id: str, **kwargs: object) -> LiteLLM_TeamTable: + database_value: Final = bool(kwargs.get("check_db_only")) + return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy( + update={ + "organization_id": None, + "rpm_limit": 7 if database_value else 2, + "metadata": {"policy_source": "database" if database_value else "cache"}, + } + ) + + 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, + team_id=team_id, + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=8 if database_value else 3), + ) + + team_loader: Final = AsyncMock(side_effect=_get_team_object) + membership_loader: Final = AsyncMock(side_effect=_get_team_membership) + monkeypatch.setattr(auth_checks, "get_team_object", team_loader) + monkeypatch.setattr(auth_checks, "get_team_membership", membership_loader) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + auth: Final = personal_key_owner_teams.model_copy(update={"requires_fresh_policy": requires_fresh_policy}) + 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): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + auth, + "srv-b-divide", + scoped_server=server, + native_tool_name="divide", + allowed_mcp_servers=[server], + ) + + assert (billed.team_rpm_limit, billed.team_member_rpm_limit, billed.team_metadata) == ( + expected_team_rpm, + expected_member_rpm, + expected_metadata, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_team_mcp_rpm_limit_charges_granting_team( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + PROXY_MaxParallelRequestsHandler_v3, + RateLimitDescriptor, + ) + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + team: Final = _personal_key_team_granting_tools(team_id, ("divide",) if team_id == "team-a" else ("multiply",)) + return team.model_copy( + update={ + "metadata": ( + {"mcp_rpm_limit": {"srv-b": 3, "srv-a": 9}} + if team_id == "team-b" + else {"mcp_rpm_limit": {"srv-b": 9}} + ) + } + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + server_b: Final = global_mcp_server_manager.registry["srv-b"] + limits: Final = await MCPRequestHandler.personal_key_team_rpm_limits(personal_key_owner_teams, server_b, "multiply") + assert limits == {"team-b": {"srv-b": 3}} + + stamped: Final = personal_key_owner_teams.model_copy(update={"mcp_source_team_rpm_limits": limits}) + descriptors: list[RateLimitDescriptor] = [] + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=MagicMock()) + limiter._add_mcp_per_team_rate_limit_descriptor(stamped, server_b.name, descriptors) + assert len(descriptors) == 1 + assert descriptors[0]["key"] == "mcp_per_team" + assert descriptors[0]["value"] == f"team-b:{server_b.name}" + assert descriptors[0]["rate_limit"]["requests_per_unit"] == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_over_budget_team_grants_no_mcp_access( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + team: Final = _personal_key_team(team_id) + return team.model_copy(update={"max_budget": 1, "spend": 5}) if team_id == "team-b" else team + + async def _spend_from_fallback(*, fallback_spend: float, **_: object) -> float: + return fallback_spend + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "get_current_spend", _spend_from_fallback) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "budget_alerts", AsyncMock()) + assert await MCPRequestHandler.get_allowed_mcp_servers(personal_key_owner_teams) == ["srv-a"] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-b", personal_key_owner_teams) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_mcp_tools_follow_the_rule_of_the_team_granting_the_server( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", personal_key_owner_teams) == ["add"] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-b", personal_key_owner_teams) is None + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-c", personal_key_owner_teams) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [()], indirect=True) +async def test_personal_key_without_teams_reaches_no_mcp_tool(personal_key_owner_teams: UserAPIKeyAuth) -> None: + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", personal_key_owner_teams) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_mcp_access_is_unchanged_when_flag_off( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + import litellm + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", False) + assert await MCPRequestHandler.get_allowed_mcp_servers(personal_key_owner_teams) == [] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", personal_key_owner_teams) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_team_key_mcp_access_ignores_its_owner_other_teams(personal_key_owner_teams: UserAPIKeyAuth) -> None: + from litellm.proxy.auth import auth_checks + + team_key: Final = personal_key_owner_teams.model_copy(update={"team_id": "team-b"}) + assert await MCPRequestHandler.get_allowed_mcp_servers(team_key) == ["srv-b"] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", team_key) is None + auth_checks.get_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [(), ("team-a",)], indirect=True) +async def test_personal_key_keeps_its_own_tool_rule_on_an_allow_all_server( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-open", personal_key_owner_teams) is None + scoped_key: Final = personal_key_owner_teams.model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id="perm-key", mcp_tool_permissions={"srv-open": ["add"]} + ) + } + ) + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-open", scoped_key) == ["add"] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-c", scoped_key) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_mcp_tools_are_denied_when_a_granting_team_cannot_be_read( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + """team-b grants srv-b, but the tool-rule read behind the source pinned to team-b fails. The fault + must deny the server's tools, not become the allow-all ``None`` of a team without a tool rule.""" + + from litellm.proxy.auth import auth_checks + + team_object_reads: Final = AsyncMock() + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + team_b_reads: Final = tuple( + call for call in team_object_reads.await_args_list if call.kwargs.get("team_id") == "team-b" + ) + if team_id == "team-b" and len(team_b_reads) == 3: + raise RuntimeError("connection reset by peer") + return _personal_key_team(team_id) + + team_object_reads.side_effect = _get_team_object + monkeypatch.setattr(auth_checks, "get_team_object", team_object_reads) + srv_b_tools: Final = await MCPRequestHandler.get_allowed_tools_for_server("srv-b", personal_key_owner_teams) + team_b_reads: Final = tuple( + call for call in team_object_reads.await_args_list if call.kwargs.get("team_id") == "team-b" + ) + assert len(team_b_reads) == 3 + assert srv_b_tools == [] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", personal_key_owner_teams) == ["add"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "personal_key_owner_teams, key_org, expected_orgs", + [ + (("team-a", "team-b"), None, ("org-team-a", "org-team-b")), + (("team-a", "team-b"), "org-team-a", ("org-team-a",)), + (("team-a", "team-b"), "org-key", ()), + ], + indirect=["personal_key_owner_teams"], +) +async def test_personal_key_sources_keep_the_key_own_org_scope( + personal_key_owner_teams: UserAPIKeyAuth, key_org: str | None, expected_orgs: tuple[str, ...] +) -> None: + """A key without an org takes each team's org. A key issued under an org keeps it and only its own org's + teams become sources, a team of another org grants nothing, so neither org ceiling is dropped.""" + key: Final = personal_key_owner_teams.model_copy(update={"org_id": key_org}) + sources: Final = await MCPRequestHandler.personal_key_team_sources(key) + assert tuple(source.org_id for source in sources) == expected_orgs + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_unreadable_team_permission_is_not_a_server_or_billing_source( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPServerAccess + from litellm.proxy.auth import auth_checks + + async def _unreadable_team(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + organization_id="org-team-a", + object_permission_id="permission-that-cannot-be-loaded", + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _unreadable_team) + + assert await MCPRequestHandler.personal_key_team_sources(personal_key_owner_teams) == () + assert await MCPRequestHandler.resolve_personal_key_team_servers(personal_key_owner_teams) == MCPServerAccess( + server_ids=(), scope="scoped" + ) + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", personal_key_owner_teams) == [] + assert await MCPRequestHandler.personal_key_attributing_source(personal_key_owner_teams, "srv-a", "add") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_team_server_access_group_read_fault_stays_scoped( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPServerAccess + from litellm.proxy.auth import auth_checks + + async def _team_with_access_group(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=team_id, access_group_ids=["unreadable-group"]) + + monkeypatch.setattr(auth_checks, "get_team_object", _team_with_access_group) + get_access_object: Final = AsyncMock(side_effect=RuntimeError("access group read failed")) + monkeypatch.setattr(auth_checks, "get_access_object", get_access_object) + + team_grants: Final = await MCPRequestHandler.personal_key_team_server_grants(personal_key_owner_teams) + assert len(team_grants) == 1 + assert team_grants[0][1] == MCPServerAccess(server_ids=(), scope="scoped") + assert await MCPRequestHandler.resolve_personal_key_team_servers(personal_key_owner_teams) == MCPServerAccess( + server_ids=(), scope="scoped" + ) + get_access_object.assert_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_scoped_bare_tool_name_bills_granting_team( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return _personal_key_team(team_id).model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + mcp_servers=["srv-a"], + mcp_tool_permissions={"srv-a": ["add"]}, + ) + } + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + + async with global_mcp_server_manager.catalog.operation(): + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + tool_name="add", + mcp_servers=("srv-a",), + ) + unscoped: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + tool_name="add", + ) + + assert billed.team_id == "team-a" + assert unscoped is personal_key_owner_teams + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_prefixed_tool_billing_resolves_server_without_loaded_tools( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server: Final = MCPServer( + server_id="math_a", + name="math_a", + server_name="math_a", + transport=MCPTransport.http, + ) + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return _personal_key_team(team_id).model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + mcp_servers=["math_a"], + mcp_tool_permissions={"math_a": ["add"]}, + ) + } + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + + assert global_mcp_server_manager.tool_name_to_mcp_server_name_mapping == {} + + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + tool_name="math_a-add", + ) + unknown_prefix: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + tool_name="unknown-add", + ) + + assert billed.team_id == "team-a" + assert unknown_prefix is personal_key_owner_teams + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_scoped_tool_billing_prefers_scoped_server( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + server_id: Final = "srv-a" if team_id == "team-a" else "srv-b" + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=f"org-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + mcp_servers=[server_id], + mcp_tool_permissions={server_id: ["search"]}, + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setitem(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, "search", "srv-a") + + billed: Final = await MCPRequestHandler.billing_auth_for_tool_call( + personal_key_owner_teams, + tool_name="search", + mcp_servers=("srv-b",), + ) + + assert billed.team_id == "team-b" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [()], indirect=True) +async def test_admin_view_personal_key_keeps_its_own_tools_on_every_server( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + """The admin-view role lists every registry server for a personal key, so its tools resolve from the + key's own rules (none here, so unrestricted) instead of being denied for lack of a team grant.""" + from litellm.proxy._types import LitellmUserRoles + + viewer: Final = personal_key_owner_teams.model_copy(update={"user_role": LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY}) + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-c", viewer) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_keyless_grant_source_of_the_owner_does_not_inherit_its_teams( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + """A source auth of an admitted subject has the owner's user_id, no team and no api_key. It must + resolve as the plain caller it is, or a token scoped to one team reaches every team's servers.""" + source: Final = UserAPIKeyAuth(user_id=personal_key_owner_teams.user_id, team_id=None) + assert await MCPRequestHandler.get_allowed_mcp_servers(source) == [] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", source) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_mcp_access_is_denied_when_owner_teams_cannot_be_read( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy.auth import auth_checks + + monkeypatch.setattr(auth_checks, "personal_key_owner_teams", AsyncMock(side_effect=RuntimeError("db down"))) + assert await MCPRequestHandler.get_allowed_mcp_servers(personal_key_owner_teams) == [] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-a", personal_key_owner_teams) == [] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv-open", personal_key_owner_teams) == [] diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d314291ef74..4896c2f1706 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -56,6 +56,7 @@ from litellm.experimental_mcp_client.client import MCPClient, PersistentMCPSessi from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( ListedToolsCaller, MCPServerManager, + global_mcp_server_manager, _deserialize_json_dict, _flow_endpoints_missing, _mcp_oauth_discovery_on_startup_enabled, @@ -67,9 +68,13 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( should_strip_caller_authorization, listed_tools_caller_for, ) +from litellm.proxy import proxy_server +from litellm.proxy.auth import auth_checks from litellm.proxy._types import ( LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, MCPApprovalStatus, MCPEnvVar, @@ -167,7 +172,13 @@ async def test_elicitation_callback_keeps_initiating_session(): from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_elicitation_callback from types import SimpleNamespace - from mcp.types import ClientCapabilities, ElicitationCapability, FormElicitationCapability, ElicitRequestFormParams, ElicitResult + from mcp.types import ( + ClientCapabilities, + ElicitationCapability, + FormElicitationCapability, + ElicitRequestFormParams, + ElicitResult, + ) from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var capabilities = ClientCapabilities(elicitation=ElicitationCapability(form=FormElicitationCapability())) @@ -183,7 +194,7 @@ async def test_elicitation_callback_keeps_initiating_session(): legacy_server.active_mcp_session_var.set(SimpleNamespace()) active_mcp_request_ctx_var.set(None) capabilities.elicitation = None - result = await callback(None, ElicitRequestFormParams(message="Confirm", requested_schema={"type":"object"})) + result = await callback(None, ElicitRequestFormParams(message="Confirm", requested_schema={"type": "object"})) assert result is accepted assert request.await_args.kwargs["related_request_id"] == "initiating-call" finally: @@ -1289,7 +1300,9 @@ class TestMCPServerManager: raw_row = MagicMock() raw_row.model_dump.return_value = row.model_dump() repository = MagicMock() - other_row: Final = row.model_copy(update={"server_id": "changed-other", "server_name": "changed_other", "auth_type": MCPAuth.none}) + other_row: Final = row.model_copy( + update={"server_id": "changed-other", "server_name": "changed_other", "auth_type": MCPAuth.none} + ) repository.table.find_many = AsyncMock(return_value=[raw_row, other_row]) async def publish_while_staged(*_args, **_kwargs): @@ -1367,7 +1380,11 @@ class TestMCPServerManager: table = SimpleNamespace( find_many=AsyncMock(return_value=[_row(cached.server_id, corrupted), _row("healthy-sibling", stored)]) ) - prisma = SimpleNamespace(db=SimpleNamespace(litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)), litellm_mcpservertable=table)) + prisma = SimpleNamespace( + db=SimpleNamespace( + litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)), litellm_mcpservertable=table + ) + ) prisma.writer_db = prisma.db monkeypatch.setattr(proxy_server, "prisma_client", prisma) @@ -2301,12 +2318,17 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_managed_scope_challenge_requires_reauthorization_without_retry(self): server = MCPServer( - server_id="managed-scope", name="managed-scope", alias="managed-scope", - url="https://up.example.com/mcp", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", + server_id="managed-scope", + name="managed-scope", + alias="managed-scope", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", ) upstream_error = self._upstream_status_error( - 403, 'Bearer error="insufficient_scope", scope="tools.write", resource_metadata="https://untrusted.test/metadata"' + 403, + 'Bearer error="insufficient_scope", scope="tools.write", resource_metadata="https://untrusted.test/metadata"', ) manager = MCPServerManager() client = AsyncMock() @@ -7390,12 +7412,21 @@ class TestMCPServerManager: caller=ListedToolsCaller(user_api_key_auth=auth), ) if managed: - manager.registry["test-server"] = manager.registry["test-server"].model_copy(update={ - "auth_type": MCPAuth.oauth2, "oauth2_flow": "authorization_code", - }) + manager.registry["test-server"] = manager.registry["test-server"].model_copy( + update={ + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + } + ) send = AsyncMock() gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/test-server") - context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + context = ServerRequestContext( + session=MagicMock(), + lifespan_context={}, + protocol_version="2025-11-25", + method="tools/call", + request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}), + ) token = active_mcp_request_ctx_var.set(context) await gate({"type": "http.response.start", "status": 200, "headers": []}) progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) @@ -7407,7 +7438,13 @@ class TestMCPServerManager: manager.create_mcp_client.return_value.call_tool.side_effect = upstream_call try: - result = await manager.call_tool(server_name="test-server", name="test_tool", arguments={}, user_api_key_auth=auth, proxy_logging_obj=logging) + result = await manager.call_tool( + server_name="test-server", + name="test_tool", + arguments={}, + user_api_key_auth=auth, + proxy_logging_obj=logging, + ) assert not result.is_error finally: gate.allow() @@ -8040,9 +8077,7 @@ class TestMCPServerManager: "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.get_mcp_jwt_signer", return_value=signer, ): - manager.record_listed_tools( - server, [MCPTool(name="turn", description="alice view", inputSchema={})], alice - ) + manager.record_listed_tools(server, [MCPTool(name="turn", description="alice view", inputSchema={})], alice) assert manager.get_listed_tool(server, "turn", bob) is None for_alice = manager.get_listed_tool(server, "turn", alice) @@ -9141,9 +9176,7 @@ class TestMCPServerTimestamps: registration_url=None, ) assert oauth_endpoints_unresolved(relay_arm) is True - assert ( - oauth_endpoints_unresolved(relay_arm.model_copy(update={"registration_url": "https://idp/reg"})) is False - ) + assert oauth_endpoints_unresolved(relay_arm.model_copy(update={"registration_url": "https://idp/reg"})) is False assert oauth_endpoints_unresolved(relay_arm.model_copy(update={"client_id": "admin-client"})) is False def test_entra_obo_without_scopes_is_unresolved(self): @@ -9405,11 +9438,13 @@ class TestMCPServerTimestamps: assert blipped.issuer == "https://idp.example.com" assert blipped.authorization_response_iss_parameter_supported is True fallback = MCPServerManager._merge_discovered_oauth_metadata( - blipped, MCPOAuthMetadata(from_origin_fallback=True), + blipped, + MCPOAuthMetadata(from_origin_fallback=True), ) assert fallback.authorization_response_iss_parameter_supported is True refreshed = MCPServerManager._merge_discovered_oauth_metadata( - fallback, MCPOAuthMetadata(discovered_issuer="https://idp.example.com"), + fallback, + MCPOAuthMetadata(discovered_issuer="https://idp.example.com"), ) assert refreshed.authorization_response_iss_parameter_supported is False @@ -9764,7 +9799,7 @@ class TestMCPServerTimestamps: "unlimited_server": { "url": "https://example.com/other-mcp", "transport": MCPTransport.http, - } + }, } await manager.load_servers_from_config(config) servers = list(manager.config_mcp_servers.values()) @@ -15213,6 +15248,457 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, assert exc_info.value.status_code == 400 +@pytest.mark.asyncio +async def test_pre_call_tool_check_stamps_personal_key_team_rpm_limits(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="srv-b", + name="srv-b", + server_name="srv-b", + transport=MCPTransport.stdio, + ) + user_api_key_auth: Final = UserAPIKeyAuth(api_key="sk-personal", user_id="alice") + expected: Final = {"team-b": {"srv-b": 3}} + proxy_logging_obj: Final = _permissive_proxy_logging() + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(manager, "check_tool_permission_for_key_team", AsyncMock(return_value=None)) + monkeypatch.setattr(MCPRequestHandler, "personal_key_team_rpm_limits", AsyncMock(return_value=expected)) + await manager.pre_call_tool_check( + name="multiply", + arguments={}, + server_name=server.name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + ) + + hook_auth: Final = proxy_logging_obj.pre_call_hook.await_args.kwargs["user_api_key_dict"] + assert hook_auth is not user_api_key_auth + assert hook_auth.mcp_source_team_rpm_limits == expected + assert user_api_key_auth.mcp_source_team_rpm_limits is None + assert proxy_logging_obj.enforce_mcp_server_rate_limits.await_args.args[0] is hook_auth + assert MCPRequestHandler.personal_key_team_rpm_limits.await_args.args == ( + user_api_key_auth, + server, + "multiply", + ) + + +@pytest.mark.asyncio +async def test_pre_call_tool_check_uses_inherited_team_auth_for_guardrails(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + server: Final = MCPServer( + server_id="srv-guardrails", + name="math_a", + server_name="math_a", + transport=MCPTransport.stdio, + ) + team: Final = LiteLLM_TeamTable( + team_id="team-a", + object_permission_id="perm-team-a", + metadata={"guardrails": ["team-guardrail"]}, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-team-a", + mcp_servers=[server.server_id], + mcp_tool_permissions={server.server_id: ["multiply"]}, + ), + ) + user_api_key_auth: Final = UserAPIKeyAuth(api_key="sk-personal", user_id="alice") + manager: Final = MCPServerManager() + proxy_logging_obj: Final = _permissive_proxy_logging() + user_api_key_cache: Final = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(return_value=None) + user_api_key_cache.async_set_cache = AsyncMock() + + async def get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + assert team_id == team.team_id + return team + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(return_value=LiteLLM_UserTable(user_id="alice", teams=["team-a"])), + ) + monkeypatch.setattr(auth_checks, "get_team_object", get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "_personal_key_team_within_budget", AsyncMock(return_value=True)) + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + monkeypatch.setattr(manager, "check_tool_permission_for_key_team", AsyncMock(return_value=None)) + + team_grants: Final = await MCPRequestHandler.personal_key_team_server_grants(user_api_key_auth) + assert team_grants[0][0].team_id == "team-a" + assert server.server_id in team_grants[0][1].server_ids + + await manager.pre_call_tool_check( + name="multiply", + arguments={}, + server_name=server.name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + guardrail_auth=None, + ) + + hook_auth: Final = proxy_logging_obj.pre_call_hook.await_args.kwargs["user_api_key_dict"] + assert hook_auth.team_id == "team-a" + assert hook_auth.team_metadata == {"guardrails": ["team-guardrail"]} + + +@pytest.mark.asyncio +async def test_pre_call_tool_check_keeps_native_name_for_team_attribution(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + server: Final = MCPServer( + server_id="math-a", + name="math_a", + alias="math_a", + server_name="math_a", + transport=MCPTransport.stdio, + ) + team_a: Final = LiteLLM_TeamTable( + team_id="team-a", + object_permission_id="perm-team-a", + metadata={"mcp_rpm_limit": {"math-a": 7}}, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-team-a", + mcp_servers=["math-a"], + mcp_tool_permissions={"math-a": ["math_a-foo"]}, + ), + ) + team_b: Final = LiteLLM_TeamTable( + team_id="team-b", + object_permission_id="perm-team-b", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-team-b", + mcp_servers=["math-a"], + mcp_tool_permissions={"math-a": ["foo"]}, + ), + ) + user_api_key_auth: Final = UserAPIKeyAuth(api_key="sk-personal", user_id="alice") + manager: Final = MCPServerManager() + proxy_logging_obj: Final = _permissive_proxy_logging() + user_api_key_cache: Final = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(return_value=None) + user_api_key_cache.async_set_cache = AsyncMock() + + async def get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return team_a if team_id == team_a.team_id else team_b + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(return_value=LiteLLM_UserTable(user_id="alice", teams=["team-a", "team-b"])), + ) + monkeypatch.setattr(auth_checks, "get_team_object", get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "_personal_key_team_within_budget", AsyncMock(return_value=True)) + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + + await manager.pre_call_tool_check( + name="math_a-foo", + arguments={}, + server_name=server.name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + guardrail_auth=None, + ) + + hook_auth: Final = proxy_logging_obj.pre_call_hook.await_args.kwargs["user_api_key_dict"] + assert hook_auth.team_id == "team-a" + assert hook_auth.mcp_source_team_rpm_limits == {"team-a": {"math-a": 7}} + + +@pytest.mark.asyncio +async def test_pre_call_tool_check_keeps_original_auth_for_team_key(monkeypatch: pytest.MonkeyPatch) -> None: + server: Final = MCPServer( + server_id="srv-team-key", + name="team-server", + server_name="team-server", + transport=MCPTransport.stdio, + ) + user_api_key_auth: Final = UserAPIKeyAuth( + api_key="sk-team", + team_id="team-a", + team_metadata={"guardrails": ["team-guardrail"]}, + ) + manager: Final = MCPServerManager() + proxy_logging_obj: Final = _permissive_proxy_logging() + monkeypatch.setattr(manager, "check_tool_permission_for_key_team", AsyncMock(return_value=None)) + + await manager.pre_call_tool_check( + name="multiply", + arguments={}, + server_name=server.name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + guardrail_auth=None, + ) + + assert proxy_logging_obj.pre_call_hook.await_args.kwargs["user_api_key_dict"] is user_api_key_auth + + +@pytest.mark.asyncio +async def test_call_tool_passes_attributed_auth_to_guardrail_hooks(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.operations import _run_post_mcp_call_guardrails + from litellm.proxy.auth import auth_checks + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable + + original_auth: Final = UserAPIKeyAuth(api_key="sk-personal", user_id="alice") + team: Final = LiteLLM_TeamTable( + team_id="team-b", + organization_id="org-team-b", + object_permission_id="perm-team-b", + metadata={"guardrails": ["team-guardrail"]}, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-team-b", + mcp_servers=["test-server"], + mcp_tool_permissions={"test-server": ["test_tool"]}, + ), + ) + manager: Final = TestMCPServerManager._manager_ready_for_call_tool( + [MCPTool(name="test_tool", inputSchema={"type": "object"})], + caller=ListedToolsCaller(user_api_key_auth=original_auth), + )[0] + server: Final = manager.registry["test-server"] + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_mcp_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + proxy_logging_obj.enforce_mcp_server_rate_limits = AsyncMock() + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr("litellm.proxy.utils.premium_user_check", MagicMock()) + monkeypatch.setattr(manager, "check_tool_permission_for_key_team", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(return_value=LiteLLM_UserTable(user_id="alice", teams=["team-b"])), + ) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=team)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "_personal_key_team_within_budget", AsyncMock(return_value=True)) + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + + await manager.call_tool( + server_name="test-server", + name="test_tool", + arguments={}, + user_api_key_auth=original_auth, + proxy_logging_obj=proxy_logging_obj, + ) + + pre_auth: Final = proxy_logging_obj.pre_call_hook.await_args.kwargs["user_api_key_dict"] + during_auth: Final = proxy_logging_obj.during_call_hook.await_args.kwargs["user_api_key_dict"] + assert pre_auth.team_id == "team-b" + assert pre_auth.team_metadata == {"guardrails": ["team-guardrail"]} + assert during_auth.team_id == "team-b" + assert during_auth.team_metadata == {"guardrails": ["team-guardrail"]} + assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["metadata"]["guardrails"] == ["team-guardrail"] + assert proxy_logging_obj.during_call_hook.await_args.kwargs["data"]["metadata"]["guardrails"] == ["team-guardrail"] + assert original_auth.team_id is None + assert original_auth.team_metadata is None + + non_inherited_auth: Final = UserAPIKeyAuth( + api_key="sk-team", + team_id="team-a", + team_metadata={"guardrails": ["team-guardrail"]}, + ) + await manager.call_tool( + server_name="test-server", + name="test_tool", + arguments={}, + user_api_key_auth=non_inherited_auth, + proxy_logging_obj=proxy_logging_obj, + ) + assert proxy_logging_obj.pre_call_hook.await_args_list[-1].kwargs["user_api_key_dict"] is non_inherited_auth + assert proxy_logging_obj.during_call_hook.await_args_list[-1].kwargs["user_api_key_dict"] is non_inherited_auth + + from mcp.types import CallToolResult + + result: Final = CallToolResult(content=[]) + non_inherited_metadata: Final = {"existing": "value"} + model_call_details: Final = {"metadata": non_inherited_metadata} + post_guardrail_auth: Final = original_auth.model_copy(update={"team_metadata": {"guardrails": ["g1"]}}) + inherited_request_data: Final = {"metadata": {"guardrails": [{"req-g": {"x": 1}}]}} + await _run_post_mcp_call_guardrails( + result=result, + litellm_logging_obj=None, + original_auth=original_auth, + guardrail_auth=post_guardrail_auth, + request_data=inherited_request_data, + ) + post_kwargs: Final = proxy_logging_obj.post_mcp_call_hook.await_args.kwargs + assert post_kwargs["user_api_key_dict"] is post_guardrail_auth + assert post_kwargs["request_data"]["metadata"]["guardrails"] == [{"req-g": {"x": 1}}, "g1"] + assert inherited_request_data["metadata"]["guardrails"] == [{"req-g": {"x": 1}}] + await _run_post_mcp_call_guardrails( + result=result, + litellm_logging_obj=MagicMock(model_call_details=model_call_details), + original_auth=non_inherited_auth, + guardrail_auth=non_inherited_auth, + request_data={"metadata": {}}, + ) + non_inherited_post_kwargs: Final = proxy_logging_obj.post_mcp_call_hook.await_args.kwargs + assert non_inherited_post_kwargs["user_api_key_dict"] is non_inherited_auth + assert non_inherited_post_kwargs["request_data"] is model_call_details + assert non_inherited_post_kwargs["request_data"]["metadata"] == {"existing": "value"} + assert model_call_details["metadata"] is non_inherited_metadata + + +@pytest.mark.asyncio +async def test_virtual_mcp_calls_use_attributed_auth_for_logging_and_hooks(monkeypatch: pytest.MonkeyPatch) -> None: + from mcp.types import CallToolResult, Tool + + from litellm.proxy._experimental.mcp_server import operations, tool_search + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable + from litellm.types.mcp_server.tool_registry import MCPTool + + original_auth: Final = UserAPIKeyAuth( + api_key="sk-personal", + user_id="owner", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="virtual-tools", + mcp_tool_search_enabled=True, + ), + ) + team: Final = LiteLLM_TeamTable( + team_id="team-a", + organization_id="org-team-a", + object_permission_id="perm-team-a", + metadata={"guardrails": ["g1"]}, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-team-a", + mcp_servers=["math-a"], + mcp_tool_permissions={"math-a": ["foo"]}, + ), + ) + server: Final = MCPServer( + server_id="math-a", + name="math_a", + server_name="math_a", + transport=MCPTransport.http, + ) + registered_tool: Final = MCPTool( + name="math_a-foo", + description="Foo", + input_schema={"type": "object"}, + handler=AsyncMock(return_value=None), + server_id="math-a", + ) + result: Final = CallToolResult(content=[]) + target_tool: Final = Tool(name="math_a-foo", inputSchema={"type": "object"}) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_mcp_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + proxy_logging_obj.enforce_mcp_server_rate_limits = AsyncMock() + + async def build_logging( + name: str, + arguments: dict[str, object], + user_api_key_auth: UserAPIKeyAuth, + logging_auth: UserAPIKeyAuth | None = None, + raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, + ) -> MagicMock: + logger_auth: Final = logging_auth or user_api_key_auth + return MagicMock(model_call_details={"metadata": {"user_api_key_team_id": logger_auth.team_id}}) + + logging_builder: Final = AsyncMock(side_effect=build_logging) + tool_call_handler: Final = AsyncMock(return_value=result) + proxy_handler: Final = AsyncMock(return_value=result) + + async def get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + assert team_id == "team-a" + return team + + async def get_user_object(user_id: str, **_: object) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id=user_id, teams=["team-a"]) + + monkeypatch.setattr(auth_checks, "get_team_object", get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_user_object", get_user_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setitem(global_mcp_server_manager.registry, "math-a", server) + monkeypatch.setitem(global_mcp_tool_registry.tools, "math_a-foo", registered_tool) + monkeypatch.setattr(operations, "_build_virtual_call_logging_obj", logging_builder) + monkeypatch.setattr(tool_search, "handle_mcp_tool_call", tool_call_handler) + monkeypatch.setattr(tool_search, "handle_mcp_proxy_tool", proxy_handler) + monkeypatch.setattr(operations, "find_mcp_proxy_tool", AsyncMock(return_value=target_tool)) + monkeypatch.setattr(operations, "_fire_mcp_tool_call_logging", AsyncMock(return_value=result)) + + await operations._dispatch_virtual_mcp_tool( + name="mcp_tool_call", + arguments={"tool_name": "math_a-foo", "arguments": {}}, + user_api_key_auth=original_auth, + client_ip=None, + ) + + assert logging_builder.await_args.kwargs["user_api_key_auth"] is original_auth + assert logging_builder.await_args.kwargs["logging_auth"].team_id == "team-a" + assert logging_builder.await_args.kwargs["logging_auth"].team_metadata == {"guardrails": ["g1"]} + assert tool_call_handler.await_args.kwargs["user_api_key_dict"] is original_auth + assert tool_call_handler.await_args.kwargs["guardrail_auth"].team_id == "team-a" + assert tool_call_handler.await_args.kwargs["guardrail_auth"].team_metadata == {"guardrails": ["g1"]} + direct_logger: Final = tool_call_handler.await_args.kwargs["litellm_logging_obj"] + assert direct_logger.model_call_details["metadata"]["user_api_key_team_id"] == "team-a" + + await operations._dispatch_virtual_mcp_tool( + name="call_tool", + arguments={"tool_id": "opaque-tool-id", "arguments": {}}, + user_api_key_auth=original_auth, + client_ip=None, + mcp_proxy_mode=True, + ) + + assert logging_builder.await_args_list[1].kwargs["user_api_key_auth"] is original_auth + assert logging_builder.await_args_list[1].kwargs["logging_auth"].team_id == "team-a" + assert logging_builder.await_args_list[1].kwargs["logging_auth"].team_metadata == {"guardrails": ["g1"]} + assert proxy_handler.await_args.kwargs["user_api_key_dict"] is original_auth + assert proxy_handler.await_args.kwargs["guardrail_auth"].team_id == "team-a" + assert proxy_handler.await_args.kwargs["guardrail_auth"].team_metadata == {"guardrails": ["g1"]} + assert proxy_handler.await_args.kwargs["resolved_tool"] is target_tool + proxy_logger: Final = operations._fire_mcp_tool_call_logging.await_args.kwargs["logging_obj"] + assert proxy_logger.model_call_details["metadata"]["user_api_key_team_id"] == "team-a" + assert original_auth.team_id is None + assert original_auth.team_metadata is None + + @pytest.mark.asyncio @pytest.mark.parametrize( ("config", "extra_headers", "expected_source", "expected_authorization"), @@ -15917,9 +16403,7 @@ async def test_discovery_cache_isolates_forwarded_credentials_and_static_auth_ca @pytest.mark.parametrize("kind", ("prompts", "resources", "templates")) @pytest.mark.parametrize("header", ("Authorization", "X-LiteLLM-API-Key")) @pytest.mark.parametrize("identified", (False, True)) -async def test_discovery_cache_isolates_keyless_admission_credentials( - kind: str, header: str, identified: bool -) -> None: +async def test_discovery_cache_isolates_keyless_admission_credentials(kind: str, header: str, identified: bool) -> None: manager: Final = MCPServerManager() upstream: Final = _DiscoveryUpstream() server: Final = _discovery_server().model_copy(update={"static_headers": {"Authorization": "Bearer upstream"}}) @@ -17283,6 +17767,7 @@ async def test_catalog_config_only_snapshot_cleans_up_after_error(monkeypatch): manager = MCPServerManager() configured = MCPServer(server_id="configured", name="configured", transport="stdio", command="echo") manager.config_mcp_servers = {configured.server_id: configured} + async def fail_operation(): async with manager.catalog.operation(): assert await manager.catalog.resolve("configured") == configured @@ -17331,14 +17816,26 @@ async def test_catalog_publishes_rediscovered_routes_despite_concurrent_owner_ch @pytest.mark.parametrize("anchored", [False, True]) async def test_catalog_reload_preserves_concurrent_config_discovery_and_routes(anchored): manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="config-race", name="config_race", transport=MCPTransport.http, - url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2, - issuer="https://issuer.example" if anchored else None, issuer_is_anchored=anchored) + server: Final = MCPServer( + server_id="config-race", + name="config_race", + transport=MCPTransport.http, + url="https://upstream.example/mcp", + auth_type=MCPAuth.oauth2, + issuer="https://issuer.example" if anchored else None, + issuer_is_anchored=anchored, + ) manager.config_mcp_servers = {server.server_id: server} manager._set_oauth_discovery_deferred(server.server_id, True) generation: Final = manager.oauth_discovery_slot(server.server_id).generation - resolved: Final = server.model_copy(update={"authorization_url": "https://issuer.example/authorize", - "token_url": "https://issuer.example/token", "scopes": ["read"], "issuer": "https://issuer.example"}) + resolved: Final = server.model_copy( + update={ + "authorization_url": "https://issuer.example/authorize", + "token_url": "https://issuer.example/token", + "scopes": ["read"], + "issuer": "https://issuer.example", + } + ) async def hydrate(target: MCPServer) -> bool: target.client_id = "persisted-client" @@ -17355,7 +17852,10 @@ async def test_catalog_reload_preserves_concurrent_config_discovery_and_routes(a prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), - patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.hydrate_config_server_dcr_client", side_effect=hydrate), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.hydrate_config_server_dcr_client", + side_effect=hydrate, + ), ): await manager.reload_servers_from_database() current: Final = manager.get_mcp_server_by_id(server.server_id) @@ -17372,14 +17872,16 @@ async def test_catalog_reload_preserves_concurrent_config_discovery_and_routes(a @pytest.mark.parametrize("change", ["credentials", "delete"]) async def test_catalog_reload_does_not_restore_replaced_config_credentials_or_routes(change): manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="config-replaced", name="config_replaced", transport=MCPTransport.http, - client_id="previous-client") + server: Final = MCPServer( + server_id="config-replaced", name="config_replaced", transport=MCPTransport.http, client_id="previous-client" + ) manager.config_mcp_servers = {server.server_id: server} manager.published_tool_routes["config_replaced-search"] = "config_replaced" async def read_rows(**_kwargs): - manager.config_mcp_servers = {} if change == "delete" else { - server.server_id: server.model_copy(update={"client_id": "new-client"})} + manager.config_mcp_servers = ( + {} if change == "delete" else {server.server_id: server.model_copy(update={"client_id": "new-client"})} + ) return [] prisma: Final = MagicMock() @@ -17402,8 +17904,13 @@ async def test_catalog_operation_retains_routes_only_for_same_configured_target( from mcp.types import Tool manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="route-race", name="route_race", transport=MCPTransport.http, - url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2) + server: Final = MCPServer( + server_id="route-race", + name="route_race", + transport=MCPTransport.http, + url="https://upstream.example/mcp", + auth_type=MCPAuth.oauth2, + ) manager.registry = {server.server_id: server} with patch("litellm.proxy.proxy_server.prisma_client", None): async with manager.catalog.operation(): @@ -17411,9 +17918,16 @@ async def test_catalog_operation_retains_routes_only_for_same_configured_target( if change == "delete": manager.registry = {} else: - manager.registry[server.server_id] = server.model_copy(update=( - {"authorization_url": "https://issuer.example/authorize", "token_url": "https://issuer.example/token"} - if change == "discovery" else {"url": "https://changed.example/mcp"})) + manager.registry[server.server_id] = server.model_copy( + update=( + { + "authorization_url": "https://issuer.example/authorize", + "token_url": "https://issuer.example/token", + } + if change == "discovery" + else {"url": "https://changed.example/mcp"} + ) + ) assert (tools[0].name in manager.published_tool_routes) is (change == "discovery") @@ -17426,8 +17940,10 @@ async def test_catalog_reload_keeps_new_route_owner_over_earlier_route(change, m manager: Final = MCPServerManager() registry: Final = tool_registry.MCPToolRegistry() monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) - servers: Final = {name: MCPServer(server_id=name, name=name, transport=MCPTransport.http, - client_id="configured-client") for name in ("old_owner", "new_owner", "concurrent_owner")} + servers: Final = { + name: MCPServer(server_id=name, name=name, transport=MCPTransport.http, client_id="configured-client") + for name in ("old_owner", "new_owner", "concurrent_owner") + } manager.config_mcp_servers = servers manager.published_tool_routes = {"shared-search": "old_owner"} if mapped else {} @@ -17453,7 +17969,9 @@ async def test_catalog_reload_keeps_new_route_owner_over_earlier_route(change, m registry.published_tools.clear() elif change == "staged_with_concurrent": manager.published_tool_routes["shared-search"] = "concurrent_owner" - registry.published_tools["shared-search"] = original_tool.model_copy(update={"handler": concurrent_handler}) + registry.published_tools["shared-search"] = original_tool.model_copy( + update={"handler": concurrent_handler} + ) manager.tool_name_to_mcp_server_name_mapping["shared-search"] = "new_owner" registry.register_tool("shared-search", "Search", {}, new_handler) else: @@ -17486,10 +18004,16 @@ async def test_catalog_observes_committed_update_and_delete_without_background_r timestamp: Final = datetime.now() row: Final = LiteLLM_MCPServerTable( - server_id="catalog-server", server_name="catalog_server", alias="catalog_server", - transport=MCPTransport.http, url="https://first.example.com/mcp", updated_at=timestamp, + server_id="catalog-server", + server_name="catalog_server", + alias="catalog_server", + transport=MCPTransport.http, + url="https://first.example.com/mcp", + updated_at=timestamp, + ) + updated: Final = row.model_copy( + update={"url": "https://second.example.com/mcp", "updated_at": timestamp + timedelta(seconds=1)} ) - updated: Final = row.model_copy(update={"url": "https://second.example.com/mcp", "updated_at": timestamp + timedelta(seconds=1)}) prisma: Final = MagicMock() prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [updated], [])) @@ -17515,8 +18039,12 @@ async def test_catalog_observes_committed_update_and_delete_without_background_r async def test_catalog_rebuilt_unchanged_server_keeps_discovered_tool_routes(): from mcp.types import Tool - row: Final = LiteLLM_MCPServerTable(server_id="unchanged-routes", alias="unchanged_routes", - url="https://upstream.example/mcp", transport=MCPTransport.http) + row: Final = LiteLLM_MCPServerTable( + server_id="unchanged-routes", + alias="unchanged_routes", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + ) manager: Final = MCPServerManager() prisma: Final = MagicMock() prisma.writer_db = prisma.db @@ -17539,9 +18067,14 @@ async def test_catalog_reload_retains_routes_discovered_for_a_late_server(change from mcp.types import Tool from litellm.proxy._experimental.mcp_server import tool_registry - row: Final = LiteLLM_MCPServerTable(server_id="late-routes", alias="late_routes", - url="https://upstream.example/mcp", transport=MCPTransport.http, updated_at=datetime.now(), - spec_path="late-openapi.json" if openapi else None) + row: Final = LiteLLM_MCPServerTable( + server_id="late-routes", + alias="late_routes", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + updated_at=datetime.now(), + spec_path="late-openapi.json" if openapi else None, + ) manager: Final = MCPServerManager() registry: Final = tool_registry.MCPToolRegistry() monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) @@ -17565,9 +18098,17 @@ async def test_catalog_reload_retains_routes_discovered_for_a_late_server(change await published.wait() if change == "deleted": return [] - return [row if change == "same" else row.model_copy(update={ - "url": "https://updated.example/mcp", "spec_path": None, - "updated_at": row.updated_at + timedelta(seconds=1)})] + return [ + row + if change == "same" + else row.model_copy( + update={ + "url": "https://updated.example/mcp", + "spec_path": None, + "updated_at": row.updated_at + timedelta(seconds=1), + } + ) + ] prisma: Final = MagicMock() prisma.writer_db = prisma.db @@ -17584,7 +18125,9 @@ async def test_catalog_reload_retains_routes_discovered_for_a_late_server(change assert manager.get_mcp_server_by_id(row.server_id) is None assert "late_routes-search" not in manager.published_tool_routes else: - assert manager.server_exposes_tool(manager.get_mcp_server_by_id(row.server_id), "late_routes-search") is (change == "same") + assert manager.server_exposes_tool(manager.get_mcp_server_by_id(row.server_id), "late_routes-search") is ( + change == "same" + ) tool: Final = registry.get_tool("late_routes-search") if openapi and change == "same": assert tool is not None @@ -17595,7 +18138,9 @@ async def test_catalog_reload_retains_routes_discovered_for_a_late_server(change @pytest.mark.asyncio @pytest.mark.parametrize("retain_operation", [False, True]) -async def test_catalog_openapi_refresh_does_not_restore_removed_operations(tmp_path, monkeypatch, respx_mock, retain_operation): +async def test_catalog_openapi_refresh_does_not_restore_removed_operations( + tmp_path, monkeypatch, respx_mock, retain_operation +): from litellm.proxy._experimental.mcp_server import tool_registry monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") @@ -17606,8 +18151,13 @@ async def test_catalog_openapi_refresh_does_not_restore_removed_operations(tmp_p paths: Final = {"/removed": {"get": {"operationId": "removed"}}, "/retained": {"get": {"operationId": "retained"}}} spec: Final = {"openapi": "3.0.0", "info": {"title": "Refresh", "version": "1"}, "paths": paths} spec_path.write_text(json.dumps(spec)) - row: Final = LiteLLM_MCPServerTable(server_id="spec-refresh", alias="spec_refresh", - url="https://upstream.example", transport=MCPTransport.http, spec_path=str(spec_path)) + row: Final = LiteLLM_MCPServerTable( + server_id="spec-refresh", + alias="spec_refresh", + url="https://upstream.example", + transport=MCPTransport.http, + spec_path=str(spec_path), + ) prisma: Final = MagicMock() prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) @@ -17647,10 +18197,15 @@ async def test_catalog_lookup_uses_one_snapshot_until_operation_finishes(): from datetime import timedelta row: Final = LiteLLM_MCPServerTable( - server_id="snapshot-server", alias="snapshot_server", transport=MCPTransport.http, - url="https://first.example.com/mcp", updated_at=datetime.now(), + server_id="snapshot-server", + alias="snapshot_server", + transport=MCPTransport.http, + url="https://first.example.com/mcp", + updated_at=datetime.now(), + ) + updated: Final = row.model_copy( + update={"url": "https://second.example.com/mcp", "updated_at": row.updated_at + timedelta(seconds=1)} ) - updated: Final = row.model_copy(update={"url": "https://second.example.com/mcp", "updated_at": row.updated_at + timedelta(seconds=1)}) prisma: Final = MagicMock() prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [updated], [updated])) @@ -17746,8 +18301,15 @@ async def test_catalog_cancelled_openapi_refresh_retains_tools_and_discovery(mon manager: Final = MCPServerManager() stamp: Final = datetime.now() - server: Final = MCPServer(server_id="staged", name="staged", transport=MCPTransport.http, - url="https://before.example/mcp", spec_path="before.json", updated_at=stamp, auth_type=MCPAuth.oauth2) + server: Final = MCPServer( + server_id="staged", + name="staged", + transport=MCPTransport.http, + url="https://before.example/mcp", + spec_path="before.json", + updated_at=stamp, + auth_type=MCPAuth.oauth2, + ) manager.registry = {server.server_id: server} manager._set_oauth_discovery_deferred(server.server_id, True) original_slot: Final = manager.oauth_discovery_slot(server.server_id) @@ -17756,8 +18318,15 @@ async def test_catalog_cancelled_openapi_refresh_retains_tools_and_discovery(mon registry: Final = tool_registry.MCPToolRegistry() registry.register_tool("staged-existing", "existing", {}, lambda: "existing") monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) - row: Final = LiteLLM_MCPServerTable(server_id=server.server_id, alias="staged", transport=MCPTransport.http, - url="https://after.example/mcp", spec_path="after.json", updated_at=stamp + timedelta(seconds=1), auth_type=MCPAuth.oauth2) + row: Final = LiteLLM_MCPServerTable( + server_id=server.server_id, + alias="staged", + transport=MCPTransport.http, + url="https://after.example/mcp", + spec_path="after.json", + updated_at=stamp + timedelta(seconds=1), + auth_type=MCPAuth.oauth2, + ) prisma: Final = MagicMock() prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) @@ -17785,8 +18354,14 @@ async def test_catalog_cancelled_openapi_refresh_retains_tools_and_discovery(mon @pytest.mark.asyncio async def test_catalog_snapshot_identity_is_independent_of_worker_oauth_discovery(): - row: Final = LiteLLM_MCPServerTable(server_id="identity-server", alias="identity_server", transport=MCPTransport.http, - url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2, updated_at=datetime.now()) + row: Final = LiteLLM_MCPServerTable( + server_id="identity-server", + alias="identity_server", + transport=MCPTransport.http, + url="https://upstream.example/mcp", + auth_type=MCPAuth.oauth2, + updated_at=datetime.now(), + ) prisma: Final = MagicMock() prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) @@ -17799,8 +18374,9 @@ async def test_catalog_snapshot_identity_is_independent_of_worker_oauth_discover ): async with manager.catalog.operation() as unresolved: assert manager.get_mcp_server_by_id(row.server_id).authorization_url is None - resolved: Final = manager.registry[row.server_id].model_copy(update={ - "authorization_url": "https://idp.example/authorize", "token_url": "https://idp.example/token"}) + resolved: Final = manager.registry[row.server_id].model_copy( + update={"authorization_url": "https://idp.example/authorize", "token_url": "https://idp.example/token"} + ) manager.registry[row.server_id] = resolved async with manager.catalog.operation() as discovered: assert manager.get_mcp_server_by_id(row.server_id).authorization_url == resolved.authorization_url @@ -17814,8 +18390,13 @@ async def test_catalog_failed_openapi_row_does_not_publish_partial_handlers(monk manager: Final = MCPServerManager() registry: Final = tool_registry.MCPToolRegistry() monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) - row: Final = LiteLLM_MCPServerTable(server_id="broken", alias="broken", transport=MCPTransport.http, - url="https://upstream.example/mcp", spec_path="broken.json") + row: Final = LiteLLM_MCPServerTable( + server_id="broken", + alias="broken", + transport=MCPTransport.http, + url="https://upstream.example/mcp", + spec_path="broken.json", + ) prisma: Final = MagicMock() prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) @@ -17856,16 +18437,22 @@ async def test_catalog_reuses_openapi_tools_until_configuration_or_background_re registry = tool_registry.MCPToolRegistry() monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) spec_path = tmp_path / "catalog.json" - spec_path.write_text(json.dumps({ - "openapi": "3.0.0", "info": {"title": "Catalog", "version": "1"}, - "paths": {"/echo": {"get": {"operationId": "echo"}}}, - })) + spec_path.write_text( + json.dumps( + { + "openapi": "3.0.0", + "info": {"title": "Catalog", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}, + } + ) + ) row = _catalog_row().model_copy(update={"spec_path": str(spec_path)}) read_rows = AsyncMock(return_value=[row]) _catalog_database(monkeypatch, read_rows) manager = MCPServerManager() with patch.object( - openapi_to_mcp_generator, "load_openapi_spec_async", + openapi_to_mcp_generator, + "load_openapi_spec_async", wraps=openapi_to_mcp_generator.load_openapi_spec_async, ) as load_spec: await manager.catalog.list() @@ -17912,10 +18499,15 @@ async def test_catalog_changed_openapi_handler_never_dispatches(monkeypatch, tmp monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) monkeypatch.setattr(operations, "global_mcp_tool_registry", registry) spec_path = tmp_path / "catalog-race.json" - spec_path.write_text(json.dumps({ - "openapi": "3.0.0", "info": {"title": "Catalog", "version": "1"}, - "paths": {"/echo": {"get": {"operationId": "echo"}}}, - })) + spec_path.write_text( + json.dumps( + { + "openapi": "3.0.0", + "info": {"title": "Catalog", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}, + } + ) + ) row = _catalog_row().model_copy(update={"spec_path": str(spec_path)}) read_rows = AsyncMock(return_value=[row]) _catalog_database(monkeypatch, read_rows) @@ -17924,9 +18516,14 @@ async def test_catalog_changed_openapi_handler_never_dispatches(monkeypatch, tmp destination = respx_mock.get("https://changed.example/echo").respond(200, text="must not execute") async with manager.catalog.operation(): admitted = manager.get_mcp_server_by_id("catalog-server") - read_rows.return_value = [row.model_copy(update={ - "url": "https://changed.example", "updated_at": datetime(2026, 1, 2), - })] + read_rows.return_value = [ + row.model_copy( + update={ + "url": "https://changed.example", + "updated_at": datetime(2026, 1, 2), + } + ) + ] await asyncio.create_task(manager.reload_servers_from_database()) with pytest.raises(HTTPException) as exc: if dispatch == "managed": @@ -17962,7 +18559,8 @@ async def test_catalog_deferred_discovery_preserves_admitted_configuration(monke entered = asyncio.Event() release = asyncio.Event() metadata = MCPOAuthMetadata( - authorization_url="https://issuer.example/authorize", token_url="https://issuer.example/token", + authorization_url="https://issuer.example/authorize", + token_url="https://issuer.example/token", ) async def discover(server): @@ -17992,15 +18590,19 @@ async def test_catalog_deferred_discovery_preserves_admitted_configuration(monke async def test_catalog_cancelled_config_hydration_preserves_published_credentials(): manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="config-hydration", name="config_hydration", transport=MCPTransport.http, - client_id="previous-client") + server: Final = MCPServer( + server_id="config-hydration", name="config_hydration", transport=MCPTransport.http, client_id="previous-client" + ) manager.config_mcp_servers = {server.server_id: server} async def cancelled_hydration(target: MCPServer) -> bool: target.client_id = "unpublished-client" raise asyncio.CancelledError() - with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.hydrate_config_server_dcr_client", side_effect=cancelled_hydration): + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.hydrate_config_server_dcr_client", + side_effect=cancelled_hydration, + ): with pytest.raises(asyncio.CancelledError): await manager.reload_servers_from_database() assert manager.get_mcp_server_by_id(server.server_id).client_id == "previous-client" @@ -18009,15 +18611,22 @@ async def test_catalog_cancelled_config_hydration_preserves_published_credential @pytest.mark.asyncio async def test_catalog_oauth_resolution_cannot_replace_an_operation_target_after_update(): manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="pinned-oauth", name="pinned_oauth", transport=MCPTransport.http, - url="https://before.example/mcp", auth_type=MCPAuth.oauth2, - authorization_url="https://before.example/authorize", token_url="https://before.example/token") + server: Final = MCPServer( + server_id="pinned-oauth", + name="pinned_oauth", + transport=MCPTransport.http, + url="https://before.example/mcp", + auth_type=MCPAuth.oauth2, + authorization_url="https://before.example/authorize", + token_url="https://before.example/token", + ) manager.registry = {server.server_id: server} with patch("litellm.proxy.proxy_server.prisma_client", None): async with manager.catalog.operation(): selected: Final = manager.get_mcp_server_by_id(server.server_id) - manager.registry[server.server_id] = server.model_copy(update={ - "url": "https://after.example/mcp", "authorization_url": "https://after.example/authorize"}) + manager.registry[server.server_id] = server.model_copy( + update={"url": "https://after.example/mcp", "authorization_url": "https://after.example/authorize"} + ) resolved: Final = await manager.ensure_oauth_metadata_discovered(selected) assert resolved.url == "https://before.example/mcp" assert resolved.authorization_url == "https://before.example/authorize" @@ -18027,8 +18636,13 @@ async def test_catalog_oauth_resolution_cannot_replace_an_operation_target_after @pytest.mark.asyncio async def test_catalog_unresolved_oauth_snapshot_fails_closed_if_target_was_replaced(): manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="replaced-oauth", name="replaced_oauth", transport=MCPTransport.http, - url="https://before.example/mcp", auth_type=MCPAuth.oauth2) + server: Final = MCPServer( + server_id="replaced-oauth", + name="replaced_oauth", + transport=MCPTransport.http, + url="https://before.example/mcp", + auth_type=MCPAuth.oauth2, + ) manager.registry = {server.server_id: server} with ( patch("litellm.proxy.proxy_server.prisma_client", None), @@ -18036,9 +18650,13 @@ async def test_catalog_unresolved_oauth_snapshot_fails_closed_if_target_was_repl ): async with manager.catalog.operation(): selected: Final = manager.get_mcp_server_by_id(server.server_id) - manager.registry[server.server_id] = server.model_copy(update={ - "url": "https://after.example/mcp", "authorization_url": "https://after.example/authorize", - "token_url": "https://after.example/token"}) + manager.registry[server.server_id] = server.model_copy( + update={ + "url": "https://after.example/mcp", + "authorization_url": "https://after.example/authorize", + "token_url": "https://after.example/token", + } + ) with pytest.raises(HTTPException) as exc: await manager.ensure_oauth_metadata_discovered(selected) assert exc.value.status_code == 503 @@ -18048,11 +18666,21 @@ async def test_catalog_unresolved_oauth_snapshot_fails_closed_if_target_was_repl @pytest.mark.asyncio async def test_catalog_unresolved_oauth_snapshot_accepts_discovery_for_the_same_target(): manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="same-oauth", name="same_oauth", transport=MCPTransport.http, - url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2) + server: Final = MCPServer( + server_id="same-oauth", + name="same_oauth", + transport=MCPTransport.http, + url="https://upstream.example/mcp", + auth_type=MCPAuth.oauth2, + ) manager.registry = {server.server_id: server} - discovered: Final = server.model_copy(update={"authorization_url": "https://issuer.example/authorize", - "token_url": "https://issuer.example/token", "scopes": ["read"]}) + discovered: Final = server.model_copy( + update={ + "authorization_url": "https://issuer.example/authorize", + "token_url": "https://issuer.example/token", + "scopes": ["read"], + } + ) with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch.object(manager, "_ensure_oauth_metadata_discovered", return_value=discovered) as discovery, @@ -18068,8 +18696,9 @@ async def test_catalog_unresolved_oauth_snapshot_accepts_discovery_for_the_same_ @pytest.mark.asyncio async def test_catalog_fresh_lookup_does_not_fall_back_to_stale_grants_when_database_fails(): manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="stale-grant", name="stale_grant", transport=MCPTransport.http, - allow_all_keys=True) + server: Final = MCPServer( + server_id="stale-grant", name="stale_grant", transport=MCPTransport.http, allow_all_keys=True + ) manager.registry = {server.server_id: server} prisma: Final = MagicMock() prisma.writer_db = prisma.db @@ -18107,7 +18736,6 @@ async def test_catalog_cancelled_registration_does_not_publish_partial_handlers( assert manager.registry == {} - def _revision_row(revision): from types import SimpleNamespace @@ -18388,9 +19016,7 @@ async def test_release_during_auth_fingerprinting_does_not_open_upstream_session _RetiringSessionStub(), ) manager.track_gateway_session(gateway_session_id) - request_task: Final = asyncio.create_task( - manager._upstream_session_for(client, server, gateway_session_headers) - ) + request_task: Final = asyncio.create_task(manager._upstream_session_for(client, server, gateway_session_headers)) try: await asyncio.wait_for(fingerprint_started.wait(), 5) @@ -18428,7 +19054,9 @@ async def test_upstream_session_is_shared_per_gateway_session_and_released_with_ first: Final = await manager._upstream_session_for(client, server, {"Mcp-Session-Id": "gw-1"}) second: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-1"}) other: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-2"}) - assert first is not None and first is second, "every call of one gateway session must share one upstream session" + assert first is not None and first is second, ( + "every call of one gateway session must share one upstream session" + ) assert other is not None and other is not first, "distinct gateway sessions must not share upstream state" manager.release_upstream_sessions("gw-1") @@ -18488,9 +19116,7 @@ async def test_upstream_sessions_are_partitioned_by_discovery_auth_fingerprint( assert first_session.closed assert len(manager._upstream_sessions) == 1 assert tuple(manager._upstream_sessions.values()) == (second_session,) - first_session_again: Final = await manager._upstream_session_for( - first_client, server, gateway_session_headers - ) + first_session_again: Final = await manager._upstream_session_for(first_client, server, gateway_session_headers) assert first_session_again is not None assert first_session_again is not first_session assert first_session_again is not second_session @@ -18541,9 +19167,7 @@ async def test_upstream_sessions_are_partitioned_by_connection_settings( assert first_session.closed assert len(manager._upstream_sessions) == 1 assert tuple(manager._upstream_sessions.values()) == (changed_session,) - first_session_again: Final = await manager._upstream_session_for( - third_client, server, gateway_session_headers - ) + first_session_again: Final = await manager._upstream_session_for(third_client, server, gateway_session_headers) assert first_session_again is not None assert first_session_again is not first_session @@ -18780,9 +19404,7 @@ async def test_upstream_session_is_not_reused_outside_an_mcp_protocol_request() manager.track_gateway_session("gw-rest") try: - session: Final = await manager._upstream_session_for( - client, server, {"mcp-session-id": "gw-rest"} - ) + session: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-rest"}) assert session is None finally: manager.release_upstream_sessions("gw-rest") @@ -19566,8 +20188,11 @@ async def test_upstream_preparation_honors_case_sensitive_extra_command(monkeypa monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true") monkeypatch.setattr(upstream, "MCP_STDIO_ALLOWED_COMMANDS", frozenset({"CustomRunner"})) server: Final = MCPServer( - server_id="custom-stdio", name="custom-stdio", transport=MCPTransport.stdio, - command="/opt/tools/CustomRunner", args=[], + server_id="custom-stdio", + name="custom-stdio", + transport=MCPTransport.stdio, + command="/opt/tools/CustomRunner", + args=[], ) client: Final = await MCPServerManager()._create_mcp_client(server) assert client.stdio_config is not None @@ -19759,10 +20384,12 @@ async def test_catalog_observes_writer_changes_while_read_replica_lags(monkeypat litellm_mcpservertable=SimpleNamespace(find_many=AsyncMock(return_value=[original])), litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=_revision_row(7))), ) - routed: Final = SimpleNamespace(**{ - name: _RoutedActions(getattr(writer, name), getattr(reader, name), lambda: True) - for name in ("litellm_mcpservertable", "litellm_config") - }) + routed: Final = SimpleNamespace( + **{ + name: _RoutedActions(getattr(writer, name), getattr(reader, name), lambda: True) + for name in ("litellm_mcpservertable", "litellm_config") + } + ) monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=routed, writer_db=writer)) manager: Final = MCPServerManager() assert (await manager.catalog.resolve(original.server_id)).name == "initial" @@ -19779,7 +20406,9 @@ async def test_catalog_observes_writer_changes_while_read_replica_lags(monkeypat @pytest.mark.asyncio async def test_catalog_revision_failure_preserves_state_and_recovers(monkeypatch): reader: Final = AsyncMock(side_effect=([_catalog_row()], [_catalog_row("updated")])) - revisions: Final = AsyncMock(side_effect=[_revision_row(7), RuntimeError("private database detail"), _revision_row(8)]) + revisions: Final = AsyncMock( + side_effect=[_revision_row(7), RuntimeError("private database detail"), _revision_row(8)] + ) _catalog_database(monkeypatch, reader, revisions) manager: Final = MCPServerManager() assert (await manager.catalog.resolve("catalog-server")).name == "initial" @@ -19905,15 +20534,20 @@ async def test_catalog_accepts_discovered_issuer_response_support( manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="discovered-issuer", name="discovered_issuer", transport=MCPTransport.http, - url="https://resource.example/mcp", auth_type=MCPAuth.oauth2, - issuer="https://issuer.example" if anchored else None, issuer_is_anchored=anchored, + server_id="discovered-issuer", + name="discovered_issuer", + transport=MCPTransport.http, + url="https://resource.example/mcp", + auth_type=MCPAuth.oauth2, + issuer="https://issuer.example" if anchored else None, + issuer_is_anchored=anchored, ) manager.registry = {server.server_id: server} manager._set_oauth_discovery_deferred(server.server_id, True) metadata: Final = MCPOAuthMetadata( discovered_issuer="https://issuer.example", - authorization_url="https://issuer.example/authorize", token_url="https://issuer.example/token", + authorization_url="https://issuer.example/authorize", + token_url="https://issuer.example/token", authorization_response_iss_parameter_supported=issuer_parameter_supported, ) discovery: Final = AsyncMock(return_value=metadata) @@ -19935,9 +20569,13 @@ async def test_catalog_rejects_a_changed_anchored_issuer_during_discovery(monkey manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="anchored-issuer", name="anchored_issuer", transport=MCPTransport.http, - url="https://resource.example/mcp", auth_type=MCPAuth.oauth2, - issuer="https://original.example", issuer_is_anchored=True, + server_id="anchored-issuer", + name="anchored_issuer", + transport=MCPTransport.http, + url="https://resource.example/mcp", + auth_type=MCPAuth.oauth2, + issuer="https://original.example", + issuer_is_anchored=True, ) manager.registry = {server.server_id: server} discovery: Final = AsyncMock() @@ -20023,13 +20661,24 @@ async def test_paginated_listing_keeps_earlier_tool_metadata_and_caller_isolatio @pytest.mark.asyncio -@pytest.mark.parametrize("request_name,result_field", [("ListPromptsRequest", "prompts"), ("ListResourcesRequest", "resources"), ("ListResourceTemplatesRequest", "resource_templates")]) +@pytest.mark.parametrize( + "request_name,result_field", + [ + ("ListPromptsRequest", "prompts"), + ("ListResourcesRequest", "resources"), + ("ListResourceTemplatesRequest", "resource_templates"), + ], +) @pytest.mark.parametrize("cursor", [None, "next"]) -async def test_disabled_stdio_catalog_is_empty_initially_and_rejects_continuation(monkeypatch, request_name, result_field, cursor): +async def test_disabled_stdio_catalog_is_empty_initially_and_rejects_continuation( + monkeypatch, request_name, result_field, cursor +): from mcp import types monkeypatch.delenv("LITELLM_ENABLE_MCP_STDIO", raising=False) - server = MCPServer(server_id="disabled", name="disabled", transport=MCPTransport.stdio, command="blocked-executable", args=[]) + server = MCPServer( + server_id="disabled", name="disabled", transport=MCPTransport.stdio, command="blocked-executable", args=[] + ) request = getattr(types, request_name)(params=types.PaginatedRequestParams(cursor=cursor)) manager = MCPServerManager() if cursor is not None: @@ -20041,7 +20690,6 @@ async def test_disabled_stdio_catalog_is_empty_initially_and_rejects_continuatio assert page.next_cursor is None - @pytest.mark.asyncio @pytest.mark.parametrize("cursor", [None, "next"]) async def test_disabled_stdio_tools_are_empty_initially_and_reject_continuation(monkeypatch, cursor): @@ -20149,8 +20797,10 @@ async def test_failed_aggregate_continuation_preserves_only_delivered_tool_metad ListToolsResult(tools=[first], next_cursor="next"), ListToolsResult(tools=[unseen], next_cursor="next" if client is clients["beta"] else None), ] + async def create_client(server, **kwargs): return clients[server.server_id] + manager.create_mcp_client = create_client monkeypatch.setattr(operations, "global_mcp_server_manager", manager) caller = UserAPIKeyAuth(api_key="owned-caller", user_id="alice") @@ -20187,7 +20837,9 @@ async def test_aggregate_publishes_complete_bare_routes_only_after_delivering_a_ clients = {server.server_id: AsyncMock() for server in servers} for name, client in clients.items(): client._last_initialize_instructions = None - client.list_tools_page.return_value = ListToolsResult(tools=[tool], next_cursor="next" if name == "beta" else None) + client.list_tools_page.return_value = ListToolsResult( + tools=[tool], next_cursor="next" if name == "beta" else None + ) async def create_client(server, **kwargs): return clients[server.server_id] @@ -20278,27 +20930,55 @@ async def test_manual_oauth_endpoints_discover_client_metadata_once( monkeypatch.setattr(manager_module, "global_mcp_server_manager", manager) configured: Final[MCPServer] if source == "config": - await manager.load_servers_from_config({"manual": { - "url": "https://up.example.com/mcp", "transport": "http", "auth_type": "oauth2", - "oauth2_flow": "authorization_code", "authorization_url": "https://up.example.com/authorize", - "token_url": "https://up.example.com/token", "scopes": ["read"], - }}) + await manager.load_servers_from_config( + { + "manual": { + "url": "https://up.example.com/mcp", + "transport": "http", + "auth_type": "oauth2", + "oauth2_flow": "authorization_code", + "authorization_url": "https://up.example.com/authorize", + "token_url": "https://up.example.com/token", + "scopes": ["read"], + } + } + ) configured = next(iter(manager.config_mcp_servers.values())) else: row: Final = LiteLLM_MCPServerTable( - server_id="manual", alias="manual", url="https://up.example.com/mcp", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", - authorization_url="https://up.example.com/authorize", token_url="https://up.example.com/token", - credentials={"scopes": ["read"]}, created_at=datetime.now(), updated_at=datetime.now(), + server_id="manual", + alias="manual", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + authorization_url="https://up.example.com/authorize", + token_url="https://up.example.com/token", + credentials={"scopes": ["read"]}, + created_at=datetime.now(), + updated_at=datetime.now(), ) built: Final = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) manager.registry[built.server_id] = built configured = built async with manager.catalog.operation(): response: Final = await discoverable_endpoints.register_client_with_server( - Request({"type": "http", "scheme": "https", "server": ("gateway.example.com", 443), - "path": "/register", "root_path": "", "headers": [], "query_string": b""}), - configured, "Gateway", None, None, None, + Request( + { + "type": "http", + "scheme": "https", + "server": ("gateway.example.com", 443), + "path": "/register", + "root_path": "", + "headers": [], + "query_string": b"", + } + ), + configured, + "Gateway", + None, + None, + None, ) body: Final = json.loads(response.body) if hasattr(response, "body") else response assert body["client_id"] == ( @@ -20326,11 +21006,19 @@ async def test_optional_client_metadata_discovery_failure_preserves_manual_endpo respx_mock.get("https://up.example.com/.well-known/oauth-authorization-server").respond(503) respx_mock.route().respond(404) manager: Final = MCPServerManager() - await manager.load_servers_from_config({"manual": { - "url": "https://up.example.com/mcp", "transport": "http", "auth_type": "oauth2", - "oauth2_flow": "authorization_code", "authorization_url": "https://up.example.com/authorize", - "token_url": "https://up.example.com/token", "scopes": ["read"], - }}) + await manager.load_servers_from_config( + { + "manual": { + "url": "https://up.example.com/mcp", + "transport": "http", + "auth_type": "oauth2", + "oauth2_flow": "authorization_code", + "authorization_url": "https://up.example.com/authorize", + "token_url": "https://up.example.com/token", + "scopes": ["read"], + } + } + ) configured: Final = next(iter(manager.config_mcp_servers.values())) resolved: Final = await manager.ensure_oauth_metadata_discovered(configured) attempts: Final = len(respx_mock.calls) @@ -20350,26 +21038,43 @@ def test_oauth_rebuild_retains_only_corroborated_cimd_capability(capability: boo from litellm.proxy._experimental.mcp_server.mcp_server_manager import carry_forward_resolved_oauth_endpoints previous: Final = MCPServer( - server_id="cimd-rebuild", name="cimd_rebuild", url="https://mcp.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.oauth2, - authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + server_id="cimd-rebuild", + name="cimd_rebuild", + url="https://mcp.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", client_id_metadata_document_supported=capability, - authorization_server_metadata={"issuer": "https://idp.example.com", "response_types_supported": ["code"], "scopes_supported": ["offline_access"]}, + authorization_server_metadata={ + "issuer": "https://idp.example.com", + "response_types_supported": ["code"], + "scopes_supported": ["offline_access"], + }, + ) + rebuilt: Final = previous.model_copy( + update={ + "client_id_metadata_document_supported": not capability if rebuild == "fresh_discovery" else None, + "authorization_url": "https://changed.example.com/authorize" + if rebuild == "repointed" + else previous.authorization_url, + "token_url": None, + "authorization_server_metadata": {"scopes_supported": []} if rebuild == "fresh_discovery" else None, + "issuer": "https://idp.example.com" if rebuild == "anchored" else None, + "issuer_is_anchored": rebuild == "anchored", + } ) - rebuilt: Final = previous.model_copy(update={ - "client_id_metadata_document_supported": not capability if rebuild == "fresh_discovery" else None, - "authorization_url": "https://changed.example.com/authorize" if rebuild == "repointed" else previous.authorization_url, - "token_url": None, - "authorization_server_metadata": {"scopes_supported": []} if rebuild == "fresh_discovery" else None, - "issuer": "https://idp.example.com" if rebuild == "anchored" else None, - "issuer_is_anchored": rebuild == "anchored", - }) carry_forward_resolved_oauth_endpoints(rebuilt, previous) - expected: Final = not capability if rebuild == "fresh_discovery" else None if rebuild in ("repointed", "anchored") else capability + expected: Final = ( + not capability if rebuild == "fresh_discovery" else None if rebuild in ("repointed", "anchored") else capability + ) assert rebuilt.client_id_metadata_document_supported is expected assert rebuilt.authorization_server_metadata == ( - {"scopes_supported": []} if rebuild == "fresh_discovery" - else previous.authorization_server_metadata if rebuild == "same" else None + {"scopes_supported": []} + if rebuild == "fresh_discovery" + else previous.authorization_server_metadata + if rebuild == "same" + else None ) @@ -20378,13 +21083,19 @@ def test_oauth_rebuild_retains_only_corroborated_cimd_capability(capability: boo async def test_repeated_stale_discovery_uses_current_callers_endpoint(endpoint: str) -> None: manager: Final = MCPServerManager() original: Final = MCPServer( - server_id="partial-replacement", name="replacement", url="https://old.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", + server_id="partial-replacement", + name="replacement", + url="https://old.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", ) replacement: Final = original.model_copy(update={endpoint: "https://new.example.com/oauth"}) manager.registry[original.server_id] = replacement resolved: Final = await manager._rejoin_oauth_metadata_discovery( - original, needed_endpoint=lambda server: getattr(server, endpoint), retry_stale=False, + original, + needed_endpoint=lambda server: getattr(server, endpoint), + retry_stale=False, ) assert resolved is replacement @@ -20394,11 +21105,19 @@ def test_openapi_extra_headers_exclude_gateway_admission_key(header_name: str) - caller_key: Final = "Bearer sk-admission-only" raw_headers: Final = {"x-litellm-api-key": caller_key, "x-tenant": "tenant-control"} server: Final = MCPServer( - server_id="header-boundary", name="header-boundary", transport=MCPTransport.http, - spec_path="/spec.yaml", auth_type=MCPAuth.none, extra_headers=[header_name, "X-Tenant"], + server_id="header-boundary", + name="header-boundary", + transport=MCPTransport.http, + spec_path="/spec.yaml", + auth_type=MCPAuth.none, + extra_headers=[header_name, "X-Tenant"], ) assert resolve_openapi_tool_auth( - server, None, None, raw_headers, UserAPIKeyAuth(api_key="sk-admission-only"), + server, + None, + None, + raw_headers, + UserAPIKeyAuth(api_key="sk-admission-only"), ) == (None, {"X-Tenant": "tenant-control"}, None) assert raw_headers == {"x-litellm-api-key": caller_key, "x-tenant": "tenant-control"} diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 66d66271f71..b8e28c3f9b8 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -1335,9 +1335,12 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): - with patch( - "litellm.proxy._experimental.mcp_server.catalog.verbose_logger", - ) as mock_logger, patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger): + with ( + patch( + "litellm.proxy._experimental.mcp_server.catalog.verbose_logger", + ) as mock_logger, + patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger), + ): # Test with server-specific auth headers mcp_server_auth_headers = { "working": {"Authorization": "Bearer working-token"}, @@ -1431,9 +1434,12 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): - with patch( - "litellm.proxy._experimental.mcp_server.catalog.verbose_logger", - ) as mock_logger, patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger): + with ( + patch( + "litellm.proxy._experimental.mcp_server.catalog.verbose_logger", + ) as mock_logger, + patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger), + ): # Test with server-specific auth headers mcp_server_auth_headers = { "failing1": {"Authorization": "Bearer failing1-token"}, @@ -3535,8 +3541,15 @@ async def test_initialize_request_with_existing_session_tracks_new_session(): patch( # test-quality-ok: registry is empty in unit tests; key owns one server "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, - return_value=[MCPServer(server_id="new-server", name="new-server", url="http://upstream/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.none)], + return_value=[ + MCPServer( + server_id="new-server", + name="new-server", + url="http://upstream/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + ], ), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", @@ -3760,10 +3773,13 @@ async def test_owner_fingerprint_hashes_custom_api_keys(): @pytest.mark.asyncio -@pytest.mark.parametrize("body", [ - b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', - b'{"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":0}}', -]) +@pytest.mark.parametrize( + "body", + [ + b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', + b'{"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":0}}', + ], +) async def test_stateful_mcp_session_owner_mismatch_returns_403(body: bytes): """ A stateful mcp-session-id is bound to its creator. A different @@ -7328,9 +7344,9 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), mcp_servers=["passthrough_server"], 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"}, + 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"}, ) assert result is None else: @@ -7340,9 +7356,9 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), mcp_servers=["passthrough_server"], 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"}, + 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"}, ) assert exc_info.value.status_code == expected_status if expected_status == 401: @@ -10710,6 +10726,7 @@ async def test_post_mcp_call_guardrails_return_the_rewritten_result(): logging_obj = _mock_mcp_logging_obj() raw_result = _call_tool_result(False, "jane@example.com") masked_result = _call_tool_result(False, "") + auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") proxy_logging_mock = _mock_mcp_proxy_logging() proxy_logging_mock.post_mcp_call_hook = AsyncMock(return_value=masked_result) @@ -10717,7 +10734,8 @@ async def test_post_mcp_call_guardrails_return_the_rewritten_result(): returned = await _run_post_mcp_call_guardrails( result=raw_result, litellm_logging_obj=logging_obj, - user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + original_auth=auth, + guardrail_auth=auth, request_data={}, ) @@ -10741,6 +10759,7 @@ async def test_post_mcp_call_guardrails_run_without_a_logging_object(): raw_result = _call_tool_result(False, "jane@example.com") masked_result = _call_tool_result(False, "") + auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") proxy_logging_mock = _mock_mcp_proxy_logging() proxy_logging_mock.post_mcp_call_hook = AsyncMock(return_value=masked_result) @@ -10748,7 +10767,8 @@ async def test_post_mcp_call_guardrails_run_without_a_logging_object(): returned = await _run_post_mcp_call_guardrails( result=raw_result, litellm_logging_obj=None, - user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + original_auth=auth, + guardrail_auth=auth, request_data={"name": "fetch_record"}, ) @@ -10775,11 +10795,56 @@ async def test_post_mcp_call_guardrails_propagate_a_block(): await _run_post_mcp_call_guardrails( result=_call_tool_result(False, "jane@example.com"), litellm_logging_obj=_mock_mcp_logging_obj(), - user_api_key_auth=None, + original_auth=None, + guardrail_auth=None, request_data={}, ) +@pytest.mark.asyncio +async def test_post_mcp_call_guardrails_preserve_and_deduplicate_mapping_entries( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._experimental.mcp_server.server import _run_post_mcp_call_guardrails + from litellm.proxy import proxy_server + + monkeypatch.setattr("litellm.proxy.utils.premium_user_check", MagicMock()) + monkeypatch.setattr(proxy_server, "premium_user", True, raising=False) + original_auth: Final = UserAPIKeyAuth(api_key="personal-key", user_id="alice") + guardrail_auth: Final = original_auth.model_copy( + update={ + "team_id": "team-a", + "team_metadata": { + "guardrails": [ + {"request-guardrail": {"team": True}}, + {"team-guardrail": {"x": 1}}, + {"team-guardrail": {"x": 2}}, + "request-guardrail", + "string-guardrail", + ] + }, + } + ) + proxy_logging_mock = _mock_mcp_proxy_logging() + proxy_logging_mock.post_mcp_call_hook = AsyncMock(return_value=_call_tool_result(False, "ok")) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock): + await _run_post_mcp_call_guardrails( + result=_call_tool_result(False, "ok"), + litellm_logging_obj=None, + original_auth=original_auth, + guardrail_auth=guardrail_auth, + request_data={"metadata": {"guardrails": [{"request-guardrail": {"request": True}}]}}, + ) + + request_data: Final = proxy_logging_mock.post_mcp_call_hook.await_args.kwargs["request_data"] + assert request_data["metadata"]["guardrails"] == [ + {"request-guardrail": {"request": True}}, + {"team-guardrail": {"x": 1}}, + "string-guardrail", + ] + + class TestListFiltersHonorThePrefixBoundary: """The listing filters compare a published (prefixed) tool name against configured entries, so they have to locate the boundary with the server's @@ -11836,8 +11901,7 @@ async def test_passthrough_probes_bind_each_token_to_its_authorized_server(allow ) assert result is None assert tuple( - (str(call.request.url), call.request.headers["Authorization"]) - for call in (*first.calls, *second.calls) + (str(call.request.url), call.request.headers["Authorization"]) for call in (*first.calls, *second.calls) ) == ( (("http://first/mcp", "Bearer first-token"), ("http://second/mcp", "Bearer second-token")) if allowed and header_shape != "non_authorization" @@ -11899,6 +11963,7 @@ async def test_streamable_http_admission_auth_sees_only_peeked_jsonrpc_method() stateless_handle.assert_not_awaited() send.assert_not_awaited() + @pytest.mark.asyncio @pytest.mark.parametrize( "payload, valid", @@ -11991,16 +12056,22 @@ async def test_progress_text_obeys_request_output_guardrails(monkeypatch, action from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks policy = ContentFilterGuardrail( - guardrail_name="progress-policy", event_hook=GuardrailEventHooks.post_mcp_call, + guardrail_name="progress-policy", + event_hook=GuardrailEventHooks.post_mcp_call, blocked_words=[BlockedWord(keyword="halfway", action=ContentFilterAction(action))], - keyword_redaction_tag="[masked]", default_on=False, + keyword_redaction_tag="[masked]", + default_on=False, ) monkeypatch.setattr(litellm, "callbacks", [policy]) monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())) delivered = AsyncMock() caller = UserAPIKeyAuth(user_id="progress-user") callback = _guarded_progress_callback( - delivered, None, caller, {"metadata": {"guardrails": ["progress-policy"] if selected else []}} + delivered, + None, + caller, + caller, + {"metadata": {"guardrails": ["progress-policy"] if selected else []}}, ) assert callback is not None await callback(0, None, message) @@ -12014,8 +12085,98 @@ async def test_progress_callback_preserves_absence_and_cancellation(): from litellm.proxy._experimental.mcp_server.operations import _guarded_progress_callback - assert _guarded_progress_callback(None, None, None, {}) is None - callback = _guarded_progress_callback(AsyncMock(side_effect=asyncio.CancelledError), None, None, {}) + assert _guarded_progress_callback(None, None, None, None, {}) is None + callback = _guarded_progress_callback(AsyncMock(side_effect=asyncio.CancelledError), None, None, None, {}) assert callback is not None with pytest.raises(asyncio.CancelledError): await callback(0, 0, None) + + +@pytest.mark.asyncio +async def test_execute_mcp_tool_derives_team_guardrail_auth_for_personal_keys( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.utils import get_server_prefix + from litellm.proxy.auth import auth_checks + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable + from litellm.proxy.utils import ProxyLogging + + user_api_key_auth: Final = UserAPIKeyAuth(api_key="sk-personal", user_id="owner") + server: Final = MCPServer( + server_id="team-server", + name="math_a", + server_name="math_a", + transport=MCPTransport.http, + url="http://localhost/mcp", + ) + team: Final = LiteLLM_TeamTable( + team_id="team-a", + organization_id="org-team-a", + object_permission_id="perm-team-a", + metadata={"guardrails": ["team-guardrail"]}, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-team-a", + mcp_servers=["team-server"], + mcp_tool_permissions={"team-server": ["test_tool"]}, + ), + ) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + + async def fake_handle_managed_mcp_tool( + *, + server_name: str, + name: str, + arguments: dict[str, object], + user_api_key_auth: UserAPIKeyAuth | None, + guardrail_auth: UserAPIKeyAuth | None, + **_: object, + ) -> CallToolResult: + await global_mcp_server_manager.pre_call_tool_check( + name=name, + arguments=arguments, + server_name=server_name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=server, + guardrail_auth=guardrail_auth, + ) + return CallToolResult(content=[], is_error=False) + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.utils.premium_user_check", MagicMock()) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(return_value=LiteLLM_UserTable(user_id="owner", teams=["team-a"])), + ) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=team)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "_personal_key_team_within_budget", AsyncMock(return_value=True)) + monkeypatch.setattr(global_mcp_server_manager, "check_tool_permission_for_key_team", AsyncMock(return_value=None)) + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + monkeypatch.setattr(mcp_operations, "_list_tools_before_first_call", AsyncMock(return_value=None)) + monkeypatch.setattr(mcp_operations, "_handle_managed_mcp_tool", fake_handle_managed_mcp_tool) + + await mcp_operations.execute_mcp_tool( + name=f"{get_server_prefix(server)}-test_tool", + arguments={}, + allowed_mcp_servers=[server], + start_time=datetime.now(), + user_api_key_auth=user_api_key_auth, + requested_server_id=server.server_id, + ) + + guardrail_auth: Final = proxy_logging_obj.pre_call_hook.await_args.kwargs["user_api_key_dict"] + assert guardrail_auth.team_id == "team-a" + assert guardrail_auth.team_metadata == {"guardrails": ["team-guardrail"]} diff --git a/tests/unit/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_rest_endpoints.py index 2a04c8cda2a..73e248ff52e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1564,7 +1564,13 @@ class TestListToolsRestAPI: monkeypatch, ): """The REST tools/list path should include tools beyond the upstream first page.""" - from mcp.types import Implementation, InitializeResult, ListToolsResult, PaginatedRequestParams, ServerCapabilities + from mcp.types import ( + Implementation, + InitializeResult, + ListToolsResult, + PaginatedRequestParams, + ServerCapabilities, + ) from mcp.types import Tool as MCPTool import litellm.experimental_mcp_client.client as mcp_client_module @@ -1604,11 +1610,13 @@ class TestListToolsRestAPI: mock_session_ctx = AsyncMock() mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock(return_value=InitializeResult( - protocol_version="2025-11-25", - capabilities=ServerCapabilities(), - server_info=Implementation(name="stub", version="1"), - )) + mock_session_instance.initialize = AsyncMock( + return_value=InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="stub", version="1"), + ) + ) mock_session_instance.list_tools.side_effect = [ ListToolsResult( tools=[ @@ -2556,10 +2564,23 @@ class TestCallToolRestAPI: assert logged_data["metadata"]["model_group"] == "MCP: demo-tool" logging_obj = failure_log.await_args.args[0] assert logging_obj.model_call_details["mcp_tool_call_metadata"] == { - "name": "demo-tool", "mcp_server_name": "Example server", + "name": "demo-tool", + "mcp_server_name": "Example server", } - async def test_executes_tool_when_allowed(self, monkeypatch): + @pytest.mark.parametrize( + ("tool_name", "native_tool_name"), + ( + ("math_a-add", "add"), + ("add", "add"), + ("math_a-math_a-foo", "math_a-foo"), + ), + ) + async def test_executes_tool_when_allowed(self, monkeypatch, tool_name, native_tool_name): + import litellm + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler, MCPServerAccess + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + async def fake_contexts(user_api_key_auth): return [user_api_key_auth] @@ -2568,8 +2589,9 @@ class TestCallToolRestAPI: class StubServer: server_id = "server-1" - alias = "server-1" - server_name = "server-1" + alias = "math_a" + server_name = "math_a" + tool_name_to_display_name = None name = "stub" allowed_tools = None mcp_info = {"server_name": "stub"} @@ -2577,6 +2599,17 @@ class TestCallToolRestAPI: auth_type = None stub_server = StubServer() + original_auth: Final = UserAPIKeyAuth(api_key="sk-personal", user_id="owner") + source_auth: Final = original_auth.model_copy(update={"team_id": "team-a", "org_id": "org-a"}) + + async def fake_attributing_source( + auth: UserAPIKeyAuth, + server_id: str, + tool_name: str, + grants: tuple[tuple[UserAPIKeyAuth, MCPServerAccess], ...] | None = None, + ) -> UserAPIKeyAuth | None: + captured["attributed_tool_name"] = tool_name + return source_auth if server_id == "server-1" and tool_name == native_tool_name else None async def fake_add_litellm_data_to_request(**kwargs): return kwargs.get("data", {}) @@ -2585,6 +2618,11 @@ class TestCallToolRestAPI: async def fake_execute_mcp_tool(**kwargs): captured.update(kwargs) + captured["executed_tool_name"] = mcp_operations.resolve_native_tool_name_for_server( + kwargs["name"], + stub_server, + kwargs["allowed_mcp_servers"], + ) return _OK_TOOL_RESULT monkeypatch.setattr( @@ -2605,6 +2643,13 @@ class TestCallToolRestAPI: lambda server_id: stub_server if server_id == "server-1" else None, raising=False, ) + monkeypatch.setattr( + global_mcp_server_manager, + "get_mcp_server_from_tool_name", + lambda tool_name: stub_server if tool_name.startswith("math_a-") else None, + ) + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(MCPRequestHandler, "personal_key_attributing_source", fake_attributing_source) monkeypatch.setattr( "litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request, @@ -2631,7 +2676,7 @@ class TestCallToolRestAPI: request_payload = { "server_id": "server-1", - "name": "demo-tool", + "name": tool_name, "arguments": {"foo": "bar"}, } request = _build_request( @@ -2642,16 +2687,84 @@ class TestCallToolRestAPI: result = await rest_endpoints.call_tool_rest_api( request, - user_api_key_dict=UserAPIKeyAuth(), + user_api_key_dict=original_auth, ) assert result == _OK_TOOL_RESULT - assert captured["name"] == "demo-tool" + assert captured["name"] == tool_name + assert captured["executed_tool_name"] == native_tool_name + assert captured["attributed_tool_name"] == native_tool_name assert captured["arguments"] == {"foo": "bar"} assert captured["allowed_mcp_servers"] == [stub_server] assert captured["oauth2_headers"] is None + assert captured["user_api_key_auth"].team_id is None + assert captured["user_api_key_auth"].org_id is None + assert captured["guardrail_auth"].team_id == "team-a" + assert captured["guardrail_auth"].org_id == "org-a" + assert fire_logging.await_args.kwargs["user_api_key_auth"] is captured["guardrail_auth"] + logging_obj = fire_logging.await_args.args[0] + assert logging_obj.model_call_details["litellm_params"]["metadata"]["user_api_key_team_id"] == "team-a" fire_logging.assert_awaited_once() + @pytest.mark.asyncio + async def test_rest_billing_uses_target_server_when_global_tool_map_is_empty( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + import litellm + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable + from litellm.proxy.auth import auth_checks + + target_server: Final = MCPServer( + server_id="server-b", + name="server-b", + server_name="server-b", + transport=MCPTransport.http, + ) + team: Final = LiteLLM_TeamTable( + team_id="team-b", + object_permission_id="perm-team-b", + organization_id="org-team-b", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-team-b", + mcp_servers=["server-b"], + mcp_tool_permissions={"server-b": ["search"]}, + ), + ) + auth: Final = UserAPIKeyAuth(api_key="sk-personal", user_id="alice") + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + assert team_id == "team-b" + return team + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(return_value=LiteLLM_UserTable(user_id="alice", teams=["team-b"])), + ) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) + monkeypatch.setitem(global_mcp_server_manager.registry, target_server.server_id, target_server) + monkeypatch.setattr(global_mcp_server_manager, "tool_name_to_mcp_server_name_mapping", {}) + + resolved_server, billed = await rest_endpoints._rest_tool_call_server_and_auth( + auth, + (target_server,), + target_server.server_id, + "search", + ) + + assert resolved_server is target_server + assert billed.team_id == "team-b" + async def test_extract_mcp_headers_scrubs_admitted_caller_credential(self) -> None: caller_key: Final = "sk-rest-caller-admission-key-123" upstream_token: Final = "Bearer unrelated-upstream-token" @@ -2736,6 +2849,7 @@ class TestCallToolRestAPI: server_id = "server-1" alias = "server-1" server_name = "server-1" + tool_name_to_display_name = None name = "stub" allowed_tools = None mcp_info = {"server_name": "stub"} @@ -2813,6 +2927,7 @@ class TestCallToolRestAPI: server_id = "server-1" alias = "server-1" server_name = "server-1" + tool_name_to_display_name = None name = "stub" allowed_tools = None mcp_info = {"server_name": "stub"} @@ -2876,6 +2991,7 @@ class TestCallToolRestAPI: server_id = "server-1" alias = "server-1" server_name = "server-1" + tool_name_to_display_name = None name = "stub" allowed_tools = None mcp_info = {"server_name": "stub"} @@ -2943,6 +3059,7 @@ class TestCallToolRestAPI: server_id = "server-1" alias = "server-1" server_name = "server-1" + tool_name_to_display_name = None name = "stub" allowed_tools = None mcp_info = {"server_name": "stub"} @@ -3036,6 +3153,7 @@ class TestCallToolRestAPI: server_id = "server-1" alias = "server-1" server_name = "server-1" + tool_name_to_display_name = None name = "stub" allowed_tools = None mcp_info = {"server_name": "stub"} @@ -3123,6 +3241,7 @@ class TestCallToolRestAPI: server_id = "server-1" alias = "server-1" server_name = "server-1" + tool_name_to_display_name = None name = "stub" allowed_tools = None mcp_info = {"server_name": "stub"} @@ -3192,6 +3311,7 @@ class TestCallToolRestAPI: server_id = "server-1" alias = "server-1" server_name = "server-1" + tool_name_to_display_name = None name = "stub" allowed_tools = None mcp_info = {"server_name": "stub"} @@ -3291,19 +3411,29 @@ class TestCallToolRestAPI: from litellm.proxy import proxy_server guardrail_error = HTTPException(status_code=400, detail={"error": "Content blocked"}) + server = MagicMock() + server.server_id = "server-1" + server.alias = "server-1" + server.server_name = "server-1" + server.short_prefix = None + server.tool_name_to_display_name = None + server.name = "stub" async def fake_add_litellm_data_to_request(**kwargs): return kwargs.get("data", {}) - async def blocking_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): - raise guardrail_error - failure_logging = AsyncMock(side_effect=RuntimeError("spend log db down")) + execute_tool = AsyncMock(side_effect=guardrail_error) + monkeypatch.setattr( + rest_endpoints, + "_resolve_allowed_mcp_servers_with_ip_filter", + AsyncMock(return_value=([server], "server-1")), + ) + monkeypatch.setattr(rest_endpoints, "execute_mcp_tool", execute_tool) monkeypatch.setattr( proxy_server, "add_litellm_data_to_request", fake_add_litellm_data_to_request, raising=False ) monkeypatch.setattr(proxy_server, "proxy_config", {}, raising=False) - monkeypatch.setattr(proxy_server.proxy_logging_obj, "pre_call_hook", blocking_pre_call_hook) monkeypatch.setattr(rest_endpoints, "fire_mcp_tool_call_failure_logging", failure_logging, raising=False) request = _build_request( @@ -3318,6 +3448,7 @@ class TestCallToolRestAPI: ) assert exc_info.value is guardrail_error + execute_tool.assert_awaited_once() failure_logging.assert_awaited_once() async def test_success_logging_cancellation_propagates(self, monkeypatch): @@ -3792,7 +3923,11 @@ class TestGetToolsForSingleServer: from litellm.proxy._experimental.mcp_server.server import MCPServer from litellm.types.mcp import MCPTransport - mock_tools = [Tool(name="tool1", description="tool1", inputSchema={}), Tool(name="tool2", description="tool2", inputSchema={}), Tool(name="tool3", description="tool3", inputSchema={})] + mock_tools = [ + Tool(name="tool1", description="tool1", inputSchema={}), + Tool(name="tool2", description="tool2", inputSchema={}), + Tool(name="tool3", description="tool3", inputSchema={}), + ] async def fake_get_tools_from_server(**kwargs): return mock_tools @@ -4531,6 +4666,7 @@ class TestToolResponseMcpInfoEnrichment: assert response.model_dump(by_alias=True, exclude={"mcp_info"}) == original assert tool.model_dump(by_alias=True) == original + class TestRestListToolsetFiltering: @pytest.mark.asyncio async def test_rest_list_filters_toolset_only_key_to_toolset_tools(self, monkeypatch): @@ -4799,8 +4935,11 @@ async def test_preview_client_honors_protocol_metadata(revision: MCPUpstreamProt from litellm.experimental_mcp_client.client import MCPClient payload: Final = NewMCPServerRequest( - server_name="preview", url="http://127.0.0.1:9/mcp", transport="http", - auth_type=MCPAuth.none, mcp_info={"protocol_version": revision}, + server_name="preview", + url="http://127.0.0.1:9/mcp", + transport="http", + auth_type=MCPAuth.none, + mcp_info={"protocol_version": revision}, ) async def inspect_client(client: MCPClient) -> dict[str, str]: @@ -4824,8 +4963,10 @@ async def test_preview_client_honors_protocol_metadata(revision: MCPUpstreamProt ), ) async def test_saved_preview_protocol_omission_and_explicit_edits( - monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuth, - metadata: dict[str, str] | None, expected: MCPUpstreamProtocol, + monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuth, + metadata: dict[str, str] | None, + expected: MCPUpstreamProtocol, ) -> None: from starlette.datastructures import Headers @@ -4834,19 +4975,29 @@ async def test_saved_preview_protocol_omission_and_explicit_edits( from litellm.proxy.management_endpoints import mcp_management_endpoints saved: Final = MCPServer( - server_id="saved-preview", name="preview", url="https://example.com/mcp", - transport="http", auth_type=auth_type, protocol_version="2025-11-25", + server_id="saved-preview", + name="preview", + url="https://example.com/mcp", + transport="http", + auth_type=auth_type, + protocol_version="2025-11-25", authentication_token="stored-token", - authorization_url="https://example.com/authorize", token_url="https://example.com/token", + authorization_url="https://example.com/authorize", + token_url="https://example.com/token", ) manager: Final = MCPServerManager() manager.registry = {saved.server_id: saved} monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) payload: Final = NewMCPServerRequest( - server_id=saved.server_id, server_name=saved.name, url=saved.url, transport="http", - auth_type=auth_type, mcp_info=metadata, - authorization_url=saved.authorization_url, token_url=saved.token_url, + server_id=saved.server_id, + server_name=saved.name, + url=saved.url, + transport="http", + auth_type=auth_type, + mcp_info=metadata, + authorization_url=saved.authorization_url, + token_url=saved.token_url, ) staged: Final = rest_endpoints._stage_server_test( payload, Headers({"x-litellm-api-key": "sk-admin", "authorization": "Bearer preview-token"}) @@ -4856,8 +5007,10 @@ async def test_saved_preview_protocol_omission_and_explicit_edits( return {"protocol_version": client.protocol_version} result: Final = await rest_endpoints._execute_with_mcp_client( - staged.request, inspect_client, - mcp_auth_header=staged.mcp_auth_header, oauth2_headers=staged.oauth2_headers, + staged.request, + inspect_client, + mcp_auth_header=staged.mcp_auth_header, + oauth2_headers=staged.oauth2_headers, ) assert result == {"protocol_version": expected} assert saved.protocol_version == "2025-11-25" @@ -4866,14 +5019,20 @@ async def test_saved_preview_protocol_omission_and_explicit_edits( @pytest.mark.asyncio @pytest.mark.parametrize("virtual", [False, True]) -@pytest.mark.parametrize("bearer,claims,expected", [ - ("sk-guardrail-admission", None, None), - ("caller.jwt.assertion", {"sub": "caller"}, "caller.jwt.assertion"), - ("upstream-token", None, "upstream-token"), -]) +@pytest.mark.parametrize( + "bearer,claims,expected", + [ + ("sk-guardrail-admission", None, None), + ("caller.jwt.assertion", {"sub": "caller"}, "caller.jwt.assertion"), + ("upstream-token", None, "upstream-token"), + ], +) async def test_rest_guardrail_bearer_excludes_admission_key( - monkeypatch: pytest.MonkeyPatch, virtual: bool, - bearer: str, claims: dict[str, str] | None, expected: str | None, + monkeypatch: pytest.MonkeyPatch, + virtual: bool, + bearer: str, + claims: dict[str, str] | None, + expected: str | None, ) -> None: import litellm from litellm.caching.dual_cache import DualCache @@ -4884,8 +5043,11 @@ async def test_rest_guardrail_bearer_excludes_admission_key( class CaptureBearer(CustomLogger): async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, - data: dict[str, object], call_type: str, + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: str, ) -> None: if call_type == "call_mcp_tool" and "incoming_bearer_token" in data: observed.put_nowait(data["incoming_bearer_token"]) @@ -4894,23 +5056,36 @@ async def test_rest_guardrail_bearer_excludes_admission_key( return "guarded-result" auth: Final = UserAPIKeyAuth( - api_key="sk-guardrail-admission", jwt_claims=claims, - object_permission={"object_permission_id": "guarded", "mcp_servers": ["guarded"], "mcp_tool_search_enabled": virtual}, + api_key="sk-guardrail-admission", + jwt_claims=claims, + object_permission={ + "object_permission_id": "guarded", + "mcp_servers": ["guarded"], + "mcp_tool_search_enabled": virtual, + }, ) managed: Final = MCPServer( - server_id="guarded", name="guarded", transport=MCPTransport.http, - spec_path="/catalog.yaml", allow_all_keys=True, + server_id="guarded", + name="guarded", + transport=MCPTransport.http, + spec_path="/catalog.yaml", + allow_all_keys=True, ) monkeypatch.setitem(rest_endpoints.global_mcp_server_manager.registry, managed.server_id, managed) monkeypatch.setattr(litellm, "callbacks", [CaptureBearer()]) global_mcp_tool_registry.register_tool( - name="guarded-echo", description="Echo", input_schema={"type": "object"}, handler=echo, + name="guarded-echo", + description="Echo", + input_schema={"type": "object"}, + handler=echo, ) - headers: Final = {"authorization": f"Bearer {bearer}", **( - {} if claims else {"x-litellm-api-key": "sk-guardrail-admission"} - )} + headers: Final = { + "authorization": f"Bearer {bearer}", + **({} if claims else {"x-litellm-api-key": "sk-guardrail-admission"}), + } payload: Final = { - "server_id": "guarded", "name": "mcp_tool_call" if virtual else "guarded-echo", + "server_id": "guarded", + "name": "mcp_tool_call" if virtual else "guarded-echo", "arguments": {"tool_name": "guarded-echo", "arguments": {}} if virtual else {}, } try: @@ -4923,7 +5098,8 @@ async def test_rest_guardrail_bearer_excludes_admission_key( assert observed.get_nowait() == expected if claims is None: list_headers: Final = rest_endpoints._extract_mcp_headers_from_request( - _build_request(headers), auth, + _build_request(headers), + auth, )[2] authorization: Final = list_headers.get("authorization", "") assert (authorization.removeprefix("Bearer ") or None) == expected diff --git a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 45d3d7fb22e..4692088e536 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -10,7 +10,7 @@ from unittest.mock import AsyncMock, patch import pytest from litellm.constants import UI_SESSION_TOKEN_TEAM_ID -from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import UI_TEAM_ID, LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( @@ -76,71 +76,39 @@ class TestAgentRequestHandler: ) # Case 1: Both key and team have agents - intersection - with patch.object( - AgentRequestHandler, "get_allowed_agents_for_key" - ) as mock_key: - with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_team" - ) as mock_team: - mock_key.return_value = RestrictedAgentAccess( - frozenset({"agent1", "agent2", "agent3"}) - ) - mock_team.return_value = RestrictedAgentAccess( - frozenset({"agent2", "agent4"}) - ) + with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + mock_key.return_value = RestrictedAgentAccess(frozenset({"agent1", "agent2", "agent3"})) + mock_team.return_value = RestrictedAgentAccess(frozenset({"agent2", "agent4"})) - result = await AgentRequestHandler.resolve_agent_access( - user_api_key_auth=mock_user_auth - ) + result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=mock_user_auth) assert result == RestrictedAgentAccess(frozenset({"agent2"})) # Case 2: Team has agents, key has none - inherit from team - with patch.object( - AgentRequestHandler, "get_allowed_agents_for_key" - ) as mock_key: - with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_team" - ) as mock_team: + with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: mock_key.return_value = UnrestrictedAgentAccess() - mock_team.return_value = RestrictedAgentAccess( - frozenset({"team_agent1", "team_agent2"}) - ) + mock_team.return_value = RestrictedAgentAccess(frozenset({"team_agent1", "team_agent2"})) - result = await AgentRequestHandler.resolve_agent_access( - user_api_key_auth=mock_user_auth - ) - assert result == RestrictedAgentAccess( - frozenset({"team_agent1", "team_agent2"}) - ) + result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=mock_user_auth) + assert result == RestrictedAgentAccess(frozenset({"team_agent1", "team_agent2"})) # Case 3: Key has agents, team has none - key restrictions stand - with patch.object( - AgentRequestHandler, "get_allowed_agents_for_key" - ) as mock_key: - with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_team" - ) as mock_team: + with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: mock_key.return_value = RestrictedAgentAccess(frozenset({"key_agent1"})) mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.resolve_agent_access( - user_api_key_auth=mock_user_auth - ) + result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=mock_user_auth) assert result == RestrictedAgentAccess(frozenset({"key_agent1"})) # Case 4: No grant anywhere - unrestricted (documented open-by-default) - with patch.object( - AgentRequestHandler, "get_allowed_agents_for_key" - ) as mock_key: - with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_team" - ) as mock_team: + with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: mock_key.return_value = UnrestrictedAgentAccess() mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.resolve_agent_access( - user_api_key_auth=mock_user_auth - ) + result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=mock_user_auth) assert result == UnrestrictedAgentAccess() async def test_agent_access_default_deny_is_opt_in(self): @@ -317,9 +285,7 @@ class TestAgentRequestHandler: """LIT-5143: a key restricted to one agent inside a team restricted to another must reach nothing. The empty intersection used to read as "no restrictions", so adding the team grant handed the key every agent on the proxy.""" - mock_user_auth: Final = UserAPIKeyAuth( - api_key="test-key", user_id="test-user", team_id="test-team" - ) + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="test-team") with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key: with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: @@ -332,9 +298,7 @@ class TestAgentRequestHandler: for agent_id in ("agent-alpha", "agent-beta", "agent-secret"): assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id=agent_id, user_api_key_auth=mock_user_auth - ) + await AgentRequestHandler.is_agent_allowed(agent_id=agent_id, user_api_key_auth=mock_user_auth) is False ), agent_id @@ -496,9 +460,7 @@ class TestAgentRequestHandler: agent_access_groups=["group-with-no-agents"], ) - with patch.object( - AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock - ) as mock_groups: + with patch.object(AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock) as mock_groups: mock_groups.return_value = [] assert await AgentRequestHandler.get_allowed_agents_for_key( @@ -506,20 +468,17 @@ class TestAgentRequestHandler: ) == RestrictedAgentAccess(frozenset()) assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id="agent-secret", user_api_key_auth=mock_user_auth - ) + await AgentRequestHandler.is_agent_allowed(agent_id="agent-secret", user_api_key_auth=mock_user_auth) is False ) - with patch.object( - AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock - ) as mock_groups: + with patch.object(AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock) as mock_groups: mock_groups.side_effect = Exception("DB Error") - assert await AgentRequestHandler.get_allowed_agents_for_key( - user_api_key_auth=mock_user_auth - ) == UnrestrictedAgentAccess() + assert ( + await AgentRequestHandler.get_allowed_agents_for_key(user_api_key_auth=mock_user_auth) + == UnrestrictedAgentAccess() + ) async def test_is_agent_allowed_respects_permissions(self): """ @@ -529,54 +488,32 @@ class TestAgentRequestHandler: mock_user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") # Agent in allowed list - should be allowed - with patch.object( - AgentRequestHandler, "resolve_agent_access" - ) as mock_get_allowed: - mock_get_allowed.return_value = RestrictedAgentAccess( - frozenset({"agent1", "agent2"}) - ) + with patch.object(AgentRequestHandler, "resolve_agent_access") as mock_get_allowed: + mock_get_allowed.return_value = RestrictedAgentAccess(frozenset({"agent1", "agent2"})) assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id="agent1", user_api_key_auth=mock_user_auth - ) - is True + await AgentRequestHandler.is_agent_allowed(agent_id="agent1", user_api_key_auth=mock_user_auth) is True ) # Agent not in allowed list - should be denied - with patch.object( - AgentRequestHandler, "resolve_agent_access" - ) as mock_get_allowed: - mock_get_allowed.return_value = RestrictedAgentAccess( - frozenset({"agent1", "agent2"}) - ) + with patch.object(AgentRequestHandler, "resolve_agent_access") as mock_get_allowed: + mock_get_allowed.return_value = RestrictedAgentAccess(frozenset({"agent1", "agent2"})) assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id="agent3", user_api_key_auth=mock_user_auth - ) - is False + await AgentRequestHandler.is_agent_allowed(agent_id="agent3", user_api_key_auth=mock_user_auth) is False ) # Restricted to nothing - should deny every agent - with patch.object( - AgentRequestHandler, "resolve_agent_access" - ) as mock_get_allowed: + with patch.object(AgentRequestHandler, "resolve_agent_access") as mock_get_allowed: mock_get_allowed.return_value = RestrictedAgentAccess(frozenset()) assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id="any_agent", user_api_key_auth=mock_user_auth - ) + await AgentRequestHandler.is_agent_allowed(agent_id="any_agent", user_api_key_auth=mock_user_auth) is False ) # Unrestricted - should allow any agent - with patch.object( - AgentRequestHandler, "resolve_agent_access" - ) as mock_get_allowed: + with patch.object(AgentRequestHandler, "resolve_agent_access") as mock_get_allowed: mock_get_allowed.return_value = UnrestrictedAgentAccess() assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id="any_agent", user_api_key_auth=mock_user_auth - ) + await AgentRequestHandler.is_agent_allowed(agent_id="any_agent", user_api_key_auth=mock_user_auth) is True ) @@ -587,9 +524,7 @@ class TestAgentRequestHandler: result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=None) assert result == UnrestrictedAgentAccess() - is_allowed = await AgentRequestHandler.is_agent_allowed( - agent_id="any_agent", user_api_key_auth=None - ) + is_allowed = await AgentRequestHandler.is_agent_allowed(agent_id="any_agent", user_api_key_auth=None) assert is_allowed is True async def test_resolve_agent_access_handles_errors_gracefully(self): @@ -605,18 +540,12 @@ class TestAgentRequestHandler: object_permission_id="test-permission", ) - with patch.object( - AgentRequestHandler, "get_allowed_agents_for_key" - ) as mock_key: - with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_team" - ) as mock_team: + with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: mock_key.side_effect = Exception("DB Error") mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.resolve_agent_access( - user_api_key_auth=mock_user_auth - ) + result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=mock_user_auth) assert result == UnrestrictedAgentAccess() async def test_accessible_agents_hides_ungranted_agents_from_non_admins(self): @@ -702,20 +631,14 @@ class TestAgentRequestHandler: access_group_ids=["ag-with-agents"], ) - with patch.object( - AgentRequestHandler, "_get_key_object_permission", return_value=None - ): + with patch.object(AgentRequestHandler, "_get_key_object_permission", return_value=None): with patch( "litellm.proxy.auth.auth_checks.get_agent_ids_from_access_groups", new_callable=AsyncMock, return_value=["agent-from-ag-1", "agent-from-ag-2"], ): - result = await AgentRequestHandler.get_allowed_agents_for_key( - user_api_key_auth=mock_user_auth - ) - assert result == RestrictedAgentAccess( - frozenset({"agent-from-ag-1", "agent-from-ag-2"}) - ) + result = await AgentRequestHandler.get_allowed_agents_for_key(user_api_key_auth=mock_user_auth) + assert result == RestrictedAgentAccess(frozenset({"agent-from-ag-1", "agent-from-ag-2"})) async def testget_allowed_agents_for_key_combines_native_and_access_groups(self): """ @@ -742,12 +665,8 @@ class TestAgentRequestHandler: new_callable=AsyncMock, return_value=["agent-from-ag"], ): - result = await AgentRequestHandler.get_allowed_agents_for_key( - user_api_key_auth=mock_user_auth - ) - assert result == RestrictedAgentAccess( - frozenset({"agent-from-ag", "native-agent-1"}) - ) + result = await AgentRequestHandler.get_allowed_agents_for_key(user_api_key_auth=mock_user_auth) + assert result == RestrictedAgentAccess(frozenset({"agent-from-ag", "native-agent-1"})) async def test_is_agent_allowed_accepts_legacy_config_agent_id_grants(self): """LIT-5144: object_permission grants stored under the pre-fix full-entry hash @@ -902,9 +821,17 @@ async def test_managed_agent_invocation_grants_intersect_verified_user_grants( assert access == RestrictedAgentAccess(frozenset({"shared"} if delegated else {"shared", "agent-only"})) target: Final = AgentResponse( - agent_id="shared", agent_name="Shared", agent_card_params={}, identity_managed=True, + agent_id="shared", + agent_name="Shared", + agent_card_params={}, + identity_managed=True, identity=AgentIdentityBinding( - agent_id="shared", provider="microsoft_entra", tenant_id="tenant", client_id="client", issuer="issuer", revision="current" + agent_id="shared", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + issuer="issuer", + revision="current", ), ) database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=target) @@ -912,7 +839,9 @@ async def test_managed_agent_invocation_grants_intersect_verified_user_grants( @pytest.mark.asyncio -@pytest.mark.parametrize("revoked", ["user", "team-member", "team-grant", "team-permission", "direct-grant", "access-group"]) +@pytest.mark.parametrize( + "revoked", ["user", "team-member", "team-grant", "team-permission", "direct-grant", "access-group"] +) async def test_delegated_grants_revoke_with_warm_user_team_and_permission_caches( monkeypatch: pytest.MonkeyPatch, revoked: str ) -> None: @@ -1003,9 +932,7 @@ async def test_strict_legacy_group_grants_ignore_stale_replica(monkeypatch: pyte frozenset({"revoked"}) ) database.writer_db.litellm_agentstable.find_many.return_value = [] - assert await AgentRequestHandler.get_allowed_agents_for_key(auth, strict=True) == RestrictedAgentAccess( - frozenset() - ) + assert await AgentRequestHandler.get_allowed_agents_for_key(auth, strict=True) == RestrictedAgentAccess(frozenset()) database.db.litellm_agentstable.find_many.assert_not_awaited() @@ -1030,9 +957,9 @@ async def test_strict_invocation_policy_outage_denies_instead_of_allowing_all( monkeypatch.setattr(proxy_server, "prisma_client", database) auth: Final = UserAPIKeyAuth( team_id="team" if team else None, - object_permission=None if team else LiteLLM_ObjectPermissionTable( - object_permission_id="grant", agent_access_groups=["group"] - ), + object_permission=None + if team + else LiteLLM_ObjectPermissionTable(object_permission_id="grant", agent_access_groups=["group"]), ) with pytest.raises(HTTPException, match="policy is unavailable") as denied: await AgentRequestHandler.resolve_key_team_agent_access(auth, strict=True) @@ -1065,9 +992,9 @@ async def test_registered_managed_target_cannot_bypass_missing_or_unavailable_po from litellm.types.agents import AgentResponse registry: Final = AgentRegistry() - registry.register_agent(AgentResponse( - agent_id="target", agent_name="Target", agent_card_params={}, identity_managed=True - )) + registry.register_agent( + AgentResponse(agent_id="target", agent_name="Target", agent_card_params={}, identity_managed=True) + ) monkeypatch.setattr(agent_registry, "global_agent_registry", registry) database: Final = MagicMock() database.writer_db.litellm_agentstable.find_unique = AsyncMock( @@ -1091,7 +1018,9 @@ async def test_delegation_without_a_verified_human_never_grants_agents(grant: bo auth: Final = UserAPIKeyAuth(agent_id="actor") auth.managed_agent_policy = AgentResponse( - agent_id="actor", agent_name="Actor", agent_card_params={}, + agent_id="actor", + agent_name="Actor", + agent_card_params={}, object_permission={"object_permission_id": "own", "agents": ["target"]} if grant else None, ) auth.managed_agent_context = ManagedAgentContext(agent_id="actor", mode="delegated") @@ -1100,7 +1029,9 @@ async def test_delegation_without_a_verified_human_never_grants_agents(grant: bo @pytest.mark.asyncio -@pytest.mark.parametrize("change", ("grant", "permission_reference", "groups", "team", "blocked", "expired", "deleted", "outage")) +@pytest.mark.parametrize( + "change", ("grant", "permission_reference", "groups", "team", "blocked", "expired", "deleted", "outage") +) async def test_managed_target_rechecks_authoritative_key_after_peer_revocation( monkeypatch: pytest.MonkeyPatch, change: str ) -> None: @@ -1112,10 +1043,22 @@ async def test_managed_target_rechecks_authoritative_key_after_peer_revocation( from litellm.types.proxy.agent_identity import AgentIdentityBinding permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["target"]) - warm: Final = UserAPIKeyAuth(api_key="a" * 64, token="a" * 64, object_permission_id="grant", object_permission=permission) + warm: Final = UserAPIKeyAuth( + api_key="a" * 64, token="a" * 64, object_permission_id="grant", object_permission=permission + ) target: Final = AgentResponse( - agent_id="target", agent_name="Target", agent_card_params={}, identity_managed=True, - identity=AgentIdentityBinding(agent_id="target", provider="microsoft_entra", tenant_id="tenant", client_id="client", issuer="issuer", revision="current"), + agent_id="target", + agent_name="Target", + agent_card_params={}, + identity_managed=True, + identity=AgentIdentityBinding( + agent_id="target", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + issuer="issuer", + revision="current", + ), ) client: Final = MagicMock() client.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=target) @@ -1126,26 +1069,38 @@ async def test_managed_target_rechecks_authoritative_key_after_peer_revocation( monkeypatch.setattr(proxy_server, "prisma_client", client) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) assert await AgentRequestHandler.is_agent_allowed("target", warm) is True - client.get_data.return_value = warm.model_copy(update={ - "object_permission": None, - "object_permission_id": "replacement" if change == "permission_reference" else "grant", - "access_group_ids": [], - "team_id": "new-team" if change == "team" else None, - "blocked": change == "blocked", - "expires": "2000-01-01T00:00:00+00:00" if change == "expired" else None, - }) - client.writer_db.litellm_objectpermissiontable.find_unique.return_value = permission.model_copy(update={"agents": []}) + client.get_data.return_value = warm.model_copy( + update={ + "object_permission": None, + "object_permission_id": "replacement" if change == "permission_reference" else "grant", + "access_group_ids": [], + "team_id": "new-team" if change == "team" else None, + "blocked": change == "blocked", + "expires": "2000-01-01T00:00:00+00:00" if change == "expired" else None, + } + ) + client.writer_db.litellm_objectpermissiontable.find_unique.return_value = permission.model_copy( + update={"agents": []} + ) if change == "team": from litellm.proxy._types import LiteLLM_TeamTable from litellm.proxy.auth import auth_checks + client.writer_db.litellm_objectpermissiontable.find_unique.return_value = permission - monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=LiteLLM_TeamTable( - team_id="new-team", object_permission=permission.model_copy(update={"agents": ["other"]}) - ))) + monkeypatch.setattr( + auth_checks, + "get_team_object", + AsyncMock( + return_value=LiteLLM_TeamTable( + team_id="new-team", object_permission=permission.model_copy(update={"agents": ["other"]}) + ) + ), + ) if change == "groups": warm.object_permission = None warm.access_group_ids = ["old-group"] from litellm.proxy.auth import auth_checks + monkeypatch.setattr(auth_checks, "get_agent_ids_from_access_groups", AsyncMock(return_value=["target"])) if change == "deleted": client.get_data.return_value = None @@ -1174,14 +1129,23 @@ async def test_managed_target_preserves_ordinary_actor_ceilings_after_key_reload from litellm.types.proxy.agent_identity import AgentIdentityBinding target: Final = AgentResponse( - agent_id="target", agent_name="Target", agent_card_params={}, identity_managed=True, + agent_id="target", + agent_name="Target", + agent_card_params={}, + identity_managed=True, identity=AgentIdentityBinding( - agent_id="target", provider="microsoft_entra", tenant_id="tenant", client_id="client", - issuer="issuer", revision="current", + agent_id="target", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + issuer="issuer", + revision="current", ), ) actor: Final = AgentResponse( - agent_id="ordinary", agent_name="Ordinary", agent_card_params={}, + agent_id="ordinary", + agent_name="Ordinary", + agent_card_params={}, access_group_ids=["actor-group"] if ceiling != "caller-team" else [], ) registry: Final = AgentRegistry() @@ -1191,18 +1155,24 @@ async def test_managed_target_preserves_ordinary_actor_ceilings_after_key_reload object_permission_id="key-grant", agents=[] if ceiling == "group-without-grant" else ["target"] ) persisted: Final = UserAPIKeyAuth( - api_key="a" * 64, agent_id="ordinary", object_permission_id="key-grant", object_permission=permission, + api_key="a" * 64, + agent_id="ordinary", + object_permission_id="key-grant", + object_permission=permission, ) auth: Final = persisted.model_copy() auth.agent_caller = AgentCaller(team_id="caller-team") if ceiling == "caller-team" else None group: Final = LiteLLM_AccessGroupTable( - access_group_id="actor-group", access_group_name="Actor group", + access_group_id="actor-group", + access_group_name="Actor group", access_agent_ids=["target"] if permitted else ["other"], ) team: Final = LiteLLM_TeamTable( - team_id="caller-team", object_permission_id="caller-grant", + team_id="caller-team", + object_permission_id="caller-grant", object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id="caller-grant", agents=["target"] if permitted else ["other"], + object_permission_id="caller-grant", + agents=["target"] if permitted else ["other"], ), ) database: Final = MagicMock() @@ -1220,7 +1190,9 @@ async def test_managed_target_preserves_ordinary_actor_ceilings_after_key_reload monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) monkeypatch.setattr(agent_registry, "global_agent_registry", registry) - assert await AgentRequestHandler.is_agent_allowed("target", auth) is (permitted and ceiling != "group-without-grant") + assert await AgentRequestHandler.is_agent_allowed("target", auth) is ( + permitted and ceiling != "group-without-grant" + ) database.get_data.assert_awaited_once() assert auth.agent_caller == (AgentCaller(team_id="caller-team") if ceiling == "caller-team" else None) @@ -1319,15 +1291,24 @@ async def test_target_authorization_uses_live_policy_despite_stale_unmanaged_reg registry.register_agent(stale) monkeypatch.setattr(agent_registry, "global_agent_registry", registry) binding: Final = AgentIdentityBinding( - agent_id="target", provider="microsoft_entra", tenant_id="tenant", client_id="client", - issuer="issuer", revision="current", + agent_id="target", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + issuer="issuer", + revision="current", + ) + current: Final = stale.model_copy( + update={ + "identity_managed": managed, + "identity": binding if managed else None, + "enabled": enabled, + } ) - current: Final = stale.model_copy(update={ - "identity_managed": managed, "identity": binding if managed else None, "enabled": enabled, - }) database: Final = MagicMock() database.writer_db.litellm_agentstable.find_unique = AsyncMock( - return_value=current, side_effect=ConnectionError("writer unavailable") if outage else None, + return_value=current, + side_effect=ConnectionError("writer unavailable") if outage else None, ) monkeypatch.setattr(proxy_server, "prisma_client", database) permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["target"]) @@ -1339,3 +1320,927 @@ async def test_target_authorization_uses_live_policy_despite_stale_unmanaged_reg assert denied.value.status_code == 503 return assert await AgentRequestHandler.is_agent_allowed("target", auth) is allowed + + +_PERSONAL_KEY_TEAM_AGENTS: Final = { + "team-a": ("agent-a",), + "team-b": ("agent-b",), + "team-x": ("agent-x",), + "team-n": ("agent-n",), + "team-open": (), + "team-grouped": (), + "team-unreadable": (), + UI_TEAM_ID: (), +} + + +@pytest.fixture +def personal_key_owner_teams(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> UserAPIKeyAuth: + """A personal key whose owner is on ``request.param`` teams; team-a grants agent-a, team-b grants + agent-b, team-open grants no agent rule at all, team-grouped declares an access group whose + lookup fails against the mocked database and team-unreadable names a permission row that did not + load.""" + from unittest.mock import MagicMock + + import litellm + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + if team_id == "team-unreadable": + return LiteLLM_TeamTable(team_id=team_id, object_permission_id=f"perm-{team_id}") + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + agent_access_groups=["grp"] if team_id == "team-grouped" else [], + ), + ) + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + prisma_client: Final = MagicMock() + prisma_client.writer_db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + user_api_key_cache: Final = MagicMock() + user_api_key_cache.async_set_cache = AsyncMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(return_value=LiteLLM_UserTable(user_id="alice", teams=list(request.param))), + ) + return UserAPIKeyAuth(api_key="sk-personal", user_id="alice", team_id=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "personal_key_owner_teams, expected", + [ + (("team-a", "team-b"), RestrictedAgentAccess(frozenset({"agent-a", "agent-b"}))), + (("team-b",), RestrictedAgentAccess(frozenset({"agent-b"}))), + ((), RestrictedAgentAccess(frozenset())), + (("team-a", "team-open"), RestrictedAgentAccess(frozenset({"agent-a"}))), + ], + indirect=["personal_key_owner_teams"], +) +async def test_personal_key_agents_are_the_union_of_its_owner_team_grants( + personal_key_owner_teams: UserAPIKeyAuth, expected: AgentAccess +) -> None: + assert await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("", "team-open")], indirect=True) +async def test_empty_owner_team_id_does_not_restrict_agent_access( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + if not team_id: + raise RuntimeError("empty team id should not be loaded") + return LiteLLM_TeamTable( + team_id=team_id, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=[], + ), + ) + + team_loader: Final = AsyncMock(side_effect=_get_team_object) + monkeypatch.setattr(auth_checks, "get_team_object", team_loader) + + owner_teams: Final = await auth_checks.personal_key_owner_teams_with_status( + valid_token=personal_key_owner_teams, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + + assert owner_teams.unreadable is False + assert tuple(team.team_id for team in owner_teams.teams) == ("team-open",) + assert await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == UnrestrictedAgentAccess() + assert all(call.kwargs["team_id"] == "team-open" for call in team_loader.await_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-open")], indirect=True) +async def test_personal_key_accessible_agents_keep_declared_grants( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.types.agents import AgentResponse + + agents: Final = ( + AgentResponse(agent_id="agent-a", agent_name="agent-a", agent_card_params={}), + AgentResponse(agent_id="agent-b", agent_name="agent-b", agent_card_params={}), + ) + accessible: Final = await accessible_agents( + personal_key_owner_teams, all_agents=agents, effective_contexts=_single_context + ) + assert tuple(agent.agent_id for agent in accessible) == ("agent-a",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-open")], indirect=True) +async def test_personal_key_strict_agent_access_keeps_declared_grants( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + assert await AgentRequestHandler.resolve_agent_access( + personal_key_owner_teams, strict=True + ) == RestrictedAgentAccess(frozenset({"agent-a"})) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-x", "team-n")], indirect=True) +async def test_personal_key_agent_access_ignores_owner_teams_in_other_org( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + organization_id: Final = "org-1" if team_id == "team-a" else "org-2" if team_id == "team-x" else None + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=organization_id, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + auth: Final = personal_key_owner_teams.model_copy(update={"org_id": "org-1"}) + assert await AgentRequestHandler.resolve_agent_access(auth) == RestrictedAgentAccess( + frozenset({"agent-a", "agent-n"}) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_agent_billing_uses_the_team_granting_the_agent( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=f"org-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(personal_key_owner_teams, "agent-b") + assert billed is not personal_key_owner_teams + assert (billed.team_id, billed.org_id) == ("team-b", "org-team-b") + assert (personal_key_owner_teams.team_id, personal_key_owner_teams.org_id) == (None, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_org_id,expected_org_id", + [("org-1", None, "org-1"), (None, "org-2", "org-2")], +) +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_agent_billing_preserves_key_organization( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, + key_org_id: str | None, + team_org_id: str | None, + expected_org_id: str, +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=team_org_id if team_id == "team-b" else "org-team-a", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + auth: Final = personal_key_owner_teams.model_copy(update={"org_id": key_org_id}) + + billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(auth, "agent-b") + + assert billed.team_id == "team-b" + assert billed.org_id == expected_org_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_agent_billing_matches_legacy_agent_ids( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy.agent_endpoints import agent_registry + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + entry: Final = { + "agent_name": "legacy-agent", + "agent_card_params": {"name": "Legacy Agent", "url": "http://localhost", "version": "1.0.0"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("legacy-agent") + assert agent is not None + legacy_id: Final = next(iter(registry.ids_for_agent(agent.agent_id) - {agent.agent_id})) + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=f"org-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=[legacy_id] if team_id == "team-b" else ["agent-a"], + ), + ) + + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(personal_key_owner_teams, agent.agent_id) + assert (billed.team_id, billed.org_id) == ("team-b", "org-team-b") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_agent_billing_inherits_adopted_organization_limits( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.proxy._types import LiteLLM_OrganizationTable, LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + organization: Final = LiteLLM_OrganizationTable( + organization_id="org-1", + organization_alias="org-one", + budget_id="budget-1", + created_by="test", + updated_by="test", + metadata={"k": "v"}, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100, rpm_limit=5, tpm_limit=50), + ) + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id="org-1" if team_id == "team-b" else "org-team-a", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + + org_loader: Final = AsyncMock(return_value=organization) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object_for_request", org_loader) + auth: Final = personal_key_owner_teams.model_copy(update={"org_id": None}) + + billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(auth, "agent-b") + + assert ( + billed.org_id, + billed.organization_rpm_limit, + billed.organization_tpm_limit, + billed.organization_metadata, + ) == ("org-1", 5, 50, {"k": "v"}) + assert billed.organization_max_budget == 100 + assert (auth.org_id, auth.organization_rpm_limit, auth.organization_tpm_limit, auth.organization_metadata) == ( + None, + None, + None, + None, + ) + org_loader.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_agent_billing_keeps_team_controls_when_org_lookup_fails( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id="org-1" if team_id == "team-b" else None, + rpm_limit=5 if team_id == "team-b" else None, + max_budget=40 if team_id == "team-b" else None, + metadata={"team_private": True} if team_id == "team-b" else None, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + + org_loader: Final = AsyncMock(side_effect=RuntimeError("org lookup failed")) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object_for_request", org_loader) + + billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(personal_key_owner_teams, "agent-b") + + assert billed is not personal_key_owner_teams + assert billed.team_id == "team-b" + assert billed.team_rpm_limit == 5 + assert billed.team_metadata == {"team_private": True} + assert billed.team_max_budget == 40 + assert billed.org_id == "org-1" + assert personal_key_owner_teams.team_id is None + assert personal_key_owner_teams.team_rpm_limit is None + org_loader.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_agent_billing_keeps_team_controls_when_membership_lookup_fails( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + organization_id=None, + rpm_limit=5 if team_id == "team-b" else None, + max_budget=40 if team_id == "team-b" else None, + metadata={"team_private": True} if team_id == "team-b" else None, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr( + auth_checks, + "get_team_membership", + AsyncMock(side_effect=RuntimeError("membership lookup failed")), + ) + + billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(personal_key_owner_teams, "agent-b") + + assert billed.team_id == "team-b" + assert billed.team_rpm_limit == 5 + assert billed.team_metadata == {"team_private": True} + assert billed.team_max_budget == 40 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "strict, expected_team_rpm, expected_member_rpm, expected_metadata", + [ + pytest.param(False, 2, 3, {"policy_source": "cache"}, id="default-uses-cache"), + pytest.param(True, 7, 8, {"policy_source": "database"}, id="strict-uses-database"), + ], +) +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_agent_billing_uses_fresh_membership_for_strict_policy( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, + strict: bool, + expected_team_rpm: int, + expected_member_rpm: int, + expected_metadata: dict[str, str], +) -> None: + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.team_membership import LiteLLM_TeamMembership + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **kwargs: object) -> LiteLLM_TeamTable: + database_value: Final = bool(kwargs.get("check_db_only")) + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=None, + rpm_limit=7 if database_value else 2, + metadata={"policy_source": "database" if database_value else "cache"}, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + + 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, + team_id=team_id, + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=8 if database_value else 3), + ) + + team_loader: Final = AsyncMock(side_effect=_get_team_object) + monkeypatch.setattr(auth_checks, "get_team_object", team_loader) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(side_effect=_get_team_membership)) + monkeypatch.setattr(auth_checks, "_check_team_member_budget", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + auth: Final = personal_key_owner_teams.model_copy(update={"requires_fresh_policy": strict}) + + billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(auth, "agent-b", strict=strict) + + assert (billed.team_rpm_limit, billed.team_member_rpm_limit, billed.team_metadata) == ( + expected_team_rpm, + expected_member_rpm, + expected_metadata, + ) + assert all(call.kwargs["check_db_only"] is strict for call in team_loader.await_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_agent_billing_carries_team_rate_limits( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.caching.caching import DualCache + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.team_membership import LiteLLM_TeamMembership + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + team_alias=f"alias-{team_id}", + organization_id=f"org-{team_id}", + rpm_limit=1 if team_id == "team-b" else None, + models=["team-only-model"] if team_id == "team-b" else [], + metadata={"team_private": True} if team_id == "team-b" else None, + object_permission_id=f"perm-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + + membership: Final = LiteLLM_TeamMembership( + user_id="alice", + team_id="team-b", + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=2), + ) + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=membership)) + + billed: Final = await AgentRequestHandler.personal_key_agent_billing_auth(personal_key_owner_teams, "agent-b") + assert (billed.team_id, billed.org_id, billed.team_rpm_limit, billed.team_member_rpm_limit) == ( + "team-b", + "org-team-b", + 1, + 2, + ) + assert billed.team_models == personal_key_owner_teams.team_models + assert billed.team_metadata == {"team_private": True} + assert billed.team_object_permission == personal_key_owner_teams.team_object_permission + assert (personal_key_owner_teams.team_id, personal_key_owner_teams.team_rpm_limit) == (None, None) + + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + descriptors: Final = limiter.create_rate_limit_descriptors( + user_api_key_dict=billed, + data={}, + rpm_limit_type=None, + tpm_limit_type=None, + model_has_failures=False, + ) + team_limits: Final = tuple( + (descriptor["key"], descriptor["value"], descriptor["rate_limit"]["requests_per_unit"]) + for descriptor in descriptors + if descriptor["key"] in {"team", "team_member"} + ) + assert team_limits == (("team", "team-b", 1), ("team_member", "team-b:alice", 2)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_agent_billing_auth_returns_non_personal_key_unchanged( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + team_key: Final = personal_key_owner_teams.model_copy(update={"team_id": "team-a"}) + assert await AgentRequestHandler.personal_key_agent_billing_auth(team_key, "agent-a") is team_key + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-open")], indirect=True) +async def test_personal_key_non_strict_agent_check_denies_ungranted_agent( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore + + monkeypatch.setattr(AgentIdentityStore, "agent", AsyncMock(return_value=None)) + assert await AgentRequestHandler.is_agent_allowed("agent-b", personal_key_owner_teams) is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strict", [False, True], ids=["non-strict", "strict"]) +@pytest.mark.parametrize( + "team_a_agents, expected", + [ + (("agent-a",), RestrictedAgentAccess(frozenset({"agent-a"}))), + ((), RestrictedAgentAccess(frozenset())), + ], + ids=["eligible-grant", "eligible-unrestricted"], +) +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_unreadable_withheld_team_only_adds_conservative_agent_restriction( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, + team_a_agents: tuple[str, ...], + expected: AgentAccess, + strict: bool, +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + if team_id == "team-b": + return LiteLLM_TeamTable(team_id=team_id, blocked=True, object_permission_id=f"perm-{team_id}") + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(team_a_agents), + agent_access_groups=[], + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + + assert ( + await AgentRequestHandler._personal_key_owner_team_agent_access(personal_key_owner_teams, strict=strict) + == expected + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_unreadable_eligible_owner_team_still_raises_in_strict_agent_access( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from fastapi import HTTPException + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=team_id, object_permission_id=f"perm-{team_id}") + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + + with pytest.raises(HTTPException, match="policy is unavailable") as denied: + await AgentRequestHandler._personal_key_owner_team_agent_access(personal_key_owner_teams, strict=True) + assert denied.value.status_code == 503 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-open")], indirect=True) +async def test_personal_key_unreadable_owner_team_does_not_make_access_unrestricted( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore + from litellm.proxy.auth import auth_checks + + get_team_object = auth_checks.get_team_object + + async def fail_team_a_lookup(team_id: str, **_: object) -> LiteLLM_TeamTable: + if team_id == "team-a": + raise RuntimeError("team lookup failed") + return await get_team_object( + team_id=team_id, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + parent_otel_span=personal_key_owner_teams.parent_otel_span, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + + monkeypatch.setattr(auth_checks, "get_team_object", fail_team_a_lookup) + monkeypatch.setattr(AgentIdentityStore, "agent", AsyncMock(return_value=None)) + + assert await AgentRequestHandler.resolve_agent_access(personal_key_owner_teams) == RestrictedAgentAccess( + frozenset() + ) + assert await AgentRequestHandler.is_agent_allowed("agent-b", personal_key_owner_teams) is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-open",)], indirect=True) +async def test_personal_key_all_unrestricted_owner_teams_remain_unrestricted( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + assert await AgentRequestHandler.resolve_agent_access(personal_key_owner_teams) == UnrestrictedAgentAccess() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", UI_TEAM_ID)], indirect=True) +async def test_personal_key_ui_team_does_not_expand_agent_access( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + assert await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == RestrictedAgentAccess( + frozenset({"agent-a"}) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_over_budget_team_grants_no_agent_access( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + team: Final = LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=f"org-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=list(_PERSONAL_KEY_TEAM_AGENTS[team_id]), + ), + ) + if team_id == "team-b": + return team.model_copy(update={"max_budget": 1, "spend": 5}) + return team + + async def _spend_from_fallback(*, fallback_spend: float, **_: object) -> float: + return fallback_spend + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "get_current_spend", _spend_from_fallback) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "budget_alerts", AsyncMock()) + assert await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == RestrictedAgentAccess( + frozenset({"agent-a"}) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", [False, True], ids=["over-budget", "blocked"]) +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-open")], indirect=True) +async def test_personal_key_withheld_restricted_owner_team_denies_all_agents( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, + blocked: bool, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + if team_id == "team-open": + return LiteLLM_TeamTable(team_id=team_id) + team: Final = LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=["agent-a"], + ), + ) + if blocked: + return team.model_copy(update={"blocked": True}) + return team.model_copy(update={"max_budget": 1, "spend": 5}) + + async def _spend_from_fallback(*, fallback_spend: float, **_: object) -> float: + return fallback_spend + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "get_current_spend", _spend_from_fallback) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "budget_alerts", AsyncMock()) + + access: Final = await AgentRequestHandler.resolve_agent_access(personal_key_owner_teams) + + assert access == RestrictedAgentAccess(frozenset()) + assert isinstance(access, RestrictedAgentAccess) + assert "agent-a" not in access.agent_ids + assert "agent-b" not in access.agent_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-x", "team-open")], indirect=True) +async def test_personal_key_unrestricted_over_budget_team_stays_unrestricted( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + if team_id == "team-open": + return LiteLLM_TeamTable(team_id=team_id) + return LiteLLM_TeamTable(team_id=team_id, max_budget=1, spend=5) + + async def _spend_from_fallback(*, fallback_spend: float, **_: object) -> float: + return fallback_spend + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "get_current_spend", _spend_from_fallback) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "budget_alerts", AsyncMock()) + + assert await AgentRequestHandler.resolve_agent_access(personal_key_owner_teams) == UnrestrictedAgentAccess() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_personal_key_window_budget_team_is_not_an_owner_grant( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + budget_limits=( + [{"budget_duration": "24h", "max_budget": 5.0, "reset_at": None}] if team_id == "team-b" else [] + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + monkeypatch.setattr(proxy_server, "get_current_spend", AsyncMock(return_value=5.0)) + allowed: Final = await auth_checks.personal_key_owner_teams( + valid_token=personal_key_owner_teams, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + assert tuple(team.team_id for team in allowed) == ("team-a",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_strict_personal_key_agent_access_reads_owner_teams_from_the_database( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + """Invocation (strict) must see a revoke at once, so the owner's teams are read with check_db_only + like the team-key branch, while listing keeps the cached read.""" + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + + seen: Final[list[bool]] = [] # mutable-ok: records each read mode + + async def _get_team_object(team_id: str, **kwargs: object) -> LiteLLM_TeamTable: + seen.append(bool(kwargs.get("check_db_only"))) + return LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", agents=["agent-a"] + ), + ) + + monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object) + await AgentRequestHandler._get_allowed_agents_for_team(personal_key_owner_teams, strict=True) + await AgentRequestHandler._get_allowed_agents_for_team(personal_key_owner_teams, strict=False) + assert seen == [True, False] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_invocation_uses_fresh_owner_team_membership( + monkeypatch: pytest.MonkeyPatch, + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore + from litellm.proxy.auth import auth_checks + + async def _get_user_object(**kwargs: object) -> LiteLLM_UserTable: + return LiteLLM_UserTable( + user_id="alice", + teams=[] if bool(kwargs.get("check_db_only")) else ["team-a"], + ) + + monkeypatch.setattr(auth_checks, "get_user_object", _get_user_object) + monkeypatch.setattr(AgentIdentityStore, "agent", AsyncMock(return_value=None)) + cached_auth: Final = personal_key_owner_teams.model_copy(update={"requires_fresh_policy": False}) + fresh_auth: Final = personal_key_owner_teams.model_copy(update={"requires_fresh_policy": True}) + + assert await AgentRequestHandler.is_agent_allowed("agent-a", fresh_auth) is False + assert await AgentRequestHandler.is_agent_allowed("agent-a", cached_auth) is True + + +@pytest.mark.asyncio +async def test_owner_team_agent_grant_lookup_returns_policy_unavailable_value() -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentPolicyUnavailable + + team: Final = LiteLLM_TeamTable(team_id="team-unreadable", object_permission_id="perm-unreadable") + + assert await AgentRequestHandler._owner_team_agent_grants(team, strict=True) == AgentPolicyUnavailable( + team_id="team-unreadable" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_keyless_grant_source_of_the_owner_does_not_inherit_its_teams( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + """The direct-grant source of an admitted subject has the owner's user_id, no team and no api_key. + It must keep its pre-flag resolution, or a delegated token scoped to team-a reaches agent-b.""" + source: Final = UserAPIKeyAuth(user_id=personal_key_owner_teams.user_id, team_id=None) + assert await AgentRequestHandler.resolve_key_team_agent_access(source) == UnrestrictedAgentAccess() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_agent_access_is_unchanged_when_flag_off( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + import litellm + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", False) + assert ( + await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == UnrestrictedAgentAccess() + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a", "team-b")], indirect=True) +async def test_team_key_agent_access_ignores_its_owner_other_teams(personal_key_owner_teams: UserAPIKeyAuth) -> None: + from litellm.proxy.auth import auth_checks + + team_key: Final = personal_key_owner_teams.model_copy(update={"team_id": "team-b"}) + assert await AgentRequestHandler.resolve_key_team_agent_access(team_key) == RestrictedAgentAccess( + frozenset({"agent-b"}) + ) + auth_checks.get_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "personal_key_owner_teams", [("team-a", "team-grouped"), ("team-a", "team-unreadable")], indirect=True +) +async def test_owner_team_whose_grants_cannot_be_read_grants_no_agent( + personal_key_owner_teams: UserAPIKeyAuth, +) -> None: + assert await AgentRequestHandler.resolve_key_team_agent_access(personal_key_owner_teams) == RestrictedAgentAccess( + frozenset({"agent-a"}) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("personal_key_owner_teams", [("team-a",)], indirect=True) +async def test_personal_key_agent_access_is_denied_when_owner_teams_cannot_be_read( + monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth +) -> None: + from litellm.proxy.auth import auth_checks + + monkeypatch.setattr( + auth_checks, "personal_key_owner_teams_with_status", AsyncMock(side_effect=RuntimeError("db down")) + ) + assert await AgentRequestHandler.resolve_agent_access(personal_key_owner_teams) == RestrictedAgentAccess( + frozenset() + ) diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index d2f2fec073c..a2448694dd7 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -418,6 +418,124 @@ def _base_patches( ] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("identity_managed", "requires_fresh_policy", "strict"), + [(True, False, True), (False, True, True), (False, False, False)], +) +async def test_invoke_agent_a2a_passes_billed_auth_to_common_processing( + monkeypatch: pytest.MonkeyPatch, + identity_managed: bool, + requires_fresh_policy: bool, + strict: bool, +) -> None: + import litellm + + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable + from litellm.proxy.agent_endpoints import agent_registry + from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + from litellm.proxy.auth import auth_checks + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + + agent: Final = _make_agent_mock() + agent.identity_managed = identity_managed + request: Final = _make_request_mock("message/send", _HELLO_MESSAGE_PARAMS) + auth: Final = UserAPIKeyAuth( + api_key="sk-personal", + user_id="alice", + team_id=None, + ) + auth.requires_fresh_policy = requires_fresh_policy + processed_metadata: Final[dict[str, object]] = {} + processed_request_data: Final[dict[str, object]] = {"metadata": processed_metadata} + + async def _common_processing( + *, user_api_key_dict: UserAPIKeyAuth, **_: object + ) -> tuple[dict[str, object], MagicMock]: + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, + data=processed_request_data, + metadata_variable_name="metadata", + ) + return {}, MagicMock() + + common_processing: Final = AsyncMock(side_effect=_common_processing) + + async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + team: Final = LiteLLM_TeamTable( + team_id=team_id, + object_permission_id=f"perm-{team_id}", + organization_id=f"org-{team_id}", + metadata={"guardrails": ["team-guardrail"]} if team_id == "team-b" else None, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id=f"perm-{team_id}", + agents=["test-agent"] if team_id == "team-b" else [], + agent_access_groups=["billing-group"] if team_id == "team-b" else [], + ), + ) + return team + + database: Final = MagicMock() + database.db.litellm_agentstable.find_many = AsyncMock(return_value=[]) + database.writer_db.litellm_agentstable.find_many = AsyncMock(return_value=[]) + team_loader: Final = AsyncMock(side_effect=_get_team_object) + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr("litellm.proxy.utils.premium_user_check", MagicMock()) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) + monkeypatch.setattr(agent_registry, "global_agent_registry", AgentRegistry()) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(return_value=LiteLLM_UserTable(user_id="alice", teams=["team-a", "team-b"])), + ) + monkeypatch.setattr(auth_checks, "get_team_object", team_loader) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + + with ExitStack() as stack: + for patcher in _base_patches(agent): + stack.enter_context(patcher) + stack.enter_context( + patch.object( + ProxyBaseLLMRequestProcessing, + "common_processing_pre_call_logic", + new=common_processing, + ) + ) + stack.enter_context(patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", False)) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response: Final = await invoke_agent_a2a( + agent_id="test-agent", + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=auth, + ) + + assert json.loads(response.body.decode())["error"]["code"] == -32603 + billed_auth: Final = common_processing.await_args.kwargs["user_api_key_dict"] + assert billed_auth is not auth + assert (billed_auth.team_id, billed_auth.org_id) == ("team-b", "org-team-b") + assert billed_auth.team_metadata == {"guardrails": ["team-guardrail"]} + assert processed_metadata["guardrails"] == ["team-guardrail"] + assert (auth.team_id, auth.org_id) == (None, None) + assert auth.team_metadata is None + assert len(team_loader.await_args_list) == 2 + assert all(call.kwargs["check_db_only"] is strict for call in team_loader.await_args_list) + if strict: + database.writer_db.litellm_agentstable.find_many.assert_awaited_once() + database.db.litellm_agentstable.find_many.assert_not_awaited() + else: + database.db.litellm_agentstable.find_many.assert_awaited_once() + database.writer_db.litellm_agentstable.find_many.assert_not_awaited() + + async def _add_proxy_data(data: dict[str, object], **kwargs: object) -> dict[str, object]: return { **data, diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 72179e661a3..c583407e61e 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -27,6 +27,7 @@ from litellm.proxy._types import ( LiteLLM_EndUserTable, LiteLLM_ObjectPermissionTable, LiteLLM_TagTable, + LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, @@ -48,6 +49,8 @@ from litellm.proxy.auth.auth_checks import ( _check_agent_access_group_model_access, check_end_user_budget, _check_team_member_budget, + personal_key_owner_teams, + personal_key_owner_teams_with_status, fetch_key_object_from_db_with_reconnect, _get_fuzzy_user_object, CallerTeamLoader, @@ -82,6 +85,7 @@ from litellm.constants import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import check_api_key_for_custom_headers_or_pass_through_endpoints +from litellm.proxy.utils import ProxyLogging from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_bearer_token, encrypt_value_helper from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -4642,6 +4646,188 @@ async def test_check_team_member_budget_dispatches_the_configured_alert_before_t assert call_info.max_budget_alert_emails == {"50": [], "100": ["finance@co.com"]} +@pytest.mark.asyncio +async def test_personal_key_owner_teams_reports_over_budget_and_blocked_teams( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.auth import auth_checks + + cache: Final = UserApiKeyCache() + logging_obj: Final = ProxyLogging(user_api_key_cache=cache) + owner: Final = LiteLLM_UserTable( + user_id="owner", + teams=["team-a", "team-b", "team-blocked", "team-other-org"], + ) + loaded_teams: Final = { + "team-a": LiteLLM_TeamTable(team_id="team-a", organization_id="org-1"), + "team-b": LiteLLM_TeamTable(team_id="team-b", organization_id="org-1"), + "team-blocked": LiteLLM_TeamTable(team_id="team-blocked", organization_id="org-1", blocked=True), + "team-other-org": LiteLLM_TeamTable(team_id="team-other-org", organization_id="org-2"), + } + memberships: Final = { + "team-a": LiteLLM_TeamMembership( + user_id="owner", + team_id="team-a", + spend=1.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=5.0), + ), + "team-b": LiteLLM_TeamMembership( + user_id="owner", + team_id="team-b", + spend=1.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ), + } + get_user = AsyncMock(return_value=owner) + get_membership = AsyncMock(side_effect=lambda *, team_id, **_: memberships[team_id]) + + async def get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable: + return loaded_teams[team_id] + + async def current_spend( + counter_key: str, + fallback_spend: float, + max_budget: float | None = None, + ) -> float: + assert counter_key.startswith("spend:team_member:") and max_budget is not None + return fallback_spend + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(auth_checks, "get_user_object", get_user) + monkeypatch.setattr(auth_checks, "get_team_object", get_team_object) + monkeypatch.setattr(auth_checks, "get_team_membership", get_membership) + monkeypatch.setattr(proxy_server, "get_current_spend", current_spend) + + result: Final = await personal_key_owner_teams_with_status( + valid_token=UserAPIKeyAuth(api_key="sk-personal", user_id="owner", org_id="org-1"), + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=logging_obj, + ) + + assert tuple(team.team_id for team in result.teams) == ("team-a",) + assert {team.team_id for team in result.withheld} == {"team-b", "team-blocked"} + assert all(team.team_id != "team-other-org" for team in (*result.teams, *result.withheld)) + assert result.unreadable is False + assert get_user.await_count == 1 + assert {call.kwargs["team_id"] for call in get_membership.await_args_list} == {"team-a", "team-b"} + + +@pytest.mark.asyncio +async def test_personal_key_owner_teams_check_db_only_refreshes_cached_owner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.auth import auth_checks + + user_id = "owner-cache-refresh" + cached_owner: Final = LiteLLM_UserTable( + user_id=user_id, + user_email="owner@example.com", + teams=["team-a"], + ) + database_owner: Final = LiteLLM_UserTable( + user_id=user_id, + user_email="owner@example.com", + teams=[], + ) + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=user_id, value=cached_owner, model_type=LiteLLM_UserTable) + logging_obj: Final = ProxyLogging(user_api_key_cache=cache) + prisma: Final = MagicMock() + prisma.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=database_owner) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-a"))) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=None)) + valid_token: Final = UserAPIKeyAuth(api_key="sk-personal", user_id=user_id) + + cached_result: Final = await personal_key_owner_teams_with_status( + valid_token=valid_token, + prisma_client=prisma, + user_api_key_cache=cache, + proxy_logging_obj=logging_obj, + ) + fresh_result: Final = await personal_key_owner_teams_with_status( + valid_token=valid_token, + prisma_client=prisma, + user_api_key_cache=cache, + proxy_logging_obj=logging_obj, + check_db_only=True, + ) + + assert tuple(team.team_id for team in cached_result.teams) == ("team-a",) + assert fresh_result.teams == () + prisma.writer_db.litellm_usertable.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_personal_key_owner_teams_check_db_only_uses_fresh_membership_budget( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.auth import auth_checks + + user_id: Final = "owner-membership-refresh" + team_id: Final = "team-membership-refresh" + cached_membership: Final = LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=1.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=5.0), + ) + database_membership: Final = LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=10.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=5.0), + ) + database_row: Final = MagicMock() + database_row.dict.return_value = database_membership.model_dump(mode="python") + prisma: Final = MagicMock() + prisma.db.litellm_teammembership.find_unique = AsyncMock() + prisma.writer_db.litellm_teammembership.find_unique = AsyncMock(return_value=database_row) + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key=auth_checks.team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + value=cached_membership, + model_type=LiteLLM_TeamMembership, + ) + logging_obj: Final = ProxyLogging(user_api_key_cache=cache) + owner: Final = LiteLLM_UserTable(user_id=user_id, teams=[team_id]) + team: Final = LiteLLM_TeamTable(team_id=team_id) + + async def current_spend( + counter_key: str, + fallback_spend: float, + max_budget: float | None = None, + ) -> float: + return fallback_spend + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=owner)) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=team)) + monkeypatch.setattr(proxy_server, "get_current_spend", current_spend) + valid_token: Final = UserAPIKeyAuth(api_key="sk-personal", user_id=user_id) + + cached_result: Final = await personal_key_owner_teams_with_status( + valid_token=valid_token, + prisma_client=prisma, + user_api_key_cache=cache, + proxy_logging_obj=logging_obj, + ) + fresh_result: Final = await personal_key_owner_teams_with_status( + valid_token=valid_token, + prisma_client=prisma, + user_api_key_cache=cache, + proxy_logging_obj=logging_obj, + check_db_only=True, + ) + + assert tuple(team.team_id for team in cached_result.teams) == (team_id,) + assert cached_result.withheld == () + assert fresh_result.teams == () + assert tuple(team.team_id for team in fresh_result.withheld) == (team_id,) + prisma.db.litellm_teammembership.find_unique.assert_not_awaited() + prisma.writer_db.litellm_teammembership.find_unique.assert_awaited_once() + + @pytest.mark.parametrize( "spend, max_budget, expect_alert", [ @@ -8422,6 +8608,84 @@ async def test_get_team_membership_db_fetch_returns_validated_membership(): assert result.spend == 1.5 +@pytest.mark.asyncio +async def test_get_team_membership_check_db_only_bypasses_stale_cache(): + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + cache: Final = UserApiKeyCache() + cache_key: Final = team_membership_reservation_cache_key(user_id="u-fresh", team_id="t-fresh") + stale_membership: Final = LiteLLM_TeamMembership(user_id="u-fresh", team_id="t-fresh", spend=1.0) + await cache.async_set_cache(key=cache_key, value=stale_membership, model_type=LiteLLM_TeamMembership) + + database_row: Final = MagicMock() + database_row.dict.return_value = {"user_id": "u-fresh", "team_id": "t-fresh", "spend": 9.0} + prisma: Final = MagicMock() + prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=database_row) + prisma.writer_db.litellm_teammembership.find_unique = AsyncMock(return_value=database_row) + + cached: Final = await get_team_membership( + user_id="u-fresh", + team_id="t-fresh", + prisma_client=prisma, + user_api_key_cache=cache, + ) + fresh: Final = await get_team_membership( + user_id="u-fresh", + team_id="t-fresh", + prisma_client=prisma, + user_api_key_cache=cache, + check_db_only=True, + ) + + assert cached is not None and cached.spend == 1.0 + assert fresh is not None and fresh.spend == 9.0 + prisma.db.litellm_teammembership.find_unique.assert_not_awaited() + prisma.writer_db.litellm_teammembership.find_unique.assert_awaited_once() + with pytest.raises(Exception, match=r"^No db connected$"): + await get_team_membership( + user_id="u-fresh", + team_id="t-fresh", + prisma_client=None, + user_api_key_cache=cache, + check_db_only=True, + ) + + +@pytest.mark.asyncio +async def test_get_team_membership_cache_miss_reads_replica_and_check_db_only_reads_writer(): + from litellm.proxy.auth.auth_checks import get_team_membership + + reader_row: Final = MagicMock() + reader_row.dict.return_value = {"user_id": "u-read", "team_id": "t-read", "spend": 2.0} + writer_row: Final = MagicMock() + writer_row.dict.return_value = {"user_id": "u-read", "team_id": "t-read", "spend": 9.0} + prisma_client: Final = MagicMock() + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=reader_row) + prisma_client.writer_db.litellm_teammembership.find_unique = AsyncMock(return_value=writer_row) + cache: Final = UserApiKeyCache() + + reader_membership: Final = await get_team_membership( + user_id="u-read", + team_id="t-read", + prisma_client=prisma_client, + user_api_key_cache=cache, + ) + writer_membership: Final = await get_team_membership( + user_id="u-read", + team_id="t-read", + prisma_client=prisma_client, + user_api_key_cache=cache, + check_db_only=True, + ) + + assert reader_membership is not None and reader_membership.spend == 2.0 + assert writer_membership is not None and writer_membership.spend == 9.0 + prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + prisma_client.writer_db.litellm_teammembership.find_unique.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_team_membership_negative_caches_a_missing_row(): """ diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 9bcf0920f3e..7af747fc26d 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3938,6 +3938,126 @@ async def test_mcp_per_team_rate_limit_is_enforced_from_team_metadata_v3() -> No await handler.enforce_mcp_server_rate_limits(second_key, server) +@pytest.mark.asyncio +async def test_inherited_mcp_team_rate_limit_deduplicates_team_and_source_descriptors_v3() -> None: + handler, _ = _make_mcp_handler() + server: Final = MCPServer( + server_id="math-a", + name="math_a", + server_name="math_a", + transport=MCPTransport.http, + ) + inherited_auth: Final = UserAPIKeyAuth( + api_key=hash_token("sk-inherited"), + team_id="team-a", + team_metadata={"mcp_rpm_limit": {"math_a": 1}}, + ).model_copy(update={"mcp_source_team_rpm_limits": {"team-a": {"math_a": 1}}}) + descriptors: Final[list[RateLimitDescriptor]] = [] + handler._add_mcp_per_key_rate_limit_descriptor(inherited_auth, "math_a", descriptors) + + assert descriptors == [] + team_descriptors: Final[list[RateLimitDescriptor]] = [] + handler._add_mcp_per_team_rate_limit_descriptor(inherited_auth, "math_a", team_descriptors) + assert [(descriptor["key"], descriptor["value"]) for descriptor in team_descriptors] == [ + ("mcp_per_team", "team-a:math_a") + ] + await handler.enforce_mcp_server_rate_limits(inherited_auth, server) + + with pytest.raises(ProxyRateLimitError, match="mcp_per_team"): + await handler.enforce_mcp_server_rate_limits(inherited_auth, server) + + +def test_inherited_mcp_team_rate_limit_uses_source_limit_when_primary_metadata_is_missing() -> None: + handler, _ = _make_mcp_handler() + auth: Final = UserAPIKeyAuth( + api_key=hash_token("sk-inherited-source-limit"), + team_id="team-a", + ).model_copy(update={"mcp_source_team_rpm_limits": {"team-a": {"math_a": 1}}}) + descriptors: Final[list[RateLimitDescriptor]] = [] + + handler._add_mcp_per_team_rate_limit_descriptor(auth, "math_a", descriptors) + + assert descriptors == [ + { + "key": "mcp_per_team", + "value": "team-a:math_a", + "rate_limit": { + "requests_per_unit": 1, + "tokens_per_unit": None, + "window_size": handler.window_size, + }, + } + ] + + +@pytest.mark.asyncio +async def test_billed_personal_key_auth_produces_team_rpm_descriptors(monkeypatch: pytest.MonkeyPatch) -> None: + from unittest.mock import AsyncMock, MagicMock + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.team_membership import LiteLLM_TeamMembership + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.auth import auth_checks + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable + + server: Final = MCPServer( + server_id="math_a", + name="math_a", + server_name="math_a", + transport=MCPTransport.http, + ) + team: Final = LiteLLM_TeamTable( + team_id="team-a", + organization_id="org-team-a", + rpm_limit=7, + object_permission_id="perm-team-a", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-team-a", + mcp_servers=["math_a"], + mcp_tool_permissions={"math_a": ["add"]}, + ), + ) + membership: Final = LiteLLM_TeamMembership( + user_id="owner", + team_id="team-a", + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=3), + ) + user_api_key_auth: Final = UserAPIKeyAuth( + api_key=hash_token("sk-billed-personal"), + user_id="owner", + ) + + monkeypatch.setattr(litellm, "personal_key_model_access_from_teams", True) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _: False) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) + monkeypatch.setattr(auth_checks, "personal_key_owner_teams", AsyncMock(return_value=(team,))) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=team)) + monkeypatch.setattr(auth_checks, "get_team_membership", AsyncMock(return_value=membership)) + monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None)) + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + + billed_auth: Final = await MCPRequestHandler.billing_auth_for_tool_call(user_api_key_auth, "math_a-add") + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + + descriptors: Final = handler.create_rate_limit_descriptors( + user_api_key_dict=billed_auth, + data={}, + rpm_limit_type=None, + tpm_limit_type=None, + model_has_failures=False, + ) + + assert [ + (descriptor["key"], descriptor["value"], descriptor["rate_limit"]["requests_per_unit"]) + for descriptor in descriptors + if descriptor["key"] in ("team", "team_member") + ] == [("team", "team-a", 7), ("team_member", "team-a:owner", 3)] + + @pytest.mark.asyncio async def test_mcp_server_rpm_is_shared_across_keys_v3() -> None: handler, _ = _make_mcp_handler() From f4347d4f5015e9cc54681b4dca2e2d60532ae414 Mon Sep 17 00:00:00 2001 From: Vansh Singh Date: Tue, 20 Oct 2026 04:32:11 +0530 Subject: [PATCH 17/18] fix(speech)!: correctly propagate spend logging and usage for Gemini TTS bridge (#44551) * fix(speech): correctly propagate spend logging and usage for Gemini TTS bridge * fix(speech): call Gemini once per TTS request and log the bridge's token usage aspeech re-ran speech whenever the first result was not a coroutine, so the speech-to-completion bridge sent every Gemini TTS request upstream twice. It now returns the sync result. The bridge's inner completion runs as an internal sub call, so only the speech call fires success callbacks, and the completion's usage rides on the audio response so the spend log and usage_object carry real token counts. Replaces the separate inner Logging object and the cost_calculator fallback, which never reached the spend row. * fix(speech): flag a cached Gemini TTS reply as a cache hit so it logs at no charge * fix(speech): price a cached Gemini TTS reply's savings from its token usage and keep its cache key * test(integration): cover Gemini TTS speech billing, logging, cache hits and deployment TPM * test(integration): start the speech TPM cell in a minute's first 10 s so the two requests never cross a minute --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/cost_calculator.py | 2 + .../speech_to_completion_bridge/handler.py | 9 +- .../transformation.py | 15 +- litellm/main.py | 8 +- litellm/types/llms/openai.py | 11 +- .../sdk/test_gemini_tts_speech_logging_sdk.py | 69 ++++ .../spend/test_gemini_tts_speech_billing.py | 334 ++++++++++++++++++ tests/unit/test_main.py | 65 +++- 8 files changed, 497 insertions(+), 16 deletions(-) create mode 100644 tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py create mode 100644 tests/integration/spend/test_gemini_tts_speech_billing.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4fcd3bde843..8c6e2fead2d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1103,6 +1103,8 @@ def get_usage_object( if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)): return None if completion_response.usage is None else _decisions_usage(completion_response.usage) + if isinstance(completion_response, HttpxBinaryResponseContent): + return completion_response.usage if usage_obj is None: return None if isinstance(usage_obj, Usage): diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index 6c33621ec89..00d112773f1 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -6,6 +6,8 @@ from typing import TYPE_CHECKING, Final from typing_extensions import TypedDict +from litellm._internal_context import internal_sub_call + if TYPE_CHECKING: from litellm import LiteLLMLoggingObj from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -111,9 +113,10 @@ class SpeechToCompletionBridgeHandler: voice=voice, ) - result: Final = completion( - **request_data, - ) + with internal_sub_call(): + result: Final = completion( + **request_data, + ) requested_response_format: Final = optional_params.get("response_format") if isinstance(result, ModelResponse): diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 2ed140c0208..9c1e5ae7a29 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -5,11 +5,12 @@ from typing import TYPE_CHECKING, Final, cast from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_param if TYPE_CHECKING: from litellm import Logging as LiteLLMLoggingObj from litellm.types.llms.openai import ChatCompletionUserMessage, HttpxBinaryResponseContent - from litellm.types.utils import ModelResponse + from litellm.types.utils import ModelResponse, Usage def _completion_response_cost(model_response: "ModelResponse") -> float | None: @@ -20,6 +21,13 @@ def _completion_response_cost(model_response: "ModelResponse") -> float | None: return response_cost if isinstance(response_cost, float) else None +def _completion_usage(model_response: "ModelResponse") -> "Usage | None": + from litellm.types.utils import Usage + + usage: Final = getattr(model_response, "usage", None) + return usage if isinstance(usage, Usage) else None + + GEMINI_TTS_CHAT_AUDIO_FORMAT: Final = "pcm16" GEMINI_TTS_RAW_RESPONSE_FORMAT: Final = "pcm" GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: Final = frozenset({"wav", GEMINI_TTS_RAW_RESPONSE_FORMAT}) @@ -179,4 +187,9 @@ class SpeechToCompletionBridgeTransformationHandler: ) binary_response: Final = HttpxBinaryResponseContent(response) binary_response.set_response_cost(_completion_response_cost(model_response)) + binary_response.usage = _completion_usage(model_response) + cached_hidden_params: Final = get_hidden_params(model_response) if served_from_cache(model_response) else None + if cached_hidden_params is not None: + set_hidden_param(binary_response, "cache_hit", True) + set_hidden_param(binary_response, "cache_key", cached_hidden_params.get("cache_key")) return binary_response diff --git a/litellm/main.py b/litellm/main.py index 38591441e01..9430af34401 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8500,14 +8500,8 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent: api_base=api_base_hint, ) - # Await normally init_response: Final = await loop.run_in_executor(None, func_with_context) - if asyncio.iscoroutine(init_response): - response = await init_response - else: - # Call the synchronous function using run_in_executor - response = await loop.run_in_executor(None, func_with_context) - return response + return await init_response if asyncio.iscoroutine(init_response) else init_response except Exception as e: custom_llm_provider = custom_llm_provider or "openai" raise exception_type( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 8cc9e912e80..8dbc0d97d54 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2,7 +2,7 @@ import builtins from collections.abc import Iterable, Mapping from enum import Enum from os import PathLike -from typing import IO, Any, Final, Generic, Literal, Optional, TypeAlias, Union +from typing import IO, TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, Union import httpx from openai import Omit @@ -96,6 +96,9 @@ from litellm.types.responses.main import ( from .base import CachedTokensDetails +if TYPE_CHECKING: + from litellm.types.utils import Usage + FileContent = IO[bytes] | bytes | PathLike FileTypes = ( @@ -119,6 +122,7 @@ class BinaryResponseSummary(TypedDict): object: ReadOnly[Literal["binary"]] content_type: ReadOnly[str | None] num_bytes: ReadOnly[int] + usage: NotRequired[ReadOnly[Mapping[str, builtins.object]]] _ResponseT = TypeVar("_ResponseT", bound=httpx.Response | SDKResponse, default=httpx.Response) @@ -126,6 +130,7 @@ _ResponseT = TypeVar("_ResponseT", bound=httpx.Response | SDKResponse, default=h class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT]): _hidden_params: dict + usage: "Usage | None" response: _ResponseT # pyright: ignore[reportIncompatibleVariableOverride] # SDK accepts both backends at runtime @property @@ -139,6 +144,7 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT def __init__(self, response: _ResponseT) -> None: super().__init__(response) # pyright: ignore[reportArgumentType] # SDK accepts both backends at runtime self._hidden_params = {} + self.usage = None @property def content_type(self) -> str | None: @@ -146,11 +152,12 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT return headers.get("content-type") def logging_summary(self) -> BinaryResponseSummary: - return { + summary: Final[BinaryResponseSummary] = { "object": "binary", "content_type": self.content_type, "num_bytes": self._num_bytes(), } + return summary if self.usage is None else {**summary, "usage": self.usage.model_dump()} def _num_bytes(self) -> int: try: diff --git a/tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py b/tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py new file mode 100644 index 00000000000..3df10da8fda --- /dev/null +++ b/tests/integration/sdk/test_gemini_tts_speech_logging_sdk.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import eventually +from integration._support.wire import wire_server +from integration.providers._gemini_tts_voice_support import ( + AUDIO_TOKENS, + BACKEND, + GEMINI_API_KEY, + PCM, + PROMPT_TOKENS, + gemini_peer, + marker, + received_text, + wav_payload, +) + +import litellm +from litellm.integrations.custom_logger import CustomLogger + +BINARY_RESPONSE: Final = "HttpxBinaryResponseContent" + + +@dataclass(frozen=True, slots=True) +class _Logged: + call_type: object + response_type: str + tokens: tuple[object, object] + + +class _Recorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self._lock: Final = threading.Lock() + self.events: tuple[_Logged, ...] = () + + def log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload: Final = kwargs.get("standard_logging_object") + logged: Final = payload if isinstance(payload, dict) else {} + event: Final = _Logged( + call_type=kwargs.get("call_type"), + response_type=type(response_obj).__name__, + tokens=(logged.get("prompt_tokens"), logged.get("completion_tokens")), + ) + with self._lock: + self.events = (*self.events, event) + + +def test_sync_speech_logs_one_success_event_carrying_the_reply_token_usage(monkeypatch: pytest.MonkeyPatch) -> None: + recorder: Final = _Recorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + with wire_server(gemini_peer) as wire: + monkeypatch.setenv("GEMINI_API_BASE", wire.url) + text: Final = marker() + audio: Final = litellm.speech(model=f"gemini/{BACKEND}", input=text, voice="Kore", api_key=GEMINI_API_KEY) + assert wav_payload(audio.content) == PCM + assert sum(1 for request in wire.drain() if received_text(request) == text) == 1 + events: Final = eventually( + lambda: recorder.events, + lambda found: any(event.response_type == BINARY_RESPONSE for event in found), + seconds=30, + ) + assert events == (_Logged("speech", BINARY_RESPONSE, (PROMPT_TOKENS, AUDIO_TOKENS)),) diff --git a/tests/integration/spend/test_gemini_tts_speech_billing.py b/tests/integration/spend/test_gemini_tts_speech_billing.py new file mode 100644 index 00000000000..f1bf3de91b3 --- /dev/null +++ b/tests/integration/spend/test_gemini_tts_speech_billing.py @@ -0,0 +1,334 @@ +from __future__ import annotations + +import os +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from datetime import UTC, datetime +from hashlib import sha256 +from pathlib import Path +from typing import Final, Literal + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, eventually, gateway_from_environment, object_value +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Wire, wire_server +from integration.providers._gemini_tts_voice_support import ( + AUDIO_TOKENS, + BACKEND, + GEMINI_API_KEY, + PCM, + PROMPT_TOKENS, + REJECTED_PREFIX, + VERTEX_LOCATION, + VERTEX_PROJECT, + gemini_peer, + marker, + received_text, + wav_payload, +) +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +Provider = Literal["gemini", "vertex_ai"] +MODELS: Final[Mapping[Provider, str]] = {"gemini": "tts-billing-gemini", "vertex_ai": "tts-billing-vertex"} +TPM_MODEL: Final = "tts-billing-tpm" +TPM_LIMIT: Final = PROMPT_TOKENS + AUDIO_TOKENS +BINARY_RESPONSE: Final = "HttpxBinaryResponseContent" +CALLBACK_LOG: Final = "TTS_BILLING_CALLBACK_LOG" +CALLBACK_MODULE: Final = """ +import json +import os + +from litellm.integrations.custom_logger import CustomLogger + + +def _event(kind, kwargs, response_obj): + payload = kwargs.get("standard_logging_object") or {} + metadata = payload.get("metadata") or {} + response = payload.get("response") + usage = response.get("usage") if isinstance(response, dict) else None + return { + "kind": kind, + "key": metadata.get("user_api_key_hash"), + "call_type": kwargs.get("call_type"), + "response_type": type(response_obj).__name__, + "response_cost": payload.get("response_cost"), + "prompt_tokens": payload.get("prompt_tokens"), + "completion_tokens": payload.get("completion_tokens"), + "cache_hit": payload.get("cache_hit"), + "saved_cache_cost": payload.get("saved_cache_cost"), + "cache_key": payload.get("cache_key"), + "logged_usage_total": usage.get("total_tokens") if isinstance(usage, dict) else None, + } + + +class SpeechEvents(CustomLogger): + def _write(self, line): + with open(os.environ["TTS_BILLING_CALLBACK_LOG"], "a") as sink: + sink.write(json.dumps(line) + "\\n") + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self._write(_event("sync_success", kwargs, response_obj)) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self._write(_event("async_success", kwargs, response_obj)) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._write(_event("sync_failure", kwargs, response_obj)) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._write(_event("async_failure", kwargs, response_obj)) + + +speech_events = SpeechEvents() +""" +SPEND_ROWS_SQL: Final = ( + 'SELECT call_type, spend, prompt_tokens, completion_tokens, cache_hit FROM "LiteLLM_SpendLogs"' + ' WHERE api_key = %s ORDER BY "startTime"' +) +KEY_SPEND_SQL: Final = 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s' + +Event = dict[str, JsonValue] + + +@dataclass(frozen=True, slots=True) +class BillingProxy: + proxy: OwnedProxy + callback_log: Path + + +@dataclass(frozen=True, slots=True) +class Sent: + status: int + body: bytes + cache_key: str + + +def _hashed(key: str) -> str: + return sha256(key.encode()).hexdigest() + + +def _model_list(token_url: str) -> list[JsonValue]: + return [ + {"model_name": MODELS["gemini"], "litellm_params": {"model": f"gemini/{BACKEND}", "api_key": GEMINI_API_KEY}}, + { + "model_name": TPM_MODEL, + "litellm_params": {"model": f"gemini/{BACKEND}", "api_key": GEMINI_API_KEY, "tpm": TPM_LIMIT}, + }, + { + "model_name": MODELS["vertex_ai"], + "litellm_params": { + "model": f"vertex_ai/{BACKEND}", + "vertex_project": VERTEX_PROJECT, + "vertex_location": VERTEX_LOCATION, + "vertex_credentials": service_account_json(VERTEX_PROJECT, token_url.rstrip("/")), + }, + }, + ] + + +def _config(directory: Path, token_url: str, *, cached: bool) -> Path: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + settings: Final = object_value(stock["litellm_settings"]) + kept: Final = {name: value for name, value in settings.items() if name not in {"cache", "cache_params"}} + cache: Final[dict[str, JsonValue]] = ( + {"cache": True, "cache_params": {**object_value(settings["cache_params"]), "namespace": marker()}} + if cached + else {"cache": False} + ) + routing: Final[dict[str, JsonValue]] = ( + {} + if cached + else { + "routing_strategy": "usage-based-routing-v2", + "redis_host": os.environ["REDIS_HOST"], + "redis_port": int(os.environ["REDIS_PORT"]), + } + ) + (directory / "tts_billing_callbacks.py").write_text(CALLBACK_MODULE) + path: Final = directory / "tts-billing.yaml" + path.write_text( + yaml.safe_dump( + { + **stock, + "model_list": _model_list(token_url), + "litellm_settings": {**kept, **cache, "callbacks": "tts_billing_callbacks.speech_events"}, + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0, **routing}, + } + ) + ) + return path + + +def _billing_proxy(wire: Wire, directory: Path, *, cached: bool) -> Iterator[BillingProxy]: + callback_log: Final = directory / "callback_events.jsonl" + callback_log.touch() + with gateway_from_environment() as gateway: + config: Final = _config(directory, gateway.upstream_url, cached=cached) + overrides: Final = {"GEMINI_API_BASE": wire.url, "VERTEXAI_API_BASE": wire.url, CALLBACK_LOG: str(callback_log)} + with owned_proxy_process(gateway, directory, overrides, config=config, workers=2) as proxy: + yield BillingProxy(proxy, callback_log) + + +@pytest.fixture(scope="module") +def tts_wire() -> Iterator[Wire]: + with wire_server(gemini_peer) as wire: + yield wire + + +@pytest.fixture(scope="module") +def uncached(tts_wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[BillingProxy]: + yield from _billing_proxy(tts_wire, tmp_path_factory.mktemp("tts-billing-uncached"), cached=False) + + +@pytest.fixture(scope="module") +def cached(tts_wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[BillingProxy]: + yield from _billing_proxy(tts_wire, tmp_path_factory.mktemp("tts-billing-cached"), cached=True) + + +def _speech(billing: BillingProxy, model: str, text: str, key: str, voice: str = "Kore") -> Sent: + response: Final = billing.proxy.gateway.request( + "POST", "/v1/audio/speech", {"model": model, "input": text, "voice": voice}, key=key + ) + return Sent(response.status_code, response.content, response.headers.get("x-litellm-cache-key", "")) + + +def _upstream_calls(wire: Wire, text: str) -> int: + return sum(1 for request in wire.drain() if received_text(request) == text) + + +def _spend_rows(key: str, count: int) -> list[dict[str, JsonValue]]: + return eventually(lambda: read_rows(SPEND_ROWS_SQL, (_hashed(key),)), lambda rows: len(rows) == count, seconds=70) + + +def _key_spend(key: str, expected: float) -> JsonValue: + rows: Final = eventually( + lambda: read_rows(KEY_SPEND_SQL, (_hashed(key),)), + lambda found: len(found) == 1 and found[0]["spend"] == pytest.approx(expected), + seconds=70, + return_last_on_timeout=True, + ) + return rows[0]["spend"] + + +def _chat_spend(billing: BillingProxy, model: str, key: str) -> float: + body: Final[dict[str, JsonValue]] = { + "model": model, + "messages": [{"role": "user", "content": marker()}], + "modalities": ["audio"], + "audio": {"voice": "Kore", "format": "pcm16"}, + } + response: Final = billing.proxy.gateway.request("POST", "/v1/chat/completions", body, key=key) + assert response.status_code == 200, response.text + row: Final = _spend_rows(key, 1)[0] + assert (row["prompt_tokens"], row["completion_tokens"]) == (PROMPT_TOKENS, AUDIO_TOKENS), row + spend: Final = row["spend"] + assert isinstance(spend, float) and spend > 0, row + return spend + + +def _events(billing: BillingProxy, key: str) -> tuple[Event, ...]: + complete_lines: Final = billing.callback_log.read_text().split("\n")[:-1] + events: Final = tuple(JSON_OBJECT.validate_json(line) for line in complete_lines) + return tuple(event for event in events if event["key"] == _hashed(key)) + + +def _speech_events(billing: BillingProxy, key: str, count: int) -> tuple[Event, ...]: + return eventually( + lambda: _events(billing, key), + lambda found: sum(1 for event in found if event["response_type"] == BINARY_RESPONSE) >= count, + seconds=70, + ) + + +def _summary(event: Event) -> tuple[JsonValue, ...]: + return (event["call_type"], event["response_type"], event["prompt_tokens"], event["completion_tokens"]) + + +@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"]) +def test_speech_calls_google_once_and_bills_and_logs_its_token_usage_once( + uncached: BillingProxy, tts_wire: Wire, provider: Provider +) -> None: + with uncached.proxy.gateway.scenario() as scenario: + speech_key: Final = scenario.key() + text: Final = marker() + sent: Final = _speech(uncached, MODELS[provider], text, speech_key) + assert sent.status == 200, sent.body[:300] + assert wav_payload(sent.body) == PCM + assert _upstream_calls(tts_wire, text) == 1 + chat_spend: Final = _chat_spend(uncached, MODELS[provider], scenario.key()) + row: Final = _spend_rows(speech_key, 1)[0] + assert row["call_type"] == "aspeech", row + assert (row["prompt_tokens"], row["completion_tokens"]) == (PROMPT_TOKENS, AUDIO_TOKENS), row + assert row["spend"] == pytest.approx(chat_spend), row + assert _key_spend(speech_key, chat_spend) == pytest.approx(chat_spend) + events: Final = _speech_events(uncached, speech_key, 1) + assert [_summary(event) for event in events] == [("aspeech", BINARY_RESPONSE, PROMPT_TOKENS, AUDIO_TOKENS)] + assert events[0]["response_cost"] == pytest.approx(chat_spend), events + assert events[0]["logged_usage_total"] == PROMPT_TOKENS + AUDIO_TOKENS, events + + +def test_cached_speech_repeat_is_a_free_cache_hit_after_a_billed_first_call( + cached: BillingProxy, tts_wire: Wire +) -> None: + model: Final = MODELS["gemini"] + with cached.proxy.gateway.scenario() as scenario: + speech_key: Final = scenario.key() + text: Final = marker() + first: Final = _speech(cached, model, text, speech_key) + repeat: Final = _speech(cached, model, text, speech_key) + assert (first.status, repeat.status) == (200, 200), (first.body[:300], repeat.body[:300]) + assert wav_payload(first.body) == PCM and repeat.body == first.body + assert _upstream_calls(tts_wire, text) == 1 + assert repeat.cache_key, "the cached repeat carried no x-litellm-cache-key header" + chat_spend: Final = _chat_spend(cached, model, scenario.key()) + billed, hit = _spend_rows(speech_key, 2) + assert billed["cache_hit"] != "True" and billed["spend"] == pytest.approx(chat_spend), billed + assert (billed["prompt_tokens"], billed["completion_tokens"]) == (PROMPT_TOKENS, AUDIO_TOKENS), billed + assert (hit["cache_hit"], hit["spend"]) == ("True", 0.0), hit + assert _key_spend(speech_key, chat_spend) == pytest.approx(chat_spend) + events: Final = _speech_events(cached, speech_key, 2) + assert len(events) == 2, events + hits: Final = tuple(event for event in events if event["cache_hit"] is True) + assert len(hits) == 1, events + assert (hits[0]["response_cost"], hits[0]["cache_key"]) == (0.0, repeat.cache_key), hits + assert hits[0]["saved_cache_cost"] == pytest.approx(chat_spend), hits + paid: Final = tuple(event for event in events if event["cache_hit"] is not True) + assert len(paid) == 1 and paid[0]["response_cost"] == pytest.approx(chat_spend), events + + +def test_speech_rejected_by_google_reaches_the_caller_and_logs_one_failure_after_one_call( + uncached: BillingProxy, tts_wire: Wire +) -> None: + with uncached.proxy.gateway.scenario() as scenario: + key: Final = scenario.key() + text: Final = marker() + sent: Final = _speech(uncached, MODELS["gemini"], text, key, voice="Custom-Voice") + assert sent.status == 400, sent.body[:300] + assert f"{REJECTED_PREFIX}Custom-Voice" in sent.body.decode(), sent.body[:300] + assert _upstream_calls(tts_wire, text) == 1 + events: Final = eventually(lambda: _events(uncached, key), lambda found: len(found) >= 2, seconds=70) + logged: Final = sorted((str(event["kind"]), event["call_type"]) for event in events) + assert logged == [("async_failure", "aspeech"), ("sync_failure", "aspeech")], events + + +def test_a_speech_reply_fills_the_deployment_tpm_limit_so_the_next_request_gets_429( + uncached: BillingProxy, tts_wire: Wire +) -> None: + eventually(lambda: datetime.now(UTC), lambda now: now.second < 10, seconds=65) + with uncached.proxy.gateway.scenario() as scenario: + key: Final = scenario.key() + first_text: Final = marker() + first: Final = _speech(uncached, TPM_MODEL, first_text, key) + assert first.status == 200, first.body[:300] + _spend_rows(key, 1) + second_text: Final = marker() + second: Final = _speech(uncached, TPM_MODEL, second_text, key) + assert second.status == 429, second.body[:300] + assert b"No deployments available" in second.body, second.body[:300] + received: Final = tuple(received_text(request) for request in tts_wire.drain()) + assert (received.count(first_text), received.count(second_text)) == (1, 0) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 63f440d4aba..aa15893e395 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4425,18 +4425,27 @@ class _RecordedSpeechSuccess: spend_metadata: Mapping[str, object] response_cost: float | None logged_response_cost: float | None + logged_tokens: tuple[object, object] + logged_cache_hit: object + logged_saved_cache_cost: object + logged_cache_key: object def _record_speech_success(payload: dict[str, object]) -> _RecordedSpeechSuccess: call_type: Final = payload.get("call_type") response_cost: Final = payload.get("response_cost") logging_payload: Final = payload.get("standard_logging_object") - logged_cost: Final = logging_payload.get("response_cost") if isinstance(logging_payload, dict) else None + logged: Final[Mapping[str, object]] = logging_payload if isinstance(logging_payload, dict) else {} + logged_cost: Final = logged.get("response_cost") return _RecordedSpeechSuccess( call_type=call_type if isinstance(call_type, str) else None, spend_metadata=get_litellm_metadata_from_kwargs(payload), response_cost=response_cost if isinstance(response_cost, float) else None, logged_response_cost=logged_cost if isinstance(logged_cost, float) else None, + logged_tokens=(logged.get("prompt_tokens"), logged.get("completion_tokens")), + logged_cache_hit=logged.get("cache_hit"), + logged_saved_cache_cost=logged.get("saved_cache_cost"), + logged_cache_key=logged.get("cache_key"), ) @@ -4445,6 +4454,11 @@ class _SuccessEventRecorder(CustomLogger): super().__init__() self.events: list[_RecordedSpeechSuccess] = [] # mutable-ok: test recorder of success-callback events + def log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events.append(_record_speech_success(kwargs)) + async def async_log_success_event( self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object ) -> None: @@ -4459,6 +4473,16 @@ async def _wait_for_success_event(recorder: _SuccessEventRecorder, call_type: st pytest.fail(f"no {call_type} success event; got {[e.call_type for e in recorder.events]}") +async def _wait_for_success_events( + recorder: _SuccessEventRecorder, call_type: str, count: int +) -> tuple[_RecordedSpeechSuccess, ...]: + for _ in range(100): + if len(events := tuple(e for e in recorder.events if e.call_type == call_type)) >= count: + return events + await asyncio.sleep(0.05) + pytest.fail(f"expected {count} {call_type} success events; got {[e.call_type for e in recorder.events]}") + + def _gemini_tts_generate_content_response() -> dict[str, object]: return { "candidates": [ @@ -4490,7 +4514,7 @@ def _gemini_tts_generate_content_response() -> dict[str, object]: @pytest.mark.asyncio -async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking( +async def test_aspeech_gemini_bridge_bills_one_call_with_proxy_metadata_and_completion_tokens( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -4510,9 +4534,11 @@ async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking( metadata={"user_api_key": "hashed-virtual-key", "user_api_key_user_id": "user-1"}, ) - assert mock_route.called + assert mock_route.call_count == 1 assert mock_route.calls.last.request.headers["x-goog-api-key"] == "fake-gemini-key" speech_event: Final = await _wait_for_success_event(recorder, call_type="aspeech") + assert recorder.events == [speech_event] + assert speech_event.logged_tokens == (5, 60) assert speech_event.spend_metadata["user_api_key"] == "hashed-virtual-key" assert speech_event.spend_metadata["user_api_key_user_id"] == "user-1" expected_prompt_cost, expected_completion_cost = litellm.cost_per_token( @@ -4525,6 +4551,39 @@ async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking( assert speech_event.logged_response_cost == pytest.approx(expected_cost) +@pytest.mark.asyncio +async def test_aspeech_gemini_bridge_logs_a_cached_reply_as_a_free_cache_hit( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + recorder: Final = _SuccessEventRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + mock_route: Final = respx_mock.post( + url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*" + ).mock(return_value=httpx.Response(200, json=_gemini_tts_generate_content_response())) + + for _ in range(2): + await litellm.aspeech( + model="gemini/gemini-2.5-flash-preview-tts", + input="cache hit check", + voice="Kore", + api_key="fake-gemini-key", + ) + + first, cached = await _wait_for_success_events(recorder, call_type="aspeech", count=2) + assert mock_route.call_count == 1 + assert len(recorder.events) == 2 + assert first.logged_cache_hit is not True + assert first.logged_response_cost is not None and first.logged_response_cost > 0 + assert cached.logged_cache_hit is True + assert cached.logged_response_cost == 0.0 + assert cached.logged_saved_cache_cost == pytest.approx(first.logged_response_cost) + assert isinstance(cached.logged_cache_key, str) and cached.logged_cache_key + + def _stream_builder_text_chunk(model: str, content: str, finish_reason: str | None = None) -> ModelResponseStream: return ModelResponseStream( id="chatcmpl-cost", From 93b79d2587d186fb7266556a34a22d03f187a3da Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 19 Oct 2026 16:26:41 -0700 Subject: [PATCH 18/18] fix(anthropic): apply tool_config cache-control points on chat and /v1/messages (#47394) * fix(anthropic): apply tool_config cache-control points on chat and /v1/messages A cache_control_injection_points entry with location tool_config was spread into the Anthropic request body on /v1/chat/completions, which Anthropic rejected with 400 Extra inputs are not permitted, and the native /v1/messages path dropped the point silently. Both paths now put the point's control on the last tool Anthropic caches and never forward the key. Co-authored-by: Tan Nguyen * refactor(anthropic): validate tool_config inputs instead of casting them Keeps the lint gates flat: the helper takes untyped tools and points and narrows them with the existing object validators, the chat transform builds the body with the marked tools instead of writing into optional_params, and the handler reads a typed view of kwargs. * fix(openai): keep tool_config cache points out of extra_body, add integration test The cache hook hands a tool_config point on for the provider step, and for providers that call through the OpenAI SDK every non-OpenAI param rides in extra_body, so OpenAI answered 400 Unrecognized request argument supplied. The new integration file drives a live proxy against scripted Anthropic, Bedrock Converse and OpenAI upstreams for /v1/chat/completions and /v1/messages. * test(integration): cover tool_config cache points on bridged /v1/messages Three new cases on /v1/messages for models without a native Anthropic messages path: OpenAI through the Responses bridge and Gemini through the chat bridge never see cache_control_injection_points in the body, and Bedrock Nova through the chat bridge gets the cachePoint on toolConfig --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Tan Nguyen --- litellm/llms/anthropic/chat/transformation.py | 15 + litellm/llms/anthropic/common_utils.py | 47 ++- .../pass_through/messages/handler.py | 10 +- .../anthropic_cache_control_hook.py | 2 +- litellm/utils.py | 5 +- .../test_tool_config_cache_point_wire.py | 312 ++++++++++++++++++ .../test_anthropic_chat_transformation.py | 96 +++++- ...erimental_pass_through_messages_handler.py | 99 ++++-- tests/unit/test_utils.py | 205 ++++++------ 9 files changed, 633 insertions(+), 158 deletions(-) create mode 100644 tests/integration/providers/test_tool_config_cache_point_wire.py diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 38af6e7998b..9624dbb582b 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -125,6 +125,7 @@ from ..common_utils import ( refusal_provider_specific_fields, requires_native_compaction_beta, strip_advisor_blocks_from_messages, + tools_marked_for_a_tool_config_point, ) if TYPE_CHECKING: @@ -2408,6 +2409,19 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Remove internal LiteLLM parameters that should not be sent to Anthropic API optional_params.pop("is_vertex_request", None) optional_params.pop("client_metadata", None) + request_params: Final = cast( # cast-ok: optional_params is the untyped dict of the shared transform signature + dict[str, object], optional_params + ) + marked_tools: Final = tools_marked_for_a_tool_config_point( + request_params.get("tools"), request_params.pop("cache_control_injection_points", None) + ) + if marked_tools is not None: + from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook + + request_litellm_params: Final = cast( # cast-ok: litellm_params is the untyped dict of the shared signature + Mapping[str, object], litellm_params + ) + AnthropicCacheControlHook.record_gateway_injection(request_litellm_params, 1) # ``top_k`` is a provider-specific kwarg that bypasses # ``map_openai_params``; gate it here, the single boundary shared by @@ -2427,6 +2441,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "model": model, "messages": anthropic_messages, **optional_params, + **({} if marked_tools is None else {"tools": list(marked_tools)}), } self._apply_output_config(data=data, model=model, optional_params=optional_params) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index dd9b22f3d5d..578329c4c11 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -50,6 +50,7 @@ from litellm.types.llms.anthropic import ( ANTHROPIC_OAUTH_TOKEN_PREFIX, ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX, ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, + ANTHROPIC_TOOL_SEARCH_TOOL_TYPES, AllAnthropicToolsValues, AnthropicMcpServerTool, AnthropicMessagesToolChoice, @@ -111,7 +112,7 @@ _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$") _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) -_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _DECODED_JSON: Final = TypeAdapter(object) @@ -160,6 +161,40 @@ def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None ) +def _tool_config_point_control(points: object) -> Mapping[object, object] | None: + point_fields: Final = tuple(validated_object_mapping(point) for point in validated_object_list(points) or ()) + tool_config_point: Final = next( + (fields for fields in point_fields if fields is not None and fields.get("location") == "tool_config"), None + ) + if tool_config_point is None: + return None + control: Final = tool_config_point.get("control") + return {"type": "ephemeral"} if control is None else validated_object_mapping(control) + + +def tools_marked_for_a_tool_config_point(tools: object, points: object) -> tuple[object, ...] | None: + """Put the ``tool_config`` point's breakpoint on the last tool Anthropic caches. + + Anthropic caches tools as one prefix ending at the marked tool, so the last one carries + the mark, skipping tool-search tools whose transform drops cache markers. A client that + already marked a tool keeps its own breakpoint. Returns ``None`` when nothing is marked. + """ + control: Final = _tool_config_point_control(points) + tool_list: Final = validated_object_list(tools) or () + tool_fields: Final = tuple(validated_object_mapping(tool) for tool in tool_list) + if control is None or any(fields is not None and "cache_control" in fields for fields in tool_fields): + return None + cacheable: Final = tuple( + (position, fields) + for position, fields in enumerate(tool_fields) + if fields is not None and fields.get("type") not in ANTHROPIC_TOOL_SEARCH_TOOL_TYPES + ) + if not cacheable: + return None + position, fields = cacheable[-1] + return (*tool_list[:position], {**fields, "cache_control": control}, *tool_list[position + 1 :]) + + def is_claude_code_user_agent(user_agent: str) -> bool: """Claude Code sends its API calls through the Anthropic SDK as `claude-cli/` and its own fetches, such as gateway model discovery, as `claude-code/`""" @@ -173,9 +208,9 @@ def validated_object_mapping(value: object) -> dict[object, object] | None: return None -def _validated_claude_code_list(value: object) -> list[object] | None: +def validated_object_list(value: object) -> Sequence[object] | None: try: - return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value) + return OBJECT_LIST_ADAPTER.validate_python(value) except ValidationError: return None @@ -202,7 +237,7 @@ def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None: if isinstance(system, str): return (system,) - blocks: Final = _validated_claude_code_list(system) + blocks: Final = validated_object_list(system) if blocks is None: return None block_mappings: Final = tuple(validated_object_mapping(block) for block in blocks) @@ -1963,8 +1998,8 @@ def tool_call_is_rebuilt_as_server_tool_use( return ( find_anthropic_server_tool_result( tool_call_id, - _validated_claude_code_list(fields.get("web_search_results")), - _validated_claude_code_list(fields.get("tool_results")), + validated_object_list(fields.get("web_search_results")), + validated_object_list(fields.get("tool_results")), ) is not None ) diff --git a/litellm/llms/anthropic/pass_through/messages/handler.py b/litellm/llms/anthropic/pass_through/messages/handler.py index 002f39c121c..d0f0fa4d901 100644 --- a/litellm/llms/anthropic/pass_through/messages/handler.py +++ b/litellm/llms/anthropic/pass_through/messages/handler.py @@ -15,10 +15,12 @@ from pydantic import TypeAdapter import litellm from litellm.anthropic_beta_headers_manager import aload_beta_headers_config +from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import ( strip_provider_specific_fields_from_anthropic_messages, + tools_marked_for_a_tool_config_point, ) from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, @@ -701,7 +703,13 @@ def anthropic_messages_handler( f"custom_llm_provider is required for Anthropic messages, passed in model={model}, custom_llm_provider={custom_llm_provider}" ) - local_vars.update(kwargs) + request_kwargs: Final = cast(dict[str, object], kwargs) # cast-ok: **kwargs of the handler is untyped + marked_tools: Final = tools_marked_for_a_tool_config_point( + local_vars.get("tools"), request_kwargs.pop("cache_control_injection_points", None) + ) + if marked_tools is not None: + AnthropicCacheControlHook.record_gateway_injection(request_kwargs, 1) + local_vars.update({**request_kwargs, **({"tools": list(marked_tools)} if marked_tools is not None else {})}) anthropic_messages_optional_request_params: Final = ( AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( params=local_vars, diff --git a/litellm/types/integrations/anthropic_cache_control_hook.py b/litellm/types/integrations/anthropic_cache_control_hook.py index 564dc1db527..636bb0b84ee 100644 --- a/litellm/types/integrations/anthropic_cache_control_hook.py +++ b/litellm/types/integrations/anthropic_cache_control_hook.py @@ -23,7 +23,7 @@ class CacheControlMessageInjectionPoint(TypedDict): class CacheControlToolConfigInjectionPoint(TypedDict): - """Type for tool_config-level injection points (Bedrock).""" + """Type for tool_config-level injection points (Anthropic and Bedrock).""" location: Literal["tool_config"] control: ChatCompletionCachedContent | None diff --git a/litellm/utils.py b/litellm/utils.py index 143c2470d33..2b7abae358b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5329,6 +5329,9 @@ def get_optional_params( EXTRA_BODY_ROUTING_KEYS: Final = frozenset({"model"}) +LITELLM_ONLY_PARAMS: Final = frozenset(("cache_control_injection_points",)) + + def add_provider_specific_params_to_optional_params( optional_params: dict, passed_params: dict, @@ -5352,7 +5355,7 @@ def add_provider_specific_params_to_optional_params( if _should_drop_param(k="extra_body", additional_drop_params=additional_drop_params) is False: extra_body: Final = dict(passed_params.pop("extra_body", None) or {}) for k in passed_params: - if k not in openai_params and passed_params[k] is not None: + if k not in openai_params and k not in LITELLM_ONLY_PARAMS and passed_params[k] is not None: extra_body[k] = passed_params[k] if not isinstance(optional_params.get("extra_body"), dict): optional_params["extra_body"] = {} diff --git a/tests/integration/providers/test_tool_config_cache_point_wire.py b/tests/integration/providers/test_tool_config_cache_point_wire.py new file mode 100644 index 00000000000..6373e122975 --- /dev/null +++ b/tests/integration/providers/test_tool_config_cache_point_wire.py @@ -0,0 +1,312 @@ +import json +import uuid +from typing import Final +from urllib.parse import unquote, urlsplit + +from integration._support.client import Gateway +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_ANTHROPIC_MODEL: Final = "claude-sonnet-4-6" +_BEDROCK_MODEL: Final = "us.anthropic.claude-sonnet-4-6" +_OPENAI_MODEL: Final = "gpt-4.1" +_GEMINI_MODEL: Final = "gemini-3.8-flash" +_NOVA_MODEL: Final = "us.amazon.nova-pro-v1:0" +_API_KEY: Final = "synthetic-provider-key" +_CONTROL: Final = {"type": "ephemeral", "ttl": "1h"} +_TOOL_CONFIG_POINT: Final = [{"location": "tool_config", "control": _CONTROL}] +_AWS: Final = { + "aws_access_key_id": "AKIAINTEGRATIONTEST", + "aws_secret_access_key": "integration-secret", + "aws_region_name": "us-east-1", +} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_TOOL_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def _openai_tools() -> list[JsonValue]: + return [ + { + "type": "function", + "function": { + "name": name, + "description": f"{name} tool", + "parameters": {"type": "object", "properties": {}}, + }, + } + for name in ("get_weather", "get_forecast") + ] + + +def _anthropic_tools() -> list[JsonValue]: + return [ + {"name": name, "description": f"{name} tool", "input_schema": {"type": "object", "properties": {}}} + for name in ("get_weather", "get_forecast") + ] + + +def _anthropic_reply(identity: str) -> bytes: + return json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": _ANTHROPIC_MODEL, + "content": [{"type": "text", "text": "done"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 3, "cache_creation_input_tokens": 10}, + } + ).encode() + + +def _converse_reply() -> bytes: + return json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "done"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 12, "outputTokens": 3, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, + } + ).encode() + + +def _assert_last_anthropic_tool_marked(body: dict[str, JsonValue]) -> None: + assert "cache_control_injection_points" not in body, body + tools: Final = _TOOL_LIST.validate_python(body["tools"]) + assert [tool["name"] for tool in tools] == ["get_weather", "get_forecast"], tools + assert "cache_control" not in tools[0], tools + assert tools[1]["cache_control"] == _CONTROL, tools + + +def test_chat_completions_tool_config_point_marks_the_last_anthropic_tool(gateway: Gateway) -> None: + identity: Final = f"anthropic-tool-config-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + _assert_last_anthropic_tool_marked(_JSON_OBJECT.validate_json(request.body)) + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_ANTHROPIC_MODEL}", + api_base=wire.url, + api_key=_API_KEY, + cache_control_injection_points=_TOOL_CONFIG_POINT, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": _openai_tools(), + }, + ) + assert response.status_code == 200, response.text + assert response.json()["usage"]["cache_creation_input_tokens"] == 10, response.text + assert len(wire.drain()) == 1 + + +def test_messages_tool_config_point_marks_the_last_anthropic_tool(gateway: Gateway) -> None: + identity: Final = f"anthropic-native-tool-config-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + _assert_last_anthropic_tool_marked(_JSON_OBJECT.validate_json(request.body)) + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_ANTHROPIC_MODEL}", + api_base=wire.url, + api_key=_API_KEY, + cache_control_injection_points=_TOOL_CONFIG_POINT, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": _anthropic_tools(), + }, + ) + assert response.status_code == 200, response.text + assert response.json()["usage"]["cache_creation_input_tokens"] == 10, response.text + assert len(wire.drain()) == 1 + + +def test_chat_completions_tool_config_point_appends_a_bedrock_converse_cache_point(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + path: Final = unquote(urlsplit(request.target).path) + assert request.method == "POST" and path == f"/model/{_BEDROCK_MODEL}/converse", request.target + body: Final = _JSON_OBJECT.validate_json(request.body) + assert "cache_control_injection_points" not in body, body + tools: Final = _TOOL_LIST.validate_python(_JSON_OBJECT.validate_python(body["toolConfig"])["tools"]) + assert [tuple(tool) for tool in tools] == [("toolSpec",), ("toolSpec",), ("cachePoint",)], tools + assert tools[2] == {"cachePoint": {"type": "default", "ttl": "1h"}}, tools + return Reply(body=_converse_reply()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/{_BEDROCK_MODEL}", + api_base=wire.url, + api_key=None, + cache_control_injection_points=_TOOL_CONFIG_POINT, + **_AWS, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": _openai_tools(), + }, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_chat_completions_tool_config_point_never_reaches_openai(gateway: Gateway) -> None: + identity: Final = f"openai-tool-config-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/chat/completions", request.target + body: Final = _JSON_OBJECT.validate_json(request.body) + assert "cache_control_injection_points" not in body, body + tools: Final = _TOOL_LIST.validate_python(body["tools"]) + assert [tuple(sorted(tool)) for tool in tools] == [("function", "type"), ("function", "type")], tools + return chat_reply(identity, _OPENAI_MODEL, "done", stream=False) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{_OPENAI_MODEL}", + api_base=wire.url, + api_key=_API_KEY, + cache_control_injection_points=_TOOL_CONFIG_POINT, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": _openai_tools(), + }, + ) + assert response.status_code == 200, response.text + assert response.json()["id"] == identity, response.text + assert len(wire.drain()) == 1 + + +def _gemini_reply() -> bytes: + return json.dumps( + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "done"}]}, "finishReason": "STOP", "index": 0} + ], + "usageMetadata": {"promptTokenCount": 12, "candidatesTokenCount": 3, "totalTokenCount": 15}, + "modelVersion": _GEMINI_MODEL, + } + ).encode() + + +def test_messages_tool_config_point_never_reaches_openai_through_the_responses_bridge(gateway: Gateway) -> None: + identity: Final = f"openai-bridge-tool-config-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + body: Final = _JSON_OBJECT.validate_json(request.body) + assert "cache_control_injection_points" not in body, body + return responses_reply(identity, _OPENAI_MODEL, "done", stream=False) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{_OPENAI_MODEL}", + api_base=wire.url, + api_key=_API_KEY, + cache_control_injection_points=_TOOL_CONFIG_POINT, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": _anthropic_tools(), + }, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_messages_tool_config_point_never_reaches_gemini_through_the_chat_bridge(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + path: Final = urlsplit(request.target).path + assert request.method == "POST" and path == f"/models/{_GEMINI_MODEL}:generateContent", request.target + body: Final = _JSON_OBJECT.validate_json(request.body) + assert "cache_control_injection_points" not in body, body + return Reply(body=_gemini_reply()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"gemini/{_GEMINI_MODEL}", + api_base=wire.url, + api_key=_API_KEY, + cache_control_injection_points=_TOOL_CONFIG_POINT, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": _anthropic_tools(), + }, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_messages_tool_config_point_appends_a_bedrock_converse_cache_point_through_the_chat_bridge( + gateway: Gateway, +) -> None: + def respond(request: Request) -> Reply: + path: Final = unquote(urlsplit(request.target).path) + assert request.method == "POST" and path == f"/model/{_NOVA_MODEL}/converse", request.target + body: Final = _JSON_OBJECT.validate_json(request.body) + assert "cache_control_injection_points" not in body, body + tools: Final = _TOOL_LIST.validate_python(_JSON_OBJECT.validate_python(body["toolConfig"])["tools"]) + assert [tuple(tool) for tool in tools] == [("toolSpec",), ("toolSpec",), ("cachePoint",)], tools + assert tools[2] == {"cachePoint": {"type": "default"}}, tools + return Reply(body=_converse_reply()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/{_NOVA_MODEL}", + api_base=wire.url, + api_key=None, + cache_control_injection_points=_TOOL_CONFIG_POINT, + **_AWS, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": _anthropic_tools(), + }, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 1308c5aa01f..ef8447712d6 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1375,7 +1375,6 @@ def test_dict_fields_take_the_json_tool_on_a_native_structured_output_model(fiel assert ([tool["name"] for tool in optional_params.get("tools", [])] == ["json_tool_call"]) is not native - @pytest.mark.parametrize( ("model", "thinking_params"), [ @@ -1388,7 +1387,13 @@ def test_dict_fields_take_the_json_tool_on_a_native_structured_output_model(fiel {"tools": [{"type": "function", "function": {"name": "lookup", "parameters": {"type": "object"}}}]}, ), ], - ids=["adaptive-thinking", "reasoning-effort", "reasoning-effort-none", "model-without-forced-tools", "caller-tools"], + ids=[ + "adaptive-thinking", + "reasoning-effort", + "reasoning-effort-none", + "model-without-forced-tools", + "caller-tools", + ], ) def test_dict_fields_keep_native_output_when_the_json_tool_would_not_be_forced(model, thinking_params): optional_params: Final = AnthropicConfig().map_openai_params( @@ -1414,6 +1419,7 @@ def test_dict_fields_keep_native_output_when_the_json_tool_would_not_be_forced(m assert "output_format" in optional_params assert "json_tool_call" not in [tool["name"] for tool in optional_params.get("tools", [])] + def test_native_structured_output_uses_bundled_capability_when_remote_map_lags( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -3704,7 +3710,9 @@ def _block_ids(message: dict[str, object], block_type: str, id_key: str) -> list return sorted(block[id_key] for block in message["content"] if block["type"] == block_type) -_SEARCH_THAT_FOUND_NOTHING: Final = _tool_search_turn([], {"type": "text", "text": "I have no tool that books flights."}) +_SEARCH_THAT_FOUND_NOTHING: Final = _tool_search_turn( + [], {"type": "text", "text": "I have no tool that books flights."} +) @pytest.mark.parametrize( @@ -9528,7 +9536,9 @@ def test_map_openai_params_runs_tool_search_on_a_custom_named_deployment_whose_m AzureAnthropicConfig().map_openai_params( non_default_params={"tools": tools}, optional_params={}, model="my-claude-deployment", drop_params=False ) - litellm.register_model({"azure_ai/my-claude-deployment": {"supports_tool_search": True, "litellm_provider": "azure_ai"}}) + litellm.register_model( + {"azure_ai/my-claude-deployment": {"supports_tool_search": True, "litellm_provider": "azure_ai"}} + ) mapped: Final = AzureAnthropicConfig().map_openai_params( non_default_params={"tools": tools}, optional_params={}, model="my-claude-deployment", drop_params=False ) @@ -9653,3 +9663,81 @@ def test_non_dict_upstream_body_still_transforms_without_refusal_fields() -> Non ) provider_specific_fields = result.choices[0].message.provider_specific_fields or {} assert "stop_details" not in provider_specific_fields + + +def _two_function_tools() -> list[dict]: + return [ + { + "type": "function", + "function": {"name": f"tool_{i}", "description": f"Tool {i}", "parameters": {"type": "object"}}, + } + for i in range(2) + ] + + +@pytest.mark.parametrize( + "control, expected_control", + [ + (None, {"type": "ephemeral"}), + ({"type": "ephemeral", "ttl": "1h"}, {"type": "ephemeral", "ttl": "1h"}), + ], +) +def test_transform_request_tool_config_point_marks_last_tool_and_drops_the_key(control, expected_control): + """LIT-9880: a tool_config injection point used to be spread into the Anthropic body, which + Anthropic rejected with 400 Extra inputs are not permitted.""" + tools = _two_function_tools() + tools_before = copy.deepcopy(tools) + point = {"location": "tool_config", **({"control": control} if control is not None else {})} + litellm_params = {"metadata": {}, "model_info": {"id": "deployment-1"}} + + data = _make_transform_request( + optional_params={"tools": tools, "cache_control_injection_points": [point]}, + litellm_params=litellm_params, + ) + + assert "cache_control_injection_points" not in data + assert data["tools"][-1]["cache_control"] == expected_control + assert "cache_control" not in data["tools"][0] + assert tools == tools_before + assert litellm_params["metadata"]["litellm_gateway_injected_cache"] == "deployment-1" + + +def test_transform_request_tool_config_point_keeps_the_clients_tool_marker(): + tools = _two_function_tools() + tools[0]["cache_control"] = {"type": "ephemeral"} + + data = _make_transform_request( + optional_params={"tools": tools, "cache_control_injection_points": [{"location": "tool_config"}]}, + litellm_params={}, + ) + + assert "cache_control_injection_points" not in data + assert data["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in data["tools"][1] + + +def test_transform_request_tool_config_point_skips_tool_search_tools(): + tool_search_type = next(iter(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES)) + tools = [*_two_function_tools(), {"type": tool_search_type, "name": "tool_search"}] + + data = _make_transform_request( + optional_params={"tools": tools, "cache_control_injection_points": [{"location": "tool_config"}]}, + litellm_params={}, + ) + + assert data["tools"][1]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in data["tools"][2] + assert "cache_control" not in data["tools"][0] + + +def test_transform_request_tool_config_point_without_tools_drops_the_key(): + litellm_params = {"metadata": {}} + + data = _make_transform_request( + optional_params={"cache_control_injection_points": [{"location": "tool_config"}]}, + litellm_params=litellm_params, + ) + + assert "cache_control_injection_points" not in data + assert "tools" not in data + assert litellm_params["metadata"] == {} diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index ffe9e46bf56..1f02b3aaa5a 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1103,7 +1103,9 @@ async def test_messages_honors_a_chat_pin_on_a_prefixed_responses_deployment( litellm.in_memory_llm_clients_cache.flush_cache() with respx.mock(assert_all_called=True) as respx_mock: - upstream: Final = respx_mock.route(host="api.perplexity.ai").mock(side_effect=httpx.ConnectError("cut at the wire")) + upstream: Final = respx_mock.route(host="api.perplexity.ai").mock( + side_effect=httpx.ConnectError("cut at the wire") + ) with pytest.raises(litellm.exceptions.InternalServerError): await litellm.anthropic.messages.acreate( max_tokens=100, @@ -1524,6 +1526,7 @@ async def test_provider_messages_api_base_env_is_not_shadowed_by_the_chat_defaul assert seen_urls == ["https://deepseek.internal.example/anthropic/v1/messages"] + @pytest.mark.asyncio async def test_anthropic_messages_forwards_safeguards_and_unknown_beta_to_anthropic(): """Shapes are what Claude Code 2.1.278 sends and api.anthropic.com returns, captured 2026-09-21.""" @@ -1638,7 +1641,9 @@ def _claude_code_auto_mode_request() -> tuple[list[dict[str, object]], list[dict return safeguards, safeguard_results -def _upstream_answering_with(safeguard_results: list[dict[str, object]], captured: dict[str, object]) -> AsyncHTTPHandler: +def _upstream_answering_with( + safeguard_results: list[dict[str, object]], captured: dict[str, object] +) -> AsyncHTTPHandler: def upstream_records_the_request(request: httpx.Request) -> httpx.Response: captured["body"] = json.loads(request.content) captured["anthropic-beta"] = request.headers.get("anthropic-beta") @@ -1664,7 +1669,9 @@ def _upstream_answering_with(safeguard_results: list[dict[str, object]], capture _CLIENT_BETA_HEADERS: Final = ( - pytest.param({"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, id="client_sends_beta"), + pytest.param( + {"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, id="client_sends_beta" + ), pytest.param({"anthropic-beta": "interleaved-thinking-2025-05-14"}, id="client_omits_beta"), pytest.param({}, id="client_sends_no_beta_header"), ) @@ -1789,7 +1796,11 @@ def _tool_use_sse(model: str) -> bytes: "delta": {"type": "input_json_delta", "partial_json": json.dumps(_CURL_TOOL_USE["input"])}, }, {"type": "content_block_stop", "index": 0}, - {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 4}}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + }, {"type": "message_stop"}, ) return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode() @@ -1957,26 +1968,20 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking(): print("Call kwargs:", json.dumps(call_kwargs, indent=2, default=str)) # Verify that litellm_metadata was passed and contains _latency_per_deployment - assert ( - "litellm_metadata" in call_kwargs - ), "litellm_metadata should be passed to anthropic_messages" + assert "litellm_metadata" in call_kwargs, "litellm_metadata should be passed to anthropic_messages" litellm_metadata = call_kwargs["litellm_metadata"] assert litellm_metadata is not None, "litellm_metadata should not be None" - assert isinstance( - litellm_metadata, dict - ), "litellm_metadata should be a dictionary" + assert isinstance(litellm_metadata, dict), "litellm_metadata should be a dictionary" # Verify _latency_per_deployment is present - assert ( - "_latency_per_deployment" in litellm_metadata - ), "litellm_metadata should contain _latency_per_deployment field" + assert "_latency_per_deployment" in litellm_metadata, ( + "litellm_metadata should contain _latency_per_deployment field" + ) # Verify the structure of _latency_per_deployment latency_per_deployment = litellm_metadata["_latency_per_deployment"] - assert isinstance( - latency_per_deployment, dict - ), "_latency_per_deployment should be a dictionary" + assert isinstance(latency_per_deployment, dict), "_latency_per_deployment should be a dictionary" print(f"✅ Latency per deployment data: {latency_per_deployment}") @@ -2134,9 +2139,7 @@ async def test_anthropic_messages_bedrock_credentials_passthrough(): when using anthropic.messages.acreate with a bedrock model """ # Mock the get_credentials method - with unittest.mock.patch.object( - BaseAWSLLM, "get_credentials" - ) as mock_get_credentials: + with unittest.mock.patch.object(BaseAWSLLM, "get_credentials") as mock_get_credentials: # Create a proper mock for credentials with the necessary attributes mock_credentials = unittest.mock.MagicMock() mock_credentials.access_key = "mock_access_key" @@ -2147,9 +2150,7 @@ async def test_anthropic_messages_bedrock_credentials_passthrough(): # We also need to mock the actual AWS request signing to avoid real API calls with unittest.mock.patch("botocore.auth.SigV4Auth.add_auth"): # Set up mock for AsyncHTTPHandler.post to avoid actual API calls - with unittest.mock.patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post" - ) as mock_post: + with unittest.mock.patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") as mock_post: # Configure mock response mock_response = unittest.mock.MagicMock() mock_response.raise_for_status = unittest.mock.MagicMock() @@ -2191,9 +2192,7 @@ async def test_anthropic_messages_bedrock_credentials_passthrough(): # Assert that our test credentials were passed correctly for param_name, param_value in aws_params.items(): - assert ( - call_args[param_name] == param_value - ), f"Parameter {param_name} was not passed correctly" + assert call_args[param_name] == param_value, f"Parameter {param_name} was not passed correctly" @pytest.mark.asyncio @@ -2221,11 +2220,8 @@ async def test_anthropic_messages_bedrock_dynamic_region(): # Patch necessary AWS components with ( unittest.mock.patch("botocore.auth.SigV4Auth.add_auth"), - unittest.mock.patch.object( - BaseAWSLLM, "get_credentials" - ) as mock_get_credentials, + unittest.mock.patch.object(BaseAWSLLM, "get_credentials") as mock_get_credentials, ): - # Setup mock credentials mock_credentials = unittest.mock.MagicMock() mock_credentials.access_key = "test_access_key" @@ -2254,9 +2250,9 @@ async def test_anthropic_messages_bedrock_dynamic_region(): # Check that the URL contains the correct region url = call_args.kwargs.get("url", "") - assert ( - f"bedrock-runtime.{test_region}.amazonaws.com" in url - ), f"URL does not contain the correct region. URL: {url}" + assert f"bedrock-runtime.{test_region}.amazonaws.com" in url, ( + f"URL does not contain the correct region. URL: {url}" + ) # Verify get_credentials was called with the correct region mock_get_credentials.assert_called_once() @@ -2294,3 +2290,42 @@ async def test_anthropic_messages_with_safeguards_and_no_classifier_keep_the_ups assert rejected.value.status_code == 400 assert "invalid_request_error" in str(rejected.value) + + +def test_native_messages_tool_config_point_marks_last_tool_and_drops_the_key(monkeypatch): + """LIT-9880 on /v1/messages: the hook wrote the tool_config point back into kwargs and the + native path dropped it silently, so the tools went out without any cache breakpoint.""" + from litellm.llms.anthropic.pass_through.messages import handler + + captured = {} + + def fake_native(**kwargs): + captured["params"] = kwargs["anthropic_messages_optional_request_params"] + captured["kwargs"] = kwargs["kwargs"] + captured["messages"] = kwargs["messages"] + return "native-passthrough" + + monkeypatch.setattr(handler.base_llm_http_handler, "anthropic_messages_handler", fake_native) + tools = [{"name": f"tool_{i}", "description": f"Tool {i}", "input_schema": {"type": "object"}} for i in range(2)] + tools_before = copy.deepcopy(tools) + + result = handler.anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "Hello"}], + model="anthropic/claude-opus-4-6", + api_key="sk-test", + tools=tools, + litellm_metadata={}, + cache_control_injection_points=[ + {"location": "message", "role": "user"}, + {"location": "tool_config", "control": {"type": "ephemeral", "ttl": "1h"}}, + ], + ) + + assert result == "native-passthrough" + assert "cache_control_injection_points" not in captured["params"] + assert "cache_control_injection_points" not in captured["kwargs"] + assert captured["params"]["tools"][-1]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + assert "cache_control" not in captured["params"]["tools"][0] + assert captured["messages"][0]["content"][-1]["cache_control"] == {"type": "ephemeral"} + assert tools == tools_before diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 69ef2caf8f4..6ffb589fd41 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1465,8 +1465,6 @@ def test_get_model_info_bedrock_inference_profile_arn_keeps_regional_row(local_m assert bare["input_cost_per_token"] == base["input_cost_per_token"] - - @pytest.mark.parametrize( "bare_key", [ @@ -2302,9 +2300,7 @@ def test_add_known_models_skips_region_scoped_bedrock_pricing_keys(monkeypatch: def test_utils_defines_tests_without_invoking_them_at_module_scope() -> None: tree: Final = ast.parse(Path(__file__).read_text()) defined: Final = frozenset( - node.name - for node in tree.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + node.name for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) ) invoked: Final = tuple( node.value.func.id @@ -3619,9 +3615,7 @@ class TestExtraBodyCannotOverrideModel: } ] - untouched = litellm.get_optional_params( - model="my-vllm-model", custom_llm_provider="hosted_vllm", tools=tools() - ) + untouched = litellm.get_optional_params(model="my-vllm-model", custom_llm_provider="hosted_vllm", tools=tools()) assert untouched["tools"][0]["function"]["custom_marker"] == "LEAK", untouched result = litellm.get_optional_params( @@ -4661,7 +4655,9 @@ def test_is_prompt_caching_valid_prompt_stops_counting_once_the_minimum_is_reach assert is_prompt_caching_valid_prompt(model="claude-opus-4-8", messages=long_prompt, min_token_count=1024) is True assert sum(counted_messages) < len(long_prompt), sum(counted_messages) - full_count = litellm.token_counter(model="claude-opus-4-8", messages=long_prompt, use_default_image_token_count=True) + full_count = litellm.token_counter( + model="claude-opus-4-8", messages=long_prompt, use_default_image_token_count=True + ) assert ( is_prompt_caching_valid_prompt(model="claude-opus-4-8", messages=long_prompt, min_token_count=full_count) is True @@ -4977,7 +4973,9 @@ class _EntryReadDuringARegistration(dict[str, object]): @pytest.mark.parametrize("walk", ["fal_ai", "fireworks_ai", "ollama", "gpt-5", "case-insensitive scan"]) -def test_request_path_cost_map_walks_survive_a_registration_mid_walk(monkeypatch: pytest.MonkeyPatch, walk: str) -> None: +def test_request_path_cost_map_walks_survive_a_registration_mid_walk( + monkeypatch: pytest.MonkeyPatch, walk: str +) -> None: from litellm.llms.fal_ai.cost_calculator import _keyed_rows from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.llms.ollama.common_utils import OllamaModelInfo @@ -5226,7 +5224,13 @@ _SUCCESS_RESPONSES_BY_CALL_TYPE: Final = ( pytest.param(EmbeddingResponse(model="text-embedding-3-small"), CallTypes.aembedding, id="embedding"), pytest.param( ResponsesAPIResponse( - id="resp_abc", created_at=1, output=[], parallel_tool_calls=False, tool_choice="auto", tools=[], model="gpt-5.6" + id="resp_abc", + created_at=1, + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + model="gpt-5.6", ), CallTypes.aresponses, id="responses", @@ -5254,7 +5258,9 @@ async def test_success_deployment_hook_raising_keeps_response_and_runs_later_hoo assert result is response assert second_hook.seen_responses == (response,) - failure_logs: Final = tuple(r for r in caplog.records if "async_post_call_success_deployment_hook error" in r.message) + failure_logs: Final = tuple( + r for r in caplog.records if "async_post_call_success_deployment_hook error" in r.message + ) assert len(failure_logs) == 1 assert "_ChatShapedSuccessDeploymentHook" in failure_logs[0].message assert str(call_type) in failure_logs[0].message @@ -5391,17 +5397,13 @@ async def test_wrapper_async_logs_converted_responses_stream_with_standard_loggi ({"metadata": None}, True, False), ], ) -def test_is_litellm_router_call_is_async_aware( - kwargs: Mapping[str, object], is_async: bool, expected: bool -) -> None: +def test_is_litellm_router_call_is_async_aware(kwargs: Mapping[str, object], is_async: bool, expected: bool) -> None: assert _is_litellm_router_call(kwargs, is_async=is_async) is expected @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) -async def test_router_aresponses_does_not_run_sdk_retries( - monkeypatch: pytest.MonkeyPatch, stream: bool -) -> None: +async def test_router_aresponses_does_not_run_sdk_retries(monkeypatch: pytest.MonkeyPatch, stream: bool) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() model_list: Final = [ @@ -5415,9 +5417,7 @@ async def test_router_aresponses_does_not_run_sdk_retries( }, } ] - router: Final = litellm.Router( - model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True - ) + router: Final = litellm.Router(model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True) try: with respx.mock(assert_all_called=True) as respx_mock: @@ -5453,9 +5453,7 @@ def test_router_responses_keeps_sdk_retries_for_sync_router_call( }, } ] - router: Final = litellm.Router( - model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True - ) + router: Final = litellm.Router(model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True) try: with respx.mock(assert_all_called=True) as respx_mock: @@ -7208,7 +7206,9 @@ def test_get_model_info_gemini(monkeypatch): assert info.get("rpm") is not None, f"{model} does not have rpm" -def test_register_model_without_a_provider_never_inherits_another_providers_entry(monkeypatch: pytest.MonkeyPatch) -> None: +def test_register_model_without_a_provider_never_inherits_another_providers_entry( + monkeypatch: pytest.MonkeyPatch, +) -> None: from litellm.utils import _invalidate_model_cost_lowercase_map key: Final = "azure/claude-sonnet-4-6" @@ -7359,7 +7359,9 @@ def test_function_setup_logs_the_search_query_edit_prompt_and_ocr_document_summa @pytest.mark.parametrize("original_function", ("atext_completion", "text_completion")) -def test_function_setup_without_a_prompt_leaves_the_missing_prompt_to_request_validation(original_function: str) -> None: +def test_function_setup_without_a_prompt_leaves_the_missing_prompt_to_request_validation( + original_function: str, +) -> None: assert _logged_request_messages(original_function, model="gpt-4o") is None @@ -7401,6 +7403,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def setup_and_teardown(): """ @@ -7414,6 +7417,7 @@ def setup_and_teardown(): loop.close() asyncio.set_event_loop(None) + MODEL: Final = "anthropic/claude-haiku-4-5" @@ -7423,6 +7427,7 @@ def test_validate_tool_choice_none(): result = validate_chat_completion_tool_choice(None, model=MODEL) assert result is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_string(): """Test that string values are returned as-is.""" @@ -7430,6 +7435,7 @@ def test_validate_tool_choice_string(): assert validate_chat_completion_tool_choice("none", model=MODEL) == "none" assert validate_chat_completion_tool_choice("required", model=MODEL) == "required" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_standard_dict(): """Test standard OpenAI format with function.""" @@ -7437,6 +7443,7 @@ def test_validate_tool_choice_standard_dict(): result = validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert result == tool_choice + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_cursor_format(): """Cursor IDE format {"type": "auto"} is unwrapped to the bare string.""" @@ -7444,6 +7451,7 @@ def test_validate_tool_choice_cursor_format(): assert validate_chat_completion_tool_choice({"type": "none"}, model=MODEL) == "none" assert validate_chat_completion_tool_choice({"type": "required"}, model=MODEL) == "required" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize( "tool_choice", @@ -7465,6 +7473,7 @@ def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): assert exc_info.value.status_code == 400 assert exc_info.value.model == MODEL + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize("tool_choice", [123, []]) def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): @@ -7475,6 +7484,7 @@ def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert exc_info.value.status_code == 400 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_without_model_is_still_a_400(): """Callers that predate the model argument keep getting a 400, with an empty model on the error.""" @@ -7483,12 +7493,14 @@ def test_validate_tool_choice_without_model_is_still_a_400(): assert exc_info.value.status_code == 400 assert exc_info.value.model == "" + @pytest.fixture() def _vcr_outcome_gate_local_testing(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -7541,6 +7553,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -7584,6 +7597,7 @@ def setup_and_teardown_local_testing(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7609,6 +7623,7 @@ def test_vertex_projects(): assert "vertex_ai_project" in optional_params assert "vertex_ai_location" in optional_params + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7624,6 +7639,7 @@ def test_bedrock_embed_v2_regular(): print(f"received optional_params: {optional_params}") assert optional_params == {"dimensions": 512} + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7642,6 +7658,7 @@ def test_bedrock_embed_v2_with_drop_params(): print(f"received optional_params: {optional_params}") assert optional_params == {"dimensions": 512, "embeddingTypes": ["binary"]} + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7663,6 +7680,7 @@ def test_openai_non_text_embedding_3_with_allowed_openai_params(): print(f"received optional_params: {optional_params}") assert optional_params.get("dimensions") == 1024 + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7689,6 +7707,7 @@ def test_openai_non_text_embedding_3_without_allowed_openai_params_raises(): finally: litellm.drop_params = prev_drop_params + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7717,6 +7736,7 @@ def test_openai_non_text_embedding_3_drop_params_per_call(): finally: litellm.drop_params = prev_drop_params + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7744,12 +7764,14 @@ def test_openai_non_text_embedding_3_drop_params_global(): finally: litellm.drop_params = prev_drop_params + @pytest.fixture() def _vcr_outcome_gate_search_tests(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.mark.usefixtures("_vcr_outcome_gate_search_tests") def test_search_tool_name_in_all_litellm_params(): """ @@ -7759,6 +7781,7 @@ def test_search_tool_name_in_all_litellm_params(): """ assert "search_tool_name" in all_litellm_params + @pytest.mark.usefixtures("_vcr_outcome_gate_search_tests") def test_filter_out_search_tool_name(): """ @@ -7785,6 +7808,7 @@ def test_filter_out_search_tool_name(): assert filtered["query"] == "latest ai developments" assert filtered["max_results"] == 5 + @pytest.mark.asyncio async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None: """Chat over the Responses bridge exits two @client wrappers with one logging object. Issue @@ -7937,9 +7961,7 @@ def test_trimming_with_system_message_within_max_tokens(): "content": "This is a medium normal message, let's say litellm is awesome.", }, ] - trimmed_messages = trim_messages( - messages, max_tokens=30, model="gpt-4-0613" - ) + trimmed_messages = trim_messages(messages, max_tokens=30, model="gpt-4-0613") assert len(trimmed_messages) == 2 assert trimmed_messages[0]["content"] == "This is a short system message" @@ -8122,9 +8144,7 @@ def test_aget_valid_models(): print(valid_models) # list of openai supported llms on litellm - expected_models = ( - litellm.open_ai_chat_completion_models | litellm.open_ai_text_completion_models - ) + expected_models = litellm.open_ai_chat_completion_models | litellm.open_ai_text_completion_models assert set(valid_models) == set(expected_models) @@ -8213,9 +8233,7 @@ def test_get_chat_completion_prompt(): prompt_variables=None, ) - assert litellm_logging_obj.messages == [ - {"role": "user", "content": updated_message} - ] + assert litellm_logging_obj.messages == [{"role": "user", "content": updated_message}] def test_redact_msgs_from_logs(): @@ -8287,9 +8305,7 @@ def test_redact_embedding_response(): litellm.turn_off_message_logging = True # Create a test EmbeddingResponse with usage data - original_usage = litellm.Usage( - prompt_tokens=10, completion_tokens=0, total_tokens=10 - ) + original_usage = litellm.Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) original_data = [ {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3, 0.4, 0.5]}, {"object": "embedding", "index": 1, "embedding": [0.6, 0.7, 0.8, 0.9, 1.0]}, @@ -8325,9 +8341,7 @@ def test_redact_embedding_response(): # Assert the redacted response preserves critical metadata assert _redacted_response_obj.usage == original_usage # usage should be preserved - assert ( - _redacted_response_obj.model == "text-embedding-3-small" - ) # model should be preserved + assert _redacted_response_obj.model == "text-embedding-3-small" # model should be preserved assert _redacted_response_obj.object == "list" # object should be preserved # Assert sensitive data is cleared @@ -8381,12 +8395,8 @@ def test_redact_msgs_from_logs_with_dynamic_params(): ) # Test Case 1: standard_callback_dynamic_params = False (or not set) - standard_callback_dynamic_params = StandardCallbackDynamicParams( - turn_off_message_logging=False - ) - litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = ( - standard_callback_dynamic_params - ) + standard_callback_dynamic_params = StandardCallbackDynamicParams(turn_off_message_logging=False) + litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = standard_callback_dynamic_params _redacted_response_obj = redact_message_input_output_from_logging( result=response_obj, model_call_details=litellm_logging_obj.model_call_details, @@ -8395,12 +8405,8 @@ def test_redact_msgs_from_logs_with_dynamic_params(): assert _redacted_response_obj.choices[0].message.content == test_content # Test Case 2: standard_callback_dynamic_params = True - standard_callback_dynamic_params = StandardCallbackDynamicParams( - turn_off_message_logging=True - ) - litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = ( - standard_callback_dynamic_params - ) + standard_callback_dynamic_params = StandardCallbackDynamicParams(turn_off_message_logging=True) + litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = standard_callback_dynamic_params _redacted_response_obj = redact_message_input_output_from_logging( result=response_obj, model_call_details=litellm_logging_obj.model_call_details, @@ -8411,9 +8417,7 @@ def test_redact_msgs_from_logs_with_dynamic_params(): # Test Case 3: standard_callback_dynamic_params does not set turn_off_message_logging # since litellm.turn_off_message_logging is True redaction should occur standard_callback_dynamic_params = StandardCallbackDynamicParams() - litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = ( - standard_callback_dynamic_params - ) + litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = standard_callback_dynamic_params _redacted_response_obj = redact_message_input_output_from_logging( result=response_obj, model_call_details=litellm_logging_obj.model_call_details, @@ -8755,10 +8759,7 @@ def test_models_by_provider(): continue elif k == "sample_spec": continue - elif ( - v["litellm_provider"] == "sagemaker" - or v["litellm_provider"] == "bedrock_converse" - ): + elif v["litellm_provider"] == "sagemaker" or v["litellm_provider"] == "bedrock_converse": continue elif v.get("mode") in ("search", "decisions", "evaluation"): continue @@ -8766,9 +8767,7 @@ def test_models_by_provider(): providers.add(v["litellm_provider"]) for provider in providers: - assert provider in models_by_provider.keys() or JSONProviderRegistry.exists( - provider - ) + assert provider in models_by_provider.keys() or JSONProviderRegistry.exists(provider) @pytest.fixture @@ -8790,16 +8789,11 @@ def restore_end_user_cost_tracking_flags(monkeypatch: pytest.MonkeyPatch) -> Non ({"user_api_key_end_user_id": "123"}, True, None), ], ) -def test_get_end_user_id_for_cost_tracking( - litellm_params, disable_end_user_cost_tracking, expected_end_user_id -): +def test_get_end_user_id_for_cost_tracking(litellm_params, disable_end_user_cost_tracking, expected_end_user_id): from litellm.utils import get_end_user_id_for_cost_tracking litellm.disable_end_user_cost_tracking = disable_end_user_cost_tracking - assert ( - get_end_user_id_for_cost_tracking(litellm_params=litellm_params) - == expected_end_user_id - ) + assert get_end_user_id_for_cost_tracking(litellm_params=litellm_params) == expected_end_user_id @pytest.mark.usefixtures("restore_end_user_cost_tracking_flags") @@ -8816,13 +8810,9 @@ def test_get_end_user_id_for_cost_tracking_prometheus_only( ): from litellm.utils import get_end_user_id_for_cost_tracking - litellm.enable_end_user_cost_tracking_prometheus_only = ( - enable_end_user_cost_tracking_prometheus_only - ) + litellm.enable_end_user_cost_tracking_prometheus_only = enable_end_user_cost_tracking_prometheus_only assert ( - get_end_user_id_for_cost_tracking( - litellm_params=litellm_params, service_type="prometheus" - ) + get_end_user_id_for_cost_tracking(litellm_params=litellm_params, service_type="prometheus") == expected_end_user_id ) @@ -8837,20 +8827,14 @@ def test_get_end_user_id_for_cost_tracking_prometheus_only( ), # Test with only litellm_metadata field (new behavior) ( - { - "litellm_metadata": { - "user_api_key_end_user_id": "user_from_litellm_metadata" - } - }, + {"litellm_metadata": {"user_api_key_end_user_id": "user_from_litellm_metadata"}}, "user_from_litellm_metadata", ), # Test with both fields - metadata should take precedence for user_api_key fields ( { "metadata": {"user_api_key_end_user_id": "user_from_metadata"}, - "litellm_metadata": { - "user_api_key_end_user_id": "user_from_litellm_metadata" - }, + "litellm_metadata": {"user_api_key_end_user_id": "user_from_litellm_metadata"}, }, "user_from_metadata", ), @@ -8866,9 +8850,7 @@ def test_get_end_user_id_for_cost_tracking_prometheus_only( ( { "metadata": {}, - "litellm_metadata": { - "user_api_key_end_user_id": "user_from_litellm_metadata" - }, + "litellm_metadata": {"user_api_key_end_user_id": "user_from_litellm_metadata"}, }, "user_from_litellm_metadata", ), @@ -8876,9 +8858,7 @@ def test_get_end_user_id_for_cost_tracking_prometheus_only( ({}, None), ], ) -def test_get_end_user_id_for_cost_tracking_metadata_handling( - litellm_params, expected_end_user_id -): +def test_get_end_user_id_for_cost_tracking_metadata_handling(litellm_params, expected_end_user_id): """ Test that get_end_user_id_for_cost_tracking correctly handles both metadata and litellm_metadata fields using the get_litellm_metadata_from_kwargs helper function. @@ -9136,9 +9116,7 @@ def test_add_custom_logger_callback_to_specific_event_e2e_failure(monkeypatch): assert len(litellm.success_callback) == curr_len_success_callback assert len(litellm.failure_callback) == curr_len_failure_callback - assert any( - isinstance(callback, OpenMeterLogger) for callback in litellm.failure_callback - ) + assert any(isinstance(callback, OpenMeterLogger) for callback in litellm.failure_callback) @pytest.mark.asyncio @@ -9235,20 +9213,14 @@ def test_validate_user_messages_invalid_content_type(): [ { "name": "default_on_guardrail", - "callbacks": [ - CustomGuardrail(guardrail_name="test_guardrail", default_on=True) - ], + "callbacks": [CustomGuardrail(guardrail_name="test_guardrail", default_on=True)], "kwargs": {"metadata": {"requester_metadata": {"guardrails": []}}}, "expected": ["test_guardrail"], }, { "name": "request_specific_guardrail", - "callbacks": [ - CustomGuardrail(guardrail_name="test_guardrail", default_on=False) - ], - "kwargs": { - "metadata": {"requester_metadata": {"guardrails": ["test_guardrail"]}} - }, + "callbacks": [CustomGuardrail(guardrail_name="test_guardrail", default_on=False)], + "kwargs": {"metadata": {"requester_metadata": {"guardrails": ["test_guardrail"]}}}, "expected": ["test_guardrail"], }, { @@ -9257,18 +9229,12 @@ def test_validate_user_messages_invalid_content_type(): CustomGuardrail(guardrail_name="default_guardrail", default_on=True), CustomGuardrail(guardrail_name="request_guardrail", default_on=False), ], - "kwargs": { - "metadata": { - "requester_metadata": {"guardrails": ["request_guardrail"]} - } - }, + "kwargs": {"metadata": {"requester_metadata": {"guardrails": ["request_guardrail"]}}}, "expected": ["default_guardrail", "request_guardrail"], }, { "name": "empty_metadata", - "callbacks": [ - CustomGuardrail(guardrail_name="test_guardrail", default_on=False) - ], + "callbacks": [CustomGuardrail(guardrail_name="test_guardrail", default_on=False)], "kwargs": {}, "expected": [], }, @@ -10074,9 +10040,7 @@ def test_pre_call_rules_see_messages_passed_positionally(monkeypatch: pytest.Mon ) with pytest.raises(litellm.APIResponseValidationError, match="failed post-call-rule check"): - litellm.completion( - "gpt-4o-mini", [{"role": "user", "content": "say something forbidden"}], api_key="sk-test" - ) + litellm.completion("gpt-4o-mini", [{"role": "user", "content": "say something forbidden"}], api_key="sk-test") allowed: Final = litellm.completion("gpt-4o-mini", [{"role": "user", "content": "say hi"}], api_key="sk-test") assert route.call_count == 1 @@ -10094,3 +10058,18 @@ def test_vertex_llama3_model_info_resolves_without_the_meta_prefix(): prefixed_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") assert litellm.get_model_info(model=short_name, custom_llm_provider="vertex_ai") == prefixed_info, model assert prefixed_info["key"] == f"vertex_ai/{model}", model + + +def test_get_optional_params_keeps_a_tool_config_cache_point_out_of_the_openai_extra_body(): + """LIT-9880: the hook hands a tool_config point on for the provider step, and OpenAI has no + step that consumes it, so it must not ride along in extra_body and 400 at OpenAI.""" + optional_params = litellm.utils.get_optional_params( + model="gpt-4.1", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + cache_control_injection_points=[{"location": "tool_config", "control": {"type": "ephemeral"}}], + extra_body={"kept": True}, + ) + + assert optional_params["extra_body"] == {"kept": True} + assert "cache_control_injection_points" not in optional_params