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/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/cost_calculator.py b/litellm/cost_calculator.py index d9e6aedf45f..a9180eeebf4 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1135,6 +1135,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/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/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/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/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/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/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/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/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/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/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/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/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..f65310a77e3 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, ) @@ -3579,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. @@ -3600,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 @@ -3615,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: @@ -4531,7 +4535,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 +5608,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: @@ -6042,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. @@ -6098,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 ... @@ -6110,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, @@ -6137,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( @@ -6902,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 = ..., @@ -6929,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 @@ -6960,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"] @@ -6971,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 90f01368127..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, @@ -3041,6 +3307,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: @@ -3099,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, @@ -3111,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: @@ -3132,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: @@ -3517,14 +3804,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: @@ -3593,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/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/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/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/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/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/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/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/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/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/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/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/proxy_server.py b/litellm/proxy/proxy_server.py index c1102cee8d1..5030b1512a6 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() @@ -20032,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/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/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/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/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/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/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/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/litellm/utils.py b/litellm/utils.py index 49e655074a8..84d91c43c59 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5358,6 +5358,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, @@ -5381,7 +5384,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/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)), 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"); 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/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/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) 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/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/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/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 {} ), 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 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, 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/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/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/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/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/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/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/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/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" 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/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(): 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/_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/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..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 @@ -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 @@ -53,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, @@ -64,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, @@ -164,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())) @@ -180,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: @@ -1286,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): @@ -1364,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) @@ -2298,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() @@ -7387,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})) @@ -7404,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() @@ -8037,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) @@ -9138,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): @@ -9402,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 @@ -9761,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()) @@ -15210,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"), @@ -15914,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"}}) @@ -17280,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 @@ -17328,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" @@ -17352,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) @@ -17369,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() @@ -17399,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(): @@ -17408,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") @@ -17423,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 {} @@ -17450,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: @@ -17483,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], [])) @@ -17512,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 @@ -17536,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) @@ -17562,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 @@ -17581,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 @@ -17592,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") @@ -17603,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]) @@ -17644,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])) @@ -17743,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) @@ -17753,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]) @@ -17782,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]) @@ -17796,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 @@ -17811,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]) @@ -17853,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() @@ -17909,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) @@ -17921,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": @@ -17959,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): @@ -17989,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" @@ -18006,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" @@ -18024,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), @@ -18033,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 @@ -18045,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, @@ -18065,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 @@ -18104,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 @@ -18385,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) @@ -18425,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") @@ -18485,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 @@ -18538,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 @@ -18777,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") @@ -19563,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 @@ -19756,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" @@ -19776,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" @@ -19902,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) @@ -19932,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() @@ -20020,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: @@ -20038,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): @@ -20146,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") @@ -20184,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] @@ -20275,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"] == ( @@ -20323,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) @@ -20347,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 ) @@ -20375,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 @@ -20391,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"} @@ -20427,3 +21149,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..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 @@ -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, @@ -1331,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"}, @@ -1427,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"}, @@ -3531,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", @@ -3756,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 @@ -5876,6 +5896,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(): """ @@ -7254,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: @@ -7266,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: @@ -10071,6 +10161,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, *, @@ -10511,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) @@ -10518,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={}, ) @@ -10542,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) @@ -10549,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"}, ) @@ -10576,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 @@ -11637,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" @@ -11700,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", @@ -11792,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) @@ -11815,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_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_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/_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 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/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/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), {}) 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) 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() 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 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/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/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/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 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.""" 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", 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/test_utils.py b/tests/unit/test_utils.py index 35326da7562..3bdee29b23c 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 @@ -4989,7 +4985,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 @@ -5238,7 +5236,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", @@ -5266,7 +5270,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 @@ -5403,17 +5409,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 = [ @@ -5427,9 +5429,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: @@ -5465,9 +5465,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: @@ -7220,7 +7218,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" @@ -7371,7 +7371,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 @@ -7413,6 +7415,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def setup_and_teardown(): """ @@ -7426,6 +7429,7 @@ def setup_and_teardown(): loop.close() asyncio.set_event_loop(None) + MODEL: Final = "anthropic/claude-haiku-4-5" @@ -7435,6 +7439,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.""" @@ -7442,6 +7447,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.""" @@ -7449,6 +7455,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.""" @@ -7456,6 +7463,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", @@ -7477,6 +7485,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): @@ -7487,6 +7496,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.""" @@ -7495,12 +7505,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(): """ @@ -7553,6 +7565,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), @@ -7596,6 +7609,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", @@ -7621,6 +7635,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", @@ -7636,6 +7651,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", @@ -7654,6 +7670,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", @@ -7675,6 +7692,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", @@ -7701,6 +7719,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", @@ -7729,6 +7748,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", @@ -7756,12 +7776,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(): """ @@ -7771,6 +7793,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(): """ @@ -7797,6 +7820,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 @@ -7949,9 +7973,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" @@ -8134,9 +8156,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) @@ -8225,9 +8245,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(): @@ -8299,9 +8317,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]}, @@ -8337,9 +8353,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 @@ -8393,12 +8407,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, @@ -8407,12 +8417,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, @@ -8423,9 +8429,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, @@ -8767,10 +8771,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 @@ -8778,9 +8779,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 @@ -8802,16 +8801,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") @@ -8828,13 +8822,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 ) @@ -8849,20 +8839,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", ), @@ -8878,9 +8862,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", ), @@ -8888,9 +8870,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. @@ -9148,9 +9128,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 @@ -9247,20 +9225,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"], }, { @@ -9269,18 +9241,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": [], }, @@ -10086,9 +10052,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 @@ -10124,3 +10088,18 @@ def test_runtime_model_cost_registration_resolves_bedrock_routing(monkeypatch: p assert model_utils.runtime_model_cost_registration("bedrock/converse/openai.test-converse", "bedrock") == rates assert model_utils.runtime_model_cost_registration(alias, "bedrock") == rates assert model_utils.runtime_model_cost_registration("bedrock/converse/openai.other-converse", "bedrock") is None + + +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 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 */