diff --git a/.github/fork-patches.txt b/.github/fork-patches.txt index 9623c3ca91db..69d14340b461 100644 --- a/.github/fork-patches.txt +++ b/.github/fork-patches.txt @@ -36,3 +36,4 @@ litellm/llms/anthropic/chat/transformation.py | _translate_legacy_thinking_for_a ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx | import { getProviderLogoAndName } from "@/components/provider_info_helpers" | Fork-only provider display-name enhancement (model group / provider columns render getProviderLogoAndName(provider).displayName instead of the raw provider slug, including a sort comparator that sorts by the rendered name). Originally lived in ui/litellm-dashboard/src/components/model_hub_table_columns.tsx; the 2026-07-17 sync's upstream refactor (#33629, shared DataTable migration) deleted that file and moved its columns to this path, dropping the enhancement along the way, so the patch was reapplied at the new location plus the sibling Providers-modal usage in AIHub/ModelHubTable.tsx (same drop-only-the-import pattern). REMOVAL CONDITION: remove once upstream ships equivalent provider display-name rendering in these files itself. ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx | import { getProviderLogoAndName } from "@/components/provider_info_helpers" | Fork-only provider display-name enhancement for the model detail modal's Providers badges (renders getProviderLogoAndName(provider).displayName instead of the raw provider slug). Sibling patch to AIHub/ModelHubTableColumns.tsx above; the 2026-07-17 sync merge kept the fork's call site but dropped only the import line, the same drop-only-the-declaration pattern documented for the anthropic legacy-thinking translation. REMOVAL CONDITION: remove once upstream ships equivalent provider display-name rendering in this file itself. .github/workflows/codspeed.yml | continue-on-error: true | rayward-external/litellm was never registered with CodSpeed (codspeed.io) — the "Run benchmarks" step's own pytest-codspeed run passes, but the upload step 401s ("Repository not found or the user does not have access to it") on every PR/push run, so continue-on-error keeps a fork-infra gap from blocking the merge gate. First observed 2026-07-16 (PR #123, the first PR whose base commit made this workflow's checks land on a PR at all). REMOVAL CONDITION: remove once an admin registers rayward-external/litellm with CodSpeed (or adds a working CODSPEED_TOKEN) and the upload step succeeds. +tests/test_litellm/proxy/test_budget_reservation.py | _arelease_max_parallel_requests_on_disconnect = AsyncMock\(\) | Upstream #33736 moved the max_parallel_requests slot release into an async proxy_logging_obj._arelease_max_parallel_requests_on_disconnect() call inside async_streaming_data_generator's disconnect cleanup, but only added coverage in the new tests/test_litellm/proxy/test_common_request_processing.py; this pre-existing sibling file's bare MagicMock() stand-ins for proxy_logging_obj were never given that attribute, so awaiting it raised "TypeError: object MagicMock can't be used in 'await' expression" in 4 tests. Both call sites are IDENTICAL to upstream (verified via `git diff upstream/litellm_internal_staging`), so this reproduces on a clean upstream checkout too — upstream's own CI misses it because the new test file doesn't exercise these older streaming-cancel fixtures. REMOVAL CONDITION: remove once upstream adds the same AsyncMock to test_budget_reservation.py's _drive_streaming_cancel/slow-path fixtures itself. diff --git a/litellm/__init__.py b/litellm/__init__.py index 6e2a03b7c7c5..2f6643c644c2 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -315,6 +315,11 @@ def _dev_env_hot_reload_enabled() -> bool: disable_add_transform_inline_image_block: bool = False disable_add_user_agent_to_request_tags: bool = False disable_anthropic_gemini_context_caching_transform: bool = False +enable_anthropic_prompt_caching: bool = os.getenv("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", "false").lower() == "true" +_anthropic_prompt_caching_ttl_env: Optional[str] = os.getenv("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL") +anthropic_prompt_caching_ttl: Optional[Literal["5m", "1h"]] = ( + "1h" if _anthropic_prompt_caching_ttl_env == "1h" else "5m" if _anthropic_prompt_caching_ttl_env == "5m" else None +) disable_vertex_batch_output_transformation: bool = False extra_spend_tag_headers: Optional[List[str]] = None in_memory_llm_clients_cache: "LLMClientCache" diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 8732f7c0176e..85273fda5181 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -412,18 +412,148 @@ def apply_to_anthropic_messages_request( return processed_messages, processed_system, remaining_points + @staticmethod + def _default_control() -> ChatCompletionCachedContent: + """Build the cache_control block for auto-injected breakpoints. + + Defaults to Anthropic's 5-minute ephemeral cache; honors the optional + ``litellm.anthropic_prompt_caching_ttl`` override ("5m" or "1h"). + """ + import litellm + + ttl = litellm.anthropic_prompt_caching_ttl + if ttl == "5m" or ttl == "1h": + return ChatCompletionCachedContent(type="ephemeral", ttl=ttl) + return ChatCompletionCachedContent(type="ephemeral") + + @staticmethod + def _request_has_cache_control( + messages: list[AllMessageValues], + system: str | list | None, + tools: list | None = None, + ) -> bool: + """Return True if the request already carries any client-supplied cache_control. + + When the client (e.g. Claude Code) already marks its own breakpoints we + stand down entirely rather than add more, per the auto-caching contract. + Tools count: they are a breakpoint the client can mark, they count toward + the provider's four-block limit, and caching only the tool definitions is + a common pattern, so injecting alongside them can exceed the cap. + """ + if any(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages): + return True + if isinstance(system, list): + if any(isinstance(block, dict) and block.get("cache_control") is not None for block in system): + return True + if tools is not None: + return any(isinstance(tool, dict) and tool.get("cache_control") is not None for tool in tools) + return False + + @staticmethod + def get_default_injection_points( + messages: list[AllMessageValues], + system: str | list | None, + model: str, + custom_llm_provider: str | None, + tools: list | None = None, + ) -> list[CacheControlInjectionPoint]: + """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on. + + Caches the system prompt and the trailing turn, so the stable prefix + (system + tools + history) is reused while the breakpoint advances with + the conversation. Returns [] (stand down) when the flag is off, the + provider does not consume cache_control breakpoints (only anthropic / + bedrock do), the model lacks prompt-caching support, or the request + already carries client-supplied cache_control. + """ + import litellm + + if litellm.enable_anthropic_prompt_caching is not True: + return [] + + provider = custom_llm_provider + if provider is None: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + ) + + try: + _, provider, _, _ = get_llm_provider(model=model) + except Exception: # noqa: BLE001 # unroutable model must never block the call, just skip auto-caching + return [] + + if provider not in ("anthropic", "bedrock"): + return [] + + from litellm.utils import supports_prompt_caching + + if not supports_prompt_caching(model=model, custom_llm_provider=provider): + return [] + + if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools): + return [] + + control = AnthropicCacheControlHook._default_control() + points: list[CacheControlInjectionPoint] = [ + CacheControlMessageInjectionPoint(location="message", role="system", index=None, control=control), + CacheControlMessageInjectionPoint(location="message", role=None, index=-1, control=control), + ] + return points + + @staticmethod + def maybe_seed_default_injection_points( + non_default_params: dict[str, Any], + messages: list[AllMessageValues], + model: str, + custom_llm_provider: str | None, + tools: list | None = None, + ) -> None: + """For /chat/completions: add default injection points to the request params. + + No-op when injection points are already configured (explicit config wins). + Seeding the param lets the existing prompt-management gate and the + AnthropicCacheControlHook run unchanged. + """ + if non_default_params.get("cache_control_injection_points"): + return + points = AnthropicCacheControlHook.get_default_injection_points( + messages=messages, + system=None, + model=model, + custom_llm_provider=custom_llm_provider, + tools=tools, + ) + if points: + non_default_params["cache_control_injection_points"] = points + @staticmethod def maybe_inject_cache_control( messages: List[Dict], system: str | list | None, kwargs: Dict[str, Any], + model: str | None = None, + custom_llm_provider: str | None = None, + tools: list[dict] | None = None, ) -> Tuple[List[Dict], str | list | None]: """Extract cache_control_injection_points from kwargs and apply if present. + When none are configured but ``litellm.enable_anthropic_prompt_caching`` + is on, synthesize default breakpoints for the native /v1/messages path. Pops the key from kwargs; if remaining (non-message) points exist they are written back so downstream transforms can handle them. """ - injection_points = kwargs.pop("cache_control_injection_points", None) + configured = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list + list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None) + ) + injection_points: list[CacheControlInjectionPoint] = configured or [] + if not injection_points and model is not None: + injection_points = AnthropicCacheControlHook.get_default_injection_points( + messages=cast(list[AllMessageValues], messages), # cast-ok: Anthropic-shaped dicts from v1/messages + system=system, + tools=tools, + model=model, + custom_llm_provider=custom_llm_provider, + ) if not injection_points: return messages, system diff --git a/litellm/integrations/opik/utils.py b/litellm/integrations/opik/utils.py index 7222c9d05021..d4850d50778c 100644 --- a/litellm/integrations/opik/utils.py +++ b/litellm/integrations/opik/utils.py @@ -1,40 +1,38 @@ import configparser import os import time +import uuid from typing import Any, Dict, Final, List, Optional, Tuple CONFIG_FILE_PATH_DEFAULT: Final[str] = "~/.opik.config" -def create_uuid7(): - ns = time.time_ns() - last = [0, 0, 0, 0] - - # Simple uuid7 implementation - sixteen_secs = 16_000_000_000 - t1, rest1 = divmod(ns, sixteen_secs) - t2, rest2 = divmod(rest1 << 16, sixteen_secs) - t3, _ = divmod(rest2 << 12, sixteen_secs) - t3 |= 7 << 12 # Put uuid version in top 4 bits, which are 0 in t3 - - # The next two bytes are an int (t4) with two bits for - # the variant 2 and a 14 bit sequence counter which increments - # if the time is unchanged. - if t1 == last[0] and t2 == last[1] and t3 == last[2]: - # Stop the seq counter wrapping past 0x3FFF. - # This won't happen in practice, but if it does, - # uuids after the 16383rd with that same timestamp - # will not longer be correctly ordered but - # are still unique due to the 6 random bytes. - if last[3] < 0x3FFF: - last[3] += 1 - else: - last[:] = (t1, t2, t3, 0) - t4 = (2 << 14) | last[3] # Put variant 0b10 in top two bits - - # Six random bytes for the lower part of the uuid - rand = os.urandom(6) - return f"{t1:>08x}-{t2:>04x}-{t3:>04x}-{t4:>04x}-{rand.hex()}" +def create_uuid7() -> str: + """Generate an RFC 9562 conformant UUIDv7 string. + + The top 48 bits encode the Unix timestamp in milliseconds. Opik's backend + validates this embedded timestamp on ingestion (it must fall within a window + around "now"), so the encoding has to be correct or trace/span batches are + rejected with HTTP 400. Implemented with the standard library only, so no + extra dependency is added to litellm. See ``opik.id_helpers`` for the + reference implementation. + """ + unix_ts_ms = int(time.time() * 1000) + + # Fill the 16-byte buffer with random data, then overwrite the structured + # parts (timestamp, version, variant) defined by the UUIDv7 layout. + uuid_bytes = bytearray(os.urandom(16)) + + # First 48 bits (6 bytes): Unix timestamp in milliseconds. + uuid_bytes[0:6] = unix_ts_ms.to_bytes(6, byteorder="big") + + # Version 7 in the top 4 bits of byte 6. + uuid_bytes[6] = 0x70 | (uuid_bytes[6] & 0x0F) + + # Variant 0b10 in the top 2 bits of byte 8. + uuid_bytes[8] = 0x80 | (uuid_bytes[8] & 0x3F) + + return str(uuid.UUID(bytes=bytes(uuid_bytes))) def _read_opik_config_file() -> Dict[str, str]: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 36d175968733..3b3c6a6ce298 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1453,6 +1453,9 @@ def _response_cost_calculator( response_cost = litellm.response_cost_calculator(**response_cost_calculator_kwargs) verbose_logger.debug(f"response_cost: {response_cost}") + additional_response_cost: object = self.model_call_details.get("additional_response_cost") + if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0: + return (response_cost or 0.0) + additional_response_cost return response_cost except Exception as e: # error calculating cost debug_info = StandardLoggingModelCostFailureDebugInformation( diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index e006662ec4d0..256fee6b1663 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -906,16 +906,17 @@ def strip_advisor_blocks_from_messages(messages: List[Any], replace_with_text: b def is_anthropic_invalid_thinking_signature_error(error_text: str) -> bool: """ - Detect Anthropic 400 when encrypted thinking signatures in history do not match - the current deployment (e.g. user rotated API key or switched model endpoint). + Detect Anthropic 400 errors caused by missing or invalid thinking signatures. - Example API message: + Known error formats: + {"message":"messages.2.content.0.thinking.signature.str: Input should be a valid string"} + messages.N.content.M.thinking.signature.str: Input should be a valid string messages.N.content.M: Invalid `signature` in `thinking` block """ if not error_text: return False lower = error_text.lower() - return "invalid" in lower and "signature" in lower and "thinking" in lower and "block" in lower + return "thinking" in lower and "signature" in lower and ("invalid" in lower or "valid string" in lower) def strip_thinking_blocks_from_anthropic_messages(messages: List[Any]) -> List[Any]: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index ebee93237666..703ccf13c273 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -237,7 +237,9 @@ async def anthropic_messages( AnthropicCacheControlHook, ) - messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(messages, system, kwargs) + messages, system = AnthropicCacheControlHook.maybe_inject_cache_control( + messages, system, kwargs, model=model, custom_llm_provider=custom_llm_provider, tools=tools + ) original_stream = stream or kwargs.get("_websearch_interception_converted_stream", False) @@ -426,7 +428,9 @@ def anthropic_messages_handler( AnthropicCacheControlHook, ) - messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(messages, system, kwargs) + messages, system = AnthropicCacheControlHook.maybe_inject_cache_control( + messages, system, kwargs, model=model, custom_llm_provider=custom_llm_provider, tools=tools + ) metadata = validate_anthropic_api_metadata(metadata) diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index ed936f6233a9..682adf5a8ffe 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -75,10 +75,23 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: model_info = get_model_info(model=base_model, custom_llm_provider="fireworks_ai") ## CALCULATE INPUT COST - - prompt_cost: float = usage["prompt_tokens"] * model_info["input_cost_per_token"] + prompt_tokens_details = usage.prompt_tokens_details + cached_tokens: int = ( + prompt_tokens_details.cached_tokens + if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None + else 0 + ) + input_cost_per_token: float = model_info["input_cost_per_token"] or 0.0 + cache_read_input_token_cost = model_info.get("cache_read_input_token_cost") + cache_read_cost_per_token: float = ( + cache_read_input_token_cost if cache_read_input_token_cost is not None else input_cost_per_token + ) + non_cached_prompt_tokens: int = max(usage.prompt_tokens - cached_tokens, 0) + + prompt_cost: float = non_cached_prompt_tokens * input_cost_per_token + cached_tokens * cache_read_cost_per_token ## CALCULATE OUTPUT COST - completion_cost = usage["completion_tokens"] * model_info["output_cost_per_token"] + output_cost_per_token: float = model_info["output_cost_per_token"] or 0.0 + completion_cost: float = usage.completion_tokens * output_cost_per_token return prompt_cost, completion_cost diff --git a/litellm/main.py b/litellm/main.py index 6fd68921fb02..3584297b35f1 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -510,6 +510,20 @@ async def acompletion( ######################################################### ######################################################### litellm_logging_obj = kwargs.get("litellm_logging_obj", None) + + from litellm.integrations.anthropic_cache_control_hook import ( + AnthropicCacheControlHook, + ) + from litellm.types.llms.openai import AllMessageValues + + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=kwargs, + messages=cast(list[AllMessageValues], messages), # cast-ok: acompletion types messages as a bare List + model=model, + custom_llm_provider=cast(Optional[str], custom_llm_provider), # cast-ok: read from untyped kwargs + tools=tools, + ) + if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( litellm_logging_obj.should_run_prompt_management_hooks( prompt_id=kwargs.get("prompt_id", None), @@ -5055,6 +5069,19 @@ def completion( # type: ignore litellm_params = {} # used to prevent unbound var errors ## PROMPT MANAGEMENT HOOKS ## + from litellm.integrations.anthropic_cache_control_hook import ( + AnthropicCacheControlHook, + ) + from litellm.types.llms.openai import AllMessageValues + + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=non_default_params, + messages=cast(list[AllMessageValues], messages), # cast-ok: completion types messages as a bare List + model=model, + custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs + tools=tools, + ) + if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( litellm_logging_obj.should_run_prompt_management_hooks( prompt_id=prompt_id, non_default_params=non_default_params diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1a24088396fd..ee996198b280 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3451,7 +3451,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2.2e-05, "output_cost_per_token": 2.64e-06, "supports_audio_input": true, @@ -3470,7 +3470,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 0.00022, "output_cost_per_token": 2.2e-05, "supports_audio_input": true, @@ -3489,7 +3489,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2.2e-05, "supported_modalities": [ @@ -4687,7 +4687,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supports_audio_input": true, @@ -4707,7 +4707,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -4739,7 +4739,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -4771,7 +4771,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supported_endpoints": [ @@ -4832,7 +4832,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 0.0002, "output_cost_per_token": 2e-05, "supports_audio_input": true, @@ -4850,7 +4850,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, "supported_modalities": [ @@ -7922,7 +7922,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2.2e-05, "output_cost_per_token": 2.64e-06, "supports_audio_input": true, @@ -7941,7 +7941,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 0.00022, "output_cost_per_token": 2.2e-05, "supports_audio_input": true, @@ -7960,7 +7960,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2.2e-05, "supported_modalities": [ @@ -22094,7 +22094,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supports_audio_input": true, @@ -22113,7 +22113,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supports_audio_input": true, @@ -22207,7 +22207,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, "supports_audio_input": true, @@ -22225,7 +22225,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, "supports_audio_input": true, @@ -22243,7 +22243,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, "supports_audio_input": true, @@ -24438,7 +24438,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -24470,7 +24470,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -24502,7 +24502,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -24535,7 +24535,7 @@ "max_input_tokens": 128000, "max_output_tokens": 32000, "max_tokens": 32000, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, "regional_processing_uplift_multiplier_eu": 1.1, @@ -24570,7 +24570,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "regional_processing_uplift_multiplier_eu": 1.1, @@ -24603,7 +24603,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supported_endpoints": [ @@ -24635,7 +24635,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -43573,7 +43573,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supported_endpoints": [ @@ -43606,7 +43606,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supported_endpoints": [ 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 421f1dcfbeab..d2f3efbc54ec 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 @@ -1220,11 +1220,29 @@ async def get_allowed_tools_for_server( global_mcp_server_manager, ) - key_tools = ( + key_direct_tools = ( 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 = (key_obj_perm.mcp_toolsets or []) if key_obj_perm else [] + key_toolset_tools = ( + (await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=key_toolset_ids)).get( + server_id + ) + if key_toolset_ids + else None + ) + + key_tools = ( + list(set(key_direct_tools or []) | set(key_toolset_tools or [])) + if key_direct_tools is not None or key_toolset_tools is not None + else None + ) team_tools = ( global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id) if team_obj_perm @@ -1430,8 +1448,18 @@ async def _get_allowed_mcp_servers_for_key( global_mcp_server_manager.expand_tool_permissions(key_object_permission.mcp_tool_permissions).keys() ) + # servers referenced by the key's toolset grants are part of the key's + # scope on every path (list, call, REST), subject to the same team/org + # ceilings as any other key-level grant + toolset_ids = key_object_permission.mcp_toolsets or [] + toolset_servers = ( + list((await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=toolset_ids)).keys()) + if toolset_ids + else [] + ) + # Combine all lists - all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers + all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers + toolset_servers return list(set(all_servers)) except Exception as e: verbose_logger.warning(f"Failed to get allowed MCP servers for key: {str(e)}") diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index 3e3e549008db..74752809e865 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -88,3 +88,19 @@ class MCPToolResultError(Exception): into two identities, breaking ``isinstance`` checks against instances created before the reload. """ + + +class MCPServerListError(Exception): + """Carrier for a classified per-server listing fault (``faults.list_outcomes.ServerListFault``). + + Raised where a server fetch used to silently return an empty tool list, so each boundary can + apply its own policy: the aggregate listing absorbs it into that server's outcome, while + single-server routes relay a truthful HTTP status instead of empty-success. The fault value is + typed as ``object`` here only to avoid a circular import with the faults package; construction + sites always pass a ``ServerListFault``. + """ + + def __init__(self, fault: object, server_name: str) -> None: + self.fault = fault + self.server_name = server_name + super().__init__(f"Listing tools from MCP server {server_name!r} failed") diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py new file mode 100644 index 000000000000..6f27c1c04722 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -0,0 +1,190 @@ +"""Per-server outcomes for the aggregate MCP tools/list fan-out. + +The aggregate listing deliberately keeps serving the healthy subset when one server fails, but a +failed server must contribute a classified outcome instead of silently shrinking the list: an empty +contribution with no signal makes a broken upstream indistinguishable from a healthy server with no +tools. Outcomes carry only machine fields (category and status code) so nothing from an upstream +body crosses the trust boundary; classification is total, so any exception out of a server fetch +becomes an outcome, never a second failure. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Literal, NamedTuple, NoReturn, TypeAlias + +import httpx +from mcp.types import Tool as MCPTool +from pydantic import BaseModel, ConfigDict +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPServerListError, + MCPUpstreamAuthError, +) + +ListFaultCategory: TypeAlias = Literal[ + "auth_required", + "forbidden", + "timeout", + "unreachable", + "upstream_error", + "internal", +] + + +class ServerListOk(BaseModel): + model_config = ConfigDict(frozen=True) + tag: Literal["ok"] = "ok" + tool_count: int + + +class ServerListFault(BaseModel): + """Why a server contributed nothing to a listing: the caller must authenticate upstream + (``auth_required``/``forbidden``), the upstream did not answer (``timeout``/``unreachable``), + the upstream answered outside its contract (``upstream_error``), or the gateway itself failed + (``internal``). ``status_code`` is the upstream HTTP status when one exists.""" + + model_config = ConfigDict(frozen=True) + tag: ListFaultCategory + status_code: int | None = None + + +ServerOutcome: TypeAlias = ServerListOk | ServerListFault + +SERVER_OUTCOMES_META_KEY = "litellm.ai/server_outcomes" +"""The tools/list result ``_meta`` key carrying per-server outcomes. Prefixed with the litellm.ai +domain per the MCP spec's ``_meta`` key format so it cannot collide with spec-reserved names.""" + + +class AggregateToolListing(NamedTuple): + tools: list[MCPTool] + outcomes: dict[str, ServerOutcome] + + +def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response]: + """Yield every ``httpx.Response`` in the exception tree (``__cause__``/``__context__``/ + ExceptionGroup members) in deliberate order, mirroring how upstream failures surface through the + MCP SDK's task groups. Explicit links come first: each node's ``raise ... from`` cause, then + group members in raise order, then the incidental ``__context__`` chain, so a response raised + while handling the real failure can never shadow one on the explicit causal chain. Consumers + apply their own predicate over the stream: selecting the first response and THEN testing it + would miss a causal auth response sitting behind an unrelated earlier one.""" + seen: set[int] = set() + stack = [exc] + while stack: + current = stack.pop() + if id(current) in seen: + continue + seen.add(id(current)) + response = getattr(current, "response", None) + if isinstance(response, httpx.Response): + yield response + if current.__context__ is not None: + stack.append(current.__context__) + exceptions = getattr(current, "exceptions", None) + if isinstance(exceptions, tuple): + stack.extend(reversed(exceptions)) + if current.__cause__ is not None: + stack.append(current.__cause__) + + +def _find_upstream_response(exc: BaseException) -> httpx.Response | None: + return next(_iter_upstream_responses(exc), None) + + +def upstream_auth_challenge(exc: BaseException) -> tuple[int, str | None] | None: + """The first upstream 401/403 in deliberate order and its ``WWW-Authenticate`` challenge, both + read from the SAME response, so the status that picks the carrier channel and the challenge that + rides with it can never come from two different responses in the tree. Non-auth responses do not + end the scan: a causal 401 behind an unrelated 5xx must still be found, or the client never + receives the challenge it needs to re-authenticate.""" + for response in _iter_upstream_responses(exc): + if response.status_code in (401, 403): + return response.status_code, response.headers.get("www-authenticate") + return None + + +def raise_classified_list_failure( + exc: BaseException, + server_name: str, + suppress_challenge: bool = False, +) -> NoReturn: + """The one place a failed server fetch chooses its carrier: an upstream 401/403 travels as + ``MCPUpstreamAuthError`` with the upstream's own challenge preserved (a challenge is only ever + fabricated at the HTTP edge, and only for a 401), everything else as ``MCPServerListError`` with + a classified fault. Every fetch site delegates here so the two channels cannot drift apart per + call site. ``suppress_challenge`` is for dcr_bridge servers, whose upstream challenge points + clients at the wrong protected-resource metadata and must never relay.""" + auth = upstream_auth_challenge(exc) + if auth is not None: + status_code, challenge = auth + raise MCPUpstreamAuthError( + status_code=status_code, + www_authenticate=None if suppress_challenge else challenge, + server_name=server_name, + ) from exc + raise MCPServerListError(classify_list_exception(exc), server_name) from exc + + +def classify_list_exception(exc: BaseException) -> ServerListFault: + """Classify a per-server listing failure into exactly one outcome. Total: an exception this + function cannot recognize is the gateway's own fault (``internal``), never a re-raise.""" + if isinstance(exc, MCPServerListError) and isinstance(exc.fault, ServerListFault): + return exc.fault + if isinstance(exc, MCPUpstreamAuthError): + tag = "forbidden" if exc.status_code == 403 else "auth_required" + return ServerListFault(tag=tag, status_code=exc.status_code) + if isinstance(exc, TimeoutError): + return ServerListFault(tag="timeout") + if isinstance(exc, ConnectionError): + return ServerListFault(tag="unreachable") + auth = upstream_auth_challenge(exc) + if auth is not None: + status_code, _ = auth + return ServerListFault( + tag="forbidden" if status_code == 403 else "auth_required", + status_code=status_code, + ) + response = _find_upstream_response(exc) + if response is not None: + return ServerListFault(tag="upstream_error", status_code=response.status_code) + if isinstance(exc, (httpx.TimeoutException,)): + return ServerListFault(tag="timeout") + if isinstance(exc, httpx.TransportError): + return ServerListFault(tag="unreachable") + return ServerListFault(tag="internal") + + +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.""" + match outcome.tag: + case "ok": + return {"status": "ok", "tool_count": outcome.tool_count} + case "auth_required" | "forbidden" | "timeout" | "unreachable" | "upstream_error" | "internal": + return { + "status": outcome.tag, + **({"http_status": outcome.status_code} if outcome.status_code is not None else {}), + } + case _: + assert_never(outcome.tag) + + +def list_fault_http_status(fault: ServerListFault) -> int: + """The truthful HTTP status for a single-upstream listing fault per RFC 9110: the upstream's own + 401/403 for auth, 504 for a timeout, 502 for an unreachable or misbehaving upstream, and 500 only + for the gateway's own failure.""" + match fault.tag: + case "auth_required": + return fault.status_code or 401 + case "forbidden": + return 403 + case "timeout": + return 504 + case "unreachable" | "upstream_error": + return 502 + case "internal": + return 500 + case _: + assert_never(fault.tag) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index f7462af1927d..5e2cec3e86db 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -50,7 +50,15 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) -from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPServerListError, + MCPUpstreamAuthError, +) +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + ServerListFault, + raise_classified_list_failure, + upstream_auth_challenge, +) from litellm.proxy._experimental.mcp_server.elicitation_handler import ( MCP_ELICITATION_AVAILABLE, ) @@ -633,49 +641,14 @@ def _caller_authorization_fans_out( def _extract_upstream_auth_failure( exc: BaseException, ) -> Optional[tuple[int, Optional[str]]]: - """Walk the exception tree looking for an HTTP 401/403 response from the - upstream MCP server. - - The MCP SDK wraps transport errors in anyio ``ExceptionGroup`` objects and - may chain through ``__cause__`` / ``__context__``. We inspect all of those - layers for an ``httpx.Response``-bearing exception (typically - ``httpx.HTTPStatusError``) and extract the status code and any upstream - ``WWW-Authenticate`` header. - - Returns ``(status_code, www_authenticate)`` on match, else ``None``. - """ - seen: set[int] = set() - stack: list[BaseException] = [exc] - while stack: - current = stack.pop() - if id(current) in seen: - continue - seen.add(id(current)) - - response = getattr(current, "response", None) - if response is not None: - status_code = getattr(response, "status_code", None) - if isinstance(status_code, int) and status_code in (401, 403): - www_authenticate: Optional[str] = None - headers = getattr(response, "headers", None) - if headers is not None: - try: - www_authenticate = headers.get("www-authenticate") - except Exception: - www_authenticate = None - return status_code, www_authenticate - - # anyio / PEP 654 ExceptionGroup - sub_exceptions = getattr(current, "exceptions", None) - if sub_exceptions: - stack.extend(sub_exceptions) + """The upstream 401/403 and its ``WWW-Authenticate`` header from the exception tree, or ``None``. - if current.__cause__ is not None: - stack.append(current.__cause__) - if current.__context__ is not None and current.__context__ is not current.__cause__: - stack.append(current.__context__) - - return None + Delegates to the shared traversal in ``faults.list_outcomes`` so every consumer (tool listing, + tool calls, the connect-time probe) selects the same response with the same deliberate order: + explicit ``raise ... from`` causes first, ExceptionGroup members in raise order, the incidental + ``__context__`` chain last. A response raised while handling the real failure can therefore never + shadow the causal one.""" + return upstream_auth_challenge(exc) def _warn_on_server_name_fields( @@ -3047,10 +3020,12 @@ async def _get_tools_from_server( server_name=server.name, ) from e verbose_logger.warning(f"Failed to get tools from server {server.name}: {str(e)}") - return [] + raise MCPServerListError(ServerListFault(tag="internal", status_code=e.status_code), server.name) from e + except MCPServerListError: + raise except Exception as e: verbose_logger.warning(f"Failed to get tools from server {server.name}: {str(e)}") - return [] + raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge) async def get_prompts_from_server( self, @@ -3682,16 +3657,17 @@ async def _fetch_tools_with_timeout( Uses anyio.fail_after() instead of asyncio.wait_for() to avoid conflicts with the MCP SDK's anyio TaskGroup. See GitHub issue #20715 for details. - An upstream HTTP 401 is converted into :class:`MCPUpstreamAuthError` - instead of being swallowed to an empty tool list, regardless of the - server's auth_type. Callers route it by surface: the single-server HTTP - routes turn it into a 401 + ``WWW-Authenticate`` challenge so standards- - compliant MCP clients trigger the upstream OAuth flow, while the - multi-server ``/mcp`` aggregator absorbs it to an empty list so one - unauthenticated server doesn't fail the whole listing. Only a 401 - (missing/invalid credential) drives the re-auth challenge; a 403 - (authenticated but forbidden, e.g. insufficient scope) is not a re-auth - signal and, like other non-auth errors, returns an empty list. + Failures never return an empty tool list. An upstream 401 or 403 raises + :class:`MCPUpstreamAuthError` carrying the upstream's own + ``WWW-Authenticate`` challenge when one was sent (a challenge is only + ever fabricated at the HTTP edge, and only for a 401: a 403 means the + caller is authenticated but not allowed, so prompting re-auth would be + wrong, while an upstream-sent 403 challenge is the RFC 6750 + insufficient_scope step-up and relays verbatim). Every other failure + raises :class:`MCPServerListError` with a classified fault. Each + boundary then applies its own policy: single-server routes relay the + truthful status, the multi-server aggregator absorbs the failure into + that server's listing outcome. Args: client: MCP client instance @@ -3705,27 +3681,18 @@ async def _fetch_tools_with_timeout( tools = await client.list_tools(raise_on_error=True) verbose_logger.debug(f"Tools from {server_name}: {tools}") return tools - except TimeoutError: + except TimeoutError as e: verbose_logger.warning(f"Timeout while listing tools from {server_name}") - return [] - except asyncio.CancelledError: + raise MCPServerListError(ServerListFault(tag="timeout"), server_name) from e + except asyncio.CancelledError as e: verbose_logger.warning(f"Task cancelled while listing tools from {server_name}") - return [] + raise MCPServerListError(ServerListFault(tag="internal"), server_name) from e except ConnectionError as e: verbose_logger.warning(f"Connection error while listing tools from {server_name}: {str(e)}") - return [] + raise MCPServerListError(ServerListFault(tag="unreachable"), server_name) from e except Exception as e: - auth_info = _extract_upstream_auth_failure(e) - if auth_info is not None and auth_info[0] == 401: - _, www_authenticate = auth_info - verbose_logger.info(f"Upstream auth failure from MCP server {server_name}: HTTP 401") - raise MCPUpstreamAuthError( - status_code=401, - www_authenticate=www_authenticate, - server_name=server_name, - ) from e verbose_logger.warning(f"Error listing tools from {server_name}: {str(e)}") - return [] + raise_classified_list_failure(e, server_name) _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 7ca4923337c0..94271c54f4b2 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -19,12 +19,20 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from litellm._logging import verbose_logger -from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPServerListError, + MCPUpstreamAuthError, +) +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + classify_list_exception, + list_fault_http_status, +) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( build_effective_auth_contexts, ) from litellm.proxy._experimental.mcp_server.utils import ( MCPMissingUserEnvVarsError, + get_server_prefix, merge_mcp_headers, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -515,20 +523,19 @@ async def _get_tools_for_single_server( # enforced even when no allowlist is set (matches the SSE/HTTP path). tools = filter_tools_by_allowed_tools(tools, server) - # Filter tools based on user_api_key_auth.object_permission.mcp_tool_permissions - # This provides per-key/team/org control over which tools can be accessed - if ( - user_api_key_auth - and user_api_key_auth.object_permission - and user_api_key_auth.object_permission.mcp_tool_permissions - ): - # Dict keys may be server_ids OR names/aliases; normalize so lookup - # by concrete server_id resolves name-keyed restrictions too. - allowed_tools_for_server = global_mcp_server_manager.expand_tool_permissions( - user_api_key_auth.object_permission.mcp_tool_permissions - ).get(server.server_id) - if allowed_tools_for_server is not None and len(allowed_tools_for_server) > 0: - # Filter tools to only include those in the allowed list + # Filter by the key's effective tool permissions through the same + # primitive the MCP protocol path uses (direct grants, toolset grants, + # and team/agent/org ceilings), so REST listing cannot drift from it + if user_api_key_auth: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + allowed_tools_for_server = await MCPRequestHandler.get_allowed_tools_for_server( + server_id=server.server_id, + user_api_key_auth=user_api_key_auth, + ) + if allowed_tools_for_server is not None: tools = [tool for tool in tools if _tool_name_matches(tool.name, allowed_tools_for_server)] return _create_tool_response_objects(tools, server) @@ -627,6 +634,16 @@ async def _list_tools_for_single_server( # matching status code and WWW-Authenticate challenge; that is what # lets standards-compliant MCP clients run the upstream OAuth flow. raise + except MCPServerListError as e: + fault = classify_list_exception(e) + verbose_logger.info(f"Listing tools from {server.name} failed with a {fault.tag} fault") + raise HTTPException( + status_code=list_fault_http_status(fault), + detail={ + "error": fault.tag, + "message": f"Failed to list tools from server {get_server_prefix(server)}", + }, + ) from e except Exception as e: verbose_logger.exception(f"Error getting tools from {server.name}: {e}") return { @@ -838,7 +855,11 @@ async def list_tool_rest_api( list_tools_result.extend(tools_result) except Exception as e: verbose_logger.exception(f"Error getting tools from {server.name}: {e}") - errors.append(f"{server.name}: {str(e)}") + errors.append( + f"{get_server_prefix(server)}: {classify_list_exception(e).tag}" + if isinstance(e, (MCPServerListError, MCPUpstreamAuthError)) + else f"{get_server_prefix(server)}: {str(e)}" + ) continue if errors and not list_tools_result: @@ -858,7 +879,10 @@ async def list_tool_rest_api( request_path=request.scope.get("_original_path") or request.url.path, ) except HTTPException as http_exc: - if http_exc.status_code == status.HTTP_404_NOT_FOUND: + if http_exc.status_code == status.HTTP_404_NOT_FOUND or server_id: + # Single-server requests relay the truthful status (a 502/504 upstream fault must + # not masquerade as a 200 empty-success body); only the multi-server aggregate + # keeps the legacy error-dict response shape below. raise # Internal access/IP 403s keep the legacy error-dict response shape # so the existing contract stays intact. diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 26322e9c58b7..a8ab0937124c 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -348,6 +348,7 @@ def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: CallToolResult, EmbeddedResource, ImageContent, + ListToolsResult, Prompt, TextContent, ) @@ -356,6 +357,14 @@ def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListOk, + ServerOutcome, + classify_list_exception, + outcome_wire_value, + ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _caller_authorization_fans_out, @@ -664,9 +673,12 @@ async def lifespan(app) -> AsyncIterator[None]: ######################################################## @server.list_tools() - async def handle_list_tools() -> List[Tool]: + async def handle_list_tools() -> "ListToolsResult | List[Tool]": """ - List all available tools. + List all available tools, with each server's listing outcome attached to the result's + ``_meta`` (SERVER_OUTCOMES_META_KEY) so a broken upstream is distinguishable from a healthy + server with no tools. Returning a ListToolsResult (rather than a bare list) makes the MCP SDK + pass the result through unwrapped, which is what lets the ``_meta`` survive to the client. Also captures the active session for propagation to callbacks. """ from mcp.server.lowlevel.server import request_ctx @@ -709,7 +721,7 @@ async def handle_list_tools() -> List[Tool]: # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") - tools = await _list_mcp_tools( + listing = await _list_mcp_tools( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, @@ -719,8 +731,15 @@ async def handle_list_tools() -> List[Tool]: log_list_tools_to_spendlogs=True, list_tools_log_source="mcp_protocol", ) - verbose_logger.info(f"MCP list_tools - Successfully returned {len(tools)} tools") - return tools + verbose_logger.info(f"MCP list_tools - Successfully returned {len(listing.tools)} tools") + if not listing.outcomes: + return listing.tools + outcome_meta = { + SERVER_OUTCOMES_META_KEY: { + key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items() + } + } + return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) except Exception as e: verbose_logger.exception(f"Error in list_tools endpoint: {str(e)}") # Return empty list instead of failing completely @@ -1746,6 +1765,13 @@ async def _gateway_initialize_instructions_request_scope( _mcp_gateway_initialize_instructions.reset(instructions_token) _mcp_gateway_server_name.reset(server_name_token) + def _aggregate_server_key(server: MCPServer) -> str: + """The client-visible key for a server in listing outcomes and spend metadata: the same + display prefix (alias, or the short prefix when that mode is enabled) the caller already + sees on the tool names. Canonical internal server names never key a caller-readable + surface; when the display naming deliberately hides them, the outcome keys must too.""" + return get_server_prefix(server) or "unknown" + async def _get_tools_from_mcp_servers( user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str], @@ -1758,7 +1784,7 @@ async def _get_tools_from_mcp_servers( litellm_trace_id: Optional[str] = None, request_tags: Optional[list[str]] = None, client_ip: Optional[str] = None, - ) -> List[MCPTool]: + ) -> AggregateToolListing: """ Helper method to fetch tools from MCP servers based on server filtering criteria. @@ -1770,10 +1796,11 @@ async def _get_tools_from_mcp_servers( oauth2_headers: Optional dict of oauth2 headers Returns: - List[MCPTool]: Combined list of tools from filtered servers + AggregateToolListing: Combined tools from filtered servers plus each server's + classified listing outcome """ if not MCP_AVAILABLE: - return [] + return AggregateToolListing(tools=[], outcomes={}) list_tools_start_time = datetime.now() litellm_logging_obj: Optional[LiteLLMLoggingObj] = None @@ -1858,10 +1885,12 @@ async def _get_tools_from_mcp_servers( async def _fetch_and_filter_server_tools( server: MCPServer, - ) -> List[MCPTool]: - """Fetch and filter tools from a single server with error handling.""" + ) -> "tuple[List[MCPTool], ServerOutcome]": + """Fetch and filter tools from a single server, classifying any failure into that + server's outcome so the aggregate can keep serving the healthy subset without a + broken server masquerading as an empty one.""" if server is None: - return [] + return [], ServerListOk(tool_count=0) server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, @@ -1931,8 +1960,8 @@ async def _fetch_and_filter_server_tools( verbose_logger.debug( f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering" ) - return filtered_tools - except MCPUpstreamAuthError: + return filtered_tools, ServerListOk(tool_count=len(filtered_tools)) + except MCPUpstreamAuthError as e: # Absorb so one unauthenticated server does not empty every other server's # tools. Surfacing the upstream 401 to the client as a re-auth challenge is # intentionally not done here: raising from this list handler cannot produce a @@ -1940,31 +1969,30 @@ async def _fetch_and_filter_server_tools( # error). Single-server routes surface it via the request-scope preemptive # check in _raise_preemptive_401_for_unauthenticated_servers instead. verbose_logger.debug(f"MCP list_tools: omitting {server.name}; it needs upstream auth") - return [] + return [], classify_list_exception(e) except Exception as e: verbose_logger.exception(f"Error getting tools from server {server.name}: {str(e)}") - return [] + return [], classify_list_exception(e) # Fetch tools from all servers in parallel tasks = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers] results = await asyncio.gather(*tasks) # Flatten results into single list - all_tools: List[MCPTool] = [tool for tools in results for tool in tools] + all_tools: List[MCPTool] = [tool for tools, _ in results for tool in tools] + server_outcomes: Dict[str, ServerOutcome] = { + _aggregate_server_key(server): outcome + for server, (_, outcome) in zip(allowed_mcp_servers, results) + if server is not None + } # If logging is enabled, enrich spend_logs_metadata with counts if litellm_logging_obj: - per_server_tool_counts: Dict[str, int] = {} - for server, server_tools in zip(allowed_mcp_servers, results): - if server is None: - continue - server_key = ( - getattr(server, "server_name", None) - or getattr(server, "alias", None) - or getattr(server, "name", None) - or "unknown" - ) - per_server_tool_counts[str(server_key)] = len(server_tools) + per_server_tool_counts: Dict[str, int] = { + _aggregate_server_key(server): len(server_tools) + for server, (server_tools, _) in zip(allowed_mcp_servers, results) + if server is not None + } metadata_dict = litellm_logging_obj.model_call_details.get("metadata") if isinstance(metadata_dict, dict): @@ -1975,6 +2003,9 @@ async def _fetch_and_filter_server_tools( spend_meta["allowed_server_count"] = len(allowed_mcp_servers) spend_meta["tool_count_total"] = len(all_tools) spend_meta["per_server_tool_counts"] = per_server_tool_counts + spend_meta["per_server_list_outcomes"] = { + key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items() + } end_time = datetime.now() try: @@ -1995,7 +2026,7 @@ async def _fetch_and_filter_server_tools( verbose_logger.info(f"Successfully fetched {len(all_tools)} tools total from all MCP servers") - return all_tools + return AggregateToolListing(tools=all_tools, outcomes=server_outcomes) except Exception as e: # Only fire failure hook if logging was requested for this list-tools execution if log_list_tools_to_spendlogs and user_api_key_auth is not None: @@ -2218,43 +2249,6 @@ async def filter_tools_by_key_team_permissions( server = global_mcp_server_manager.get_mcp_server_by_id(server_id) return [t for t in tools if strip_known_server_prefix(t.name, server) in allowed_tool_names] - async def _merge_toolset_permissions( - user_api_key_auth: Optional[UserAPIKeyAuth], - ) -> Optional[UserAPIKeyAuth]: - """ - Resolve mcp_toolsets on the key's object_permission into tool-level permissions - and merge them (union) into object_permission.mcp_tool_permissions. - - Returns the (possibly mutated copy of) user_api_key_auth. - """ - if user_api_key_auth is None: - return None - op = user_api_key_auth.object_permission - if op is None: - return user_api_key_auth - toolset_ids = getattr(op, "mcp_toolsets", None) or [] - if not toolset_ids: - return user_api_key_auth - - toolset_perms = await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=toolset_ids) - if not toolset_perms: - return user_api_key_auth - - # Merge toolset_perms into existing mcp_tool_permissions (union) - existing = dict(op.mcp_tool_permissions or {}) - for server_id, tool_names in toolset_perms.items(): - existing_tools = existing.get(server_id, []) - merged = list(set(existing_tools) | set(tool_names)) - existing[server_id] = merged - - # Build updated object_permission with merged tool permissions and server IDs. - # Union the toolset's server IDs into mcp_servers so downstream server-level - # filtering doesn't silently drop servers that the toolset references but that - # aren't already in the key's explicit mcp_servers list. - merged_servers = list(set(op.mcp_servers or []) | set(existing.keys())) - updated_op = op.model_copy(update={"mcp_servers": merged_servers, "mcp_tool_permissions": existing}) - return user_api_key_auth.model_copy(update={"object_permission": updated_op}) - async def _list_mcp_tools( user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, @@ -2265,7 +2259,7 @@ async def _list_mcp_tools( log_list_tools_to_spendlogs: bool = False, list_tools_log_source: Optional[str] = None, client_ip: Optional[str] = None, - ) -> List[MCPTool]: + ) -> AggregateToolListing: """ List all available MCP tools. @@ -2277,19 +2271,14 @@ async def _list_mcp_tools( client_ip: Client IP for IP-based server access control Returns: - List[MCPTool]: Combined list of tools from all accessible servers + AggregateToolListing: Combined tools from all accessible servers plus each server's + classified listing outcome """ if not MCP_AVAILABLE: - return [] - - # Resolve toolset permissions and merge into the key's object_permission - # so that the existing filter_tools_by_key_team_permissions logic picks them up. - user_api_key_auth = await _merge_toolset_permissions(user_api_key_auth) + return AggregateToolListing(tools=[], outcomes={}) - # Get tools from managed MCP servers with error handling - managed_tools = [] try: - managed_tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, @@ -2300,12 +2289,12 @@ async def _list_mcp_tools( list_tools_log_source=list_tools_log_source, client_ip=client_ip, ) - verbose_logger.debug(f"Successfully fetched {len(managed_tools)} tools from managed MCP servers") + verbose_logger.debug(f"Successfully fetched {len(listing.tools)} tools from managed MCP servers") + return listing except Exception as e: verbose_logger.exception(f"Error getting tools from managed MCP servers: {str(e)}") - # Continue with empty managed tools list instead of failing completely - - return managed_tools + # Continue with an empty listing instead of failing completely + return AggregateToolListing(tools=[], outcomes={}) async def _list_mcp_prompts( user_api_key_auth: Optional[UserAPIKeyAuth] = None, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index fa57a2b3eb21..2f6b54a264aa 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -91,7 +91,7 @@ async def handle_mcp_tool_search( from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools - mcp_tools = await _list_mcp_tools( + mcp_listing = await _list_mcp_tools( user_api_key_auth=user_api_key_dict, mcp_servers=mcp_servers, client_ip=client_ip, @@ -100,6 +100,7 @@ async def handle_mcp_tool_search( oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) + mcp_tools = mcp_listing.tools tools = [ { "name": t.name, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c1cf0d90a15d..a104db18b3d8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -527,6 +527,7 @@ async def common_checks( request_headers=_safe_get_request_headers(request=request), request_query_params=_safe_get_request_query_params(request=request), llm_router=llm_router, + request=request, ) if route in MODEL_DISCOVERY_ROUTES: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index a610e44e69cb..38900260c98f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -14,6 +14,9 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_ENDPOINT_MARKER, +) from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS from litellm.types.utils import CustomPricingLiteLLMParams @@ -1482,13 +1485,50 @@ def _format_model_candidates( return candidates +def _request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool: + """Whether FastAPI resolved this request to a user-defined pass-through handler. + + Reads the marker set by ``create_pass_through_route`` off the dispatched endpoint + (``request.scope["endpoint"]``). Because routing has already run by the time auth + dependencies execute, this reflects the handler that actually serves the request: + a custom path colliding with a built-in route resolves to the built-in handler, + which carries no marker, so model-access checks are never wrongly skipped. + """ + if request is None: + return False + scope = getattr(request, "scope", None) + if not isinstance(scope, dict): + return False + endpoint = scope.get("endpoint") + # Identity check against True (not truthiness): the marker is set to the literal + # True, and this keeps a spec'd Mock request (whose attribute access yields truthy + # child mocks) from being misread as a pass-through dispatch. + return getattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is True + + def get_model_from_request( request_data: dict, route: str, request_headers: Optional[Mapping[str, Any]] = None, request_query_params: Optional[Mapping[str, Any]] = None, llm_router: Optional[Router] = None, + request: Request | None = None, ) -> Optional[Union[str, List[str]]]: + """Resolve the model(s) a request targets, for model-access and budget checks. + + Returns ``None`` when the request was dispatched to a user-defined pass-through + endpoint: its body is forwarded verbatim to the configured upstream, so a + ``model`` field there names an upstream model, not a LiteLLM-managed one, and + enforcing key/team model allowlists against it would reject valid requests. The + check reads the FastAPI-resolved endpoint (``request.scope["endpoint"]``), not the + request path, so a custom path that collides with a built-in route never + suppresses model-access checks: on a collision the built-in handler is dispatched + and does not carry the marker. Built-in provider passthrough routes + (``/vertex_ai``, ``/gemini``, ...) are separate handlers and keep model enforcement. + """ + if _request_dispatched_to_pass_through_endpoint(request): + return None + candidates = _extract_model_candidates_from_request( request_data=request_data, route=route, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4d07d4c043cc..1a1b355cb171 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -162,6 +162,7 @@ def _get_model_from_request_context( request_headers=_safe_get_request_headers(request=request), request_query_params=_safe_get_request_query_params(request=request), llm_router=llm_router, + request=request, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 02bb66388ca3..c7c9397d8509 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -151,6 +151,80 @@ async def _record_streaming_client_disconnect_if_needed( return True +def _deferred_stream_logging_is_armed(request_data: dict) -> bool: + logging_obj = request_data.get("litellm_logging_obj") + if logging_obj is None: + return False + return ( + getattr(logging_obj, "_on_deferred_stream_complete", None) is not None + and getattr(logging_obj, "_deferred_stream_complete_args", None) is not None + ) + + +async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, response: object) -> bool: + """ + A client disconnect throws GeneratorExit/CancelledError into the streaming + generator, so neither the success nor the failure logging callback fires + and the chunks already streamed (plus any sub-call cost folded into the + logging object) would never reach spend tracking. Assemble the partial + response from the wrapper's collected chunks and dispatch success logging + for it; dispatch_success_handlers dedups against a natural end-of-stream + dispatch via has_dispatched_final_stream_success. + + Awaited directly by the shielded cleanup rather than scheduled with + create_task: the client is already gone so the extra latency is harmless, + and an unrooted task could be garbage-collected before it bills. + + Returns True when a disconnect-time success event owns the request's + max_parallel_requests slot release (one was dispatched here, or one had + already been dispatched for this stream), so the caller can skip the + explicit slot release and avoid a double release. Returns False when no + success event fired (logging disabled, nothing streamed, or assembly + failed) and the caller must release the slot itself. + """ + if litellm.disable_streaming_logging is True: + return False + logging_obj = request_data.get("litellm_logging_obj") + if not isinstance(logging_obj, LiteLLMLoggingObj): + return False + if logging_obj.model_call_details.get("has_dispatched_final_stream_success"): + # A natural end-of-stream success event already fired and released the + # slot; do not bill again, and let the caller skip the slot release. + return True + chunks: object = getattr(response, "chunks", None) + if not isinstance(chunks, list) or not chunks: + return False + verbose_proxy_logger.debug( + "Billing partial streamed spend for %s chunks after client disconnect, litellm_call_id=%s", + len(chunks), + request_data.get("litellm_call_id"), + ) + messages: object = getattr(response, "messages", None) + try: + partial_response = litellm.stream_chunk_builder( + chunks=chunks, + messages=messages if isinstance(messages, list) else None, + logging_obj=logging_obj, + ) + except Exception as e: # noqa: BLE001 # partial billing is best-effort; never break stream teardown + verbose_proxy_logger.debug("Failed to assemble partial streamed response for disconnect billing: %s", e) + return False + if partial_response is None: + return False + try: + await logging_obj.dispatch_success_handlers( + partial_response, + cache_hit=False, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + except Exception as e: # noqa: BLE001 # partial billing is best-effort; never break stream teardown + verbose_proxy_logger.debug("Failed to dispatch disconnect billing event: %s", e) + return False + return True + + async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None: pending_tasks = [task for task in tasks if not task.done()] for task in pending_tasks: @@ -2575,6 +2649,8 @@ async def _finalize_streaming_generator_cleanup( response: Any, stream_completed: bool = False, client_disconnected: bool = False, + user_api_key_dict: UserAPIKeyAuth | None = None, + proxy_logging_obj: ProxyLogging | None = None, ) -> None: with anyio.CancelScope(shield=True): should_record_client_disconnect = client_disconnected or (not stream_completed) @@ -2586,7 +2662,28 @@ async def _finalize_streaming_generator_cleanup( client_disconnected, ) if recorded_client_disconnect: + deferred_stream_logging_armed = _deferred_stream_logging_is_armed(request_data) ProxyLogging._fire_deferred_stream_logging(request_data) + # A disconnect-time success event (the deferred-guardrail flush + # above, or the partial-spend billing below) releases the + # request's max_parallel_requests slot through the limiter's + # own success callback. Release the slot explicitly only when + # no such event fires, so exactly one release happens; two + # concurrent releases would race and double-decrement under the + # limiter's in-memory fallback. + success_event_owns_slot_release = deferred_stream_logging_armed + if not deferred_stream_logging_armed: + success_event_owns_slot_release = await _bill_partial_streamed_spend_on_disconnect( + request_data, response + ) + if ( + not success_event_owns_slot_release + and proxy_logging_obj is not None + and user_api_key_dict is not None + ): + await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect( + user_api_key_dict, request_data + ) if hasattr(response, "aclose"): try: @@ -2675,12 +2772,13 @@ async def async_streaming_data_generator( except (asyncio.CancelledError, GeneratorExit): # Client disconnected mid-stream. CancelledError / GeneratorExit # are BaseException and bypass the success/failure logging - # callbacks that release the pre-call max_parallel_requests +1; - # release it here. This is the outermost generator Starlette closes - # on disconnect, so the nested iterator hook (which only sees - # GeneratorExit on GC) cannot own the refund. + # callbacks that release the pre-call max_parallel_requests +1. + # Flag the disconnect; the shielded cleanup in `finally` owns the + # slot release so it can coordinate with disconnect-time success + # billing and release exactly once. This is the outermost generator + # Starlette closes on disconnect, so the nested iterator hook (which + # only sees GeneratorExit on GC) cannot own the refund. if not stream_completed: - proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict) client_disconnected = True if not delivered_chunk: from litellm.proxy.spend_tracking.budget_reservation import ( @@ -2723,6 +2821,8 @@ async def async_streaming_data_generator( response=response, stream_completed=stream_completed, client_disconnected=client_disconnected, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/__init__.py new file mode 100644 index 000000000000..0fc74ddec93f --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/__init__.py @@ -0,0 +1,50 @@ +""" +Author: Madan Singhal +Date: 23/06/26 + +""" + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .singulr import SingulrGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", +): + import litellm + + _cb = SingulrGuardrail( + singulr_api_base=getattr(litellm_params, "singulr_api_base", None) or litellm_params.api_base, + singulr_api_key=getattr(litellm_params, "singulr_api_key", None) or litellm_params.api_key, + singulr_application_id=getattr(litellm_params, "singulr_application_id", None), + singulr_guardrail_id=getattr(litellm_params, "singulr_guardrail_id", None), + block_on_error=getattr(litellm_params, "block_on_error", None), + timeout=litellm_params.timeout, + guardrail_name=guardrail.get( + "guardrail_name", + "", + ), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback( + _cb, + ) + + return _cb + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.SINGULR.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.SINGULR.value: SingulrGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py new file mode 100644 index 000000000000..36a09a4ea254 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -0,0 +1,216 @@ +import os +from typing import Any +from urllib.parse import urlparse + +import httpx +import pydantic + +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.base import ( + GuardrailConfigModel, +) +from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( + SingulrGuardrailPayload, + SingulrGuardrailRequest, + SingulrGuardrailResponse, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +_DEFAULT_API_BASE = "http://localhost:8003" +_GUARD_ENDPOINT = "/api/v1/ai-gateway/litellm" +_DEFAULT_TIMEOUT = 30.0 + + +class SingulrGuardrail(CustomGuardrail): + def __init__( + self, + singulr_api_key: str | None = None, + singulr_api_base: str | None = None, + singulr_application_id: str | None = None, + singulr_guardrail_id: str | None = None, + block_on_error: bool | None = None, + timeout: float | None = None, + **kwargs: Any, + ) -> None: + self.singulr_api_key = singulr_api_key or os.environ.get("SINGULR_API_KEY") + self.singulr_api_base = (singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).rstrip( + "/" + ) + parsed = urlparse(self.singulr_api_base) + if parsed.scheme == "http" and parsed.hostname not in ( + "localhost", + "127.0.0.1", + ): + raise ValueError( + f"Singulr: api_base {self.singulr_api_base} uses plain HTTP for a " + "non-local endpoint. Guardrail payloads contain the API token, full " + "conversation content, and the guardrail decision, so this endpoint " + "must use HTTPS." + ) + + self.singulr_application_id = singulr_application_id or os.environ.get("SINGULR_ENFORCEMENT_ENTITY_ID") + self.singulr_guardrail_id = singulr_guardrail_id or os.environ.get("SINGULR_GUARDRAIL_ID") + + if block_on_error is None: + env = os.environ.get("SINGULR_BLOCK_ON_ERROR", "true") + self.block_on_error = env.lower() in ("true", "1", "yes") + else: + self.block_on_error = block_on_error + + self.timeout = _DEFAULT_TIMEOUT if timeout is None else timeout + + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + ) + + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + super().__init__(**kwargs) + + @staticmethod + def get_config_model() -> type["GuardrailConfigModel"] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( + SingulrGuardrailConfigModel, + ) + + return SingulrGuardrailConfigModel + + def _build_payload( + self, + request_data: dict[str, Any], + inputs: GenericGuardrailAPIInputs, + input_type: str, + ) -> dict[str, Any]: + if not request_data: + texts = inputs.get("texts", []) + + payload = SingulrGuardrailPayload( + input_type=input_type, + is_playground_request=True, + playground_text=texts[0] if texts else None, + ) + else: + response = request_data.get("response") + singulr_req_object = SingulrGuardrailRequest( + model=request_data.get("model"), + messages=request_data.get("messages"), + tools=request_data.get("tools"), + model_response=response.model_dump(mode="json") if input_type == "response" and response else None, + litellm_metadata=request_data.get("litellm_metadata"), + ) + payload = SingulrGuardrailPayload( + litellm_call_id=request_data.get("litellm_call_id"), + request_data=singulr_req_object, + input_type=input_type, + ) + + return payload.model_dump(mode="json") + + def _build_headers(self) -> dict[str, str]: + return dict( + (header, value) + for header, value in ( + ("Content-Type", "application/json"), + ("X-Singulr-Gateway-Token", self.singulr_api_key), + ( + "X-Singulr-Enforcement-Entity-Id", + self.singulr_application_id or "", + ), + ("X-Singulr-Guardrail-Id", self.singulr_guardrail_id or ""), + ) + if value + ) + + async def _call_api(self, payload: dict[str, Any]) -> SingulrGuardrailResponse | None: + endpoint = f"{self.singulr_api_base}{_GUARD_ENDPOINT}" + verbose_proxy_logger.debug("Singulr: %s", endpoint) + + try: + response = await self.async_handler.post( + url=endpoint, + headers=self._build_headers(), + json=payload, + timeout=self.timeout, + ) + response.raise_for_status() + result = SingulrGuardrailResponse.model_validate(response.json()) + verbose_proxy_logger.debug("Singulr: result=%s", result) + return result + + except httpx.HTTPStatusError as exc: + verbose_proxy_logger.error( + "Singulr API returned HTTP %s: %s", + exc.response.status_code, + str(exc), + ) + if self.block_on_error: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=(f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}"), + ) from exc + return None + + except httpx.TransportError as exc: + verbose_proxy_logger.error("Singulr API unreachable: %s", str(exc)) + if self.block_on_error: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Singulr API unreachable (block_on_error=True): {exc}", + ) from exc + return None + + except (ValueError, pydantic.ValidationError) as exc: + verbose_proxy_logger.error("Singulr API returned an invalid response: %s", str(exc)) + if self.block_on_error: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Singulr API returned an invalid response: {exc}", + ) from exc + return None + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: str, + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + payload = self._build_payload(request_data, inputs, input_type) + if not payload: + return inputs + + result = await self._call_api(payload) + if result is None: + return inputs + + verbose_proxy_logger.debug( + "Singulr: should_block=%s blocking_due_to=%s", + result.should_block, + result.blocking_due_to, + ) + + if result.should_block: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Blocked by Singulr: {result.blocking_due_to or 'unknown'}", + ) + + return inputs diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d60c17c744f1..22ea9fe176a9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -7,6 +7,7 @@ import asyncio import binascii import os +import uuid from datetime import datetime from typing import ( TYPE_CHECKING, @@ -185,6 +186,69 @@ return results """ +PARALLEL_ACQUIRE_SCRIPT = """ +-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge. +-- Each gauge key is a sorted set of per-request slot ids scored by acquire +-- time (Redis server clock). In-flight requests are counted by ZCARD after +-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM +-- counters the gauge is never reset while requests are in flight, a +-- rejected request never occupies a slot, and a slot leaked by a crashed +-- worker self-heals after the slot TTL even under continuous traffic. +-- +-- KEYS: one gauge zset key per descriptor. +-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id). +-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +for i = 1, #KEYS do + local limit = tonumber(ARGV[(i - 1) * 3 + 1]) + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl) + local in_flight = redis.call('ZCARD', KEYS[i]) + if in_flight + 1 > limit then + return { 1, i, in_flight, limit } + end +end +local results = { 0 } +for i = 1, #KEYS do + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + local slot_id = ARGV[(i - 1) * 3 + 3] + redis.call('ZADD', KEYS[i], now, slot_id) + redis.call('EXPIRE', KEYS[i], slot_ttl) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +""" + +PARALLEL_RELEASE_SCRIPT = """ +-- Release one slot per gauge key by removing this request's slot id. +-- ZREM of an absent member (or key) is a no-op, so a release without a +-- matching acquire (proxy-side rejection, double-fired callback, slot +-- already expired) can never free a slot owned by another request. +-- KEYS: gauge zset keys. ARGV: per-key slot_id. +-- Returns the remaining in-flight count per key. +local results = {} +for i = 1, #KEYS do + redis.call('ZREM', KEYS[i], ARGV[i]) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +""" + +PARALLEL_COUNT_SCRIPT = """ +-- Read the current in-flight count per gauge key (prunes expired slots +-- first so leaked slots do not inflate the reading). +-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local results = {} +for i = 1, #KEYS do + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i])) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +""" + TOKEN_INCREMENT_SCRIPT = """ local results = {} @@ -248,6 +312,19 @@ # mirror ``x-ratelimit-*`` headers into the SLP. Streaming exits # common_request_processing before ``async_post_call_success_hook`` runs. RATE_LIMIT_RESPONSE_KEY = "_litellm_proxy_rate_limit_response" +# Holds the acquisition the pre-call hook made for this request: the slot id +# plus the gauge counter keys it was registered under. The success/failure +# callbacks release only this exact acquisition: those callbacks also fire +# for requests rejected at pre-call (which never acquired a slot), and an +# id-less release would free a slot still owned by another in-flight request +# — every rejection would then raise effective concurrency above the +# configured limit. +MAX_PARALLEL_SLOT_ACQUIRED_KEY = "_litellm_max_parallel_slot_acquired" +# How long an acquired slot counts toward the in-flight total before it is +# considered leaked (worker crashed without any release callback firing) and +# pruned. Also the longest request duration the gauge can track: a request +# running longer than this stops occupying its slot. +PARALLEL_REQUEST_SLOT_TTL_SECONDS = 3600 # Stash keys live ONLY in metadata channels — never at the top level of the # request body. Top-level keys are forwarded as body params to upstream # providers, which reject unknown fields with 400/429 errors. @@ -258,6 +335,7 @@ TPM_RESERVATION_RELEASED_KEY, RATE_LIMIT_DESCRIPTORS_KEY, RATE_LIMIT_RESPONSE_KEY, + MAX_PARALLEL_SLOT_ACQUIRED_KEY, ) @@ -274,6 +352,17 @@ class RateLimitDescriptor(TypedDict): rate_limit: Optional[RateLimitDescriptorRateLimitObject] +class ParallelRequestGauge(TypedDict): + counter_key: str + limit: int + descriptor_key: str + + +class ParallelSlotAcquisition(TypedDict): + slot_id: str + counter_keys: list[str] + + class RateLimitStatus(TypedDict): code: str current_limit: int @@ -310,10 +399,22 @@ def __init__( self.check_and_increment_by_n_script = ( self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT) ) + self.parallel_acquire_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + PARALLEL_ACQUIRE_SCRIPT + ) + self.parallel_release_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + PARALLEL_RELEASE_SCRIPT + ) + self.parallel_count_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + PARALLEL_COUNT_SCRIPT + ) else: self.batch_rate_limiter_script = None self.token_increment_script = None self.check_and_increment_by_n_script = None + self.parallel_acquire_script = None + self.parallel_release_script = None + self.parallel_count_script = None self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) @@ -559,7 +660,6 @@ def is_cache_list_over_limit( counter_key = keys_to_fetch[i + 1] counter_value = cache_values[i + 1] requests_limit = key_metadata[window_key]["requests_limit"] - max_parallel_requests_limit = key_metadata[window_key]["max_parallel_requests_limit"] tokens_limit = key_metadata[window_key]["tokens_limit"] # Determine which limit to use for current_limit and limit_remaining @@ -568,9 +668,6 @@ def is_cache_list_over_limit( if counter_key.endswith(":requests"): current_limit = requests_limit rate_limit_type = "requests" - elif counter_key.endswith(":max_parallel_requests"): - current_limit = max_parallel_requests_limit - rate_limit_type = "max_parallel_requests" elif counter_key.endswith(":tokens"): current_limit = tokens_limit rate_limit_type = "tokens" @@ -694,6 +791,7 @@ async def should_rate_limit( parent_otel_span: Optional[Span] = None, read_only: bool = False, skip_tpm_check: bool = False, + parallel_slot_id: str | None = None, ) -> RateLimitResponse: """ Check if any of the rate limit descriptors should be rate limited. @@ -710,15 +808,122 @@ async def should_rate_limit( ``reserve_tpm_tokens`` reservation path should set this to avoid the +1-per-key Lua / in-memory increment double-charging the tokens counter. + + ``max_parallel_requests`` descriptors are enforced by the dedicated + concurrency-gauge path (``_check_parallel_request_gauges``), never by + the windowed counters. The gauge phase must stay AFTER the windowed + check so a windowed rejection never strands an acquired slot; the + reverse order would leak one gauge slot per RPM/TPM rejection. + ``parallel_slot_id`` names the slot an admission registers; callers + that enforce (not read_only) should pass the id they will later + release with — when omitted, a generated slot id is used and the slot + can only be reclaimed by TTL expiry. """ current_time = self._get_current_time() now = current_time.timestamp() now_int = int(now) # Convert to integer for Redis Lua script - # Collect all keys and their metadata upfront + keys_to_fetch, key_metadata, gauges = self._collect_windowed_keys_and_gauges( + descriptors=descriptors, + skip_tpm_check=skip_tpm_check, + ) + + windowed_response = RateLimitResponse(overall_code="OK", statuses=[]) + if keys_to_fetch: + ## CHECK IN-MEMORY CACHE + cache_values = await self.internal_usage_cache.async_batch_get_cache( + keys=keys_to_fetch, + parent_otel_span=parent_otel_span, + local_only=True, + ) + + if cache_values is not None: + rate_limit_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata) + if rate_limit_response["overall_code"] == "OVER_LIMIT": + return rate_limit_response + + ## IF under limit in-memory, check Redis + if read_only: + # READ-ONLY MODE: Just read current values without incrementing + cache_values = await self.internal_usage_cache.async_batch_get_cache( + keys=keys_to_fetch, + parent_otel_span=parent_otel_span, + local_only=False, # Check Redis too + ) + + # For keys that don't exist yet, set them to 0 + if cache_values is None: + cache_values = [] + for _ in keys_to_fetch: + cache_values.append(str(now_int) if _.endswith(":window") else 0) + elif self.batch_rate_limiter_script is not None: + # NORMAL MODE: Increment counters in Redis + # Group keys by hash tag for Redis cluster compatibility + cache_values = await self._execute_redis_batch_rate_limiter_script( + keys_to_fetch=keys_to_fetch, + now_int=now_int, + ) + + # update in-memory cache with new values + for i in range(0, len(cache_values), 2): + window_key = keys_to_fetch[i] + counter_key = keys_to_fetch[i + 1] + window_value = cache_values[i] + counter_value = cache_values[i + 1] + await self.internal_usage_cache.async_set_cache( + key=counter_key, + value=counter_value, + ttl=self.window_size, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + await self.internal_usage_cache.async_set_cache( + key=window_key, + value=window_value, + ttl=self.window_size, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + else: + # NORMAL MODE: In-memory sliding window (no Redis) + cache_values = await self.in_memory_cache_sliding_window( + keys=keys_to_fetch, + now_int=now_int, + window_size=self.window_size, + ) + + windowed_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata) + if windowed_response["overall_code"] == "OVER_LIMIT": + return windowed_response + + if not gauges: + return windowed_response + + gauge_response = await self._check_parallel_request_gauges( + gauges=gauges, + slot_id=parallel_slot_id or uuid.uuid4().hex, + parent_otel_span=parent_otel_span, + read_only=read_only, + ) + return RateLimitResponse( + overall_code=gauge_response["overall_code"], + statuses=[*windowed_response["statuses"], *gauge_response["statuses"]], + ) + + def _collect_windowed_keys_and_gauges( + self, + descriptors: list[RateLimitDescriptor], + skip_tpm_check: bool, + ) -> tuple[list[str], dict[str, dict[str, Any]], list[ParallelRequestGauge]]: + """ + Split descriptors into the windowed (window_key, counter_key) fetch + list with its per-window metadata, and the concurrency gauges for + descriptors carrying a max_parallel_requests limit. + """ keys_to_fetch: List[str] = [] - key_metadata = {} # Store metadata for each key + key_metadata: dict[str, dict[str, Any]] = {} + gauges: list[ParallelRequestGauge] = [] for descriptor in descriptors: descriptor_key = descriptor["key"] descriptor_value = descriptor["value"] @@ -732,6 +937,17 @@ async def should_rate_limit( window_key = f"{{{descriptor_key}:{descriptor_value}}}:window" + if max_parallel_requests_limit is not None: + gauges.append( + ParallelRequestGauge( + counter_key=self.create_rate_limit_keys( + descriptor_key, descriptor_value, "max_parallel_requests" + ), + limit=int(max_parallel_requests_limit), + descriptor_key=descriptor_key, + ) + ) + rate_limit_set = False if requests_limit is not None: rpm_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, "requests") @@ -741,12 +957,6 @@ async def should_rate_limit( tpm_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, "tokens") keys_to_fetch.extend([window_key, tpm_key]) rate_limit_set = True - if max_parallel_requests_limit is not None: - max_parallel_requests_key = self.create_rate_limit_keys( - descriptor_key, descriptor_value, "max_parallel_requests" - ) - keys_to_fetch.extend([window_key, max_parallel_requests_key]) - rate_limit_set = True if not rate_limit_set: continue @@ -754,77 +964,252 @@ async def should_rate_limit( key_metadata[window_key] = { "requests_limit": (int(requests_limit) if requests_limit is not None else None), "tokens_limit": int(tokens_limit) if tokens_limit is not None else None, - "max_parallel_requests_limit": ( - int(max_parallel_requests_limit) if max_parallel_requests_limit is not None else None - ), "window_size": int(window_size), "descriptor_key": descriptor_key, } + return keys_to_fetch, key_metadata, gauges + + def _gauge_status(self, gauge: ParallelRequestGauge, in_flight: int, code: str) -> RateLimitStatus: + return RateLimitStatus( + code=code, + current_limit=gauge["limit"], + limit_remaining=max(0, gauge["limit"] - in_flight), + rate_limit_type="max_parallel_requests", + descriptor_key=gauge["descriptor_key"], + ) + + def _gauge_in_flight_from_cache_value(self, raw_value: Any) -> int: + """ + In-flight count from a cached gauge value: a dict of slot_id -> + acquire timestamp when the in-memory registry is authoritative, or + the mirrored integer count from the last Redis script result. + """ + if raw_value is None: + return 0 + if isinstance(raw_value, dict): + cutoff = self._get_current_time().timestamp() - PARALLEL_REQUEST_SLOT_TTL_SECONDS + return sum(1 for ts in raw_value.values() if isinstance(ts, (int, float)) and ts >= cutoff) + return max(0, int(raw_value)) + + async def _check_parallel_request_gauges( + self, + gauges: list[ParallelRequestGauge], + slot_id: str, + parent_otel_span: Span | None = None, + read_only: bool = False, + ) -> RateLimitResponse: + """ + Enforce max_parallel_requests as a concurrency gauge over a per-slot + registry: each admitted request registers ``slot_id`` with its + acquire time, and admission requires in_flight + 1 <= limit over the + unexpired slots. Unlike the windowed RPM/TPM counters, the gauge is + never reset while requests are in flight, a rejected request never + occupies a slot, and a slot leaked by a crashed worker is pruned + after PARALLEL_REQUEST_SLOT_TTL_SECONDS even under continuous + traffic. Releases remove exactly this request's slot id, so a + double-fired or unmatched release can never free another request's + slot. + """ + gauge_keys = [gauge["counter_key"] for gauge in gauges] - ## CHECK IN-MEMORY CACHE - cache_values = await self.internal_usage_cache.async_batch_get_cache( - keys=keys_to_fetch, + if read_only: + if self.parallel_count_script is not None: + try: + raw_counts = await self.parallel_count_script( + keys=gauge_keys, + args=[PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges], + ) + counts = [max(0, int(value)) for value in raw_counts] + except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500 + verbose_proxy_logger.warning(f"parallel_count_script failed, using local mirror: {str(e)}") + counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) + else: + counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) + statuses = [] + overall_code = "OK" + for gauge, in_flight in zip(gauges, counts): + code = "OVER_LIMIT" if in_flight >= gauge["limit"] else "OK" + if code == "OVER_LIMIT": + overall_code = "OVER_LIMIT" + statuses.append(self._gauge_status(gauge, in_flight, code)) + return RateLimitResponse(overall_code=overall_code, statuses=statuses) + + local_counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) + for gauge, in_flight in zip(gauges, local_counts): + if in_flight >= gauge["limit"]: + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[self._gauge_status(gauge, in_flight, "OVER_LIMIT")], + ) + + if self.parallel_acquire_script is not None: + try: + raw = await self.parallel_acquire_script( + keys=gauge_keys, + args=[ + arg for gauge in gauges for arg in (gauge["limit"], PARALLEL_REQUEST_SLOT_TTL_SECONDS, slot_id) + ], + ) + except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500 + verbose_proxy_logger.warning( + f"parallel_acquire_script failed, falling back to in-memory gauge: {str(e)}" + ) + async with self._check_and_increment_lock: + return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span) + if int(raw[0]) == 1: + gauge = gauges[int(raw[1]) - 1] + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[self._gauge_status(gauge, int(raw[2]), "OVER_LIMIT")], + ) + statuses = [] + for gauge, in_flight in zip(gauges, raw[1:]): + await self.internal_usage_cache.async_set_cache( + key=gauge["counter_key"], + value=int(in_flight), + ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + statuses.append(self._gauge_status(gauge, int(in_flight), "OK")) + return RateLimitResponse(overall_code="OK", statuses=statuses) + + async with self._check_and_increment_lock: + return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span) + + async def _read_local_gauge_counts( + self, + gauge_keys: list[str], + parent_otel_span: Span | None = None, + ) -> list[int]: + values = await self.internal_usage_cache.async_batch_get_cache( + keys=gauge_keys, parent_otel_span=parent_otel_span, local_only=True, ) + if values is None: + return [0 for _ in gauge_keys] + return [self._gauge_in_flight_from_cache_value(value) for value in values] + + async def _acquire_parallel_slots_in_memory( + self, + gauges: list[ParallelRequestGauge], + slot_id: str, + parent_otel_span: Span | None = None, + ) -> RateLimitResponse: + """ + All-or-nothing in-memory slot-registry acquire. Caller holds the lock. + + A cached dict is the authoritative in-memory registry. A cached + integer is the count mirrored from the last successful Redis script + call: when Redis fails over to this path, that mirror still counts + the slots in flight on the Redis side, so it is carried forward as + an integer counter (not discarded as an empty registry, which would + briefly double the admitted concurrency during a Redis outage). + """ + now = self._get_current_time().timestamp() + cutoff = now - PARALLEL_REQUEST_SLOT_TTL_SECONDS + states: list[tuple[dict[str, float] | None, int]] = [] + for gauge in gauges: + raw_value = await self.internal_usage_cache.async_get_cache( + key=gauge["counter_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + if isinstance(raw_value, dict): + registry: dict[str, float] | None = { + key: float(ts) for key, ts in raw_value.items() if isinstance(ts, (int, float)) and ts >= cutoff + } + in_flight = len(registry or {}) + elif raw_value is None: + registry = {} + in_flight = 0 + else: + registry = None + in_flight = max(0, int(raw_value)) + if in_flight + 1 > gauge["limit"]: + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[self._gauge_status(gauge, in_flight, "OVER_LIMIT")], + ) + states.append((registry, in_flight)) - if cache_values is not None: - rate_limit_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata) - if rate_limit_response["overall_code"] == "OVER_LIMIT": - return rate_limit_response + statuses = [] + for gauge, (registry, in_flight) in zip(gauges, states): + new_value: Union[dict[str, float], int] = ( + {**registry, slot_id: now} if registry is not None else in_flight + 1 + ) + await self.internal_usage_cache.async_set_cache( + key=gauge["counter_key"], + value=new_value, + ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + statuses.append(self._gauge_status(gauge, in_flight + 1, "OK")) + return RateLimitResponse(overall_code="OK", statuses=statuses) - ## IF under limit in-memory, check Redis - if read_only: - # READ-ONLY MODE: Just read current values without incrementing - cache_values = await self.internal_usage_cache.async_batch_get_cache( - keys=keys_to_fetch, - parent_otel_span=parent_otel_span, - local_only=False, # Check Redis too - ) - - # For keys that don't exist yet, set them to 0 - if cache_values is None: - cache_values = [] - for _ in keys_to_fetch: - cache_values.append(str(now_int) if _.endswith(":window") else 0) - elif self.batch_rate_limiter_script is not None: - # NORMAL MODE: Increment counters in Redis - # Group keys by hash tag for Redis cluster compatibility - cache_values = await self._execute_redis_batch_rate_limiter_script( - keys_to_fetch=keys_to_fetch, - now_int=now_int, - ) - - # update in-memory cache with new values - for i in range(0, len(cache_values), 2): - window_key = keys_to_fetch[i] - counter_key = keys_to_fetch[i + 1] - window_value = cache_values[i] - counter_value = cache_values[i + 1] - await self.internal_usage_cache.async_set_cache( + async def _release_parallel_request_slots( + self, + acquisition: ParallelSlotAcquisition, + parent_otel_span: Span | None = None, + ) -> None: + """ + Release the max_parallel_requests slots acquired at pre-call by + removing this request's slot id from every gauge it was registered + under. Removing an absent slot id is a no-op, so a release without a + matching acquire or a double-fired release can never free another + request's slot. The in-memory fallback decrements integer mirror + values (floored at 0) because the mirror carries no per-slot ids. + """ + counter_keys = acquisition["counter_keys"] + slot_id = acquisition["slot_id"] + if not counter_keys or not slot_id: + return + if self.parallel_release_script is not None: + try: + raw = await self.parallel_release_script( + keys=counter_keys, + args=[slot_id for _ in counter_keys], + ) + for counter_key, remaining in zip(counter_keys, raw): + await self.internal_usage_cache.async_set_cache( + key=counter_key, + value=max(0, int(remaining)), + ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + return + except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the in-memory release, never a 500 + verbose_proxy_logger.warning( + f"parallel_release_script failed, falling back to in-memory release: {str(e)}" + ) + + async with self._check_and_increment_lock: + for counter_key in counter_keys: + raw_value = await self.internal_usage_cache.async_get_cache( key=counter_key, - value=counter_value, - ttl=self.window_size, litellm_parent_otel_span=parent_otel_span, local_only=True, ) + if isinstance(raw_value, dict): + if slot_id not in raw_value: + continue + new_value: Union[dict[str, float], int] = { + key: ts for key, ts in raw_value.items() if key != slot_id + } + elif raw_value is None: + continue + else: + new_value = max(0, int(raw_value) - 1) await self.internal_usage_cache.async_set_cache( - key=window_key, - value=window_value, - ttl=self.window_size, + key=counter_key, + value=new_value, + ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, litellm_parent_otel_span=parent_otel_span, local_only=True, ) - else: - # NORMAL MODE: In-memory sliding window (no Redis) - cache_values = await self.in_memory_cache_sliding_window( - keys=keys_to_fetch, - now_int=now_int, - window_size=self.window_size, - ) - - rate_limit_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata) - return rate_limit_response async def atomic_check_and_increment_by_n( self, @@ -2027,10 +2412,18 @@ async def async_pre_call_hook( # shrinking the effective TPM budget by N and causing # false-positive 429s under bursts. When reservation is disabled, # this pass enforces TPM directly from the post-call counters. + parallel_counter_keys = [ + self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") + for d in descriptors + if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None + ] + parallel_slot_id = uuid.uuid4().hex if parallel_counter_keys else None + response = await self.should_rate_limit( descriptors=descriptors, parent_otel_span=user_api_key_dict.parent_otel_span, skip_tpm_check=self.tpm_reservation_enabled, + parallel_slot_id=parallel_slot_id, ) if response["overall_code"] == "OVER_LIMIT": @@ -2049,6 +2442,15 @@ async def async_pre_call_hook( key=RATE_LIMIT_RESPONSE_KEY, value=response, ) + if parallel_slot_id is not None: + self._stash_value_in_metadata_channels( + data=data, + key=MAX_PARALLEL_SLOT_ACQUIRED_KEY, + value={ + "slot_id": parallel_slot_id, + "counter_keys": parallel_counter_keys, + }, + ) # ---------------------------------------------------------------- # TPM token reservation @@ -2108,6 +2510,13 @@ async def async_pre_call_hook( ) if tpm_response["overall_code"] == "OVER_LIMIT": + acquisition = self._get_parallel_slot_acquisition(kwargs=data) + if acquisition is not None: + await self._release_parallel_request_slots( + acquisition=acquisition, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + self._clear_parallel_slot_marker(data) self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, @@ -2480,6 +2889,50 @@ def _is_reservation_released( """True if a prior callback already refunded this request's reservation.""" return bool(cls._lookup_stashed_value(kwargs, standard_logging_metadata, TPM_RESERVATION_RELEASED_KEY)) + @classmethod + def _get_parallel_slot_acquisition( + cls, + kwargs: Any, + standard_logging_metadata: dict[str, Any] | None = None, + ) -> ParallelSlotAcquisition | None: + """The slot acquisition this request's pre-call hook made, if any.""" + candidate = cls._lookup_stashed_value(kwargs, standard_logging_metadata, MAX_PARALLEL_SLOT_ACQUIRED_KEY) + if not isinstance(candidate, dict): + return None + slot_id = candidate.get("slot_id") + counter_keys = candidate.get("counter_keys") + if not isinstance(slot_id, str) or not slot_id: + return None + if not isinstance(counter_keys, list) or not counter_keys: + return None + if not all(isinstance(key, str) and key for key in counter_keys): + return None + return ParallelSlotAcquisition(slot_id=slot_id, counter_keys=counter_keys) + + @staticmethod + def _clear_parallel_slot_marker(data: Any) -> None: + """ + Remove the acquired-slot marker from every metadata channel a sibling + callback might read, so one release per acquire is an invariant even + when multiple callbacks fire for the same request. + """ + if not isinstance(data, dict): + return + for channel in ("metadata", "litellm_metadata"): + channel_dict = data.get(channel) + if isinstance(channel_dict, dict): + channel_dict.pop(MAX_PARALLEL_SLOT_ACQUIRED_KEY, None) + litellm_params = data.get("litellm_params") + if isinstance(litellm_params, dict): + lp_metadata = litellm_params.get("metadata") + if isinstance(lp_metadata, dict): + lp_metadata.pop(MAX_PARALLEL_SLOT_ACQUIRED_KEY, None) + slo = data.get("standard_logging_object") + if isinstance(slo, dict): + slo_meta = slo.get("metadata") + if isinstance(slo_meta, dict): + slo_meta.pop(MAX_PARALLEL_SLOT_ACQUIRED_KEY, None) + @staticmethod def _mark_reservation_released(data: Any) -> None: """ @@ -2621,7 +3074,6 @@ def _build_success_event_pipeline_operations( standard_logging_object = kwargs.get("standard_logging_object") or {} standard_logging_metadata = standard_logging_object.get("metadata") or {} - user_api_key = standard_logging_metadata.get("user_api_key_hash") model_group = get_model_group_from_litellm_kwargs(kwargs) # Get total tokens from response @@ -2658,20 +3110,6 @@ def _build_success_event_pipeline_operations( pipeline_operations: List[RedisPipelineIncrementOperation] = [] - # max_parallel_requests is its own counter (api-key only) — always decrement. - if user_api_key: - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=self.create_rate_limit_keys( - key="api_key", - value=user_api_key, - rate_limit_type="max_parallel_requests", - ), - increment_value=-1, - ttl=self.window_size, - ) - ) - # ---------------------------------------------------------------- # TPM reconciliation # Per-scope behavior: @@ -2719,6 +3157,19 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti try: verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") + standard_logging_object = kwargs.get("standard_logging_object") or {} + standard_logging_metadata = standard_logging_object.get("metadata") or {} + acquisition = self._get_parallel_slot_acquisition( + kwargs=kwargs, + standard_logging_metadata=standard_logging_metadata, + ) + if acquisition is not None: + await self._release_parallel_request_slots( + acquisition=acquisition, + parent_otel_span=litellm_parent_otel_span, + ) + self._clear_parallel_slot_marker(kwargs) + pipeline_operations = self._build_success_event_pipeline_operations( kwargs=kwargs, response_obj=response_obj, @@ -2855,22 +3306,19 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(kwargs) standard_logging_object = kwargs.get("standard_logging_object") or {} standard_logging_metadata = standard_logging_object.get("metadata") or {} - user_api_key = standard_logging_metadata.get("user_api_key_hash") pipeline_operations: List[RedisPipelineIncrementOperation] = [] - if user_api_key: - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=self.create_rate_limit_keys( - key="api_key", - value=user_api_key, - rate_limit_type="max_parallel_requests", - ), - increment_value=-1, - ttl=self.window_size, - ) + acquisition = self._get_parallel_slot_acquisition( + kwargs=kwargs, + standard_logging_metadata=standard_logging_metadata, + ) + if acquisition is not None: + await self._release_parallel_request_slots( + acquisition=acquisition, + parent_otel_span=litellm_parent_otel_span, ) + self._clear_parallel_slot_marker(kwargs) # Skip the reservation refund if async_post_call_failure_hook # already released it (proxy-level rejection that also bubbles up @@ -2920,40 +3368,35 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti except Exception as e: verbose_proxy_logger.exception(f"Error in rate limit failure event: {str(e)}") - async def async_release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None: + async def async_release_max_parallel_requests_on_disconnect( + self, + user_api_key_dict: UserAPIKeyAuth, + request_data: dict | None = None, + ) -> None: """ Release the api-key ``max_parallel_requests`` slot that - ``async_pre_call_hook`` reserved, for a request that ended without + ``async_pre_call_hook`` acquired, for a request that ended without either logging callback firing. - The +1 is normally undone by ``async_log_success_event`` (natural + The slot is normally released by ``async_log_success_event`` (natural stream completion) or ``async_log_failure_event`` (LLM error). When a client cancels a stream mid-flight, the cancellation surfaces as ``asyncio.CancelledError`` / ``GeneratorExit`` and neither callback - runs, so without this the counter leaks one slot per cancelled stream - until the key wedges at its limit. - """ - if not user_api_key_dict.api_key or user_api_key_dict.max_parallel_requests is None: + runs, so without this the slot leaks per cancelled stream until its + TTL prunes it. ``request_data`` carries the stashed acquisition; + its presence (not the key object's current max_parallel_requests + configuration, which can change mid-request) decides whether there + is anything to release. + """ + acquisition = self._get_parallel_slot_acquisition(kwargs=request_data) + if acquisition is None: return - await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=[ - RedisPipelineIncrementOperation( - key=self.create_rate_limit_keys( - key="api_key", - value=user_api_key_dict.api_key, - rate_limit_type="max_parallel_requests", - ), - increment_value=-1, - # Refresh the window TTL on the decrement, matching the - # failure path. max_parallel_requests is a concurrency - # gauge, not a rolling-window count, so the key must - # outlive in-flight requests rather than expire mid-stream. - ttl=self.window_size, - ) - ], - litellm_parent_otel_span=None, + await self._release_parallel_request_slots( + acquisition=acquisition, + parent_otel_span=None, ) + self._clear_parallel_slot_marker(request_data) async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ @@ -3002,17 +3445,29 @@ async def async_post_call_failure_hook( traceback_str: Optional[str] = None, ) -> None: """ - Release any TPM reservation when the request is rejected after the - pre-call hook reserved tokens but before the LLM call ran (e.g. a - downstream guardrail/auth hook raised). Without this, those - reservations are stranded — async_log_failure_event is a litellm - completion-level callback and never fires for proxy-side rejections. - - Idempotent via TPM_RESERVATION_RELEASED_KEY: if both this hook and + Release the parallel-request slot and any TPM reservation when the + request is rejected after the pre-call hook acquired them but before + the LLM call ran (e.g. a downstream guardrail/auth hook raised). + Without this, those resources are stranded — async_log_failure_event + is a litellm completion-level callback and never fires for proxy-side + rejections, so a leaked slot would occupy the gauge for the full + PARALLEL_REQUEST_SLOT_TTL_SECONDS. + + Idempotent: the slot release clears the acquisition marker (and slot + removal is a no-op ZREM on a second run), and the TPM refund is + guarded by TPM_RESERVATION_RELEASED_KEY — if both this hook and async_log_failure_event end up running in the same flow, only the - first refund applies. + first release/refund applies. """ try: + acquisition = self._get_parallel_slot_acquisition(kwargs=request_data) + if acquisition is not None: + await self._release_parallel_request_slots( + acquisition=acquisition, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + self._clear_parallel_slot_marker(request_data) + if self._is_reservation_released(kwargs=request_data): return reserved_tokens = self._get_reserved_tokens_from_kwargs(kwargs=request_data) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d920ee474cc0..d50db8324ef8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -723,12 +723,13 @@ async def get_mcp_tools( """ from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools - tools = await _list_mcp_tools( + listing = await _list_mcp_tools( user_api_key_auth=user_api_key_dict, mcp_auth_header=None, mcp_servers=None, mcp_server_auth_headers=None, ) + tools = listing.tools dumped_tools = [dict(tool) for tool in tools] return {"tools": dumped_tools} diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 2aff663038b7..acb2e50c79bd 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -68,6 +68,7 @@ from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + LITELLM_PASS_THROUGH_ENDPOINT_MARKER, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, EndpointType, PassthroughStandardLoggingPayload, @@ -1771,6 +1772,7 @@ async def endpoint_func( # type: ignore if hasattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY): delattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY) + setattr(endpoint_func, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) return endpoint_func diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 7a3b0d5a7276..492131811120 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -802,6 +802,19 @@ def _maybe_setup_prometheus_multiproc_dir( ), envvar="MAX_REQUESTS_BEFORE_RESTART_JITTER", ) +@click.option( + "--limit_concurrency", + default=None, + type=click.IntRange(min=1), + help=( + "Set uvicorn's concurrency limit. Uvicorn counts both active tasks and " + "accepted connections and returns HTTP 503 after the limit is reached. " + "Idle connections can consume capacity, so use upstream connection/header " + "timeouts and per-client connection limits. Only applies to uvicorn " + "(ignored under --run_gunicorn / --run_hypercorn / --run_granian)." + ), + envvar="LIMIT_CONCURRENCY", +) @click.option( "--enforce_prisma_migration_check", is_flag=True, @@ -870,6 +883,7 @@ def run_server( timeout_worker_healthcheck, max_requests_before_restart, max_requests_before_restart_jitter: Optional[int], + limit_concurrency: Optional[int], enforce_prisma_migration_check: bool, use_v2_migration_resolver: bool, reload: bool, @@ -1243,6 +1257,8 @@ def run_server( if max_requests_before_restart is not None: uvicorn_args["limit_max_requests"] = max_requests_before_restart if run_gunicorn is False and run_hypercorn is False and run_granian is False: + if limit_concurrency is not None: + uvicorn_args["limit_concurrency"] = limit_concurrency if max_requests_before_restart_jitter is not None: ProxyInitializationHelpers._apply_uvicorn_max_requests_jitter( uvicorn_args=uvicorn_args, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dcde9a27ec01..8936f6e9ca99 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1076,9 +1076,10 @@ async def _run_pw_migration(): # lazily by the flusher on first tick (see `_state_loaded` flag) so # hot-reloaded routers also get their persisted priors. if llm_router is not None and getattr(llm_router, "adaptive_routers", None): - for _ar in llm_router.adaptive_routers.values(): - await _ar.load_state_from_db(prisma_client) - _ar._state_loaded = True + for _tagged_routers in llm_router.adaptive_routers.values(): + for _tagged in _tagged_routers: + await _tagged.strategy.load_state_from_db(prisma_client) + _tagged.strategy._state_loaded = True asyncio.create_task(_adaptive_router_flusher_loop()) ## [Optional] Initialize dd tracer @@ -3248,16 +3249,18 @@ async def _adaptive_router_flusher_loop(): adaptive_routers = getattr(llm_router, "adaptive_routers", None) or {} if not adaptive_routers or prisma_client is None: continue - for ar in adaptive_routers.values(): - # Lazy state load: covers adaptive routers registered via - # `/config/reload` after proxy boot. - if not getattr(ar, "_state_loaded", False): - try: - await ar.load_state_from_db(prisma_client) - finally: - ar._state_loaded = True - await ar.queue.flush_state_to_db(prisma_client) - await ar.queue.flush_session_to_db(prisma_client) + for tagged_routers in adaptive_routers.values(): + for tagged in tagged_routers: + ar = tagged.strategy + # Lazy state load: covers adaptive routers registered via + # `/config/reload` after proxy boot. + if not getattr(ar, "_state_loaded", False): + try: + await ar.load_state_from_db(prisma_client) + finally: + ar._state_loaded = True + await ar.queue.flush_state_to_db(prisma_client) + await ar.queue.flush_session_to_db(prisma_client) except asyncio.CancelledError: raise except Exception: @@ -3705,22 +3708,22 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: litellm_config_cache.redis_cache = redis_cache -def resolve_complexity_router_plugins( - model_name: str, - complexity_router_config: dict, +def resolve_routing_plugins( + plugin_paths: list, config_file_path: str | None, -) -> None: + source_label: str, +) -> list: """ - Resolves `complexity_router_config["plugins"]` dotted-path strings to live - instances via `get_instance_fn` (the same convention `litellm_settings.callbacks` - uses), in place. Raises at config-load time if a path resolves to something that - doesn't implement `RoutingPlugin`, rather than deferring to a confusing - `AttributeError` on the first request that reaches the plugin pipeline. + Resolves a list of routing-plugin entries to live `RoutingPlugin` instances. + Each string entry is resolved through `get_instance_fn` (the same dotted-path + convention `litellm_settings.callbacks` uses, which resolves both local module + files next to the config and modules installed as Python packages); non-string + entries are assumed to already be instances and passed through. Raises at + config-load time if any entry resolves to something that doesn't implement + `RoutingPlugin`, rather than deferring to a confusing `AttributeError` on the + first request that reaches the plugin pipeline. `source_label` names the config + key being resolved so the error points the operator at the right place. """ - plugin_paths = complexity_router_config.get("plugins") - if not isinstance(plugin_paths, list): - return - resolved_plugins = [ get_instance_fn(value=plugin_path, config_file_path=config_file_path) if isinstance(plugin_path, str) @@ -3736,12 +3739,31 @@ def resolve_complexity_router_plugins( getattr(resolved_plugin, "run", None) ): raise ValueError( - f"complexity_router_config.plugins entry {plugin_path!r} on model {model_name!r} " - f"resolved to {resolved_plugin!r}, which does not implement the RoutingPlugin " - "interface (an async `run(context)` method). Fix the referenced module before " - "starting the proxy." + f"{source_label} entry {plugin_path!r} resolved to {resolved_plugin!r}, which does " + "not implement the RoutingPlugin interface (an async `run(context)` method). Fix the " + "referenced module before starting the proxy." ) - complexity_router_config["plugins"] = resolved_plugins + return resolved_plugins + + +def resolve_complexity_router_plugins( + model_name: str, + complexity_router_config: dict, + config_file_path: str | None, +) -> None: + """ + Resolves `complexity_router_config["plugins"]` dotted-path strings to live + instances in place, via `resolve_routing_plugins`. + """ + plugin_paths = complexity_router_config.get("plugins") + if not isinstance(plugin_paths, list): + return + + complexity_router_config["plugins"] = resolve_routing_plugins( + plugin_paths=plugin_paths, + config_file_path=config_file_path, + source_label=f"complexity_router_config.plugins on model {model_name!r}", + ) class ProxyConfig: @@ -4871,6 +4893,12 @@ async def load_config(self, router: Optional[litellm.Router], config_file_path: for k, v in router_settings.items(): if k in available_args: + if k == "plugins" and isinstance(v, list): + v = resolve_routing_plugins( + plugin_paths=v, + config_file_path=config_file_path, + source_label="router_settings.plugins", + ) router_params[k] = v elif k in {"health_check_interval", "health_check_concurrency"}: raise ValueError( @@ -7373,12 +7401,13 @@ async def async_data_generator( except (asyncio.CancelledError, GeneratorExit): # Client disconnected mid-stream. CancelledError / GeneratorExit are # BaseException, so they bypass the success/failure logging callbacks - # that normally release the pre-call max_parallel_requests +1; release - # it here. This is the outermost generator Starlette closes on + # that normally release the pre-call max_parallel_requests +1. Flag the + # disconnect; the shielded cleanup in `finally` owns the slot release + # so it can coordinate with disconnect-time success billing and release + # exactly once. This is the outermost generator Starlette closes on # disconnect, so it fires reliably regardless of needs_iterator_wrap # (a nested iterator hook would only see GeneratorExit on GC). if not stream_completed: - proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict) client_disconnected = True raise except Exception as e: @@ -7424,6 +7453,8 @@ async def async_data_generator( response=response, stream_completed=stream_completed, client_disconnected=client_disconnected, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) @@ -16010,7 +16041,11 @@ async def get_adaptive_router_state( status_code=404, detail={"error": "No adaptive_router is configured on this proxy."}, ) - snapshots = [await ar.get_state_snapshot() for ar in llm_router.adaptive_routers.values()] + snapshots = [ + await tagged.strategy.get_state_snapshot() + for tagged_routers in llm_router.adaptive_routers.values() + for tagged in tagged_routers + ] return {"routers": snapshots} diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index f7f6adaa8a22..27ffc49901b3 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -11,14 +11,16 @@ import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status -from fastapi.responses import ORJSONResponse +from fastapi.responses import ORJSONResponse, StreamingResponse import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import * from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -604,6 +606,7 @@ async def rag_query( general_settings, llm_router, proxy_config, + select_data_generator, version, ) @@ -673,6 +676,31 @@ async def rag_query( **request_data, ) + hidden_params = getattr(response, "_hidden_params", {}) or {} + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=hidden_params.get("litellm_call_id", None) or "", + model_id=hidden_params.get("model_id", None) or "", + cache_key=hidden_params.get("cache_key", None) or "", + api_base=hidden_params.get("api_base", None) or "", + version=version, + response_cost=hidden_params.get("response_cost", None), + request_data=request_data, + ) + + if isinstance(response, CustomStreamWrapper): + return StreamingResponse( + select_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + request=request, + ), + media_type="text/event-stream", + headers=custom_headers, + ) + + fastapi_response.headers.update(custom_headers) return response except HTTPException: diff --git a/litellm/proxy/types_utils/utils.py b/litellm/proxy/types_utils/utils.py index e2206541fc6c..8d7aedce4d0e 100644 --- a/litellm/proxy/types_utils/utils.py +++ b/litellm/proxy/types_utils/utils.py @@ -34,16 +34,12 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any: module_name = ".".join(parts[:-1]) instance_name = parts[-1] - # If config_file_path is provided, use it to determine the module spec and load the module + module_file_path = None if config_file_path is not None: directory = os.path.dirname(config_file_path) - module_file_path = os.path.join(directory, *module_name.split(".")) - module_file_path += ".py" - - # Check if the file exists before trying to load it - if not os.path.exists(module_file_path): - raise ImportError(f"Could not find module file {module_file_path}") + module_file_path = os.path.join(directory, *module_name.split(".")) + ".py" + if module_file_path is not None and os.path.exists(module_file_path): spec = importlib.util.spec_from_file_location(module_name, module_file_path) # type: ignore if spec is None: raise ImportError(f"Could not find a module specification for {module_file_path}") @@ -52,7 +48,6 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any: raise ImportError(f"Could not find a module loader for {module_file_path}") spec.loader.exec_module(module) # type: ignore else: - # Dynamically import the module module = importlib.import_module(module_name) # Get the instance from the module diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9f36e7293303..48164ce913a5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2583,35 +2583,30 @@ def _fire_deferred_stream_logging(request_data: dict) -> None: logging_obj._deferred_stream_complete_args = None asyncio.create_task(_deferred_cb(*_args)) - def _release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None: + async def _arelease_max_parallel_requests_on_disconnect( + self, + user_api_key_dict: UserAPIKeyAuth, + request_data: dict | None = None, + ) -> None: """ Release the api-key max_parallel_requests slot when a streaming - response is cancelled mid-flight (client disconnect). Neither the - success nor failure logging callback fires on the resulting - CancelledError / GeneratorExit, so the pre-call +1 would otherwise - leak. - - Must be called from the outermost streaming generator (the one - Starlette drives and closes on disconnect). A nested iterator-hook - generator only receives GeneratorExit when it is garbage collected, - which is non-deterministic, so the refund cannot live there. - - Scheduled fire-and-forget (no await) because awaiting is not - permitted while unwinding a GeneratorExit. + response is cancelled mid-flight (client disconnect) and no logging + callback fired for it. Neither the success nor failure callback runs on + the resulting CancelledError / GeneratorExit, so the pre-call +1 would + otherwise leak. + + Awaited from the shielded streaming cleanup rather than scheduled + fire-and-forget, so the caller can make it the single owner of the + release: when a disconnect-time success event does fire (partial-spend + billing or a deferred-guardrail flush), that event's own limiter + callback releases the slot and this is not called at all. Two + concurrent releases of the same acquisition would otherwise race and + double-decrement under the limiter's in-memory fallback. """ limiter = self.get_proxy_hook("parallel_request_limiter") if not isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3): return - try: - asyncio.create_task(limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)) - except RuntimeError: - # No running event loop (e.g. interpreter/loop shutdown); the - # counter's window TTL will reclaim the slot. - verbose_proxy_logger.warning( - "parallel_request_limiter_v3: could not schedule " - "max_parallel_requests release on disconnect; no running " - "event loop. Slot will be reclaimed when its window TTL expires" - ) + await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict, request_data) def _init_response_taking_too_long_task(self, data: Optional[dict] = None): """ diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 6b5f087f902c..29891ccfd247 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -11,12 +11,14 @@ import asyncio import contextvars +from contextlib import contextmanager from functools import partial from typing import ( TYPE_CHECKING, Any, Coroutine, Dict, + Iterator, List, Optional, Tuple, @@ -27,6 +29,9 @@ import httpx import litellm +from litellm._internal_context import is_internal_call +from litellm.cost_calculator import vector_store_search_cost +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion @@ -188,6 +193,25 @@ async def aingest( ) +@contextmanager +def _suppressed_sub_call_billing() -> Iterator[None]: + """ + Suppress a sub-call's own billing event so the parent aquery event bills it. + + Every suppressed sub-call's cost must be folded into the parent event: + into the response's hidden response_cost on the non-streaming path, or via + the logging object's additional_response_cost on the streaming path (the + streamed cost is computed from assembled chunks after this pipeline + returns, so there is no response object to fold into here). + """ + previous = is_internal_call.get() + is_internal_call.set(True) + try: + yield + finally: + is_internal_call.set(previous) + + async def _execute_query_pipeline( model: str, messages: List[Any], @@ -209,27 +233,46 @@ async def _execute_query_pipeline( raise ValueError("No query found in messages for RAG query") # 2. Search vector store - search_response = await litellm.vector_stores.asearch( - vector_store_id=retrieval_config["vector_store_id"], - query=query_text, - max_num_results=retrieval_config.get("top_k", 10), - custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"), - **kwargs, - ) + with _suppressed_sub_call_billing(): + search_response = await litellm.vector_stores.asearch( + vector_store_id=retrieval_config["vector_store_id"], + query=query_text, + max_num_results=retrieval_config.get("top_k", 10), + custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"), + **kwargs, + ) + + search_provider = retrieval_config.get("custom_llm_provider", "openai") + try: + search_cost = sum( + vector_store_search_cost( + model=search_provider if "/" in search_provider else None, + custom_llm_provider=search_provider, + response=search_response, + ) + ) + except Exception: # noqa: BLE001 - cost accounting must never break the query path + search_cost = 0.0 rerank_response = None + rerank_cost = 0.0 context_chunks = search_response.get("data", []) # 3. Optional rerank if rerank and rerank.get("enabled"): documents = RAGQuery.extract_documents_from_search(search_response) if documents: - rerank_response = await litellm.arerank( - model=rerank["model"], - query=query_text, - documents=documents, - top_n=rerank.get("top_n", 5), - ) + with _suppressed_sub_call_billing(): + rerank_response = await litellm.arerank( + model=rerank["model"], + query=query_text, + documents=documents, + top_n=rerank.get("top_n", 5), + ) + rerank_hidden_params = getattr(rerank_response, "_hidden_params", None) + if isinstance(rerank_hidden_params, dict): + rerank_response_cost: float | None = rerank_hidden_params.get("response_cost") + rerank_cost = rerank_response_cost or 0.0 context_chunks = RAGQuery.get_top_chunks_from_rerank(search_response, rerank_response) # 4. Build context message and call completion @@ -237,28 +280,40 @@ async def _execute_query_pipeline( modified_messages = messages[:-1] + [context_message] + [messages[-1]] # Use router if available to properly resolve virtual model names - if router is not None: - response = await router.acompletion( - model=model, - messages=modified_messages, - stream=stream, - **kwargs, - ) - else: - response = await litellm.acompletion( - model=model, - messages=modified_messages, - stream=stream, - **kwargs, - ) + with _suppressed_sub_call_billing(): + if router is not None: + response = await router.acompletion( + model=model, + messages=modified_messages, + stream=stream, + **kwargs, + ) + else: + response = await litellm.acompletion( + model=model, + messages=modified_messages, + stream=stream, + **kwargs, + ) # 5. Attach search results to response + sub_call_cost = search_cost + rerank_cost if not stream and isinstance(response, ModelResponse): response = RAGQuery.add_search_results_to_response( response=response, search_results=search_response, rerank_results=rerank_response, ) + if sub_call_cost > 0: + hidden_params = getattr(response, "_hidden_params", None) + if isinstance(hidden_params, dict): + completion_response_cost: float | None = hidden_params.get("response_cost") + if completion_response_cost is not None: + hidden_params["response_cost"] = completion_response_cost + sub_call_cost + elif sub_call_cost > 0: + logging_obj: object = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.model_call_details["additional_response_cost"] = sub_call_cost return response # type: ignore[return-value] diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index e03f0296109f..392bb7bcab27 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -260,7 +260,7 @@ async def _get_mcp_tools_from_manager( # names), so use None and let the auth object's mcp_servers do the filtering. effective_server_filter = None if resolved_toolset_ids else (resolved_mcp_servers or None) - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=effective_server_filter, @@ -270,6 +270,7 @@ async def _get_mcp_tools_from_manager( litellm_trace_id=litellm_trace_id, request_tags=request_tags, ) + tools = listing.tools allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) allowed_mcp_servers = global_mcp_server_manager.get_mcp_servers_from_ids( # type: ignore[attr-defined] diff --git a/litellm/router.py b/litellm/router.py index 5f2d347e9533..330733fd0dbc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -34,6 +34,7 @@ Optional, Set, Tuple, + TypeVar, Union, cast, ) @@ -87,7 +88,11 @@ from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.router_strategy.simple_shuffle import simple_shuffle -from litellm.router_strategy.tag_based_routing import get_deployments_for_tag +from litellm.router_strategy.tag_based_routing import ( + _get_tags_from_request_kwargs, + get_deployments_for_tag, + is_valid_deployment_tag, +) from litellm.router_utils.add_retry_fallback_headers import ( _HiddenParamsHost, add_fallback_headers_to_response, @@ -176,6 +181,7 @@ MockRouterTestingParams, ModelGroupInfo, OptionalPreCallChecks, + PreRoutingStrategy, RetryPolicy, RouterCacheEnum, RouterGeneralSettings, @@ -187,6 +193,7 @@ RoutingPlugin, RoutingStrategy, SearchToolTypedDict, + TaggedPreRoutingStrategy, ) from litellm.types.services import ServiceTypes from litellm.types.utils import ( @@ -261,6 +268,9 @@ def _cost_value_as_float(value: Union[str, int, float, None]) -> float | None: return None +_PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") + + class RoutingArgs(enum.Enum): ttl = 60 # 1min (RPM/TPM expire key) @@ -752,10 +762,10 @@ def __init__( self.provider_default_deployment_ids: List[str] = [] self.pattern_router = PatternMatchRouter() self.team_pattern_routers: Dict[str, PatternMatchRouter] = {} # {"TEAM_ID": PatternMatchRouter} - self.auto_routers: Dict[str, "AutoRouter"] = {} - self.complexity_routers: Dict[str, "ComplexityRouter"] = {} - self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {} - self.quality_routers: Dict[str, "QualityRouter"] = {} + self.auto_routers: dict[str, list[TaggedPreRoutingStrategy["AutoRouter"]]] = {} + self.complexity_routers: dict[str, list[TaggedPreRoutingStrategy["ComplexityRouter"]]] = {} + self.adaptive_routers: dict[str, list[TaggedPreRoutingStrategy["AdaptiveRouter"]]] = {} + self.quality_routers: dict[str, list[TaggedPreRoutingStrategy["QualityRouter"]]] = {} self.routing_plugins: list[RoutingPlugin] = list(plugins) if plugins else [] # Initialize model_group_alias early since it's used in set_model_list @@ -2302,6 +2312,9 @@ def __init__(self, async_generator: AsyncGenerator): logging_obj=model_response.logging_obj, ) self._async_generator = async_generator + inner_chunks: object = getattr(model_response, "chunks", None) + if isinstance(inner_chunks, list): + self.chunks = inner_chunks # Preserve hidden params (including litellm_overhead_time_ms) from original response if hasattr(model_response, "_hidden_params"): self._hidden_params = model_response._hidden_params.copy() @@ -7938,6 +7951,11 @@ def _is_auto_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: return True return False + @staticmethod + def _deployment_tags(deployment: Deployment) -> tuple[str, ...]: + """Deployment tags used to disambiguate strategy registries keyed by model_name.""" + return tuple(deployment.litellm_params.tags or ()) + def init_auto_router_deployment(self, deployment: Deployment): """ Initialize the auto-router deployment. @@ -7973,11 +7991,12 @@ def init_auto_router_deployment(self, deployment: Deployment): embedding_model=embedding_model, litellm_router_instance=self, ) - if deployment.model_name in self.auto_routers: - raise ValueError( - f"Auto-router deployment {deployment.model_name} already exists. Please use a different model name." - ) - self.auto_routers[deployment.model_name] = autor_router + self._register_pre_routing_strategy( + registry=self.auto_routers, + deployment=deployment, + strategy=autor_router, + strategy_label="Auto-router", + ) def _is_complexity_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """ @@ -8028,20 +8047,54 @@ def init_complexity_router_deployment(self, deployment: Deployment): litellm_router_instance=self, complexity_router_config=complexity_router_config, ) - if deployment.model_name in self.complexity_routers: - raise ValueError( - f"Complexity-router deployment {deployment.model_name} already exists. Please use a different model name." - ) - self.complexity_routers[deployment.model_name] = complexity_router + self._register_pre_routing_strategy( + registry=self.complexity_routers, + deployment=deployment, + strategy=complexity_router, + strategy_label="Complexity-router", + ) def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" return litellm_params.model.startswith("auto_router/adaptive_router") + @staticmethod + def _has_registered_strategy( + registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], + model_name: str, + tags: tuple[str, ...], + ) -> bool: + """True when a strategy for this (model_name, tags) pair is already registered.""" + return any(existing.tags == tags for existing in registry.get(model_name, [])) + + def _register_pre_routing_strategy( + self, + registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], + deployment: Deployment, + strategy: _PreRoutingStrategyT, + strategy_label: str, + ) -> None: + """ + Register `strategy` under `deployment.model_name`, scoped by its tags. + Reusing a `model_name` is allowed when tags differ; a repeat of the same + (model_name, tags) pair is a misconfiguration and is rejected. + """ + tags = self._deployment_tags(deployment) + if self._has_registered_strategy(registry, deployment.model_name, tags): + raise ValueError( + f"{strategy_label} deployment {deployment.model_name} with tags {list(tags)} already exists. " + "Please use a different model name or set different tags." + ) + registry[deployment.model_name] = [ + *registry.get(deployment.model_name, []), + TaggedPreRoutingStrategy(tags=tags, strategy=strategy), + ] + def _finalize_adaptive_router_if_configured(self) -> None: """Locate every adaptive-router deployment in the finalized model_list and build an AdaptiveRouter for each. Safe no-op when none are configured. - Idempotent: skips any deployment whose model_name is already initialized.""" + Idempotent: skips any deployment whose (model_name, tags) pair is already + initialized, so hot-reloads don't rebuild routers that would lose state.""" # Drop any adaptive-router hooks left over from a previous Router # instance (e.g. after `/config/reload` replaced `llm_router`). Without # this, stale AdaptiveRouterPostCallHook callbacks from the old Router @@ -8064,23 +8117,31 @@ def _finalize_adaptive_router_if_configured(self) -> None: litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)), model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info), ) - if model_name in self.adaptive_routers: + if self._has_registered_strategy(self.adaptive_routers, model_name, self._deployment_tags(deployment)): continue self.init_adaptive_router_deployment(deployment=deployment) - for model_name, complexity_router in self.complexity_routers.items(): - if not complexity_router.config.adaptive or model_name in self.adaptive_routers: - continue - adaptive_router = complexity_router._ensure_adaptive_router() - if adaptive_router is not None: - self.adaptive_routers[model_name] = adaptive_router + for model_name, tagged_complexity_routers in self.complexity_routers.items(): + for tagged in tagged_complexity_routers: + complexity_router = tagged.strategy + if not complexity_router.config.adaptive: + continue + if self._has_registered_strategy(self.adaptive_routers, model_name, tagged.tags): + continue + adaptive_router = complexity_router._ensure_adaptive_router() + if adaptive_router is not None: + self.adaptive_routers[model_name] = [ + *self.adaptive_routers.get(model_name, []), + TaggedPreRoutingStrategy(tags=tagged.tags, strategy=adaptive_router), + ] for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): litellm.logging_callback_manager.remove_callback_from_all_lists(callback) - for adaptive_router in self.adaptive_routers.values(): - litellm.logging_callback_manager.add_litellm_callback( - AdaptiveRouterPostCallHook(adaptive_router=adaptive_router) - ) + for tagged_adaptive_routers in self.adaptive_routers.values(): + for tagged in tagged_adaptive_routers: + litellm.logging_callback_manager.add_litellm_callback( + AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) + ) def init_adaptive_router_deployment(self, deployment: Deployment) -> None: """ @@ -8133,18 +8194,18 @@ def init_adaptive_router_deployment(self, deployment: Deployment) -> None: if cost is not None: model_to_cost[name] = float(cost) - if deployment.model_name in self.adaptive_routers: - raise ValueError( - f"Adaptive-router deployment {deployment.model_name} already exists. Please use a different model name." - ) - adaptive_router = AdaptiveRouter( router_name=deployment.model_name, config=config, model_to_prefs=model_to_prefs, model_to_cost=model_to_cost, ) - self.adaptive_routers[deployment.model_name] = adaptive_router + self._register_pre_routing_strategy( + registry=self.adaptive_routers, + deployment=deployment, + strategy=adaptive_router, + strategy_label="Adaptive-router", + ) litellm.logging_callback_manager.add_litellm_callback( AdaptiveRouterPostCallHook(adaptive_router=adaptive_router) ) @@ -8196,11 +8257,12 @@ def init_quality_router_deployment(self, deployment: Deployment): litellm_router_instance=self, quality_router_config=quality_router_config, ) - if deployment.model_name in self.quality_routers: - raise ValueError( - f"Quality-router deployment {deployment.model_name} already exists. Please use a different model name." - ) - self.quality_routers[deployment.model_name] = quality_router + self._register_pre_routing_strategy( + registry=self.quality_routers, + deployment=deployment, + strategy=quality_router, + strategy_label="Quality-router", + ) def deployment_is_active_for_environment(self, deployment: Deployment) -> bool: """ @@ -11193,6 +11255,35 @@ def _filter_by_routing_plugin_candidates( return filtered + def _select_pre_routing_strategy(self, model: str, request_kwargs: Dict) -> "PreRoutingStrategy | None": + """ + Resolve the pre-routing strategy for `model`, disambiguating deployments + that share a `model_name` by matching the request's tags against each + registered strategy's tags before falling back to the first registered. + """ + candidates: list[TaggedPreRoutingStrategy[PreRoutingStrategy]] = [ + *self.auto_routers.get(model, []), + *self.complexity_routers.get(model, []), + *self.adaptive_routers.get(model, []), + *self.quality_routers.get(model, []), + ] + if not candidates: + return None + if len(candidates) == 1: + return candidates[0].strategy + + request_tags = _get_tags_from_request_kwargs(request_kwargs) + if request_tags: + for tagged in candidates: + if tagged.tags and is_valid_deployment_tag( + list(tagged.tags), request_tags, self.tag_filtering_match_any + ): + return tagged.strategy + for tagged in candidates: + if "default" in tagged.tags: + return tagged.strategy + return candidates[0].strategy + async def async_pre_routing_hook( self, model: str, @@ -11215,12 +11306,7 @@ async def async_pre_routing_hook( if self.routing_plugins: await self._run_routing_plugins(model=model, request_kwargs=request_kwargs, messages=messages) - router_strategy = ( - self.auto_routers.get(model) - or self.complexity_routers.get(model) - or self.adaptive_routers.get(model) - or self.quality_routers.get(model) - ) + router_strategy = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs) if router_strategy is None: return None diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index fa6f14e9b26d..695d8b8aeaae 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -28,6 +28,7 @@ from .config import ( DEFAULT_CODE_KEYWORDS, + DEFAULT_ESCALATION_KEYWORDS, DEFAULT_REASONING_KEYWORDS, DEFAULT_SIMPLE_KEYWORDS, DEFAULT_TECHNICAL_KEYWORDS, @@ -173,6 +174,11 @@ def __init__( self.config.custom_technical_keywords, ) self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS + self.escalation_keywords = ( + self.config.escalation_keywords + if self.config.escalation_keywords is not None + else DEFAULT_ESCALATION_KEYWORDS + ) # Lazily built on first semantic request and cached for reuse (route # embeddings are static, only the prompt is embedded per request). The lock @@ -668,6 +674,53 @@ def _soft_floor_pick( } return best_model + def _escalation_triggered(self, user_message: str) -> bool: + """Whether the prompt asks to escalate to a stronger model. + + Matching is a case-sensitive substring test so the default "LITELLM ESCALATE" + only fires on the deliberate, shouted form and not on incidental lowercase + mentions of the word (e.g. "how do I escalate this ticket"). + """ + if not self.escalation_keywords: + return False + return any(keyword in user_message for keyword in self.escalation_keywords) + + def _tier_for_model(self, model: str) -> ComplexityTier | None: + """Return the most-severe configured tier whose pool contains this model.""" + pools = self._tier_pools() + matched = tuple(ComplexityTier(tier_name) for tier_name, models in pools.items() if model in models) + if not matched: + return None + return max(matched, key=TIER_SEVERITY_ORDER.index) + + def _escalate_tier(self, tier: ComplexityTier) -> ComplexityTier: + """Bump a tier one step up to the next-higher configured tier. + + Returns the input tier unchanged when it is already the highest configured + tier, so escalation can never route below the model the user would otherwise + have received. + """ + configured = frozenset(self.config.tiers) + current_index = TIER_SEVERITY_ORDER.index(tier) + higher_tiers = tuple( + candidate for candidate in TIER_SEVERITY_ORDER[current_index + 1 :] if candidate.value in configured + ) + return higher_tiers[0] if higher_tiers else tier + + def _escalated_pin(self, pinned_model: str) -> str | None: + """Bump a session's pinned model to the next-higher configured tier. + + Returns None when the pin no longer maps to any configured tier, signalling + a full reclassification instead. + """ + pinned_tier = self._tier_for_model(pinned_model) + if pinned_tier is None: + return None + escalated_tier = self._escalate_tier(pinned_tier) + if escalated_tier == pinned_tier: + return pinned_model + return self.get_model_for_tier(escalated_tier) + def _lexical_tier_override(self, user_message: str) -> ComplexityTier | None: """When keyword_tier_rules match literally, the most-severe matched tier wins. @@ -910,29 +963,41 @@ async def async_pre_routing_hook( if cache_key is not None: pinned_model = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) if isinstance(pinned_model, str): - # Refresh the TTL on every hit so an active session doesn't lose its - # pin mid-conversation just because it outlives the original write. - await self.litellm_router_instance.cache.async_set_cache( - key=cache_key, - value=pinned_model, - ttl=self.config.session_affinity_ttl_seconds, - ) - if self.config.adaptive: - from litellm.router_strategy.adaptive_router.config import ( - ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + routed_model: str | None = pinned_model + if self.escalation_keywords: + resolved_messages = self._resolve_messages(messages, request_kwargs) + user_message = ( + self._extract_user_message_and_system_prompt(resolved_messages)[0] + if resolved_messages + else None + ) + if user_message is not None and self._escalation_triggered(user_message): + routed_model = self._escalated_pin(pinned_model) + if routed_model is not None: + # Refresh the TTL on every hit so an active session doesn't lose its + # pin mid-conversation just because it outlives the original write. + await self.litellm_router_instance.cache.async_set_cache( + key=cache_key, + value=routed_model, + ttl=self.config.session_affinity_ttl_seconds, + ) + if self.config.adaptive: + from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ) + + kwargs_metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(kwargs_metadata, dict): + kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = routed_model + cause = "session_affinity_escalation" if routed_model != pinned_model else "session_affinity_pin" + verbose_router_logger.info( + f"ComplexityRouter: routing decision cause={cause}, routed_model={routed_model}" + ) + has_original_messages = messages is not None and len(messages) > 0 + return PreRoutingHookResponse( + model=routed_model, + messages=messages if has_original_messages else None, ) - - kwargs_metadata = request_kwargs.setdefault("metadata", {}) - if isinstance(kwargs_metadata, dict): - kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = pinned_model - verbose_router_logger.info( - f"ComplexityRouter: routing decision cause=session_affinity_pin, routed_model={pinned_model}" - ) - has_original_messages = messages is not None and len(messages) > 0 - return PreRoutingHookResponse( - model=pinned_model, - messages=messages if has_original_messages else None, - ) response = await self._classify_and_route( model=model, @@ -1004,13 +1069,17 @@ async def _classify_and_route( messages=messages if has_original_messages else None, ) + escalate = self._escalation_triggered(user_message) + override_tier = await self._resolve_keyword_tier_override(user_message, request_kwargs) if override_tier is not None: - routed_model = await self._pick_model_for_tier(override_tier, messages, resolved_messages, request_kwargs) - cause = "semantic_keyword_match" if self.config.semantic_keyword_matching else "literal_keyword_match" + routed_tier = self._escalate_tier(override_tier) if escalate else override_tier + routed_model = await self._pick_model_for_tier(routed_tier, messages, resolved_messages, request_kwargs) + base_cause = "semantic_keyword_match" if self.config.semantic_keyword_matching else "literal_keyword_match" + cause = f"{base_cause}+escalation" if escalate else base_cause verbose_router_logger.info( f"ComplexityRouter: routing decision cause={cause}, " - f"tier={override_tier.value}, routed_model={routed_model}" + f"tier={routed_tier.value}, routed_model={routed_model}" ) return PreRoutingHookResponse( model=routed_model, @@ -1018,6 +1087,9 @@ async def _classify_and_route( ) tier, score, signals = await self.aclassify(user_message, system_prompt, request_kwargs) + if escalate: + tier = self._escalate_tier(tier) + signals = [*signals, "escalation"] if self.config.adaptive: routed_model = self._soft_floor_pick(tier, user_message, request_kwargs) adaptive = self._ensure_adaptive_router() diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 1f9847989700..17c2c287dde2 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -162,6 +162,9 @@ def _normalize_keywords(self) -> "KeywordTierRule": # Note: "async", "kubernetes", "docker" are in DEFAULT_CODE_KEYWORDS ] +DEFAULT_ESCALATION_KEYWORDS: list[str] = ["LITELLM ESCALATE"] + + DEFAULT_SIMPLE_KEYWORDS: list[str] = [ "what is", "what's", @@ -339,6 +342,16 @@ class ComplexityRouterConfig(BaseModel): ), ) + escalation_keywords: list[str] | None = Field( + default=None, + description=( + "Case-sensitive phrases a user can include to force a bump to the next-higher " + "complexity tier when they aren't satisfied with results (they can force a stronger " + "model, but not choose which one). Defaults to ['LITELLM ESCALATE'] when unset; " + "set to an empty list to disable." + ), + ) + # Deterministic keyword -> tier overrides, evaluated before weighted scoring keyword_tier_rules: list[KeywordTierRule] | None = Field( default=None, @@ -400,6 +413,13 @@ def _coerce_tier_values(cls, value: object) -> object: coerced[key] = item return coerced + @field_validator("escalation_keywords") + @classmethod + def _normalize_escalation_keywords(cls, value: list[str] | None) -> list[str] | None: + if value is None: + return None + return [stripped for keyword in value if (stripped := keyword.strip())] + @model_validator(mode="after") def _validate_llm_classifier_config(self) -> "ComplexityRouterConfig": if self.classifier_type == "llm" and self.classifier_llm_config is None: diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 3dda4e3990cd..86e69467dbf6 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -53,6 +53,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import ( CiscoAIDefenseGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( + SingulrGuardrailConfigModel, +) from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( HeadroomGuardrailConfigModel, ) @@ -125,6 +128,7 @@ class SupportedGuardrailIntegrations(Enum): RUBRIK = "rubrik" VIGIL_GUARD = "vigil_guard" REPELLOAI = "repelloai" + SINGULR = "singulr" HEADROOM = "headroom" COMPRESR = "compresr" @@ -932,6 +936,7 @@ class LitellmParams( HiddenlayerGuardrailConfigModel, QostodianNexusConfigModel, VigilGuardGuardrailConfigModel, + SingulrGuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") mode: Union[str, List[str], Mode] = Field( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index daac1e4506f4..9f689a2dd318 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -529,6 +529,7 @@ class ChatCompletionDeltaToolCallChunk(TypedDict, total=False): class ChatCompletionCachedContent(TypedDict): type: Literal["ephemeral"] + ttl: NotRequired[Literal["5m", "1h"]] class ChatCompletionThinkingBlock(TypedDict, total=False): diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 3524a7eb7f72..098e99fe1984 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -11,6 +11,14 @@ # exact byte/string body, such as AWS SigV4-signed requests. LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY = "litellm_pass_through_raw_body" +# Attribute set on the FastAPI endpoint function of every user-defined pass-through +# route. Auth reads it off the dispatched endpoint (``request.scope["endpoint"]``) to +# decide whether a request body ``model`` names an upstream model rather than a +# LiteLLM-managed one. Keying off the resolved endpoint (not the request path) means a +# custom path that collides with a built-in route never suppresses model-access checks: +# on a collision FastAPI dispatches the built-in handler, which does not carry this flag. +LITELLM_PASS_THROUGH_ENDPOINT_MARKER = "__litellm_pass_through_endpoint__" + class EndpointType(str, Enum): VERTEX_AI = "vertex-ai" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/singulr.py b/litellm/types/proxy/guardrails/guardrail_hooks/singulr.py new file mode 100644 index 000000000000..62d3b8653ef1 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/singulr.py @@ -0,0 +1,63 @@ +from typing import Any, Optional + +from pydantic import BaseModel, Field + +from .base import GuardrailConfigModel + + +class SingulrGuardrailRequest(BaseModel): + model: Optional[str] = None + messages: Optional[list[dict[str, Any]]] = None + tools: Optional[list[dict[str, Any]]] = None + model_response: Optional[dict[str, Any]] = None + litellm_metadata: Optional[dict[str, Any]] = None + + +class SingulrGuardrailPayload(BaseModel): + litellm_call_id: Optional[str] = None + request_data: Optional[SingulrGuardrailRequest] = None + input_type: str + is_playground_request: Optional[bool] = None + playground_text: Optional[str] = None + + +class SingulrGuardrailResponse(BaseModel): + """Response returned by the Singulr guardrail API.""" + + should_block: bool = False + blocking_due_to: Optional[str] = None + + +class SingulrGuardrailConfigModel(GuardrailConfigModel): + singulr_api_key: Optional[str] = Field( + default=None, + description="The Singulr API key. Generate API key from Singulr Platform.", + ) + + singulr_api_base: Optional[str] = Field( + default=None, + description="The Singulr API base URL. Get base URL from Singulr Platform.", + ) + + singulr_application_id: Optional[str] = Field( + default=None, + description="The Singulr application ID. Get application ID from Singulr Platform.", + ) + + singulr_guardrail_id: Optional[str] = Field( + default=None, + description="The Singulr Guardrail ID. Get guardrail ID from Singulr Platform.", + ) + + block_on_error: Optional[bool] = Field( + default=None, + description=( + "Whether to block requests when the Singulr Guardrails API is unavailable " + "or returns an error. If enabled, requests fail closed. " + "If disabled, requests continue without guardrail enforcement (fail open)." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Singulr" diff --git a/litellm/types/router.py b/litellm/types/router.py index 69a8ca9f19e2..28e4a8272e87 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -5,7 +5,18 @@ import datetime import enum from dataclasses import dataclass -from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hints +from typing import ( + Any, + Dict, + Generic, + List, + Literal, + Optional, + Tuple, + TypeVar, + Union, + get_type_hints, +) import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -830,6 +841,31 @@ class PreRoutingHookResponse(BaseModel): messages: Optional[List[Dict[str, Any]]] +_PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True) + + +@dataclass(frozen=True, slots=True) +class TaggedPreRoutingStrategy(Generic[_PreRoutingStrategyT_co]): + """A pre-routing strategy paired with the deployment `tags` it was registered under.""" + + tags: tuple[str, ...] + strategy: _PreRoutingStrategyT_co + + +@runtime_checkable +class PreRoutingStrategy(Protocol): + """Structural interface shared by the auto / complexity / adaptive / quality routers.""" + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: dict[str, Any], + messages: list[dict[str, Any]] | None = None, + input: "str | list[Any] | None" = None, + specific_deployment: bool | None = False, + ) -> "PreRoutingHookResponse | None": ... + + class RoutingContext(BaseModel): """ Passed through a Router's `plugins` pipeline before the routing decision is made. diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 04f1ff68c5d2..ec8a9336ca73 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -266,6 +266,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): "audio_transcription", "responses", "ocr", + "realtime", ] ] tpm: Optional[int] @@ -402,6 +403,11 @@ class CallTypes(str, Enum): vector_store_search = "vector_store_search" avector_store_search = "avector_store_search" + ingest = "ingest" + aingest = "aingest" + query = "query" + aquery = "aquery" + ######################################################### # Container Call Types ######################################################### diff --git a/litellm/utils.py b/litellm/utils.py index e19d2b36a52a..174bed09396f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3198,6 +3198,12 @@ def _check_valid_arg(supported_params: Optional[list]): non_default_params=non_default_params, optional_params={}, kwargs=kwargs ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "gemini": + # OpenAI SDKs (and litellm's own client) send encoding_format="float" + # by default; float lists are exactly what the vertex API returns, so + # the param is a no-op — don't reject the provider default. Other + # values (e.g. "base64") stay on the unsupported-param path below. + if non_default_params.get("encoding_format") == "float": + non_default_params.pop("encoding_format") supported_params = get_supported_openai_params( model=model, custom_llm_provider="vertex_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ffbc0dcd098e..b1a87c444c89 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3451,7 +3451,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2.2e-05, "output_cost_per_token": 2.64e-06, "supports_audio_input": true, @@ -3470,7 +3470,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 0.00022, "output_cost_per_token": 2.2e-05, "supports_audio_input": true, @@ -3489,7 +3489,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2.2e-05, "supported_modalities": [ @@ -4687,7 +4687,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supports_audio_input": true, @@ -4707,7 +4707,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -4739,7 +4739,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -4771,7 +4771,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supported_endpoints": [ @@ -4832,7 +4832,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 0.0002, "output_cost_per_token": 2e-05, "supports_audio_input": true, @@ -4850,7 +4850,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, "supported_modalities": [ @@ -7922,7 +7922,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2.2e-05, "output_cost_per_token": 2.64e-06, "supports_audio_input": true, @@ -7941,7 +7941,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 0.00022, "output_cost_per_token": 2.2e-05, "supports_audio_input": true, @@ -7960,7 +7960,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2.2e-05, "supported_modalities": [ @@ -22169,7 +22169,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supports_audio_input": true, @@ -22188,7 +22188,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supports_audio_input": true, @@ -22282,7 +22282,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, "supports_audio_input": true, @@ -22300,7 +22300,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, "supports_audio_input": true, @@ -22318,7 +22318,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, "supports_audio_input": true, @@ -24513,7 +24513,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -24545,7 +24545,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -24577,7 +24577,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -24610,7 +24610,7 @@ "max_input_tokens": 128000, "max_output_tokens": 32000, "max_tokens": 32000, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 2.4e-05, "regional_processing_uplift_multiplier_eu": 1.1, @@ -24645,7 +24645,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "regional_processing_uplift_multiplier_eu": 1.1, @@ -24678,7 +24678,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supported_endpoints": [ @@ -24710,7 +24710,7 @@ "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, "output_cost_per_token": 1.6e-05, "supported_endpoints": [ @@ -43694,7 +43694,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supported_endpoints": [ @@ -43727,7 +43727,7 @@ "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, - "mode": "chat", + "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, "supported_endpoints": [ diff --git a/router_plugins.json b/router_plugins.json new file mode 100644 index 000000000000..ffcddf89fd16 --- /dev/null +++ b/router_plugins.json @@ -0,0 +1,28 @@ +[ + { + "name": "TEMPLATE: copy this block for a new plugin, then delete this entry", + "description": "One line on what the plugin does and the routing signal it publishes.", + "author": "Plugin author's name.", + "repo": "https://github.com// (public source repository).", + "commit": "Full 40-char git SHA to pin when the plugin is not yet on PyPI; omit once 'pypi' is set.", + "version": "Plugin release version, e.g. 1.0.0.", + "pypi": "PyPI spec pinned to a version, e.g. my-plugin==1.0.0, or null if unpublished.", + "litellm_version": "Minimum compatible litellm version, e.g. >=1.94.0.", + "entrypoint": "Dotted import path to the plugin instance, e.g. my_plugin.plugin.instance.", + "license": "SPDX license id, e.g. MIT.", + "tags": ["searchable", "keywords"] + }, + { + "name": "language-detector", + "description": "Detects the user's language and publishes a routing signal.", + "author": "Jean Nuñez", + "repo": "https://github.com/jeann2013/language-detector", + "commit": "9e712819269173fc25a16f59ca3e9890f7864ac1", + "version": "1.0.0", + "pypi": null, + "litellm_version": ">=1.94.0", + "entrypoint": "litellm_plugin_language_detector.plugin.language_detector_plugin", + "license": "MIT", + "tags": ["language", "classification", "routing"] + } +] diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0e1eafb51965..67e9f4f78a7d 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -13,6 +13,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `realtime/` - realtime websocket sessions, including the pipecat audio path - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials; also the dashboard UI behavior on top of them, driven through the proxy-served UI at /ui with playwright (optional dep behind importorskip) +- `mcp/` - the MCP server surface over api_key auth: an admin registers an upstream MCP server through the management API and grants keys access via `object_permission.mcp_servers`, then the suite asserts tool listing and calling honor that permission (a key without the grant sees none of the server's tools and is refused a `tools/call` with a 403) - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) @@ -51,7 +52,7 @@ The shape is layered so tests stay declarative Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `Gateway` and adds suite-specific routes. Cleanup runs through that same `Gateway`, so whatever keys or customers your test creates get torn down by the `resources` fixture -Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The skip-vs-fail split is deliberate: a test marked `e2e` skips when no proxy answers its liveness probe, but once a request reaches the proxy any wrong behavior is a hard failure, never a skip +Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure coverage of the harness itself carries no marker and runs regardless. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache @@ -131,7 +132,7 @@ Quota Management - behavior features (entity- or config-driven caps and their ac quota_management... behavior : ratelimit | budget | spend_tracking variant : rpm | tpm | priority_generous | priority_strict - key | internal_user | end_user | organization | team_member | tag + key | internal_user | end_user | organization | team | team_member | tag | model_max | soft | key_multi_window | team_multi_window | fallback | spend_counter chat_completions | stream | embeddings | cache_hit | key_rollup @@ -139,7 +140,8 @@ quota_management... | spend_calculate | pagination assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking - | isolates_per_model | routes_to_fallback | reseed_matches_db | logs_cost | zero_cost + | isolates_per_model | isolates_per_member | enforced_across_keys | routes_to_fallback + | reseed_matches_db | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows | writes_failure_row | returns_cost | keeps_total e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] @@ -173,7 +175,7 @@ other... ``` ## Hard Rules -- no monkeypatching, mock tests or unit tests of any kind. if a contributor asks you to write an end to end test, do NOT stage a unit test with it. if you find a product gap, call it out in the PR description +- no monkeypatching or mock tests, and never substitute a unit test for e2e feature coverage: a product feature is proven end to end against a live proxy, not with a unit test. if a contributor asks you to write an end to end test, do NOT stage a unit test of the feature with it; if you find a product gap, call it out in the PR description. tests that cover the harness itself are the exception and are allowed (for example `coverage_registry/test_collector.py`, which unit-tests the coverage collector): they carry no `e2e` marker, exercise harness plumbing rather than a product feature, and run whether or not a proxy is up - use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want. diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 2082f2c9de42..555ac0482e2b 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -54,7 +54,7 @@ The suites run against a live proxy, so bring one up first. `docker-compose.yml` docker compose down -v ``` -Tests marked `@pytest.mark.e2e` skip when no proxy answers `/health/liveliness`, so a run that reports everything skipped means the stack isn't up, not that anything passed +Tests marked `@pytest.mark.e2e` hard-fail when no proxy answers `/health/liveliness`, so a run that goes red with `No live proxy` at setup means the stack isn't up; they never skip for a missing proxy, so an absent stack can't be mistaken for a pass ## What a complete test looks like @@ -132,7 +132,7 @@ The shape is layered so tests stay declarative Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `Gateway` and adds suite-specific routes. Cleanup runs through that same `Gateway`, so whatever keys or customers your test creates get torn down by the `resources` fixture -Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The skip-vs-fail split is deliberate: a test marked `e2e` skips when no proxy answers its liveness probe, but once a request reaches the proxy any wrong behavior is a hard failure, never a skip +Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure coverage of the harness itself carries no marker and runs regardless. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache diff --git a/tests/e2e/access_control/conftest.py b/tests/e2e/access_control/conftest.py index 9f4a00fe06fc..b5681ff76adc 100644 --- a/tests/e2e/access_control/conftest.py +++ b/tests/e2e/access_control/conftest.py @@ -1,4 +1,4 @@ -"""Access-control suite client fixture; lifecycle/skip/marker live in the parent conftest.""" +"""Access-control suite client fixture; lifecycle/liveness gate/marker live in the parent conftest.""" import pytest diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py index 2c6070c437af..d3b6d42bc24e 100644 --- a/tests/e2e/batches/conftest.py +++ b/tests/e2e/batches/conftest.py @@ -1,6 +1,6 @@ """Batches suite's `client` fixture. -The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. BatchClient holds the shared Gateway, so the `resources` fixture cleans up keys through it; tests register file deletes and batch cancels via `resources.defer(...)`. diff --git a/tests/e2e/bob_the_builder.py b/tests/e2e/bob_the_builder.py deleted file mode 100644 index 18aff2edc989..000000000000 --- a/tests/e2e/bob_the_builder.py +++ /dev/null @@ -1,247 +0,0 @@ -"""Bob the builder: on a red e2e run, ask Devin to fix the failing tests. - -Wired as a ``pytest_sessionfinish`` step (see ``conftest.py``). When the run went -red and remediation is enabled, it hands the failing tests plus their captured -tracebacks to Devin *through the LiteLLM proxy's own MCP gateway* -- the same -gateway + master key the suite already uses -- so Devin files a Linear ticket per -failure and opens fix PRs. Nothing new ships in the runner pod: the proxy already -registers the ``devin`` MCP server and holds ``DEVIN_API_KEY``, injecting it -upstream, so this process only needs the proxy key it always has. - -Opt-in via ``E2E_DEVIN_REMEDIATION=1`` so a normal local ``pytest tests/e2e`` run -never spawns a Devin session. ``DEVIN_DRY_RUN=1`` prints the prompt it would send -and makes no call. Everything is best-effort: any error here is logged and -swallowed so the run's exit status still reflects the tests, not remediation. -""" - -from __future__ import annotations - -import hashlib -import os -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from pathlib import Path -from typing import Protocol, cast - -import pytest -from pydantic import BaseModel, ConfigDict - -from e2e_config import MASTER_KEY, PROXY_BASE_URL -from e2e_http import Success -from transport import HttpTransport - -REMEDIATION_ENV = "E2E_DEVIN_REMEDIATION" -_LIST_PATH = "/mcp-rest/tools/list" -_CALL_PATH = "/mcp-rest/tools/call" - - -@dataclass(frozen=True, slots=True) -class Failure: - """One failed test: its pytest node id and the captured failure text.""" - - nodeid: str - detail: str - - -@dataclass(frozen=True, slots=True) -class Config: - server: str - create_tool: str - linear_team: str - target_repo: str - target_ref: str - max_failures: int - max_detail_chars: int - tags: tuple[str, ...] - dry_run: bool - - -class _NoParams(BaseModel): - pass - - -class _McpToolInfo(BaseModel): - model_config = ConfigDict(extra="allow") - server_name: str | None = None - alias: str | None = None - - -class _McpTool(BaseModel): - model_config = ConfigDict(extra="allow") - name: str - mcp_info: _McpToolInfo | None = None - - -class _McpToolsList(BaseModel): - model_config = ConfigDict(extra="allow") - tools: tuple[_McpTool, ...] = () - - -class _DevinSessionArgs(BaseModel): - prompt: str - title: str - tags: list[str] - - -class _ToolCallBody(BaseModel): - name: str - arguments: _DevinSessionArgs - - -class _ToolCallResult(BaseModel): - model_config = ConfigDict(extra="allow") - - -class _Report(Protocol): - @property - def nodeid(self) -> str: ... - - @property - def longreprtext(self) -> str: ... - - -class _TerminalReporter(Protocol): - stats: Mapping[str, Sequence[_Report]] - - -def _env(name: str, default: str) -> str: - value = os.environ.get(name, "").strip() - return value or default - - -def load_config() -> Config: - raw_tags = _env("DEVIN_TAGS", "e2e,stage") - return Config( - server=_env("DEVIN_MCP_SERVER", "devin"), - create_tool=_env("DEVIN_SESSION_TOOL", "devin_session_create"), - linear_team=_env("DEVIN_LINEAR_TEAM", "LIT"), - target_repo=_env("DEVIN_TARGET_REPO", "BerriAI/litellm"), - target_ref=_env("DEVIN_TARGET_REF", "litellm_internal_staging"), - max_failures=int(_env("DEVIN_MAX_FAILURES", "50")), - max_detail_chars=int(_env("DEVIN_MAX_DETAIL_CHARS", "3000")), - tags=tuple(t.strip() for t in raw_tags.split(",") if t.strip()), - dry_run=_env("DEVIN_DRY_RUN", "0") == "1", - ) - - -def collect_failures(session: pytest.Session, max_detail_chars: int) -> tuple[Failure, ...]: - """Pull the failed and errored tests (with their tracebacks) off the run's - terminal reporter. Returns empty when nothing failed or the reporter is - absent (e.g. a skipped, proxy-less session).""" - plugin: object = session.config.pluginmanager.getplugin("terminalreporter") - if plugin is None: - return () - reporter = cast(_TerminalReporter, plugin) - reports = (*reporter.stats.get("failed", ()), *reporter.stats.get("error", ())) - return tuple( - Failure(nodeid=r.nodeid, detail=r.longreprtext.strip()[-max_detail_chars:]) for r in reports - ) - - -def dedup_tag(failures: tuple[Failure, ...]) -> str: - """Stable short tag identifying this exact set of failing tests, so repeated - nightly runs on the same failures reference one body of work.""" - joined = "\n".join(sorted(f.nodeid for f in failures)) - return "e2e-fail-" + hashlib.sha256(joined.encode()).hexdigest()[:12] - - -def _revision() -> str: - for candidate in (Path(__file__).parent / ".litellm-revision", Path("/app/e2e/.litellm-revision")): - try: - return candidate.read_text(encoding="utf-8").strip() - except OSError: - continue - return _env("E2E_REVISION", "unknown") - - -def build_prompt(cfg: Config, failures: tuple[Failure, ...], tag: str) -> str: - shown = failures[: cfg.max_failures] - header = ( - f"The LiteLLM end-to-end suite failed on the " - f"{_env('E2E_ENVIRONMENT', 'stage')} proxy. Source repo {cfg.target_repo} " - f"at revision {_revision()} (branch {cfg.target_ref}). {len(failures)} " - f"test(s) failed" - + (f"; the first {len(shown)} are shown" if len(shown) < len(failures) else "") - + ".\n\n" - ) - task = ( - "For each failing test below:\n" - f"1. Open a Linear ticket under the {cfg.linear_team} team describing the " - "failure (test id, the assertion/error, likely cause), unless an open " - "ticket for that same test already exists -- do not create duplicates.\n" - f"2. Fix it in {cfg.target_repo}, branching off {cfg.target_ref} and " - "following the repo's CONTRIBUTING and CLAUDE.md conventions (meaningful " - "regression coverage, conventional commits, run the suite locally), then " - "open a PR that references the Linear ticket.\n" - "3. Prefer one focused PR per failing test; if several share a root cause, " - "group them and say so.\n" - f"Before starting, search existing sessions/PRs tagged '{tag}' or " - "referencing these test ids and continue that work instead of restarting.\n\n" - "Failing tests and their captured output:\n" - ) - blocks = [f"### {i}. {f.nodeid}\n```\n{f.detail}\n```\n" for i, f in enumerate(shown, start=1)] - return header + task + "\n".join(blocks) - - -def _resolve_tool_name(transport: HttpTransport, cfg: Config) -> str | None: - """Find Devin's create-session tool on the gateway. The proxy prefixes tools - with the server alias, so match by suffix and (when present) the owning - server.""" - result = transport.get( - _LIST_PATH, headers=transport.master, params=_NoParams(), response_type=_McpToolsList - ) - if not isinstance(result, Success): - print(f"bob_the_builder: could not list gateway MCP tools: {result}") - return None - for tool in result.data.tools: - owner = tool.mcp_info.server_name or tool.mcp_info.alias if tool.mcp_info else None - if (owner is None or owner == cfg.server) and ( - tool.name == cfg.create_tool or tool.name.endswith(cfg.create_tool) - ): - return tool.name - print( - f"bob_the_builder: no '{cfg.create_tool}' tool for server '{cfg.server}' on the gateway; " - f"saw {[t.name for t in result.data.tools]}" - ) - return None - - -def remediate(session: pytest.Session) -> None: - """Entry point called from ``pytest_sessionfinish``. No-op unless remediation - is enabled and the run actually had failures.""" - if os.environ.get(REMEDIATION_ENV) != "1": - return - cfg = load_config() - failures = collect_failures(session, cfg.max_detail_chars) - if not failures: - return - - tag = dedup_tag(failures) - title = f"Fix {len(failures)} failing LiteLLM e2e test(s) [{tag}]" - prompt = build_prompt(cfg, failures, tag) - args = _DevinSessionArgs(prompt=prompt, title=title, tags=[*cfg.tags, tag]) - - if cfg.dry_run: - print("bob_the_builder: DRY RUN -- would create a Devin session:") - print(f" server : {cfg.server}\n tool : {cfg.create_tool}\n title : {title}") - print(f" tags : {args.tags}\n---- prompt ----\n{prompt}") - return - - try: - transport = HttpTransport(base_url=PROXY_BASE_URL, master_key=MASTER_KEY) - tool_name = _resolve_tool_name(transport, cfg) - if tool_name is None: - return - result = transport.post( - _CALL_PATH, - headers=transport.master, - json=_ToolCallBody(name=tool_name, arguments=args), - response_type=_ToolCallResult, - ) - if isinstance(result, Success): - print(f"bob_the_builder: created Devin session for {len(failures)} failure(s) [{tag}]") - print(result.data.model_dump_json()) - else: - print(f"bob_the_builder: Devin session call failed: {result}") - except Exception as exc: # noqa: BLE001 - remediation must never fail the run - print(f"bob_the_builder: remediation error (ignored): {exc}") diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 3aec104c8617..88a9deecb7e6 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -15,14 +15,14 @@ import functools import sys -from collections.abc import Generator, Iterator +from collections.abc import Iterator from pathlib import Path import pytest import requests from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL -from e2e_result_reporter import covers_from_item, format_e2e_result_line, result_from_pytest +from junit_properties import attach_result_properties from lifecycle import GatewayProvider, ResourceManager @@ -40,6 +40,17 @@ def pytest_configure(config: pytest.Config) -> None: ) +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + """Attach the two custom signals (suite package and covered cell ids) to every + test's user_properties so the standard JUnit report (`--junitxml`) records them + as `` entries, on every outcome including skips and setup errors. + Downstream (Loki/Grafana) reads outcome and duration from the standard report + and these properties for package rollups and coverage drill-down. See + junit_properties.py.""" + for item in items: + attach_result_properties(item) + + def _liveness_reason(label: str, base_url: str) -> str | None: """None if `base_url` answers its liveness probe, else a failure reason.""" try: @@ -86,30 +97,6 @@ def pytest_runtest_call(item: pytest.Item) -> None: item.session.stash[_E2E_TEST_RAN] = True -@pytest.hookimpl(wrapper=True, tryfirst=True) -def pytest_runtest_makereport( - item: pytest.Item, call: pytest.CallInfo[object] -) -> Generator[None, pytest.TestReport, pytest.TestReport]: - """Emit one structured E2E_RESULT line per finished test for Loki/Grafana. - - Status-history panels should aggregate by package (and optional covers), not - scrape pytest progress basenames. See e2e_result_reporter.py. - """ - report = yield - result = result_from_pytest( - nodeid=str(report.nodeid), - when=str(report.when), - failed=bool(report.failed), - skipped=bool(report.skipped), - passed=bool(report.passed), - duration_seconds=float(report.duration), - covers=covers_from_item(item), - ) - if result is not None: - print(format_e2e_result_line(result), flush=True) - return report - - def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: """Once the whole e2e session is done (all suites), truncate the spend logs so the DB doesn't accumulate test rows. Sessions where no e2e test body ran leave @@ -132,13 +119,6 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: if spend_dir in sys.path: sys.path.remove(spend_dir) - try: - from bob_the_builder import remediate - - remediate(session) - except Exception as exc: # noqa: BLE001 - remediation is best-effort - print(f"devin remediation best-effort failed: {exc}") - @pytest.fixture def resources(client: GatewayProvider) -> Iterator[ResourceManager]: diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index fac266149f44..0d61f48703d6 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -7,10 +7,13 @@ - {id: quota_management.ratelimit.priority_generous.picks_under_tpm, module: quota_management, tier: P1, behavior: ratelimit, variant: priority_generous, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "dynamic_rate_limiter_v3.py:36-52", rationale: "Generous mode (<80% sat) allows priority borrowing"} - {id: quota_management.ratelimit.priority_strict.picks_under_tpm, module: quota_management, tier: P1, behavior: ratelimit, variant: priority_strict, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "dynamic_rate_limiter_v3.py:53-71", rationale: "Strict mode (>=80% sat) enforces priority fairness"} - {id: quota_management.budget.key.blocks_over_limit, module: quota_management, tier: P0, behavior: budget, variant: key, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A key's max_budget blocks further paid calls once spend crosses it"} +- {id: quota_management.budget.team.blocks_over_limit, module: quota_management, tier: P0, behavior: budget, variant: team, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A team's max_budget blocks every key on the team once combined spend crosses it, including keys that spent nothing themselves"} - {id: quota_management.budget.internal_user.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An internal user's max_budget governs personal keys"} +- {id: quota_management.budget.internal_user.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An internal user's max_budget governs every personal key it owns; a second untouched key is blocked once the shared user budget is exhausted"} - {id: quota_management.budget.end_user.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: end_user, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A customer (end-user) max_budget blocks calls attributed via user="} - {id: quota_management.budget.organization.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An organization's max_budget blocks keys under its teams"} - {id: quota_management.budget.team_member.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A member's per-team budget blocks independently of the team budget"} +- {id: quota_management.budget.team_member.isolates_per_member, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [isolates_per_member], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "One team member's exhausted per-team budget does not block a different member on the same team"} - {id: quota_management.budget.tag.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: tag, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "router_strategy/budget_limiter.py", rationale: "Proxy-level tag budgets block tagged requests at the cap"} - {id: quota_management.budget.model_max.isolates_per_model, module: quota_management, tier: P1, behavior: budget, variant: model_max, assertions: [isolates_per_model], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "model_max_budget caps one model without touching a sibling's budget"} - {id: quota_management.budget.soft.alerts_without_blocking, module: quota_management, tier: P1, behavior: budget, variant: soft, assertions: [alerts_without_blocking], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "soft_budget alerts but never blocks traffic"} diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index a117cbd570d9..29d54b011be2 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -1,5 +1,7 @@ # local setup to run e2e tests configs: + mcp_upstream_server: + file: ../mcp_tests/mcp_e2e_upstream_server.py litellm_config: content: | general_settings: @@ -131,7 +133,27 @@ services: target: /app/config.yaml command: ["--config", "/app/config.yaml", "--port", "4000"] -# throwaway db +# deterministic self-hosted upstream MCP server (FastMCP add/multiply over +# streamable-http), reachable by the litellm container at mcp-upstream:8090/mcp. +# Not a depends_on of litellm on purpose: only the mcp suite needs it, and it +# boots long before the proxy is live, so it must not gate the other suites' +# stack. The suite registers it through /v1/mcp/server at test time. + mcp-upstream: + image: ghcr.io/berriai/litellm:main-latest + entrypoint: ["python3", "/app/mcp_upstream_server.py"] + environment: + MCP_HOST: 0.0.0.0 + MCP_PORT: "8090" + configs: + - source: mcp_upstream_server + target: /app/mcp_upstream_server.py + healthcheck: + test: ["CMD", "python3", "-c", "import socket; socket.create_connection(('127.0.0.1', 8090), 2).close()"] + interval: 3s + timeout: 3s + retries: 40 + +# throwaway db db: image: postgres:16 environment: diff --git a/tests/e2e/e2e_result_reporter.py b/tests/e2e/e2e_result_reporter.py deleted file mode 100644 index 22f7581818fe..000000000000 --- a/tests/e2e/e2e_result_reporter.py +++ /dev/null @@ -1,144 +0,0 @@ -"""Structured e2e result lines for Loki / Grafana status history. - -Pytest progress lines are a bad dashboard source: they only expose file basenames, -break under quiet modes, and force status-history rows to explode with suite growth. - -Each finished test emits one logfmt line: - - E2E_RESULT package=logging file=test_langfuse_e2e.py outcome=failed - duration_ms=1234 node_id=logging/test_langfuse_e2e.py::TestX::test_y - covers=logging.langfuse.team.success - -Grafana package status-history queries max(fail) by package over E2E_RESULT lines. -Drill-down uses node_id / covers in Explore, not status-history cardinality. -""" - -from __future__ import annotations - -from collections.abc import Iterable, Sequence -from dataclasses import dataclass -from pathlib import Path -from typing import Literal, Protocol, runtime_checkable - -Outcome = Literal["passed", "failed", "error", "skipped"] - - -@dataclass(frozen=True, slots=True) -class E2EResult: - package: str - file: str - outcome: Outcome - duration_ms: int - node_id: str - covers: tuple[str, ...] - - -@runtime_checkable -class _MarkerArgs(Protocol): - args: Sequence[object] - - -@runtime_checkable -class _ItemWithCovers(Protocol): - def iter_markers(self, name: str) -> Iterable[object]: ... - - -def package_from_nodeid(nodeid: str) -> str: - """Top-level suite package under tests/e2e/, or 'root' for top-level files. - - Pytest nodeids are relative to the invocation cwd. Repo-root runs look like - `tests/e2e/logging/...`; suite-cwd runs look like `logging/...`. Strip the - `tests/e2e` prefix so package is the suite dir either way. - """ - path_part = nodeid.split("::", 1)[0].replace("\\", "/") - parts = tuple(p for p in path_part.split("/") if p and p != ".") - if len(parts) >= 3 and parts[0] == "tests" and parts[1] == "e2e": - parts = parts[2:] - if len(parts) <= 1: - return "root" - return parts[0] - - -def file_from_nodeid(nodeid: str) -> str: - path_part = nodeid.split("::", 1)[0].replace("\\", "/") - return Path(path_part).name - - -def covers_from_item(item: object) -> tuple[str, ...]: - """Read @pytest.mark.covers cell ids from a pytest Item.""" - if not isinstance(item, _ItemWithCovers): - return () - return tuple( - dict.fromkeys( - arg - for marker in item.iter_markers(name="covers") - if isinstance(marker, _MarkerArgs) - for arg in marker.args - if isinstance(arg, str) and arg - ) - ) - - -def outcome_from_report(when: str, failed: bool, skipped: bool, passed: bool) -> Outcome | None: - """Map pytest TestReport fields to a terminal outcome. None if not final.""" - if when == "setup" and skipped: - return "skipped" - if when == "setup" and failed: - return "error" - if when != "call": - return None - if skipped: - return "skipped" - if failed: - return "failed" - if passed: - return "passed" - return "failed" - - -def _logfmt_escape(value: str) -> str: - if value == "": - return '""' - needs_quote = any(ch.isspace() or ch in "\"=\\" for ch in value) - if not needs_quote: - return value - escaped = value.replace("\\", "\\\\").replace('"', '\\"') - return f'"{escaped}"' - - -def format_e2e_result_line(result: E2EResult) -> str: - covers = ",".join(result.covers) - fields = ( - ("package", result.package), - ("file", result.file), - ("outcome", result.outcome), - ("duration_ms", str(result.duration_ms)), - ("node_id", result.node_id), - ("covers", covers), - ) - body = " ".join(f"{key}={_logfmt_escape(value)}" for key, value in fields) - return f"E2E_RESULT {body}" - - -def result_from_pytest( - *, - nodeid: str, - when: str, - failed: bool, - skipped: bool, - passed: bool, - duration_seconds: float, - covers: tuple[str, ...] = (), -) -> E2EResult | None: - outcome = outcome_from_report(when=when, failed=failed, skipped=skipped, passed=passed) - if outcome is None: - return None - duration_ms = max(0, int(round(duration_seconds * 1000))) - return E2EResult( - package=package_from_nodeid(nodeid), - file=file_from_nodeid(nodeid), - outcome=outcome, - duration_ms=duration_ms, - node_id=nodeid, - covers=covers, - ) diff --git a/tests/e2e/grafana/status_history_panels.md b/tests/e2e/grafana/status_history_panels.md deleted file mode 100644 index f8cda509c631..000000000000 --- a/tests/e2e/grafana/status_history_panels.md +++ /dev/null @@ -1,66 +0,0 @@ -# Grafana: package status history for e2e - -Dashboard: [LiteLLM E2E](https://berriai.grafana.net/d/mup2cfn/litellm-e2e) (`mup2cfn`). - -The old **test suite status history** panel scraped pytest progress lines and -grouped by **file basename** (`test_foo.py`). That does not scale: multi-class -files collapse to one bit, and full `node_id` cardinality melts status-history. - -## Emitter - -After each test finishes, `tests/e2e/conftest.py` prints one logfmt line: - -``` -E2E_RESULT package=logging file=test_langfuse_e2e.py outcome=failed duration_ms=1500 node_id="logging/..." covers=cell.id -``` - -## Panel: package status history (replace panel 11) - -**Type:** Status history -**Interval:** 15m (or 1h for multi-day ranges) -**Description:** Per top-level package under `tests/e2e/`: red if any test failed or errored in the bucket. - -```logql -max by (package) ( - max_over_time( - {service_name="litellm-e2e"} - |= "E2E_RESULT" - | logfmt - | outcome != "" - | label_format result=`{{ if or (eq .outcome "failed") (eq .outcome "error") }}1{{ else }}0{{ end }}` - | unwrap result - [$__interval] - ) -) -``` - -Value mappings: `0` → Pass (green), `1` → Fail (red). - -If `service_name` is missing on older scrapes, use: - -```logql -{cluster="berrie-litellm-stage", pod=~"litellm-e2e-.+"} -``` - -instead of `{service_name="litellm-e2e"}`. - -## Panel: failed tests (logs drill-down) - -```logql -{service_name="litellm-e2e"} |= "E2E_RESULT" | logfmt | outcome=~"failed|error" -``` - -Show fields: `package`, `file`, `node_id`, `covers`, `duration_ms`. - -## Panel (optional): filter by package variable - -Dashboard variable `package` (custom or from label_values on E2E_RESULT): - -```logql -{service_name="litellm-e2e"} |= "E2E_RESULT" | logfmt | package=`$package` | outcome=~"failed|error" -``` - -## Do not - -- Put full `node_id` as the status-history series key (cardinality). -- Rely on `::S+ PASSED` progress regex as the primary signal once E2E_RESULT is live. diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py new file mode 100644 index 000000000000..e4f59f5c4d21 --- /dev/null +++ b/tests/e2e/junit_properties.py @@ -0,0 +1,59 @@ +"""Custom per-test signals for the standard JUnit reporter. + +The e2e suite ships results to Loki/Grafana from a standard pytest JUnit report +(`--junitxml=e2e-report.xml`), not a bespoke log line. JUnit already records +outcome, duration, and node id for every ``; the only signals it cannot +derive on its own are the normalized suite package and the coverage-registry cell +ids a test covers. Those ride along as JUnit `` entries via each item's +`user_properties`, attached in `conftest.py::pytest_collection_modifyitems`. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import pytest + + +def package_from_nodeid(nodeid: str) -> str: + """Top-level suite package under tests/e2e/, or 'root' for top-level files. + + Pytest nodeids are relative to the invocation cwd. Repo-root runs look like + `tests/e2e/logging/...`; suite-cwd runs look like `logging/...`. Strip the + `tests/e2e` prefix so package is the suite dir either way. + """ + path_part = nodeid.split("::", 1)[0].replace("\\", "/") + raw = tuple(p for p in path_part.split("/") if p and p != ".") + parts = raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw + if len(parts) <= 1: + return "root" + return parts[0] + + +def dedupe_covers(marker_args: Iterable[tuple[object, ...]]) -> tuple[str, ...]: + """Flatten @pytest.mark.covers arg lists into unique, order-preserving cell + ids, dropping anything that is not a non-empty string.""" + return tuple(dict.fromkeys(arg for args in marker_args for arg in args if isinstance(arg, str) and arg)) + + +def covers_from_item(item: pytest.Item) -> tuple[str, ...]: + """Read @pytest.mark.covers cell ids off a pytest Item, order-preserving.""" + return dedupe_covers(marker.args for marker in item.iter_markers(name="covers")) + + +def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]: + """The custom signals a standard reporter cannot derive: the normalized suite + package and the comma-joined coverage-registry cell ids this test covers.""" + return ( + ("package", package_from_nodeid(item.nodeid)), + ("covers", ",".join(covers_from_item(item))), + ) + + +def attach_result_properties(item: pytest.Item) -> None: + """Attach result_properties to an item's user_properties, idempotently: a + second call is a no-op, so a collection that runs the hook more than once + never emits duplicate entries.""" + if any(name == "package" for name, _ in item.user_properties): + return + item.user_properties.extend(result_properties(item)) diff --git a/tests/e2e/llm_translation/conftest.py b/tests/e2e/llm_translation/conftest.py index 2a87ef7259dc..5258b751a8c5 100644 --- a/tests/e2e/llm_translation/conftest.py +++ b/tests/e2e/llm_translation/conftest.py @@ -1,6 +1,6 @@ """LLM-translation suite's `client` fixture. -The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. PassthroughClient holds the shared Gateway, so the `resources` fixture cleans up keys this suite creates. """ diff --git a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md index 4795c3b9f54e..bae858d50afa 100644 --- a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md @@ -49,10 +49,10 @@ kept commented out in `PROVIDERS` until they pass end-to-end here; re-enable the uncommenting their entry. Every provider is provisioned and asserted; the suite never skips a provider. Per -`tests/e2e/CLAUDE.md` the only sanctioned skip is the whole-suite proxy-liveness -skip, so a provider whose credentials or upstream realtime model are missing on the -gateway is a hard failure, not a skip. Give the gateway each provider's credentials -to turn its tests green. +`tests/e2e/CLAUDE.md` there is no sanctioned skip: the whole-suite proxy-liveness +probe hard-fails when no proxy answers, and a provider whose credentials or upstream +realtime model are missing on the gateway is likewise a hard failure, not a skip. +Give the gateway each provider's credentials to turn its tests green. ## Running @@ -63,5 +63,5 @@ the deployments itself), then uv run pytest tests/e2e/llm_translation/realtime/ -v ``` -The whole suite skips only when no proxy answers `GET /health/liveliness` at +The whole suite hard-fails at setup when no proxy answers `GET /health/liveliness` at `LITELLM_PROXY_URL` (default `http://localhost:4000`). diff --git a/tests/e2e/llm_translation/realtime/conftest.py b/tests/e2e/llm_translation/realtime/conftest.py index 15cd789664eb..8e6e596bcd32 100644 --- a/tests/e2e/llm_translation/realtime/conftest.py +++ b/tests/e2e/llm_translation/realtime/conftest.py @@ -1,6 +1,6 @@ """Realtime suite's `client` and `realtime_models` fixtures. -The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. RealtimeClient holds the shared Gateway, so the `resources` fixture cleans up keys this suite creates. diff --git a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py index 6aaffdd208ef..f99fa8d86b32 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py @@ -6,10 +6,10 @@ reconstructed transcript and usage, and a full tool-call round-trip (call -> tool result -> a follow-up response that uses the result). -One GA-speaking client validates every provider; only the model alias changes. A -provider whose realtime alias is not configured on the proxy skips (skip on -environment); once it is configured, a protocol failure is a hard failure. See -REALTIME_COVERAGE_MATRIX.md. +One GA-speaking client validates every provider; only the model alias changes. +Every provider is provisioned at session start, so a missing realtime alias is a +hard failure, not a skip; once configured, a protocol failure is likewise a hard +failure. See REALTIME_COVERAGE_MATRIX.md. """ import pytest diff --git a/tests/e2e/llm_translation/test_ocr_rust_e2e.py b/tests/e2e/llm_translation/test_ocr_rust_e2e.py index 921010e5eaef..e735d9c01b5f 100644 --- a/tests/e2e/llm_translation/test_ocr_rust_e2e.py +++ b/tests/e2e/llm_translation/test_ocr_rust_e2e.py @@ -7,9 +7,9 @@ rather than another inline body. Start the proxy with the Rust OCR path enabled: Each case creates its deployment, drives a real /v1/ocr call, and asserts a -well-formed OCR document comes back. Per the e2e "skip on environment, fail on -behavior" rule, a case skips when no proxy answers but fails (never skips) once a -request reaches it: the proxy fetches each provider's referenced secrets, so a +well-formed OCR document comes back. Per the e2e hard-fail contract, a case +fails when no proxy answers and also fails once a request reaches it: the proxy +fetches each provider's referenced secrets, so a missing credential surfaces as a live provider error rather than silent green. """ diff --git a/tests/e2e/management/conftest.py b/tests/e2e/management/conftest.py index 18da1305c135..4f2dc874a33d 100644 --- a/tests/e2e/management/conftest.py +++ b/tests/e2e/management/conftest.py @@ -1,6 +1,6 @@ """Management suite fixtures: the client plus a logged-in dashboard page. -Lifecycle/skip/marker live in the parent conftest. The browser fixtures drive +Lifecycle/liveness gate/marker live in the parent conftest. The browser fixtures drive the dashboard the proxy serves at /ui, so browser tests exercise exactly what an end user sees. playwright is an optional dependency loaded behind importorskip inside the fixture, so the API tests in this suite collect and run without it: diff --git a/tests/e2e/mcp/conftest.py b/tests/e2e/mcp/conftest.py new file mode 100644 index 000000000000..77fef5747062 --- /dev/null +++ b/tests/e2e/mcp/conftest.py @@ -0,0 +1,16 @@ +"""MCP suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness handling, and the +`e2e`/`covers` markers live in the parent tests/e2e/conftest.py. McpClient holds +the shared Gateway, so the `resources` fixture tears down whatever this suite +creates (keys via the Gateway, MCP servers via the deferred cleanups). +""" + +import pytest + +from mcp_client import McpClient, build_client + + +@pytest.fixture(scope="session") +def client() -> McpClient: + return build_client() diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py new file mode 100644 index 000000000000..a1dac3fdac40 --- /dev/null +++ b/tests/e2e/mcp/mcp_client.py @@ -0,0 +1,153 @@ +"""Client for the MCP e2e suite: admin server registration plus the api_key tool +surface. + +An admin registers an upstream MCP server through the management API +(`/v1/mcp/server`, persisted in the DB) and grants a virtual key access to it via +`object_permission.mcp_servers`. Keys then reach the server through the REST bridge +the proxy exposes for api_key auth (`/mcp-rest/tools/list`, `/mcp-rest/tools/call`), +which `user_api_key_auth` gates the same way the JSON-RPC `/mcp` surface does. The +request/response bodies are co-located here because only this suite speaks MCP. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from pydantic import BaseModel, ConfigDict, Field, RootModel + +from e2e_gateway import Gateway, build_gateway +from e2e_http import Headers, NoBody, Result, unwrap +from models import KeyGenerateBody, ObjectPermission + + +class ApiKeyHeaders(Headers): + x_litellm_api_key: str = Field(serialization_alias="x-litellm-api-key") + + +class McpServerNewBody(BaseModel): + server_name: str + alias: str + url: str + transport: str = "http" + + +class McpServerNewResponse(BaseModel): + server_id: str + + +class McpServerRow(BaseModel): + server_id: str + alias: str | None = None + url: str | None = None + + +class McpServersListResponse(RootModel[list[McpServerRow]]): + pass + + +class McpToolMcpInfo(BaseModel): + server_id: str | None = None + alias: str | None = None + + +class McpToolEntry(BaseModel): + name: str + description: str | None = None + mcp_info: McpToolMcpInfo | None = None + + +class McpToolsListResponse(BaseModel): + tools: list[McpToolEntry] = [] + error: str | None = None + message: str | None = None + + def tool_names_for_server(self, server_id: str) -> frozenset[str]: + return frozenset( + tool.name + for tool in self.tools + if tool.mcp_info is not None and tool.mcp_info.server_id == server_id + ) + + +class McpCallToolBody(BaseModel): + name: str + arguments: dict[str, int] + server_id: str + + +class McpCallContent(BaseModel): + type: str | None = None + text: str | None = None + + +class McpCallToolResponse(BaseModel): + model_config = ConfigDict(populate_by_name=True) + content: list[McpCallContent] = [] + is_error: bool | None = Field(default=None, alias="isError") + + @property + def first_text(self) -> str | None: + return self.content[0].text if self.content else None + + +@dataclass(frozen=True, slots=True) +class McpClient: + gateway: Gateway + + def register_server(self, *, server_name: str, alias: str, url: str) -> str: + return unwrap( + self.gateway.transport.post( + "/v1/mcp/server", + headers=self.gateway.transport.master, + json=McpServerNewBody(server_name=server_name, alias=alias, url=url), + response_type=McpServerNewResponse, + ) + ).server_id + + def delete_server(self, server_id: str) -> None: + _ = self.gateway.transport.delete( + f"/v1/mcp/server/{server_id}", + headers=self.gateway.transport.master, + json=NoBody(), + response_type=NoBody, + ) + + def registered_servers(self) -> list[McpServerRow]: + return unwrap( + self.gateway.transport.get( + "/v1/mcp/server", + headers=self.gateway.transport.master, + params=NoBody(), + response_type=McpServersListResponse, + ) + ).root + + def generate_key(self, *, user_id: str, mcp_servers: list[str] | None) -> str: + object_permission = ( + ObjectPermission(mcp_servers=mcp_servers) if mcp_servers is not None else None + ) + return self.gateway.generate_key( + KeyGenerateBody(models=[], user_id=user_id, object_permission=object_permission) + ) + + def list_tools(self, key: str) -> Result[McpToolsListResponse]: + return self.gateway.transport.get( + "/mcp-rest/tools/list", + headers=ApiKeyHeaders(x_litellm_api_key=key), + params=NoBody(), + response_type=McpToolsListResponse, + ) + + def call_tool( + self, key: str, *, server_id: str, name: str, arguments: dict[str, int] + ) -> Result[McpCallToolResponse]: + return self.gateway.transport.post( + "/mcp-rest/tools/call", + headers=ApiKeyHeaders(x_litellm_api_key=key), + json=McpCallToolBody(name=name, arguments=arguments, server_id=server_id), + response_type=McpCallToolResponse, + ) + + +def build_client() -> McpClient: + return McpClient(gateway=build_gateway()) diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py new file mode 100644 index 000000000000..eaa49af5b698 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -0,0 +1,103 @@ +"""Live e2e: a virtual key without MCP access is denied an MCP server's tools. + +An admin registers an upstream MCP server through the management API (persisted in +the DB, picked up without a restart) and queues its deletion. Two keys are created +against that one server: one granted access through `object_permission.mcp_servers` +and one with no MCP grant at all. The permitted key is the control that proves the +upstream is alive and the tool is callable, so a failure on the denied key is an +authorization denial rather than a dead server. The denied key must then see none +of the server's tools on `tools/list` and must be refused with a 403 on +`tools/call`. + +Both the recorded state (the server is registered; the permitted key resolves its +tools) and the enforced behavior (the unpermitted key sees nothing and is blocked) +are asserted, so a regression that leaks tools to an ungranted key or drops the +call-time permission check fails here. +""" + +import os + +import pytest + +from e2e_config import unique_marker +from e2e_http import UnknownApiError, unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient + +pytestmark = pytest.mark.e2e + +MCP_UPSTREAM_URL = os.environ.get("E2E_MCP_UPSTREAM_URL", "http://mcp-upstream:8090/mcp") +MATH_TOOLS = frozenset({"add", "multiply"}) + + +def _register_math_server(client: McpClient, resources: ResourceManager) -> str: + name = f"e2e_math_{unique_marker()}" + server_id = client.register_server(server_name=name, alias=name, url=MCP_UPSTREAM_URL) + resources.defer(lambda: client.delete_server(server_id)) + return server_id + + +def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str] | None) -> str: + label = "allowed" if mcp_servers else "denied" + key = client.generate_key(user_id=f"e2e-mcp-{label}-{unique_marker()}", mcp_servers=mcp_servers) + resources.defer(lambda: client.gateway.delete_key(key)) + return key + + +def _assert_registered(client: McpClient, server_id: str) -> None: + registered = {row.server_id for row in client.registered_servers()} + assert server_id in registered, f"registered server {server_id} absent from /v1/mcp/server: {registered}" + + +class TestMcpKeyWithoutAccessIsDenied: + @pytest.mark.covers("mcp.list_tools.api_key.denied_without_permission") + def test_list_tools_denied_without_permission( + self, client: McpClient, resources: ResourceManager + ) -> None: + server_id = _register_math_server(client, resources) + _assert_registered(client, server_id) + + permitted_key = _key(client, resources, mcp_servers=[server_id]) + denied_key = _key(client, resources, mcp_servers=None) + + permitted_tools = unwrap(client.list_tools(permitted_key)).tool_names_for_server(server_id) + assert MATH_TOOLS <= permitted_tools, ( + f"granted key did not see the server's tools (upstream dead or grant not applied): " + f"{permitted_tools}" + ) + + denied_tools = unwrap(client.list_tools(denied_key)).tool_names_for_server(server_id) + assert denied_tools == frozenset(), ( + f"ungranted key saw the server's tools; tools/list leaked across the permission " + f"boundary: {denied_tools}" + ) + + @pytest.mark.covers("mcp.call_tool.api_key.denied_without_permission") + def test_call_tool_denied_without_permission( + self, client: McpClient, resources: ResourceManager + ) -> None: + server_id = _register_math_server(client, resources) + _assert_registered(client, server_id) + + permitted_key = _key(client, resources, mcp_servers=[server_id]) + denied_key = _key(client, resources, mcp_servers=None) + + permitted_tools = unwrap(client.list_tools(permitted_key)).tool_names_for_server(server_id) + assert "add" in permitted_tools, ( + f"granted key did not discover the add tool (upstream dead or grant not applied): " + f"{permitted_tools}" + ) + + permitted_call = unwrap( + client.call_tool(permitted_key, server_id=server_id, name="add", arguments={"a": 3, "b": 4}) + ) + assert permitted_call.is_error is not True, f"granted key's tool call errored: {permitted_call}" + assert permitted_call.first_text == "7", ( + f"granted key's add(3, 4) did not return 7 (upstream not reachable): {permitted_call}" + ) + + match client.call_tool(denied_key, server_id=server_id, name="add", arguments={"a": 3, "b": 4}): + case UnknownApiError(status_code=403, body=body): + assert "access_denied" in body, f"403 was not an MCP access denial: {body}" + case other: + pytest.fail(f"ungranted key's tool call was not refused with 403 access_denied: {other}") diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 82c276d0b641..39832d1a17f3 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -39,6 +39,10 @@ class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None +class ObjectPermission(BaseModel): + mcp_servers: list[str] | None = None + + class KeyGenerateBody(BaseModel): models: list[str] = [] duration: str | None = None @@ -57,6 +61,7 @@ class KeyGenerateBody(BaseModel): rpm_limit: int | None = None allowed_routes: list[str] | None = None metadata: KeyMetadata | None = None + object_permission: ObjectPermission | None = None class KeyGenerateResponse(BaseModel): diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 7b37c3af98e5..01e4d63c1c3e 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -39,6 +39,19 @@ class UserNewResponse(BaseModel): user_id: str +class UserInfoParams(BaseModel): + user_id: str + + +class UserInfoRow(BaseModel): + spend: float | None = None + max_budget: float | None = None + + +class UserInfoResponse(BaseModel): + user_info: UserInfoRow | None = None + + class UserDeleteBody(BaseModel): user_ids: list[str] @@ -262,6 +275,19 @@ def delete_user(self, user_id: str) -> None: response_type=NoBody, ) + def user_info(self, user_id: str) -> UserInfoRow | None: + result = self.gateway.transport.get( + "/user/info", + headers=self.gateway.transport.master, + params=UserInfoParams(user_id=user_id), + response_type=UserInfoResponse, + ) + match result: + case Success(data=data): + return data.user_info + case _: + return None + # ---- customer / end-user ------------------------------------------- def create_customer(self, customer_id: str, *, max_budget: float) -> str: diff --git a/tests/e2e/quota_management/budgets/conftest.py b/tests/e2e/quota_management/budgets/conftest.py index 236822f43097..4299d2ffd49e 100644 --- a/tests/e2e/quota_management/budgets/conftest.py +++ b/tests/e2e/quota_management/budgets/conftest.py @@ -1,6 +1,6 @@ """Budgets suite's `client` fixture. -The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. BudgetClient holds the shared Gateway, so the `resources` fixture cleans up keys through it; tests register entity deletes via `resources.defer(...)`. diff --git a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py index dbe1cfa4ea80..0b8adfc47ae4 100644 --- a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py @@ -4,7 +4,7 @@ the budgeted entity + a key, run() drives spend until a `budget_exceeded` block, teardown() deletes everything init() created (always runs, even on failure/skip). Covers the entities with no prior live coverage - internal user, end-user, -organization, team member. See BUDGET_TEST_COVERAGE_MATRIX.md. +organization, team member - plus key and team. See BUDGET_TEST_COVERAGE_MATRIX.md. A non-budget error fails hard (never a skip); if calls never get blocked, budget enforcement is broken -> fail. @@ -18,16 +18,17 @@ from budget_client import BudgetClient, is_budget_block from e2e_config import unique_marker -from e2e_http import require_successful_call +from e2e_http import StreamingResponse, require_successful_call from lifecycle import run_case pytestmark = pytest.mark.e2e -def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> None: - """Send paid calls until the entity's budget blocks one. Key/user/org/member - block within a couple calls off real-time reservation counters; the end-user - budget enforces off table spend that lands on the batch write, so it takes a - few more. A non-budget error fails hard (never a skip).""" +def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> StreamingResponse: + """Send paid calls until the entity's budget blocks one; return the blocked + response so callers can assert on its shape. Key/user/org/member block within + a couple calls off real-time reservation counters; the end-user budget + enforces off table spend that lands on the batch write, so it takes a few + more. A non-budget error fails hard (never a skip).""" for _ in range(40): result = client.chat( key, @@ -37,7 +38,7 @@ def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> user=user or None, ) if is_budget_block(result): - return + return result require_successful_call(result) time.sleep(2) pytest.fail("budget never enforced within the call budget") @@ -69,10 +70,53 @@ def teardown(self) -> None: class KeyBudgetCase(_BudgetCase): + """A bare key (no team_id / user_id) carrying its own max_budget, so only the + key-level budget can be the thing that blocks. The refusal must be a 429 + budget_exceeded; any other error already fails via _assert_budget_blocks.""" + def init(self) -> None: self.key = self.client.generate_key(max_budget=3e-6) self._undo.append(lambda: self.client.delete_key(self.key)) + def run(self) -> None: + blocked = _assert_budget_blocks(self.client, self.key) + assert blocked.status_code == 429, ( + f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" + ) + + +class TeamBudgetCase(_BudgetCase): + """An admin caps a whole team: two keys under a tiny-budget team, neither with + a key-level budget. Key A is driven until the team cap blocks it; key B's very + first call must then be refused too, proving the cap sits on the team, not the + key that spent. Both refusals must be 429 budget_exceeded.""" + + def init(self) -> None: + team_id = self.client.create_team( + alias=f"e2e-budget-team-{unique_marker()}", max_budget=3e-6 + ) + self._undo.append(lambda: self.client.delete_team(team_id)) + self.key = self.client.generate_key(team_id=team_id) + self._undo.append(lambda: self.client.delete_key(self.key)) + self._sibling_key = self.client.generate_key(team_id=team_id) + self._undo.append(lambda: self.client.delete_key(self._sibling_key)) + + def run(self) -> None: + blocked = _assert_budget_blocks(self.client, self.key) + assert blocked.status_code == 429, ( + f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" + ) + sibling = self.client.chat( + self._sibling_key, + "claude-haiku-4-5", + f"spend {unique_marker()}", + max_tokens=16, + ) + assert is_budget_block(sibling) and sibling.status_code == 429, ( + f"a sibling key on the capped team must get the same 429 budget_exceeded, " + f"got {sibling.status_code}: {sibling.body[:200]}" + ) + class InternalUserBudgetCase(_BudgetCase): def init(self) -> None: @@ -97,20 +141,31 @@ def run(self) -> None: class OrganizationBudgetCase(_BudgetCase): + """Org carries the tiny budget; the team under it and the key carry none, so + the org is the only entity that can block (the historically weak link). The + refusal must be a 429 budget_exceeded that names the org as the blocker.""" + def init(self) -> None: - # Org carries the tiny budget; the team under it has none, so a block here - # is org-level enforcement (the historically weak link). - org_id = self.client.create_org( + self._org_id = self.client.create_org( max_budget=3e-6, alias=f"e2e-budget-org-{unique_marker()}" ) - self._undo.append(lambda: self.client.delete_org(org_id)) + self._undo.append(lambda: self.client.delete_org(self._org_id)) team_id = self.client.create_team( - alias=f"e2e-budget-team-{unique_marker()}", organization_id=org_id + alias=f"e2e-budget-team-{unique_marker()}", organization_id=self._org_id ) self._undo.append(lambda: self.client.delete_team(team_id)) self.key = self.client.generate_key(team_id=team_id) self._undo.append(lambda: self.client.delete_key(self.key)) + def run(self) -> None: + blocked = _assert_budget_blocks(self.client, self.key) + assert blocked.status_code == 429, ( + f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" + ) + assert f"Organization={self._org_id}" in blocked.body, ( + f"refusal must name the org as the blocker, got: {blocked.body[:200]}" + ) + class TeamMemberBudgetCase(_BudgetCase): def init(self) -> None: @@ -138,6 +193,10 @@ def _case_id(case_cls: Type[_BudgetCase]) -> str: KeyBudgetCase, marks=pytest.mark.covers("quota_management.budget.key.blocks_over_limit"), ), + pytest.param( + TeamBudgetCase, + marks=pytest.mark.covers("quota_management.budget.team.blocks_over_limit"), + ), pytest.param( InternalUserBudgetCase, marks=pytest.mark.covers("quota_management.budget.internal_user.blocks_over_limit"), diff --git a/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py b/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py new file mode 100644 index 000000000000..87855a9a1c10 --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py @@ -0,0 +1,119 @@ +"""Live e2e: per-team-member budgets are enforced independently between members. + +Two members share one team that has a large team budget. The tight member is capped +at a tiny per-team budget and spends past it; the roomy member has plenty of room. +Once the tight member is blocked with budget_exceeded, the roomy member still serves +on the same team, its calls land in the spend logs under its own user id, and the +tight member stays blocked. A shared or leaky member counter would either block the +roomy member too or let the tight member back through once its peer spent. +""" + +import time +from collections.abc import Iterator +from dataclasses import dataclass + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import Success, require_successful_call +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage + +pytestmark = pytest.mark.e2e + +MODEL = "gpt-5.5" +TEAM_BUDGET = 100.0 +TIGHT_MEMBER_BUDGET = 3e-6 +ROOMY_MEMBER_BUDGET = 100.0 +ROOMY_BURST = 3 + + +@dataclass(frozen=True, slots=True) +class _Pair: + team_id: str + tight_user_id: str + roomy_user_id: str + tight_key: str + roomy_key: str + + +@pytest.fixture(scope="class") +def pair(client: BudgetClient) -> Iterator[_Pair]: + """One team with a large budget and two members on it: a tight member capped at + a tiny per-team budget and a roomy member with headroom, each with their own key. + Shared across the class and torn down LIFO best-effort when it finishes.""" + resources = ResourceManager(client=client.gateway) + try: + marker = unique_marker() + team_id = client.create_team(alias=f"e2e-member-iso-{marker}", max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_team(team_id)) + tight_user = client.create_user(max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_user(tight_user)) + roomy_user = client.create_user(max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_user(roomy_user)) + client.add_team_member(team_id, tight_user, max_budget_in_team=TIGHT_MEMBER_BUDGET) + client.add_team_member(team_id, roomy_user, max_budget_in_team=ROOMY_MEMBER_BUDGET) + tight_key = client.generate_key(team_id=team_id, user_id=tight_user) + resources.defer(lambda: client.delete_key(tight_key)) + roomy_key = client.generate_key(team_id=team_id, user_id=roomy_user) + resources.defer(lambda: client.delete_key(roomy_key)) + yield _Pair( + team_id=team_id, + tight_user_id=tight_user, + roomy_user_id=roomy_user, + tight_key=tight_key, + roomy_key=roomy_key, + ) + finally: + resources.teardown() + + +def _roomy_send(client: BudgetClient, key: str) -> str: + """One roomy-member call that must go through; returns its request id.""" + match client.gateway.chat( + key, + ChatBody( + model=MODEL, + messages=[ChatMessage(role="user", content=f"roomy {unique_marker()}")], + max_tokens=16, + ), + ): + case Success(data=response): + assert response.id is not None, "roomy member call returned no id" + return response.id + case other: + pytest.fail(f"roomy member call failed while a peer was over budget: {other}") + + +class TestTeamMemberBudgetIsolation: + @pytest.mark.covers("quota_management.budget.team_member.isolates_per_member") + def test_blocked_member_does_not_block_peer(self, client: BudgetClient, pair: _Pair) -> None: + blocked = False + for _ in range(40): + result = client.chat(pair.tight_key, MODEL, f"tight {unique_marker()}", max_tokens=16) + if is_budget_block(result): + blocked = True + break + require_successful_call(result) + time.sleep(2) + assert blocked, "tight member's per-team budget never enforced" + + sent = frozenset(_roomy_send(client, pair.roomy_key) for _ in range(ROOMY_BURST)) + + assert is_budget_block( + client.chat(pair.tight_key, MODEL, f"tight {unique_marker()}", max_tokens=16) + ), "tight member stopped being blocked once the peer spent" + + rows = client.gateway.poll_logs_for_key( + pair.roomy_key, predicate=lambda rs: bool(sent & {r.request_id for r in rs}) + ) + logged = [row for row in rows if row.request_id in sent] + assert logged, "none of the roomy member's calls reached the spend logs" + for row in logged: + assert row.user == pair.roomy_user_id, ( + f"roomy call {row.request_id} logged under user {row.user}, not {pair.roomy_user_id}" + ) + assert row.team_id == pair.team_id, ( + f"roomy call {row.request_id} logged under team {row.team_id}, not {pair.team_id}" + ) diff --git a/tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py b/tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py new file mode 100644 index 000000000000..4dc7a2df647f --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_user_budget_across_keys_e2e.py @@ -0,0 +1,79 @@ +"""Live e2e: a per-user max_budget is enforced across ALL of that user's keys. + +An internal user's budget governs every personal key it owns, not only the one +that happened to spend it down. One user with a tiny max_budget owns two keys: +driving the first key to a budget_exceeded block then makes a fresh, untouched +second key of the same user (which carries no budget of its own, so nothing but the +shared user budget can block it) reject the same way, and the user's recorded spend +has crossed the cap. A key-scoped-only budget would leave the second key serving. +""" + +import time + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +MODEL = "gpt-5.5" +TINY_CAP = 3e-6 +RECORDED_SPEND_DEADLINE_SECONDS = 90 +SECOND_KEY_BLOCK_ATTEMPTS = 6 + + +def _call(client: BudgetClient, key: str) -> StreamingResponse: + return client.chat(key, MODEL, f"across {unique_marker()}", max_tokens=16) + + +def _drive_to_block(client: BudgetClient, key: str, subject: str) -> None: + for _ in range(40): + result = _call(client, key) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(2) + pytest.fail(f"user budget never enforced on {subject} within the call budget") + + +def _expect_prompt_block(client: BudgetClient, key: str, subject: str) -> None: + """The shared user budget is already exhausted before this key makes a single + call, so a key with no budget of its own must be rejected promptly. The small + bounded retry only absorbs spend-propagation lag between the two keys; it is far + below the spend a key-scoped budget would need to accumulate to block itself, so + a block here can only come from the shared user budget.""" + for _ in range(SECOND_KEY_BLOCK_ATTEMPTS): + result = _call(client, key) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(2) + pytest.fail( + f"{subject} was not blocked by the shared user budget within {SECOND_KEY_BLOCK_ATTEMPTS} calls" + ) + + +class TestUserBudgetAcrossKeys: + @pytest.mark.covers("quota_management.budget.internal_user.enforced_across_keys") + def test_user_budget_blocks_a_second_key(self, client: BudgetClient, resources: ResourceManager) -> None: + user_id = client.create_user(max_budget=TINY_CAP) + resources.defer(lambda: client.delete_user(user_id)) + + first_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(first_key)) + second_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(second_key)) + + _drive_to_block(client, first_key, "the first key") + _expect_prompt_block(client, second_key, "the second key") + + deadline = time.monotonic() + RECORDED_SPEND_DEADLINE_SECONDS + while time.monotonic() < deadline: + info = client.user_info(user_id) + if info is not None and (info.spend or 0.0) >= TINY_CAP: + return + time.sleep(5) + pytest.fail(f"user spend never reached the {TINY_CAP} cap in the recorded state") diff --git a/tests/e2e/quota_management/ratelimit/conftest.py b/tests/e2e/quota_management/ratelimit/conftest.py index 4a5a73bb5e4b..59dee5e65b38 100644 --- a/tests/e2e/quota_management/ratelimit/conftest.py +++ b/tests/e2e/quota_management/ratelimit/conftest.py @@ -1,6 +1,6 @@ """Quota-management suite's `client` fixture. -The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. QuotaClient holds the shared Gateway, so the `resources` fixture cleans up keys through it. """ diff --git a/tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md b/tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md index 062ef8d73da5..6baebc4c28c5 100644 --- a/tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md +++ b/tests/e2e/quota_management/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md @@ -80,5 +80,5 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`. `proxy_batch_write_at` (~60s) means rows land late; every read polls to a deadline. Fresh scoped key per test (isolation, xdist-safe, cleaned up). Assert invariants (`spend > 0`, `total == prompt + completion`, aggregate == sum), not literal -$/token values, so pricing drift is not a failure. Skip on environment (no proxy / -no provider key), fail on behavior (a real 2xx call with a wrong/missing row). +$/token values, so pricing drift is not a failure. Hard-fail when no proxy +answers, fail on behavior (a real 2xx call with a wrong/missing row). diff --git a/tests/e2e/quota_management/spend_tracking/conftest.py b/tests/e2e/quota_management/spend_tracking/conftest.py index 0e80764236bd..434af15b1821 100644 --- a/tests/e2e/quota_management/spend_tracking/conftest.py +++ b/tests/e2e/quota_management/spend_tracking/conftest.py @@ -1,6 +1,6 @@ """Spend-tracking suite's `client` fixture and driver-model registration. -The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. SpendClient exposes the shared Gateway (GatewayProvider), so the `resources` fixture cleans up keys and customers this suite creates. diff --git a/tests/e2e/router/conftest.py b/tests/e2e/router/conftest.py index 046cdd80c2ba..344d8ab5c137 100644 --- a/tests/e2e/router/conftest.py +++ b/tests/e2e/router/conftest.py @@ -1,6 +1,6 @@ """Router suite's `client` fixture. -The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. ComplexityRouterClient holds the shared Gateway, so the `resources` fixture cleans up keys this suite creates. diff --git a/tests/local_testing/test_anthropic_prompt_caching.py b/tests/local_testing/test_anthropic_prompt_caching.py index ff89c3845e4d..ef374de5e2ae 100644 --- a/tests/local_testing/test_anthropic_prompt_caching.py +++ b/tests/local_testing/test_anthropic_prompt_caching.py @@ -172,7 +172,7 @@ def anthropic_messages(): "content": [ { "type": "text", - "text": "Here is the full text of a complex legal agreement" * 400, + "text": "Here is the full text of a complex legal agreement" * 500, "cache_control": {"type": "ephemeral"}, } ], diff --git a/tests/mcp_tests/mcp_e2e_upstream_server.py b/tests/mcp_tests/mcp_e2e_upstream_server.py new file mode 100644 index 000000000000..28fb08464812 --- /dev/null +++ b/tests/mcp_tests/mcp_e2e_upstream_server.py @@ -0,0 +1,40 @@ +"""Deterministic upstream MCP server for the mcp e2e suite. + +A tiny FastMCP server exposing `add` and `multiply` over streamable-http so the +suite has a self-hosted, offline upstream to register and exercise. DNS-rebinding +protection is turned off because the litellm container reaches this over the +compose network by service name (`mcp-upstream:8090`), not localhost, and the +stack is an isolated throwaway. Bind host/port come from MCP_HOST/MCP_PORT. +""" + +import os + +from mcp.server.fastmcp import FastMCP +from mcp.server.transport_security import TransportSecuritySettings + +mcp: FastMCP = FastMCP( + "e2e-math", + host=os.getenv("MCP_HOST", "0.0.0.0"), + port=int(os.getenv("MCP_PORT", "8090")), + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), +) + + +@mcp.tool() +def add(a: int, b: int) -> int: + """Add two integers""" + return a + b + + +@mcp.tool() +def multiply(a: int, b: int) -> int: + """Multiply two integers""" + return a * b + + +def main() -> None: + mcp.run(transport="streamable-http") + + +if __name__ == "__main__": + main() diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 204367921b59..560766eedd5a 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -934,8 +934,8 @@ def mock_get_server_by_id(server_id): mcp_auth_header=mock_auth_header, mcp_servers=["server1"], ) - assert len(result) == 1, "Should only return tools from server1" - assert result[0].name == "tool1", "Should return tool from server1" + assert len(result.tools) == 1, "Should only return tools from server1" + assert result.tools[0].name == "tool1", "Should return tool from server1" # Test Case 2: Without specific MCP servers # Create a different mock manager for the second test case @@ -977,9 +977,9 @@ async def mock_get_tools_side_effect( mcp_auth_header=mock_auth_header, mcp_servers=None, ) - assert len(result) == 2, "Should return tools from all servers" + assert len(result.tools) == 2, "Should return tools from all servers" assert ( - result[0].name == "tool1" and result[1].name == "tool2" + result.tools[0].name == "tool1" and result.tools[1].name == "tool2" ), "Should return tools from all servers" # @@ -1014,8 +1014,8 @@ async def mock_get_tools_side_effect( mcp_auth_header=mock_auth_header, mcp_servers=["group-a"], ) - assert len(result) == 1, "Should only return tools from server3" - assert result[0].name == "tool1", "Should return tool from server1" + assert len(result.tools) == 1, "Should only return tools from server3" + assert result.tools[0].name == "tool1", "Should return tool from server1" except AssertionError as e: pytest.fail(f"Test failed: {str(e)}") @@ -2435,11 +2435,12 @@ def mock_client_constructor(*args, **kwargs): mock_client_constructor, ): # Call _get_tools_from_mcp_servers which should apply the filtering - filtered_tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=mock_user_auth, mcp_auth_header="Bearer test_token", mcp_servers=None, # Get from all servers ) + filtered_tools = listing.tools # Verify that only allowed tools are returned assert ( @@ -2548,11 +2549,12 @@ def mock_client_constructor(*args, **kwargs): mock_client_constructor, ): # Call _get_tools_from_mcp_servers which should apply the filtering - filtered_tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=mock_user_auth, mcp_auth_header="Bearer test_token", mcp_servers=None, # Get from all servers ) + filtered_tools = listing.tools # Verify that only safe tools are returned (dangerous tools filtered out) assert ( @@ -2649,11 +2651,12 @@ def mock_client_constructor(*args, **kwargs): mock_client_constructor, ): # Call _get_tools_from_mcp_servers which should apply the filtering - filtered_tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=mock_user_auth, mcp_auth_header="Bearer test_token", mcp_servers=None, # Get from all servers ) + filtered_tools = listing.tools # Should return all tools when no restrictions assert ( diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 848a6c28a57d..a969d21a6819 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1820,7 +1820,7 @@ def test_init_auto_router_deployment_success(mock_auto_router, model_list): # Verify the auto-router was added to the router's auto_routers dict assert "test-auto-router" in router.auto_routers - assert router.auto_routers["test-auto-router"] == mock_auto_router_instance + assert router.auto_routers["test-auto-router"][0].strategy == mock_auto_router_instance @patch("litellm.router_strategy.auto_router.auto_router.AutoRouter") @@ -1833,7 +1833,11 @@ def test_init_auto_router_deployment_duplicate_model_name(mock_auto_router, mode mock_auto_router.return_value = mock_auto_router_instance # Add an existing auto-router - router.auto_routers["test-auto-router"] = mock_auto_router_instance + from litellm.types.router import TaggedPreRoutingStrategy + + router.auto_routers["test-auto-router"] = [ + TaggedPreRoutingStrategy(tags=(), strategy=mock_auto_router_instance) + ] # Try to add another auto-router with the same name litellm_params = LiteLLM_Params( @@ -1849,7 +1853,7 @@ def test_init_auto_router_deployment_duplicate_model_name(mock_auto_router, mode ) with pytest.raises( - ValueError, match="Auto-router deployment test-auto-router already exists" + ValueError, match="Auto-router deployment test-auto-router with tags .* already exists" ): router.init_auto_router_deployment(deployment) diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 2021e7836897..e45c2f4615e9 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -2,7 +2,9 @@ import datetime import json import os +import subprocess import sys +import textwrap import unittest from typing import List, Optional, Tuple from unittest.mock import ANY, MagicMock, Mock, patch @@ -1624,3 +1626,242 @@ def test_existing_cache_control_counted_toward_limit(self): sys_blocks = sum(1 for b in (result_sys or []) if isinstance(b, dict) and b.get("cache_control") is not None) total_blocks = sys_blocks + sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in result_msgs) assert total_blocks <= 4 + + +class TestEnableAnthropicPromptCaching: + """Auto-injected default breakpoints via litellm.enable_anthropic_prompt_caching.""" + + MESSAGES: List[AllMessageValues] = [ + {"role": "system", "content": "a long system prompt"}, + {"role": "user", "content": "first turn"}, + {"role": "assistant", "content": "a reply"}, + {"role": "user", "content": "latest turn"}, + ] + + def _points(self, model="claude-sonnet-4-5", provider="anthropic", messages=None, system=None, tools=None): + return AnthropicCacheControlHook.get_default_injection_points( + messages=copy.deepcopy(self.MESSAGES) if messages is None else messages, + system=system, + model=model, + custom_llm_provider=provider, + tools=tools, + ) + + def test_disabled_by_default(self): + assert litellm.enable_anthropic_prompt_caching is False + assert self._points() == [] + + def test_injects_system_and_trailing_turn(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert self._points() == [ + {"location": "message", "role": "system", "index": None, "control": {"type": "ephemeral"}}, + {"location": "message", "role": None, "index": -1, "control": {"type": "ephemeral"}}, + ] + + def test_bedrock_claude_is_injected(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + points = self._points(model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", provider="bedrock") + assert [p["index"] for p in points] == [None, -1] + + @pytest.mark.parametrize("model, provider", [("gpt-4o", "openai"), ("gemini-2.0-flash", "gemini")]) + def test_non_anthropic_providers_never_injected(self, monkeypatch, model, provider): + """These report supports_prompt_caching=True but never consume cache_control markers.""" + from litellm.utils import supports_prompt_caching + + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert supports_prompt_caching(model=model, custom_llm_provider=provider) is True + assert self._points(model=model, provider=provider) == [] + + def test_model_without_caching_support_not_injected(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert self._points(model="anthropic.claude-3-5-sonnet-20240620-v1:0", provider="bedrock") == [] + + def test_stands_down_when_client_sent_cache_control(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = [ + {"role": "system", "content": [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": "latest turn"}, + ] + assert self._points(messages=messages) == [] + + def test_stands_down_when_system_block_has_cache_control(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + system = [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}] + assert self._points(messages=[{"role": "user", "content": "hi"}], system=system) == [] + + @staticmethod + def _tools(count: int, cached: bool) -> List[dict]: + tool: dict = {"type": "function", "function": {"name": "t", "description": "d", "parameters": {}}} + if cached: + tool["cache_control"] = {"type": "ephemeral"} + return [{**tool, "function": {**tool["function"], "name": f"t{i}"}} for i in range(count)] + + def test_stands_down_when_only_tools_carry_cache_control(self, monkeypatch): + """Caching just the tool definitions is a normal client pattern, and those + breakpoints count toward the provider's four-block limit. Three of them plus + our two would be five, which Anthropic rejects outright.""" + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert self._points(tools=self._tools(3, cached=True)) == [] + + def test_injects_when_tools_carry_no_cache_control(self, monkeypatch): + """Tools alone must not suppress injection; only client-marked ones do.""" + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert [p["index"] for p in self._points(tools=self._tools(3, cached=False))] == [None, -1] + + @pytest.mark.parametrize("tools", [None, []]) + def test_absent_tools_do_not_suppress_injection(self, monkeypatch, tools): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert [p["index"] for p in self._points(tools=tools)] == [None, -1] + + def test_seed_stands_down_when_only_tools_carry_cache_control(self, monkeypatch): + """Same guard on the /chat/completions seeding path.""" + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + params: dict = {} + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=params, + messages=copy.deepcopy(self.MESSAGES), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=self._tools(3, cached=True), + ) + assert "cache_control_injection_points" not in params + + def test_v1_messages_stands_down_when_only_tools_carry_cache_control(self, monkeypatch): + """Same guard on the /v1/messages path, where tools reach the hook directly.""" + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + result_msgs, result_sys = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(messages), + "sys", + {}, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=self._tools(3, cached=True), + ) + assert result_sys == "sys" + assert result_msgs == messages + + def test_default_ttl_is_anthropics_five_minute_cache(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert all(p["control"] == {"type": "ephemeral"} for p in self._points()) + + @pytest.mark.parametrize("ttl", ["5m", "1h"]) + def test_ttl_override_applied(self, monkeypatch, ttl): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", ttl) + assert all(p["control"] == {"type": "ephemeral", "ttl": ttl} for p in self._points()) + + def test_seed_does_not_override_configured_points(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + configured = [{"location": "message", "role": "user", "index": 0}] + params = {"cache_control_injection_points": configured} + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=params, + messages=copy.deepcopy(self.MESSAGES), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + assert params["cache_control_injection_points"] is configured + + def test_seed_adds_defaults_when_enabled(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + params: dict = {} + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=params, + messages=copy.deepcopy(self.MESSAGES), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + assert [p["index"] for p in params["cache_control_injection_points"]] == [None, -1] + + def test_seed_is_noop_when_disabled(self): + params: dict = {} + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=params, + messages=copy.deepcopy(self.MESSAGES), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + assert params == {} + + def test_v1_messages_applies_defaults_end_to_end(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = [ + {"role": "user", "content": [{"type": "text", "text": "first"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "reply"}]}, + {"role": "user", "content": [{"type": "text", "text": "latest"}]}, + ] + result_msgs, result_sys = AnthropicCacheControlHook.maybe_inject_cache_control( + messages, + "a system prompt", + {}, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert result_sys == [{"type": "text", "text": "a system prompt", "cache_control": {"type": "ephemeral"}}] + assert result_msgs[-1]["content"][-1]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in result_msgs[0]["content"][-1] + + def test_v1_messages_is_noop_when_disabled(self): + messages = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + result_msgs, result_sys = AnthropicCacheControlHook.maybe_inject_cache_control( + messages, + "sys", + {}, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert result_sys == "sys" + assert result_msgs == messages + + +class TestAnthropicPromptCachingEnvVars: + """Both settings are read from the environment at import, so an admin can enable + auto-caching without a config file. Each case re-imports litellm in a subprocess + so the env is read fresh without contaminating this process's module graph. + """ + + @staticmethod + def _import_litellm_with_env(env_override: dict) -> Tuple[bool, Optional[str]]: + env = os.environ.copy() + env.pop("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", None) + env.pop("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL", None) + env.update(env_override) + script = textwrap.dedent( + """ + import json, litellm + print(json.dumps([litellm.enable_anthropic_prompt_caching, litellm.anthropic_prompt_caching_ttl])) + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, env=env, timeout=300 + ) + assert result.returncode == 0, result.stderr + enabled, ttl = json.loads(result.stdout.strip().splitlines()[-1]) + return enabled, ttl + + def test_unset_env_leaves_auto_caching_off(self): + assert self._import_litellm_with_env({}) == (False, None) + + @pytest.mark.parametrize("value", ["true", "True", "TRUE"]) + def test_env_enables_auto_caching_case_insensitively(self, value): + enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value}) + assert enabled is True + + @pytest.mark.parametrize("value", ["false", "0", "yes", ""]) + def test_env_only_enables_on_true(self, value): + enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value}) + assert enabled is False + + @pytest.mark.parametrize("value", ["5m", "1h"]) + def test_ttl_env_is_applied(self, value): + _, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value}) + assert ttl == value + + @pytest.mark.parametrize("value", ["10m", "1H", "3600", "ephemeral"]) + def test_unsupported_ttl_env_falls_back_to_provider_default(self, value): + """An unparseable TTL must fall back to Anthropic's 5m default, never reach the provider verbatim.""" + _, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value}) + assert ttl is None diff --git a/tests/test_litellm/integrations/test_opik_utils.py b/tests/test_litellm/integrations/test_opik_utils.py new file mode 100644 index 000000000000..a4250acf1dc2 --- /dev/null +++ b/tests/test_litellm/integrations/test_opik_utils.py @@ -0,0 +1,29 @@ +"""Unit tests for the native Opik integration's UUIDv7 id generation.""" + +import uuid +from datetime import datetime, timezone +from unittest.mock import patch + +from litellm.integrations.opik.utils import create_uuid7 + + +def _timestamp_ms(uuid_str: str) -> int: + """Return the unix-ms timestamp encoded in a UUIDv7's top 48 bits.""" + return uuid.UUID(uuid_str).int >> 80 + + +def test_create_uuid7_is_valid_version_7_uuid(): + parsed = uuid.UUID(create_uuid7()) + assert parsed.version == 7 + assert parsed.variant == uuid.RFC_4122 + + +def test_create_uuid7_encodes_timestamp_in_milliseconds(): + fixed = datetime(2026, 6, 24, 10, 0, 0, tzinfo=timezone.utc) + + with patch( + "litellm.integrations.opik.utils.time.time", return_value=fixed.timestamp() + ): + value = create_uuid7() + + assert _timestamp_ms(value) == int(fixed.timestamp() * 1000) diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 3c410cf84df3..6ab0f2c08abe 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1261,6 +1261,23 @@ def test_is_anthropic_invalid_thinking_signature_error_positive(self): ) assert is_anthropic_invalid_thinking_signature_error(raw) is True + def test_is_anthropic_invalid_thinking_signature_error_positive_bedrock(self): + from litellm.llms.anthropic.common_utils import ( + is_anthropic_invalid_thinking_signature_error, + ) + + # Real user-reported Bedrock scenario + raw = '{"message":"messages.2.content.0.thinking.signature.str: Input should be a valid string"}' + assert is_anthropic_invalid_thinking_signature_error(raw) is True + + def test_is_anthropic_invalid_thinking_signature_error_positive_vertex(self): + from litellm.llms.anthropic.common_utils import ( + is_anthropic_invalid_thinking_signature_error, + ) + + raw = "messages.4.content.1.thinking.signature.str: Input should be a valid string" + assert is_anthropic_invalid_thinking_signature_error(raw) is True + def test_is_anthropic_invalid_thinking_signature_error_negative(self): from litellm.llms.anthropic.common_utils import ( is_anthropic_invalid_thinking_signature_error, @@ -1271,6 +1288,11 @@ def test_is_anthropic_invalid_thinking_signature_error_negative(self): is_anthropic_invalid_thinking_signature_error("rate limit exceeded") is False ) + assert ( + is_anthropic_invalid_thinking_signature_error("invalid_request_error: model not found") + is False + ) + assert is_anthropic_invalid_thinking_signature_error("thinking signature is malformed") is False def test_strip_thinking_blocks_from_anthropic_messages(self): from litellm.llms.anthropic.common_utils import ( diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py new file mode 100644 index 000000000000..99dcaa36c756 --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -0,0 +1,66 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.fireworks_ai.cost_calculator import cost_per_token +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + +MODEL = "accounts/fireworks/models/glm-5p2" +INPUT_COST = 1.4e-06 +CACHE_READ_COST = 2.6e-07 +OUTPUT_COST = 4.4e-06 + + +def _usage(prompt_tokens: int, cached_tokens: int, completion_tokens: int) -> Usage: + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens), + ) + + +def test_cached_prompt_tokens_billed_at_cache_read_rate(): + prompt_tokens = 7036 + cached_tokens = 7020 + completion_tokens = 8 + + prompt_cost, completion_cost = cost_per_token( + model=MODEL, usage=_usage(prompt_tokens, cached_tokens, completion_tokens) + ) + + expected_prompt_cost = (prompt_tokens - cached_tokens) * INPUT_COST + cached_tokens * CACHE_READ_COST + assert prompt_cost == pytest.approx(expected_prompt_cost) + assert completion_cost == pytest.approx(completion_tokens * OUTPUT_COST) + + full_rate_cost = prompt_tokens * INPUT_COST + assert prompt_cost < full_rate_cost + + +def test_warm_call_cheaper_than_cold_call(): + prompt_tokens = 7036 + completion_tokens = 8 + + cold_prompt_cost, _ = cost_per_token( + model=MODEL, usage=_usage(prompt_tokens, 16, completion_tokens) + ) + warm_prompt_cost, _ = cost_per_token( + model=MODEL, usage=_usage(prompt_tokens, 7020, completion_tokens) + ) + + assert warm_prompt_cost < cold_prompt_cost + + +def test_no_cached_tokens_matches_full_input_rate(): + prompt_tokens = 100 + completion_tokens = 10 + + prompt_cost, completion_cost = cost_per_token( + model=MODEL, usage=_usage(prompt_tokens, 0, completion_tokens) + ) + + assert prompt_cost == pytest.approx(prompt_tokens * INPUT_COST) + assert completion_cost == pytest.approx(completion_tokens * OUTPUT_COST) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 6f132aaae9cd..9375f7481c8e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -301,6 +301,187 @@ async def test_get_allowed_mcp_servers_for_key_returns_sentinel_marker(self): assert result == [SpecialMCPServerNames.no_mcp_servers.value] + def _toolset_only_object_permission(self, toolset_ids): + key_object_permission = MagicMock() + key_object_permission.mcp_servers = [] + key_object_permission.mcp_access_groups = [] + key_object_permission.mcp_tool_permissions = None + key_object_permission.mcp_toolsets = toolset_ids + return key_object_permission + + def _mock_manager_with_toolsets(self, toolset_perms): + mock_manager = MagicMock() + mock_manager.expand_permission_list = MagicMock(side_effect=lambda servers: servers) + mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {}) + mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value=toolset_perms) + return mock_manager + + async def test_get_allowed_mcp_servers_for_key_includes_toolset_servers(self): + """A key granted only mcp_toolsets must reach the toolset's servers on + every path (list, call, REST); regression for the list-ok/call-403 bug""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + key_object_permission = self._toolset_only_object_permission(["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with ( + patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth) + + assert result == ["server-a"] + mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with(toolset_ids=["toolset-1"]) + + async def test_get_allowed_mcp_servers_for_key_skips_toolset_resolution_when_none_granted(self): + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + key_object_permission = self._toolset_only_object_permission([]) + key_object_permission.mcp_servers = ["server-direct"] + mock_manager = self._mock_manager_with_toolsets({}) + + with ( + patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth) + + assert result == ["server-direct"] + mock_manager.resolve_toolset_tool_permissions.assert_not_awaited() + + async def test_get_allowed_mcp_servers_toolset_only_key_end_to_end_inheritance(self): + """The full get_allowed_mcp_servers flow (key/team inheritance, no team + restriction) surfaces toolset-granted servers for a toolset-only key""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + key_object_permission = self._toolset_only_object_permission(["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with ( + patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[])), + patch.object(MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == ["server-a"] + + async def test_toolset_servers_stay_capped_by_team_ceiling(self): + """Toolset grants expand the KEY's scope, which the team ceiling still + intersects; a toolset must never grant a server the team does not allow. + Pins that toolset expansion lives in the intersected key scope, not the + additive access-group path""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="test-team") + key_object_permission = self._toolset_only_object_permission(["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets( + {"server-in-team": ["lookup_status"], "server-outside-team": ["other_tool"]} + ) + + with ( + patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object(MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[])), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + AsyncMock(return_value=["server-in-team", "server-unrelated"]), + ), + patch.object(MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[])), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == ["server-in-team"] + + async def test_get_allowed_tools_for_server_unions_toolset_and_direct_tools(self): + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + key_object_permission = self._toolset_only_object_permission(["toolset-1"]) + key_object_permission.mcp_tool_permissions = {"server-a": ["direct_tool"]} + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with ( + patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + ): + result = await MCPRequestHandler.get_allowed_tools_for_server( + server_id="server-a", + user_api_key_auth=user_api_key_auth, + ) + + assert result is not None + assert set(result) == {"direct_tool", "lookup_status"} + + async def test_get_allowed_tools_for_server_toolset_only_key_restricts_to_toolset_tools(self): + """A toolset grant must RESTRICT the server's tools, not fall through to + the allow-all default; otherwise merging servers alone would over-grant + every tool on a toolset-referenced server""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + key_object_permission = self._toolset_only_object_permission(["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with ( + patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + ): + allowed = await MCPRequestHandler.get_allowed_tools_for_server( + server_id="server-a", + user_api_key_auth=user_api_key_auth, + ) + is_granted_tool_allowed = await MCPRequestHandler.is_tool_allowed_for_server( + tool_name="lookup_status", + server_id="server-a", + user_api_key_auth=user_api_key_auth, + ) + is_other_tool_allowed = await MCPRequestHandler.is_tool_allowed_for_server( + tool_name="delete_everything", + server_id="server-a", + user_api_key_auth=user_api_key_auth, + ) + + assert allowed == ["lookup_status"] + assert is_granted_tool_allowed is True + assert is_other_tool_allowed is False + + async def test_get_allowed_tools_for_server_without_restrictions_stays_allow_all(self): + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + key_object_permission = self._toolset_only_object_permission([]) + mock_manager = self._mock_manager_with_toolsets({}) + + with ( + patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + ): + result = await MCPRequestHandler.get_allowed_tools_for_server( + server_id="server-a", + user_api_key_auth=user_api_key_auth, + ) + + assert result is None + async def test_permission_inheritance_edge_cases(self): """Test edge cases in permission inheritance""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py new file mode 100644 index 000000000000..cb27e992ecbd --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -0,0 +1,240 @@ +"""Classification and rendering matrix for per-server tools/list outcomes: every failure mode maps +to exactly one category, wire values never carry upstream prose, and single-upstream HTTP statuses +stay truthful to who failed.""" + +import httpx +import pytest + +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPServerListError, + MCPUpstreamAuthError, +) +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + ServerListFault, + ServerListOk, + classify_list_exception, + list_fault_http_status, + outcome_wire_value, +) + + +def test_carried_fault_passes_through(): + fault = ServerListFault(tag="timeout") + assert classify_list_exception(MCPServerListError(fault, "srv")) is fault + + +def test_upstream_auth_error_maps_to_auth_required_and_forbidden(): + assert classify_list_exception(MCPUpstreamAuthError(401, None, "srv")).tag == "auth_required" + assert classify_list_exception(MCPUpstreamAuthError(403, None, "srv")).tag == "forbidden" + + +def test_timeout_and_connection_errors_classify_without_status(): + assert classify_list_exception(TimeoutError()).tag == "timeout" + assert classify_list_exception(ConnectionError()).tag == "unreachable" + + +def test_embedded_upstream_response_status_wins(): + response = httpx.Response(503, request=httpx.Request("POST", "https://mcp.example.com/mcp")) + exc = httpx.HTTPStatusError("boom", request=response.request, response=response) + wrapped = RuntimeError("wrapper") + wrapped.__cause__ = exc + fault = classify_list_exception(wrapped) + assert fault.tag == "upstream_error" + assert fault.status_code == 503 + + +def test_embedded_401_classifies_auth_required(): + response = httpx.Response(401, request=httpx.Request("POST", "https://mcp.example.com/mcp")) + exc = httpx.HTTPStatusError("no", request=response.request, response=response) + assert classify_list_exception(exc).tag == "auth_required" + + +def test_context_response_does_not_shadow_the_causal_chain_response(): + real_response = httpx.Response(401, request=httpx.Request("POST", "https://mcp.example.com/mcp")) + real = httpx.HTTPStatusError("upstream rejected", request=real_response.request, response=real_response) + incidental_response = httpx.Response(500, request=httpx.Request("POST", "https://hooks.example.com/log")) + incidental = httpx.HTTPStatusError( + "logging hook failed", request=incidental_response.request, response=incidental_response + ) + wrapper = RuntimeError("wrapper") + wrapper.__cause__ = real + wrapper.__context__ = incidental + fault = classify_list_exception(wrapper) + assert fault.tag == "auth_required" + assert fault.status_code == 401 + + +def test_exception_group_members_are_searched_in_raise_order(): + first_response = httpx.Response(502, request=httpx.Request("POST", "https://mcp.example.com/mcp")) + first = httpx.HTTPStatusError("first", request=first_response.request, response=first_response) + second_response = httpx.Response(503, request=httpx.Request("POST", "https://mcp.example.com/mcp")) + second = httpx.HTTPStatusError("second", request=second_response.request, response=second_response) + fault = classify_list_exception(BaseExceptionGroup("task group", [first, second])) + assert fault.tag == "upstream_error" + assert fault.status_code == 502 + + +def test_unknown_exception_is_internal(): + assert classify_list_exception(ValueError("who knows")).tag == "internal" + + +def test_wire_value_carries_no_prose(): + fault = ServerListFault(tag="upstream_error", status_code=500) + assert outcome_wire_value(fault) == {"status": "upstream_error", "http_status": 500} + assert outcome_wire_value(ServerListOk(tool_count=7)) == {"status": "ok", "tool_count": 7} + assert outcome_wire_value(ServerListFault(tag="timeout")) == {"status": "timeout"} + + +@pytest.mark.parametrize( + "tag,status_code,expected", + [ + ("auth_required", 401, 401), + ("auth_required", None, 401), + ("forbidden", 403, 403), + ("timeout", None, 504), + ("unreachable", None, 502), + ("upstream_error", 500, 502), + ("internal", None, 500), + ], +) +def test_single_upstream_http_status_is_truthful(tag, status_code, expected): + assert list_fault_http_status(ServerListFault(tag=tag, status_code=status_code)) == expected + + +@pytest.mark.asyncio +async def test_cancelled_fetch_is_a_classified_fault_not_a_healthy_empty_server(): + """A cancelled per-server fetch must not masquerade as ok(tool_count=0): cancellation was already + suppressed before the outcome plumbing existed, so it stays suppressed, but as an internal fault + the outcome reporting can see.""" + import asyncio + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + manager = MCPServerManager() + client = MagicMock() + client.list_tools = AsyncMock(side_effect=asyncio.CancelledError()) + + with pytest.raises(MCPServerListError) as exc_info: + await manager._fetch_tools_with_timeout(client, "cancelled_srv") + + assert exc_info.value.fault.tag == "internal" + + +def test_auth_challenge_and_status_come_from_the_causal_response(): + """An incidental 403 raised while handling the causal 401 (context chain) must not shadow it: + the carrier channel and the challenge both derive from the response on the explicit causal + chain, so the caller is challenged to authenticate rather than told it is forbidden.""" + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import upstream_auth_challenge + + causal = httpx.HTTPStatusError( + "auth", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response( + 401, + headers={"www-authenticate": 'Bearer resource_metadata="https://mcp.example.com/.well-known"'}, + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + ), + ) + incidental = httpx.HTTPStatusError( + "hook", + request=httpx.Request("POST", "https://hook.example.com/log"), + response=httpx.Response(403, request=httpx.Request("POST", "https://hook.example.com/log")), + ) + wrapper = RuntimeError("fetch failed") + wrapper.__cause__ = causal + wrapper.__context__ = incidental + + result = upstream_auth_challenge(wrapper) + assert result is not None + status_code, challenge = result + assert status_code == 401 + assert challenge == 'Bearer resource_metadata="https://mcp.example.com/.well-known"' + + +def test_raise_classified_list_failure_routes_auth_to_upstream_auth_error(): + """The single choice-point sends 401/403 through MCPUpstreamAuthError with the upstream's own + challenge and everything else through MCPServerListError, so fetch sites cannot drift.""" + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import raise_classified_list_failure + + auth_exc = httpx.HTTPStatusError( + "auth", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response( + 401, + headers={"www-authenticate": "Bearer realm=x"}, + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + ), + ) + with pytest.raises(MCPUpstreamAuthError) as auth_info: + raise_classified_list_failure(auth_exc, "srv") + assert auth_info.value.status_code == 401 + assert auth_info.value.www_authenticate == "Bearer realm=x" + + with pytest.raises(MCPServerListError) as fault_info: + raise_classified_list_failure(RuntimeError("boom"), "srv") + assert fault_info.value.fault.tag == "internal" + + +def test_causal_auth_behind_unrelated_response_is_still_found(): + """The auth scan must not end at the first response of any status: a causal 401 sitting deeper + in the tree than an unrelated 5xx (retry attempts, multi-stream task groups) must still surface + with its challenge, or the client is told upstream_error and never re-authenticates.""" + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import upstream_auth_challenge + + deep_auth = httpx.HTTPStatusError( + "auth", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response( + 401, + headers={"www-authenticate": "Bearer realm=upstream"}, + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + ), + ) + earlier_5xx = httpx.HTTPStatusError( + "flaky attempt", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response(500, request=httpx.Request("POST", "https://mcp.example.com/mcp")), + ) + earlier_5xx.__cause__ = deep_auth + wrapper = RuntimeError("fetch failed") + wrapper.__cause__ = earlier_5xx + + result = upstream_auth_challenge(wrapper) + assert result is not None + assert result == (401, "Bearer realm=upstream") + + +def test_classification_agrees_with_auth_scan_on_nested_auth(): + """classify_list_exception derives its auth arm from the same scan as the carrier choice-point, + so a nested 401 behind a 5xx classifies auth_required, never upstream_error(500).""" + deep_auth = httpx.HTTPStatusError( + "auth", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response(401, request=httpx.Request("POST", "https://mcp.example.com/mcp")), + ) + earlier_5xx = httpx.HTTPStatusError( + "flaky attempt", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response(500, request=httpx.Request("POST", "https://mcp.example.com/mcp")), + ) + earlier_5xx.__cause__ = deep_auth + wrapper = RuntimeError("fetch failed") + wrapper.__cause__ = earlier_5xx + + fault = classify_list_exception(wrapper) + assert fault.tag == "auth_required" + assert fault.status_code == 401 + + +def test_pure_non_auth_response_still_classifies_upstream_error(): + """With no auth response anywhere in the tree, the first response in deliberate order still + drives the generic upstream_error classification.""" + exc = httpx.HTTPStatusError( + "boom", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response(502, request=httpx.Request("POST", "https://mcp.example.com/mcp")), + ) + fault = classify_list_exception(exc) + assert fault.tag == "upstream_error" + assert fault.status_code == 502 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 16c36af5156d..fdc77d19d73c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -273,6 +273,7 @@ def _http_server(server_id: str, name: str, **kwargs) -> MCPServer: return MCPServer( server_id=server_id, name=name, + alias=name, url=f"https://{name}/mcp", transport=MCPTransport.http, **kwargs, @@ -284,7 +285,8 @@ async def test_aggregate_list_tools_absorbs_one_unauthenticated_server(): """Regression: across the aggregate (/mcp), a delegate/passthrough server that raises MCPUpstreamAuthError must not empty every other server's tools. Re-raising it on the aggregate path (introduced with the passthrough feature) zeroed the whole list because the - fan-out gather propagated it.""" + fan-out gather propagated it. The failed server now contributes an "auth_required" outcome + instead of vanishing, so it stays distinguishable from a healthy server with no tools.""" from unittest.mock import patch from mcp.types import Tool as MCPTool @@ -312,22 +314,24 @@ async def fake_get_tools(server, **kwargs): ), patch.object( mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) ): - tools = await mcp_server._get_tools_from_mcp_servers( + listing = await mcp_server._get_tools_from_mcp_servers( user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"), mcp_auth_header=None, mcp_servers=None, ) - assert [t.name for t in tools] == ["working_docs-read"] + assert [t.name for t in listing.tools] == ["working_docs-read"] + assert listing.outcomes["delegate_docs"].tag == "auth_required" + assert listing.outcomes["working_docs"].tag == "ok" @pytest.mark.asyncio async def test_single_server_route_also_absorbs_upstream_auth_error(): """A single-server route (//mcp) absorbs an upstream-auth error just like the aggregate: - the failing server is omitted (empty list) rather than re-raised. Surfacing it to the client as a - 401 + WWW-Authenticate challenge cannot be done from this list handler — the MCP session manager - serializes a raise into a JSON-RPC error, not an HTTP 401 — so re-auth surfacing is handled by a - request-scope preemptive check, tracked separately.""" + the failing server contributes no tools and an "auth_required" outcome rather than re-raising. + Surfacing it to the client as a 401 + WWW-Authenticate challenge cannot be done from this list + handler — the MCP session manager serializes a raise into a JSON-RPC error, not an HTTP 401 — so + re-auth surfacing is handled by a request-scope preemptive check, tracked separately.""" from unittest.mock import patch from litellm.proxy._experimental.mcp_server import server as mcp_server @@ -351,12 +355,13 @@ async def fake_get_tools(server, **kwargs): ), patch.object( mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) ): - tools = await mcp_server._get_tools_from_mcp_servers( + listing = await mcp_server._get_tools_from_mcp_servers( user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"), mcp_auth_header=None, mcp_servers=["delegate_docs"], ) - assert tools == [] + assert listing.tools == [] + assert listing.outcomes["delegate_docs"].tag == "auth_required" finally: _mcp_gateway_server_name.reset(token) @@ -388,10 +393,11 @@ async def fake_get_tools(server, **kwargs): mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) ): # Aggregate route: no explicit server filter, even though only one server is accessible. - tools = await mcp_server._get_tools_from_mcp_servers( + listing = await mcp_server._get_tools_from_mcp_servers( user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"), mcp_auth_header=None, mcp_servers=None, ) - assert tools == [] + assert listing.tools == [] + assert listing.outcomes["delegate_docs"].tag == "auth_required" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 7e59904a39c4..ae4f12fc1e17 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1094,8 +1094,10 @@ async def mock_get_tools_from_server( ) # Verify that tools from the working server are returned - assert len(result) == 1 - assert result[0].name == "working_tool_1" + assert len(result.tools) == 1 + assert result.tools[0].name == "working_tool_1" + assert result.outcomes["working"].tag == "ok" + assert result.outcomes["failing"].tag == "internal" # Verify failure logging mock_logger.exception.assert_any_call( @@ -1188,7 +1190,9 @@ async def mock_get_tools_from_server( ) # Verify that empty list is returned - assert len(result) == 0 + assert len(result.tools) == 0 + assert result.outcomes["failing1"].tag == "internal" + assert result.outcomes["failing2"].tag == "internal" # Verify failure logging for both servers mock_logger.exception.assert_any_call( @@ -3074,7 +3078,7 @@ async def mock_get_tools_from_server( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", mock_manager, ): - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=None, @@ -3082,8 +3086,8 @@ async def mock_get_tools_from_server( ) # Server prefix is always added regardless of number of allowed servers - assert len(tools) == 1 - assert tools[0].name == "zapier-toolA" + assert len(listing.tools) == 1 + assert listing.tools[0].name == "zapier-toolA" @pytest.mark.asyncio @@ -3153,7 +3157,7 @@ async def mock_get_tools_from_server( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", mock_manager, ): - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=None, @@ -3161,7 +3165,7 @@ async def mock_get_tools_from_server( ) # Should be prefixed since multiple servers are allowed - names = sorted([t.name for t in tools]) + names = sorted([t.name for t in listing.tools]) assert names == ["jira-toolA", "zapier-toolA"] @@ -3437,7 +3441,7 @@ async def mock_get_tools_from_server( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", mock_manager, ): - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=None, @@ -3445,8 +3449,8 @@ async def mock_get_tools_from_server( ) # Should only return tool1 and tool2 - assert len(tools) == 2 - tool_names = sorted([t.name for t in tools]) + assert len(listing.tools) == 2 + tool_names = sorted([t.name for t in listing.tools]) assert tool_names == ["tool1", "tool2"] @@ -3553,7 +3557,7 @@ async def mock_get_tools_from_server( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_team_object_permission", AsyncMock(return_value=team_object_permission), ): - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=None, @@ -3561,8 +3565,8 @@ async def mock_get_tools_from_server( ) # Should only return tool2 and tool3 (intersection of key and team permissions) - assert len(tools) == 2 - tool_names = sorted([t.name for t in tools]) + assert len(listing.tools) == 2 + tool_names = sorted([t.name for t in listing.tools]) assert tool_names == ["tool2", "tool3"] @@ -3640,7 +3644,7 @@ async def mock_get_tools_from_server( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", mock_manager, ): - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=None, @@ -3648,8 +3652,8 @@ async def mock_get_tools_from_server( ) # Should return all tools when no restrictions - assert len(tools) == 3 - tool_names = sorted([t.name for t in tools]) + assert len(listing.tools) == 3 + tool_names = sorted([t.name for t in listing.tools]) assert tool_names == ["tool1", "tool2", "tool3"] @@ -3746,7 +3750,7 @@ async def mock_get_tools_from_server( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", mock_manager, ): - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=None, @@ -3754,8 +3758,8 @@ async def mock_get_tools_from_server( ) # Should only return the 2 tools that match (after stripping prefix) - assert len(tools) == 2 - tool_names = sorted([t.name for t in tools]) + assert len(listing.tools) == 2 + tool_names = sorted([t.name for t in listing.tools]) # Tools still have prefixes in the output, but were filtered correctly assert tool_names == [ "GITMCP-fetch_litellm_documentation", @@ -4278,7 +4282,7 @@ def _capture_function_setup(*_args, **kwargs): ): mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1]) - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_auth, mcp_auth_header=None, mcp_servers=["server_a"], @@ -4288,7 +4292,7 @@ def _capture_function_setup(*_args, **kwargs): request_tags=["team-a"], ) - assert tools == [tool_1] + assert listing.tools == [tool_1] dummy_logging_obj.async_success_handler.assert_awaited_once() assert dummy_logging_obj.async_success_handler.await_args.kwargs["result"] == [tool_1.model_dump(mode="json")] assert function_setup_kwargs["metadata"]["tags"] == ["team-a"] @@ -4297,6 +4301,7 @@ def _capture_function_setup(*_args, **kwargs): assert spend_meta["tool_count_total"] == 1 assert spend_meta["allowed_server_count"] == 1 assert spend_meta["per_server_tool_counts"]["server_a"] == 1 + assert spend_meta["per_server_list_outcomes"] == {"server_a": {"status": "ok", "tool_count": 1}} @pytest.mark.asyncio @@ -4359,7 +4364,7 @@ async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fai ): mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1]) - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_auth, mcp_auth_header=None, mcp_servers=["server_a"], @@ -4368,7 +4373,7 @@ async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fai list_tools_log_source="mcp_protocol", ) - assert tools == [tool_1] + assert listing.tools == [tool_1] dummy_logging_obj.async_success_handler.assert_awaited_once() @@ -4665,7 +4670,7 @@ async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token(): ): mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1]) - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_auth, mcp_auth_header=None, mcp_servers=["atlassian_test"], @@ -4681,7 +4686,7 @@ async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token(): call_kwargs = mock_manager._get_tools_from_server.await_args.kwargs assert call_kwargs["extra_headers"] == {"Authorization": f"Bearer {STORED_TOKEN}"} - assert tools == [tool_1] + assert listing.tools == [tool_1] # --------------------------------------------------------------------------- @@ -5207,7 +5212,7 @@ async def capture_extra_headers(*args, **kwargs): mock_manager.filter_server_ids_by_ip_with_info = MagicMock(return_value=(["legacy-m2m-id"], 0)) mock_manager._get_tools_from_server = AsyncMock(side_effect=capture_extra_headers) - tools = await _get_tools_from_mcp_servers( + listing = await _get_tools_from_mcp_servers( user_api_key_auth=user_auth, mcp_auth_header=None, mcp_servers=["legacy_m2m"], @@ -5222,7 +5227,7 @@ async def capture_extra_headers(*args, **kwargs): "P1 security issue: caller's Authorization header was forwarded to M2M server. " "Expected None, got: " + str(captured_extra_headers) ) - assert tools == [tool_1] + assert listing.tools == [tool_1] @pytest.mark.asyncio @@ -7439,6 +7444,150 @@ async def test_call_mcp_tool_skips_failure_hook_for_upstream_auth_error(): proxy_logging_mock.post_call_failure_hook.assert_not_awaited() +@pytest.mark.asyncio +async def test_aggregate_listing_reports_per_server_outcomes(): + """A failed server must contribute a classified outcome, not just silently shrink the list: + without the outcome a broken upstream is indistinguishable from a healthy server with no tools.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + set_auth_context, + ) + except ImportError: + pytest.skip("MCP server not available") + + from litellm.proxy._experimental.mcp_server.exceptions import MCPServerListError + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListFault + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + set_auth_context(user_api_key_auth) + + working_server = MagicMock() + working_server.name = "working_server" + working_server.alias = "working" + working_server.allowed_tools = None + working_server.disallowed_tools = None + working_server.server_id = "working_server" + working_server.server_name = "working_server" + working_server.auth_type = None + working_server.extra_headers = None + + broken_server = MagicMock() + broken_server.name = "broken_server" + broken_server.alias = "broken" + broken_server.allowed_tools = None + broken_server.disallowed_tools = None + broken_server.server_id = "broken_server" + broken_server.server_name = "broken_server" + broken_server.auth_type = None + broken_server.extra_headers = None + + mock_manager = MagicMock() + mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["working_server", "broken_server"]) + mock_manager.get_mcp_server_by_id = lambda server_id: ( + working_server if server_id == "working_server" else broken_server + ) + mock_manager.filter_server_ids_by_ip_with_info = lambda server_ids, client_ip: (server_ids, 0) + + async def mock_get_tools_from_server(server, **kwargs): + if server.name == "working_server": + tool1 = MagicMock() + tool1.name = "working_tool_1" + tool1.description = "Working tool 1" + tool1.inputSchema = {} + return [tool1] + raise MCPServerListError(ServerListFault(tag="upstream_error", status_code=500), server.name) + + mock_manager._get_tools_from_server = mock_get_tools_from_server + + with patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + mock_manager, + ): + listing = await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=None, + mcp_servers=["working_server", "broken_server"], + mcp_server_auth_headers=None, + ) + + assert [tool.name for tool in listing.tools] == ["working_tool_1"] + assert listing.outcomes["working"].tag == "ok" + assert listing.outcomes["working"].tool_count == 1 + assert listing.outcomes["broken"].tag == "upstream_error" + assert listing.outcomes["broken"].status_code == 500 + assert "working_server" not in listing.outcomes + assert "broken_server" not in listing.outcomes + + +@pytest.mark.asyncio +async def test_outcome_keys_use_display_prefix_never_canonical_names(): + """Outcome keys are client-visible and must use the same display naming (alias or short prefix) + the caller already sees on tool names: keying them by canonical server_name would let any + authenticated caller enumerate internal server names the alias scheme deliberately hides.""" + try: + from litellm.proxy._experimental.mcp_server.server import _aggregate_server_key + except ImportError: + pytest.skip("MCP server not available") + + server = MagicMock() + server.alias = "public-alias" + server.server_name = "internal-canonical-name" + server.name = "internal-canonical-name" + server.short_prefix = None + server.server_id = "srv-1" + + key = _aggregate_server_key(server) + assert key == "public-alias" + assert "internal-canonical-name" not in key + + +@pytest.mark.asyncio +async def test_handle_list_tools_attaches_outcome_meta(): + """The protocol handler returns a ListToolsResult whose _meta carries the per-server outcomes, + so MCP clients can tell a degraded listing from a genuinely empty one.""" + try: + from litellm.proxy._experimental.mcp_server.server import handle_list_tools + except ImportError: + pytest.skip("MCP server not available") + + from mcp.types import ListToolsResult, Tool + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListFault, + ServerListOk, + ) + + tool = Tool(name="t1", inputSchema={"type": "object"}) + listing = AggregateToolListing( + tools=[tool], + outcomes={"healthy": ServerListOk(tool_count=1), "broken": ServerListFault(tag="unreachable")}, + ) + + async def fake_auth_context(): + return (None, None, None, None, None, None, None) + + 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.server._list_mcp_tools", + new=AsyncMock(return_value=listing), + ), + ): + result = await handle_list_tools() + + assert isinstance(result, ListToolsResult) + wire = result.model_dump(by_alias=True) + outcomes_meta = wire["_meta"][SERVER_OUTCOMES_META_KEY] + assert outcomes_meta["healthy"] == {"status": "ok", "tool_count": 1} + assert outcomes_meta["broken"] == {"status": "unreachable"} + + def _make_oauth2_server( alias: str, *, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 80bf08a5ebae..55b6bbbdbc29 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -11,7 +11,11 @@ import pytest from fastapi import HTTPException -from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPServerListError, + MCPUpstreamAuthError, +) +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListFault # Add the parent directory to the path so we can import litellm sys.path.insert(0, "../../../../../") @@ -904,9 +908,12 @@ async def test_list_surfaces_resolver_401_as_upstream_auth_error(self): assert exc_info.value.www_authenticate == challenge @pytest.mark.asyncio - async def test_list_absorbs_non_auth_httpexception(self): - """A non-auth HTTP error (e.g. 412 no endpoint, 503 IdP down) must stay absorbed to [] so one - misconfigured/unavailable server does not blank the whole aggregate listing.""" + async def test_list_surfaces_non_auth_httpexception_as_internal_fault(self): + """A non-auth HTTP error (e.g. 412 no endpoint, 503 IdP down) now raises MCPServerListError + with an "internal" fault carrying the status code instead of absorbing to []: the silent + empty list made a misconfigured/unavailable server indistinguishable from a healthy server + with no tools. The aggregate absorbs it into that server's outcome, so one broken server + still does not blank the whole aggregate listing.""" server = MCPServer( server_id="te-412", name="te-412-server", @@ -921,10 +928,10 @@ async def test_list_absorbs_non_auth_httpexception(self): manager._create_mcp_client = AsyncMock( side_effect=HTTPException(status_code=412, detail="token exchange endpoint is not configured") ) - result = await manager._get_tools_from_server( - server=server, oauth2_headers={"Authorization": "Bearer subj-jwt"} - ) - assert result == [] + with pytest.raises(MCPServerListError) as exc_info: + await manager._get_tools_from_server(server=server, oauth2_headers={"Authorization": "Bearer subj-jwt"}) + assert exc_info.value.fault == ServerListFault(tag="internal", status_code=412) + assert exc_info.value.server_name == "te-412-server" def _upstream_status_error(self, status_code: int, www_authenticate: Optional[str] = None) -> httpx.HTTPStatusError: """Build an httpx.HTTPStatusError shaped like the one the MCP SDK surfaces for an upstream @@ -7751,14 +7758,16 @@ def _upstream_status_error(status_code: int, challenge: str) -> httpx.HTTPStatus class TestMCPToolsListAuthSurfacing: - """Regression: MCP tools/list 401 auth failures must surface as MCPUpstreamAuthError. + """Regression: MCP tools/list failures must surface as typed exceptions, never a silent []. Previously a missing/expired per-user OAuth token, or an upstream 401 for any non-carveout auth_type, was swallowed to an empty tool list, so a single-server client saw a 200 with no tools instead of a 401 challenge. The listing helpers - now raise MCPUpstreamAuthError on a 401 regardless of auth_type; the single-server - routes turn it into a 401 + WWW-Authenticate while the aggregator absorbs it to an - empty list. Only a 401 challenges; a 403 (forbidden) degrades like any other error. + now raise MCPUpstreamAuthError on an upstream 401 or 403 and MCPServerListError + with a classified fault for every other failure; single-server routes relay a + truthful HTTP status while the aggregator absorbs each failure into that + server's outcome, so a broken upstream is never indistinguishable from a + healthy server with no tools. """ @pytest.mark.asyncio @@ -7780,25 +7789,38 @@ async def test_fetch_tools_with_timeout_surfaces_upstream_401(self): assert exc_info.value.server_name == "static-key-server" @pytest.mark.asyncio - async def test_fetch_tools_with_timeout_absorbs_upstream_403(self): - """Only a 401 drives the re-auth challenge. A 403 (authenticated but - forbidden, e.g. insufficient scope) is not a re-auth signal, so even - with a WWW-Authenticate header it degrades to an empty list rather than - surfacing a challenge.""" + async def test_fetch_tools_with_timeout_surfaces_upstream_403(self): + """An upstream 403 (authenticated but forbidden, e.g. insufficient scope) now raises + MCPUpstreamAuthError instead of absorbing to []: the silent empty list made a forbidden + upstream indistinguishable from a healthy server with no tools. The upstream + WWW-Authenticate is preserved so single-server routes can relay the real challenge.""" manager = MCPServerManager() challenge = 'Bearer error="insufficient_scope", scope="read:tools"' client = MagicMock() client.list_tools = AsyncMock(side_effect=_upstream_status_error(403, challenge)) - assert await manager._fetch_tools_with_timeout(client, "forbidden-server") == [] + with pytest.raises(MCPUpstreamAuthError) as exc_info: + await manager._fetch_tools_with_timeout(client, "forbidden-server") + + assert exc_info.value.status_code == 403 + assert exc_info.value.www_authenticate == challenge + assert exc_info.value.server_name == "forbidden-server" @pytest.mark.asyncio - async def test_fetch_tools_with_timeout_returns_empty_on_non_auth_error(self): + async def test_fetch_tools_with_timeout_raises_classified_fault_on_non_auth_error(self): + """A non-auth listing failure now raises MCPServerListError carrying a classified fault + instead of absorbing to []: the silent empty list made a broken upstream indistinguishable + from a healthy server with no tools. An unrecognized exception classifies as the gateway's + own fault ("internal").""" manager = MCPServerManager() client = MagicMock() client.list_tools = AsyncMock(side_effect=RuntimeError("upstream 500")) - assert await manager._fetch_tools_with_timeout(client, "srv") == [] + with pytest.raises(MCPServerListError) as exc_info: + await manager._fetch_tools_with_timeout(client, "srv") + + assert exc_info.value.fault == ServerListFault(tag="internal") + assert exc_info.value.server_name == "srv" @pytest.mark.asyncio async def test_get_tools_from_server_surfaces_unusable_user_token(self): @@ -7825,9 +7847,12 @@ async def test_get_tools_from_server_surfaces_unusable_user_token(self): assert exc_info.value.server_name == "oauth-srv" @pytest.mark.asyncio - async def test_get_tools_from_server_absorbs_non_challenge_http_error(self): - """A non-auth HTTPException (500) stays absorbed so one misconfigured server cannot blank - the listing; 401/403 are the challenge-class statuses routed to MCPUpstreamAuthError.""" + async def test_get_tools_from_server_surfaces_non_challenge_http_error_as_internal_fault(self): + """A non-auth HTTPException (500) now raises MCPServerListError with an "internal" fault + carrying the status code instead of absorbing to []: the silent empty list made a + misconfigured server indistinguishable from a healthy server with no tools. The aggregate + absorbs it into that server's outcome; single-server routes relay a truthful status. + 401/403 remain the challenge-class statuses routed to MCPUpstreamAuthError.""" manager = MCPServerManager() server = MCPServer(server_id="stdio-srv", name="stdio-srv", transport=MCPTransport.http) manager._create_mcp_client = AsyncMock( @@ -7837,7 +7862,83 @@ async def test_get_tools_from_server_absorbs_non_challenge_http_error(self): ) ) - assert await manager._get_tools_from_server(server) == [] + with pytest.raises(MCPServerListError) as exc_info: + await manager._get_tools_from_server(server) + + assert exc_info.value.fault == ServerListFault(tag="internal", status_code=500) + assert exc_info.value.server_name == "stdio-srv" + + @pytest.mark.asyncio + async def test_get_tools_from_server_generic_arm_extracts_nested_auth_challenge(self): + """A 401 buried in the exception tree at client-build time must travel the same channel as + one raised during the fetch: MCPUpstreamAuthError with the upstream's own challenge. Before + the shared choice-point it classified into a challenge-less fault, so single-server routes + answered 401 without the WWW-Authenticate the client needs to start the OAuth flow.""" + import httpx + + from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPUpstreamAuthError, + ) + + manager = MCPServerManager() + server = MCPServer(server_id="nested-srv", name="nested-srv", transport=MCPTransport.http) + causal = httpx.HTTPStatusError( + "auth", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response( + 401, + headers={"www-authenticate": "Bearer realm=upstream"}, + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + ), + ) + wrapper = RuntimeError("client build failed") + wrapper.__cause__ = causal + manager._create_mcp_client = AsyncMock(side_effect=wrapper) + + with pytest.raises(MCPUpstreamAuthError) as exc_info: + await manager._get_tools_from_server(server) + + assert exc_info.value.status_code == 401 + assert exc_info.value.www_authenticate == "Bearer realm=upstream" + + @pytest.mark.asyncio + async def test_get_tools_from_server_generic_arm_strips_challenge_for_dcr_bridge(self): + """The dcr_bridge challenge suppression must hold on the generic arm too, not only when the + fetch itself raised MCPUpstreamAuthError: a bridge client following the upstream challenge + would fail the RFC 9728 resource match against the gateway URL it dialed.""" + import httpx + + from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPUpstreamAuthError, + ) + from litellm.types.mcp import MCPAuth + + manager = MCPServerManager() + bridge_server = MCPServer( + server_id="bridge-nested", + name="bridge-nested", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + dcr_bridge=True, + ) + causal = httpx.HTTPStatusError( + "auth", + request=httpx.Request("POST", "https://upstream.example/mcp"), + response=httpx.Response( + 401, + headers={"www-authenticate": 'Bearer resource_metadata="https://upstream.example/.wk"'}, + request=httpx.Request("POST", "https://upstream.example/mcp"), + ), + ) + wrapper = RuntimeError("client build failed") + wrapper.__cause__ = causal + manager._create_mcp_client = AsyncMock(side_effect=wrapper) + + with pytest.raises(MCPUpstreamAuthError) as exc_info: + await manager._get_tools_from_server(bridge_server) + + assert exc_info.value.status_code == 401 + assert exc_info.value.www_authenticate is None @pytest.mark.asyncio async def test_get_tools_from_server_suppresses_upstream_challenge_for_dcr_bridge(self): @@ -8496,3 +8597,34 @@ def test_build_mcp_server_table_carries_null_oauth2_flow(): table = manager._build_mcp_server_table(server) assert table.oauth2_flow is None + + +@pytest.mark.asyncio +async def test_resolve_toolset_tool_permissions_single_db_fetch_across_checks(): + """The server-level and tool-level permission primitives each resolve the + key's toolsets during one request; the shared cache must dedupe the DB + fetch so the request costs a single toolset query however many checks run""" + from litellm.caching.caching import DualCache + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + manager = MCPServerManager() + toolset = MagicMock() + toolset.tools = [{"server_id": "server-a", "tool_name": "lookup_status"}] + list_toolsets_mock = AsyncMock(return_value=[toolset]) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.toolset_db.list_mcp_toolsets", + list_toolsets_mock, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + ): + first = await manager.resolve_toolset_tool_permissions(toolset_ids=["ts-1"]) + second = await manager.resolve_toolset_tool_permissions(toolset_ids=["ts-1"]) + + assert first == {"server-a": ["lookup_status"]} + assert second == first + list_toolsets_mock.assert_awaited_once() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index b8f0b205831d..1da44029b5c2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -16,6 +16,7 @@ import pytest from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.proxy._experimental.mcp_server.tool_search import ( MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, @@ -381,7 +382,7 @@ async def test_mcp_tool_search_call_returns_tool_defs(self) -> None: with patch( "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", new_callable=AsyncMock, - return_value=[mock_tool], + return_value=AggregateToolListing(tools=[mock_tool], outcomes={}), ): result = await self._get_call_fn()( request=request, @@ -508,7 +509,7 @@ async def test_mcp_tool_search_forwards_client_ip_for_ip_filtering(self) -> None patch( "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", new_callable=AsyncMock, - return_value=[], + return_value=AggregateToolListing(tools=[], outcomes={}), ) as mock_list, ): await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 99e051823611..d4ba66c43811 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -690,16 +690,16 @@ async def fake_get_allowed_mcp_servers(*args, **kwargs): ) request = _build_request(path="/mcp-rest/tools/list", method="GET") - result = await rest_endpoints.list_tool_rest_api( - request, - server_id="server-1", - user_api_key_dict=UserAPIKeyAuth(), - ) + with pytest.raises(HTTPException) as exc_info: + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + user_api_key_dict=UserAPIKeyAuth(), + ) - assert result["tools"] == [] - assert result["error"] == "unexpected_error" - assert "access_denied" in result["message"] - assert "server server-1" in result["message"] + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "access_denied" + assert "server-1" in exc_info.value.detail["message"] async def test_lists_tools_for_allowed_server(self, monkeypatch): async def fake_contexts(user_api_key_auth): @@ -911,6 +911,64 @@ async def fake_get_tools(*args, **kwargs): assert exc_info.value.status_code == upstream_status assert exc_info.value.headers == {"www-authenticate": challenge} + async def test_single_server_upstream_fault_surfaces_truthful_status(self, monkeypatch): + """A single-server listing whose upstream breaks (5xx, timeout, unreachable) must answer + with the truthful gateway status instead of masking the failure as an empty-success + {"tools": [], "error": null} body a caller cannot distinguish from a toolless server.""" + from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPServerListError, + ) + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + ServerListFault, + ) + + class StubServer: + alias = "server-1" + server_name = "server-1" + name = "flaky" + allowed_tools = None + mcp_info = {"server_name": "flaky"} + available_on_public_internet = True + + stub_server = StubServer() + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + async def fake_get_tools(*args, **kwargs): + raise MCPServerListError(ServerListFault(tag="upstream_error", status_code=503), "flaky") + + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + with pytest.raises(HTTPException) as exc_info: + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.detail["error"] == "upstream_error" + assert "server-1" in exc_info.value.detail["message"] + assert "flaky" not in exc_info.value.detail["message"] + async def test_aggregate_list_absorbs_one_server_auth_failure(self, monkeypatch): """The multi-server aggregate listing degrades a server whose upstream rejects auth to an empty contribution and still returns the healthy @@ -1108,15 +1166,15 @@ async def fake_get_allowed_mcp_servers(*args, **kwargs): ) request = _build_request(path="/mcp-rest/tools/list", method="GET") - result = await rest_endpoints.list_tool_rest_api( - request, - server_id="restricted-server", - user_api_key_dict=UserAPIKeyAuth(), - ) + with pytest.raises(HTTPException) as exc_info: + await rest_endpoints.list_tool_rest_api( + request, + server_id="restricted-server", + user_api_key_dict=UserAPIKeyAuth(), + ) - assert result["tools"] == [] - assert result["error"] == "unexpected_error" - assert "access_denied" in result["message"] + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "access_denied" async def test_mcp_server_name_query_param_resolves_to_server(self, monkeypatch): """mcp_server_name is a name-based alias for server_id: it should @@ -2654,3 +2712,74 @@ def test_alias_none_is_explicit_in_mcp_info(self): "server_id": "server-uuid", "alias": None, } + + +class TestRestListToolsetFiltering: + @pytest.mark.asyncio + async def test_rest_list_filters_toolset_only_key_to_toolset_tools(self, monkeypatch): + """A toolset-only key reaching a toolset server via REST list must see + only the toolset's tools; the raw catalog leaked every tool on the + server when the filter read object_permission directly instead of the + shared toolset-aware primitive""" + from unittest.mock import patch + + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + stub_server = MCPServer( + server_id="server-a", + name="stubtools", + transport=MCPTransport.http, + ) + stub_server.alias = "stubtools" + stub_server.server_name = "stubtools" + stub_server.allowed_tools = None + stub_server.disallowed_tools = None + stub_server.mcp_info = {"server_name": "stubtools"} + + upstream_tools = [ + MCPTool(name="lookup_status", inputSchema={"type": "object"}), + MCPTool(name="delete_everything", inputSchema={"type": "object"}), + ] + + key_object_permission = MagicMock() + key_object_permission.mcp_servers = [] + key_object_permission.mcp_access_groups = [] + key_object_permission.mcp_tool_permissions = None + key_object_permission.mcp_toolsets = ["toolset-1"] + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + mock_manager = MagicMock() + mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {}) + mock_manager.resolve_toolset_tool_permissions = AsyncMock( + return_value={"server-a": ["lookup_status"]} + ) + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + AsyncMock(return_value=upstream_tools), + ) + + with ( + patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission), + patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + ): + result = await rest_endpoints._get_tools_for_single_server( + server=stub_server, + server_auth_header=None, + raw_headers=None, + user_api_key_auth=user_auth, + ) + + assert [tool.name for tool in result] == ["lookup_status"] diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index cc4a7d5bfb4e..2da645bf4e1a 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2047,6 +2047,88 @@ async def test_common_checks_metadata_route_keeps_key_tags_out_of_provider_metad assert "metadata" not in request_body +def _pass_through_request() -> "Request": + """A Request whose FastAPI-resolved endpoint carries the pass-through marker, + i.e. the request was dispatched to a user-defined pass-through handler.""" + from fastapi import Request + + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_ENDPOINT_MARKER, + ) + + def pass_through_endpoint(): + ... + + setattr(pass_through_endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) + return Request(scope={"type": "http", "headers": [], "endpoint": pass_through_endpoint}) + + +def _builtin_request() -> "Request": + """A Request dispatched to a built-in (non-pass-through) handler, e.g. what a + custom path colliding with a core route actually resolves to.""" + from fastapi import Request + + def chat_completions(): + ... + + return Request(scope={"type": "http", "headers": [], "endpoint": chat_completions}) + + +@pytest.mark.asyncio +async def test_common_checks_auth_enforced_pass_through_ignores_upstream_model(): + """An auth-enforced (`auth: true`) user-defined pass-through endpoint must + authenticate the key but forward the body unchanged; a body `model` naming an + upstream-only model must not be rejected against the team/key model allowlist + when the request was dispatched to the pass-through handler. The same body on a + request dispatched to a built-in handler (e.g. a path collision) must still be + enforced.""" + from litellm.proxy.auth.auth_checks import common_checks + + team_object = LiteLLM_TeamTable(team_id="team-1", models=["gpt-4o"]) + valid_token = UserAPIKeyAuth( + token="test-token", + team_id="team-1", + models=[], + metadata={"allowed_passthrough_routes": ["/my-custom-endpoint"]}, + ) + + with patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={}, + ): + result = await common_checks( + request_body={"model": "upstream-special-model", "prompt": "hi"}, + team_object=team_object, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/my-custom-endpoint", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=_pass_through_request(), + ) + assert result is True + + with pytest.raises(ProxyException) as exc_info: + await common_checks( + request_body={"model": "upstream-special-model", "prompt": "hi"}, + team_object=team_object, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=_builtin_request(), + ) + assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied + + @pytest.mark.asyncio async def test_virtual_key_soft_budget_check_with_user_obj(): """Test _virtual_key_soft_budget_check includes user_email when user_obj is provided""" diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 17ff700791f9..b5d8727f7e62 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import pytest +from fastapi import Request from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( @@ -331,6 +332,70 @@ def test_should_fall_back_to_body_when_no_standard_header(self): assert result == "body-user" +def _request_dispatched_to(endpoint) -> Request: + """Build a minimal Request whose FastAPI-resolved endpoint is ``endpoint``, + mirroring what Starlette sets in ``scope`` once routing has matched.""" + return Request(scope={"type": "http", "headers": [], "endpoint": endpoint}) + + +def _pass_through_endpoint(): + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_ENDPOINT_MARKER, + ) + + def endpoint(): # stand-in for create_pass_through_route's handler + ... + + setattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) + return endpoint + + +def test_get_model_from_request_skips_pass_through_dispatched_request(): + """When FastAPI dispatched the request to a user-defined pass-through handler, + the body `model` names an upstream model and must not be treated as a LiteLLM + model for allowlist/budget enforcement.""" + assert ( + get_model_from_request( + request_data={"model": "upstream-special-model"}, + route="/my-custom-endpoint", + request=_request_dispatched_to(_pass_through_endpoint()), + ) + is None + ) + + +def test_get_model_from_request_enforces_when_builtin_handler_dispatched(): + """A custom pass-through path that collides with a built-in route resolves to the + built-in handler (no marker), so the body `model` must still be extracted and + enforced. Same request path as above, but dispatched to a non-pass-through + endpoint: the model must NOT be suppressed.""" + + def builtin_chat_completions(): + ... + + assert ( + get_model_from_request( + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + request=_request_dispatched_to(builtin_chat_completions), + ) + == "gpt-4o" + ) + + +def test_get_model_from_request_no_request_extracts_model(): + """Callers without a request object (e.g. budget reservation) still extract the + model; the pass-through suppression only applies to a dispatched pass-through + handler.""" + assert ( + get_model_from_request( + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + ) + == "gpt-4o" + ) + + def test_get_model_from_request_supports_google_model_names_with_slashes(): assert ( get_model_from_request( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py new file mode 100644 index 000000000000..14d8e90e0277 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py @@ -0,0 +1,550 @@ +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy.guardrails.guardrail_hooks.singulr.singulr import SingulrGuardrail +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( + SingulrGuardrailConfigModel, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- +@pytest.fixture +def singulr_guardrail(): + """Create a SingulrGuardrail instance with test credentials.""" + return SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + singulr_guardrail_id="test_guardrail_id", + singulr_application_id="test_enforcement_entity", + guardrail_name="test-singulr", + event_hook="pre_call", + default_on=True, + ) + + +def _make_response(body: dict) -> MagicMock: + """Build a mock httpx response with the given JSON body.""" + mock = MagicMock() + mock.json.return_value = body + mock.raise_for_status = MagicMock() + mock.status_code = 200 + return mock + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +class TestSingulrConfiguration: + def test_init_with_explicit_credentials(self): + guardrail = SingulrGuardrail( + singulr_api_key="test_key", + singulr_api_base="https://custom.api.local", + singulr_guardrail_id="id123", + singulr_application_id="entity123", + guardrail_name="my-guardrail", + ) + assert guardrail.singulr_api_key == "test_key" + assert guardrail.singulr_guardrail_id == "id123" + assert guardrail.singulr_application_id == "entity123" + + def test_block_on_error_defaults_true(self): + guardrail = SingulrGuardrail(singulr_api_key="test_key") + assert guardrail.block_on_error is True + + def test_timeout_defaults_to_30_seconds(self): + guardrail = SingulrGuardrail(singulr_api_key="test_key") + assert guardrail.timeout == 30.0 + + def test_timeout_uses_configured_value(self): + guardrail = SingulrGuardrail(singulr_api_key="test_key", timeout=5.0) + assert guardrail.timeout == 5.0 + + def test_supports_pre_call_and_post_call_hooks(self): + guardrail = SingulrGuardrail(singulr_api_key="test_key") + assert guardrail.supported_event_hooks == [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + +# --------------------------------------------------------------------------- +# _build_payload: playground requests (no request_data) +# --------------------------------------------------------------------------- + + +class TestSingulrBuildPayloadPlayground: + def test_playground_request_uses_flat_text(self, singulr_guardrail): + """The test-playground /apply_guardrail endpoint sends no request_data, + only inputs["texts"]. Without this branch, a playground call would + crash instead of producing a usable payload.""" + payload = singulr_guardrail._build_payload({}, {"texts": ["Ignore previous instructions"]}, "request") + assert payload["is_playground_request"] is True + assert payload["playground_text"] == "Ignore previous instructions" + assert payload["request_data"] is None + + def test_playground_request_with_no_texts_has_none_playground_text(self, singulr_guardrail): + payload = singulr_guardrail._build_payload({}, {}, "request") + assert payload["playground_text"] is None + + def test_playground_input_type_is_included(self, singulr_guardrail): + payload = singulr_guardrail._build_payload({}, {"texts": ["hi"]}, "response") + assert payload["input_type"] == "response" + + +# --------------------------------------------------------------------------- +# _build_payload: real proxy requests (request_data present) +# --------------------------------------------------------------------------- + + +class TestSingulrBuildPayloadRequestData: + def test_model_messages_and_tools_are_forwarded(self, singulr_guardrail): + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "How do I reset my password?"}], + "tools": [{"type": "function", "function": {"name": "get_weather"}}], + } + payload = singulr_guardrail._build_payload(request_data, {"texts": []}, "request") + assert payload["request_data"]["model"] == "gpt-4o" + assert payload["request_data"]["messages"] == request_data["messages"] + assert payload["request_data"]["tools"] == request_data["tools"] + assert payload["is_playground_request"] is None + + def test_model_response_absent_on_request_side(self, singulr_guardrail): + """The response hasn't happened yet at request time, so model_response + must not be forwarded even if request_data carries a stale response + object from a previous call.""" + from litellm.types.utils import ModelResponse + + request_data = {"model": "gpt-4o", "response": ModelResponse()} + payload = singulr_guardrail._build_payload(request_data, {"texts": []}, "request") + assert payload["request_data"]["model_response"] is None + + def test_model_response_is_forwarded_and_json_serializable(self, singulr_guardrail): + """Regression: request_data["response"] is a ModelResponse (pydantic) + object containing nested non-JSON-safe values (e.g. a `created` + unix timestamp is fine, but nested pydantic submodels are not plain + dicts). Without mode="json" on both the inner and outer dumps, this + payload cannot be sent via httpx's json= kwarg.""" + import json as _json + + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + response = ModelResponse( + choices=[Choices(message=Message(role="assistant", content="Go to settings."))], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + request_data = {"model": "gpt-4o", "response": response} + payload = singulr_guardrail._build_payload(request_data, {"texts": ["Go to settings."]}, "response") + + # Must not raise - this is what httpx's json= kwarg effectively does. + serialized = _json.dumps(payload) + assert "Go to settings." in serialized + assert payload["request_data"]["model_response"]["choices"][0]["message"]["content"] == "Go to settings." + + def test_model_requested_tool_calls_are_forwarded_in_model_response(self, singulr_guardrail): + """Tool calls the model requests arrive inside response.choices[].message.tool_calls. + They must survive the dump so Singulr can inspect what tools the + model is trying to invoke.""" + from litellm.types.utils import Choices, Message, ModelResponse + + response = ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_current_time", "arguments": "{}"}, + } + ], + ) + ) + ], + ) + request_data = {"model": "gpt-4o", "response": response} + payload = singulr_guardrail._build_payload(request_data, {"texts": []}, "response") + + tool_calls = payload["request_data"]["model_response"]["choices"][0]["message"]["tool_calls"] + assert tool_calls[0]["function"]["name"] == "get_current_time" + + def test_litellm_metadata_is_forwarded(self, singulr_guardrail): + request_data = {"model": "gpt-4o", "litellm_metadata": {"user_api_key_hash": "abc123"}} + payload = singulr_guardrail._build_payload(request_data, {"texts": []}, "request") + assert payload["request_data"]["litellm_metadata"] == {"user_api_key_hash": "abc123"} + + def test_internal_logging_object_is_not_forwarded(self, singulr_guardrail): + """Regression: request_data can carry internal proxy objects (e.g. the + Logging instance) that aren't JSON-serializable at all. _build_payload + must only pull known request/response fields out of request_data, + not dump it wholesale, or this crashes on every real proxy call.""" + import json as _json + + class _NotSerializable: + pass + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "litellm_logging_obj": _NotSerializable(), + } + payload = singulr_guardrail._build_payload(request_data, {"texts": ["hi"]}, "request") + + # Must not raise. + _json.dumps(payload) + assert "litellm_logging_obj" not in payload["request_data"] + + +# --------------------------------------------------------------------------- +# Allow / block decisions +# --------------------------------------------------------------------------- + + +class TestSingulrAllowAction: + @pytest.mark.asyncio + async def test_allow_returns_inputs_unchanged(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + inputs = {"texts": ["How do I reset my password?"]} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + result = await singulr_guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + + +class TestSingulrBlockAction: + @pytest.mark.asyncio + async def test_block_raises_guardrail_exception(self, singulr_guardrail): + """Regression: a should_block=True response must stop the request + instead of silently letting it through.""" + resp = _make_response( + { + "should_block": True, + "blocking_due_to": "PII Information detected", + } + ) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + with pytest.raises(GuardrailRaisedException) as exc_info: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["My SSN is 123-45-6789"]}, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert "PII Information detected" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_block_without_reason_uses_unknown_placeholder(self, singulr_guardrail): + resp = _make_response({"should_block": True}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + with pytest.raises(GuardrailRaisedException, match="unknown"): + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + + +# --------------------------------------------------------------------------- +# HTTP call wiring (endpoint, timeout, headers) +# --------------------------------------------------------------------------- + + +class TestSingulrRequestWiring: + @pytest.mark.asyncio + async def test_sends_configured_timeout(self): + """litellm_params.timeout must reach the httpx call so operators can + tighten or loosen the latency budget instead of being stuck with a + hardcoded 30s regardless of configuration.""" + guardrail = SingulrGuardrail( + singulr_api_key="test_key", + singulr_api_base="https://api.test.singulr.ai", + timeout=5.0, + ) + resp = _make_response({"should_block": False}) + with patch.object(guardrail.async_handler, "post", return_value=resp) as mock_post: + await guardrail.apply_guardrail( + inputs={"texts": ["test"]}, + request_data={}, + input_type="request", + ) + assert mock_post.call_args.kwargs["timeout"] == 5.0 + + +class TestSingulrBuildHeaders: + def test_content_type_always_present(self, singulr_guardrail): + assert singulr_guardrail._build_headers()["Content-Type"] == "application/json" + + def test_all_optional_headers_included_when_set(self, singulr_guardrail): + headers = singulr_guardrail._build_headers() + assert headers["X-Singulr-Gateway-Token"] == "test_token_1234" + assert headers["X-Singulr-Enforcement-Entity-Id"] == "test_enforcement_entity" + assert headers["X-Singulr-Guardrail-Id"] == "test_guardrail_id" + + def test_optional_headers_absent_when_unset(self): + guardrail = SingulrGuardrail(guardrail_name="bare") + headers = guardrail._build_headers() + assert "X-Singulr-Gateway-Token" not in headers + assert "X-Singulr-Enforcement-Entity-Id" not in headers + assert "X-Singulr-Guardrail-Id" not in headers + + +# --------------------------------------------------------------------------- +# Non-JSON / malformed response handling +# --------------------------------------------------------------------------- + + +class TestSingulrInvalidResponse: + @pytest.mark.asyncio + async def test_non_json_response_block_on_error_false_returns_inputs(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=False, + ) + mock_resp = MagicMock() + mock_resp.raise_for_status = MagicMock() + mock_resp.json.side_effect = ValueError("No JSON object could be decoded") + + inputs = {"texts": ["test"]} + with patch.object(guardrail.async_handler, "post", return_value=mock_resp): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + assert result is inputs + + @pytest.mark.asyncio + async def test_non_json_response_block_on_error_true_raises(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=True, + ) + mock_resp = MagicMock() + mock_resp.raise_for_status = MagicMock() + mock_resp.json.side_effect = ValueError("No JSON object could be decoded") + + with patch.object(guardrail.async_handler, "post", return_value=mock_resp): + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail( + inputs={"texts": ["test"]}, + request_data={}, + input_type="request", + ) + + @pytest.mark.asyncio + async def test_response_missing_expected_fields_block_on_error_true_raises(self): + """Regression: a response body that fails SingulrGuardrailResponse + validation (e.g. should_block is a string, not a bool) must raise + GuardrailRaisedException instead of letting pydantic.ValidationError + propagate unhandled.""" + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=True, + ) + resp = _make_response({"should_block": "not-a-bool"}) + with patch.object(guardrail.async_handler, "post", return_value=resp): + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail( + inputs={"texts": ["test"]}, + request_data={}, + input_type="request", + ) + + +# --------------------------------------------------------------------------- +# Transport error handling +# --------------------------------------------------------------------------- + + +class TestSingulrTransportError: + @pytest.mark.asyncio + async def test_remote_protocol_error_block_on_error_false_returns_inputs(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=False, + ) + inputs = {"texts": ["test"]} + with patch.object( + guardrail.async_handler, + "post", + side_effect=httpx.RemoteProtocolError("malformed HTTP response"), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + assert result is inputs + + @pytest.mark.asyncio + async def test_remote_protocol_error_block_on_error_true_raises(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=True, + ) + with patch.object( + guardrail.async_handler, + "post", + side_effect=httpx.RemoteProtocolError("malformed HTTP response"), + ): + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail( + inputs={"texts": ["test"]}, + request_data={}, + input_type="request", + ) + + +# --------------------------------------------------------------------------- +# HTTP status error handling +# --------------------------------------------------------------------------- + + +class TestSingulrHttpStatusError: + @pytest.mark.asyncio + async def test_http_error_message_names_status_code_not_unreachable(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=True, + ) + mock_response = MagicMock() + mock_response.status_code = 403 + mock_response.text = "Forbidden" + exc = httpx.HTTPStatusError("403 Forbidden", request=MagicMock(), response=mock_response) + mock_response.raise_for_status.side_effect = exc + + with patch.object(guardrail.async_handler, "post", return_value=mock_response): + with pytest.raises(GuardrailRaisedException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["test"]}, + request_data={}, + input_type="request", + ) + msg = str(exc_info.value) + assert "403" in msg + assert "unreachable" not in msg.lower() + + @pytest.mark.asyncio + async def test_http_error_block_on_error_false_returns_inputs(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=False, + ) + mock_response = MagicMock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + exc = httpx.HTTPStatusError("500", request=MagicMock(), response=mock_response) + mock_response.raise_for_status.side_effect = exc + + inputs = {"texts": ["test"]} + with patch.object(guardrail.async_handler, "post", return_value=mock_response): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + assert result is inputs + + +# --------------------------------------------------------------------------- +# Config model +# --------------------------------------------------------------------------- + + +class TestSingulrConfigModel: + def test_ui_friendly_name(self): + assert SingulrGuardrailConfigModel.ui_friendly_name() == "Singulr" + + +# --------------------------------------------------------------------------- +# Initializer and registry +# --------------------------------------------------------------------------- + + +class TestSingulrInitializer: + def test_guardrail_initializer_registry_has_entry(self): + from litellm.proxy.guardrails.guardrail_hooks.singulr import ( + initialize_guardrail, + ) + + assert callable(initialize_guardrail) + + def test_initialize_guardrail_reads_singulr_prefixed_fields(self): + """Regression: the UI config form (and YAML config) populate the + singulr_-prefixed fields declared on SingulrGuardrailConfigModel, not + the generic api_base/api_key fields. initialize_guardrail must read + those, or a UI-configured singulr_api_base is silently ignored and + the guardrail falls back to the localhost default.""" + from litellm.proxy.guardrails.guardrail_hooks.singulr import ( + initialize_guardrail, + ) + from litellm.types.guardrails import Guardrail, LitellmParams + + litellm_params = LitellmParams( + guardrail="singulr", + mode="pre_call", + singulr_api_base="https://configured.singulr.ai", + singulr_api_key="configured_key", + singulr_application_id="configured_app_id", + singulr_guardrail_id="configured_guardrail_id", + ) + guardrail: Guardrail = { + "guardrail_name": "test-singulr", + "litellm_params": litellm_params, + } + + cb = initialize_guardrail(litellm_params, guardrail) + + assert cb.singulr_application_id == "configured_app_id" + assert cb.singulr_guardrail_id == "configured_guardrail_id" + + def test_initialize_guardrail_wires_timeout(self): + """BaseLitellmParams.timeout exists so operators can override the + per-request latency budget. initialize_guardrail must forward it to + SingulrGuardrail instead of leaving every deployment stuck on the + hardcoded default regardless of configuration.""" + from litellm.proxy.guardrails.guardrail_hooks.singulr import ( + initialize_guardrail, + ) + from litellm.types.guardrails import Guardrail, LitellmParams + + litellm_params = LitellmParams( + guardrail="singulr", + mode="pre_call", + singulr_api_key="configured_key", + timeout=12.5, + ) + guardrail: Guardrail = { + "guardrail_name": "test-singulr", + "litellm_params": litellm_params, + } + + cb = initialize_guardrail(litellm_params, guardrail) + + assert cb.timeout == 12.5 diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index e7d2909263af..c76e1a60afd0 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -17,6 +17,10 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + MAX_PARALLEL_SLOT_ACQUIRED_KEY, + PARALLEL_REQUEST_SLOT_TTL_SECONDS, +) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) @@ -566,10 +570,9 @@ async def mock_increment_pipeline(increment_list, **kwargs): # Verify that the correct token count was used based on the rate limit type assert ( - len(captured_operations) == 2 - ), "Should have 2 operations: max_parallel_requests decrement and TPM increment" + len(captured_operations) == 1 + ), "Should have 1 operation: the TPM increment (parallel slots are released via the gauge, not the pipeline)" - # Find the TPM increment operation (not the max_parallel_requests decrement) tpm_operation = None for op in captured_operations: if op["key"].endswith(":tokens"): @@ -655,7 +658,10 @@ async def mock_increment_pipeline(increment_list, **kwargs): @pytest.mark.asyncio async def test_async_log_failure_event_v3(): """ - Simple test for async_log_failure_event - should decrement max_parallel_requests by 1 + async_log_failure_event releases exactly this request's slot id: the + first release removes it, and repeated or unknown-slot releases are + no-ops that can never free another request's slot (releasing more than + was acquired is what previously let concurrency exceed the limit). """ _api_key = "sk-12345" _api_key = hash_token(_api_key) @@ -663,33 +669,246 @@ async def test_async_log_failure_event_v3(): parallel_request_handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" - # Mock kwargs with user_api_key via standard_logging_object - mock_kwargs = { - "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}} - } + await _seed_max_parallel_requests_slots(local_cache, counter_key, ["slot-a", "slot-b"]) - # Capture pipeline operations - captured_ops = [] + def kwargs_with_slot(slot_id): + return { + "metadata": { + MAX_PARALLEL_SLOT_ACQUIRED_KEY: { + "slot_id": slot_id, + "counter_keys": [counter_key], + } + }, + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + } - async def mock_pipeline(increment_list, **kwargs): - captured_ops.extend(increment_list) + async def in_flight(): + return parallel_request_handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) - parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( - mock_pipeline + await parallel_request_handler.async_log_failure_event( + kwargs=kwargs_with_slot("slot-a"), response_obj=None, start_time=None, end_time=None ) + assert await in_flight() == 1 + + for slot_id in ("slot-a", "slot-unknown", "slot-a"): + await parallel_request_handler.async_log_failure_event( + kwargs=kwargs_with_slot(slot_id), response_obj=None, start_time=None, end_time=None + ) + assert await in_flight() == 1 - # Call async_log_failure_event await parallel_request_handler.async_log_failure_event( - kwargs=mock_kwargs, response_obj=None, start_time=None, end_time=None + kwargs=kwargs_with_slot("slot-b"), response_obj=None, start_time=None, end_time=None + ) + assert await in_flight() == 0 + + +@pytest.mark.asyncio +async def test_failure_event_without_acquired_slot_does_not_release_v3(): + """ + Failure callbacks also fire for requests rejected at pre-call, which never + acquired a parallel slot. Releasing on those frees a slot still owned by + another in-flight request, so every 429 would raise effective concurrency + above the configured limit. Without the acquired-slot marker the gauge + must stay untouched. + """ + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + await _seed_max_parallel_requests_slots( + local_cache, counter_key, ["slot-a", "slot-b", "slot-c"] + ) + + await handler.async_log_failure_event( + kwargs={ + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}} + }, + response_obj=None, + start_time=None, + end_time=None, + ) + assert ( + handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) + == 3 + ) + + +@pytest.mark.asyncio +async def test_max_parallel_requests_not_reset_by_window_roll_v3(): + """ + max_parallel_requests is a concurrency gauge, not a windowed counter: the + rate-limit window rolling over must not reset it while requests are still + in flight. Previously the gauge shared the sliding-window reset with + RPM/TPM, so every window roll forgot all in-flight requests and admitted + a fresh batch of `limit` on top of what was still running. + """ + controller = TimeController() + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=controller.now, + ) + _api_key = hash_token("sk-12345") + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=2) + + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + controller.advance(handler.window_size + 1) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert "max_parallel_requests" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_rejected_request_does_not_consume_parallel_slot_v3(): + """ + A 429-rejected request must not occupy a parallel-request slot: nothing + ever releases a slot for a request that was never admitted, so the old + increment-then-check behavior wedged the gauge above the limit and + rejected requests that should have been admitted after a release. + """ + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + _api_key = hash_token("sk-12345") + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + + admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=admitted_data, + call_type="", + ) + acquisition = admitted_data["metadata"][MAX_PARALLEL_SLOT_ACQUIRED_KEY] + assert isinstance(acquisition, dict) + assert isinstance(acquisition["slot_id"], str) and acquisition["slot_id"] + assert acquisition["counter_keys"] == [f"{{api_key:{_api_key}}}:max_parallel_requests"] + + for _ in range(3): + with pytest.raises(HTTPException): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + await handler.async_log_failure_event( + kwargs={ + "metadata": {MAX_PARALLEL_SLOT_ACQUIRED_KEY: acquisition}, + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + +@pytest.mark.asyncio +async def test_parallel_gauge_uses_atomic_redis_script_v3(): + """ + With Redis available, gauge admission goes through the atomic + check-and-acquire script (limit, slot TTL, and this request's slot id as + args), the returned in-flight count is mirrored into the local cache, + and an over-limit script result maps to a 429 without occupying a slot. + """ + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) ) + _api_key = hash_token("sk-12345") + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + captured_calls = [] - # Verify correct operation was created - assert len(captured_ops) == 1 - op = captured_ops[0] - assert op["key"] == f"{{api_key:{_api_key}}}:max_parallel_requests" - assert op["increment_value"] == -1 - assert op["ttl"] == 60 # default window size + async def fake_acquire(keys, args): + captured_calls.append((list(keys), list(args))) + return [0, 3] + + handler.parallel_acquire_script = fake_acquire + + data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type="", + ) + stashed_acquisition = data["metadata"][MAX_PARALLEL_SLOT_ACQUIRED_KEY] + assert isinstance(stashed_acquisition, dict) + stashed_slot_id = stashed_acquisition["slot_id"] + assert isinstance(stashed_slot_id, str) and stashed_slot_id + assert stashed_acquisition["counter_keys"] == [counter_key] + assert captured_calls == [ + ([counter_key], [5, PARALLEL_REQUEST_SLOT_TTL_SECONDS, stashed_slot_id]) + ] + assert ( + await handler.internal_usage_cache.async_get_cache( + key=counter_key, litellm_parent_otel_span=None, local_only=True + ) + == 3 + ) + gauge_statuses = [ + s + for s in data["litellm_proxy_rate_limit_response"]["statuses"] + if s["rate_limit_type"] == "max_parallel_requests" + ] + assert gauge_statuses == [ + { + "code": "OK", + "current_limit": 5, + "limit_remaining": 2, + "rate_limit_type": "max_parallel_requests", + "descriptor_key": "api_key", + } + ] + + async def fake_acquire_over_limit(keys, args): + return [1, 1, 5, 5] + + handler.parallel_acquire_script = fake_acquire_over_limit + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert "max_parallel_requests" in exc_info.value.detail @pytest.mark.asyncio @@ -3227,27 +3446,28 @@ def test_get_key_mcp_rpm_limit_precedence(): assert get_team_mcp_rpm_limit(none_set) is None -async def _seed_max_parallel_requests_counter( - dual_cache: DualCache, counter_key: str, window_size: int +_TEST_SLOT_ID = "slot-disconnect-test" + + +async def _seed_max_parallel_requests_slots( + dual_cache: DualCache, counter_key: str, slot_ids: List[str] ) -> None: - await dual_cache.async_increment_cache_pipeline( - increment_list=[ - RedisPipelineIncrementOperation( - key=counter_key, increment_value=1, ttl=window_size - ) - ] + await dual_cache.async_set_cache( + key=counter_key, + value={slot_id: time.time() for slot_id in slot_ids}, + local_only=True, ) async def _build_seeded_limiter(): - """Build a v3 limiter whose api-key counter already holds the pre-call +1.""" + """Build a v3 limiter whose api-key slot registry already holds the pre-call slot.""" api_key = hash_token("sk-disconnect") cache = DualCache() limiter = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(cache) ) counter_key = f"{{api_key:{api_key}}}:max_parallel_requests" - await _seed_max_parallel_requests_counter(cache, counter_key, limiter.window_size) + await _seed_max_parallel_requests_slots(cache, counter_key, [_TEST_SLOT_ID]) user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) return limiter, cache, counter_key, user_api_key_dict @@ -3286,14 +3506,370 @@ async def test_release_max_parallel_requests_on_disconnect_v3(): user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=2) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" - await _seed_max_parallel_requests_counter( - local_cache, counter_key, handler.window_size + await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID]) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 1 + + await handler.async_release_max_parallel_requests_on_disconnect( + user_api_key_dict, + request_data={ + "metadata": { + MAX_PARALLEL_SLOT_ACQUIRED_KEY: { + "slot_id": _TEST_SLOT_ID, + "counter_keys": [counter_key], + } + } + }, ) - assert await local_cache.async_get_cache(key=counter_key) == 1 - await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 - assert await local_cache.async_get_cache(key=counter_key) == 0 + +@pytest.mark.asyncio +async def test_release_on_disconnect_works_when_key_config_changed_v3(): + """ + The disconnect release must be driven by the stashed acquisition, not the + key object's current max_parallel_requests configuration: if the limit is + cleared on the key while a request is in flight, the acquired slot still + has to be released or it lingers until TTL pruning. + """ + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID]) + + await handler.async_release_max_parallel_requests_on_disconnect( + UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=None), + request_data={ + "metadata": { + MAX_PARALLEL_SLOT_ACQUIRED_KEY: { + "slot_id": _TEST_SLOT_ID, + "counter_keys": [counter_key], + } + } + }, + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_releases_parallel_slot_v3(): + """ + A proxy-level rejection raised by a downstream hook after the rate + limiter's pre-call hook acquired a slot (guardrail, budget check) must + release that slot via async_post_call_failure_hook: + async_log_failure_event never fires for proxy-side rejections, so + without this the slot lingers for the full slot TTL and moderate + rejection rates wedge the key at its limit. The release must also be + idempotent with a later failure callback in the same flow. + """ + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=admitted_data, + call_type="", + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 1 + + await handler.async_post_call_failure_hook( + request_data=admitted_data, + original_exception=Exception("guardrail rejected the request"), + user_api_key_dict=user_api_key_dict, + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 + + await handler.async_log_failure_event( + kwargs={ + "metadata": admitted_data["metadata"], + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + +@pytest.mark.asyncio +async def test_success_event_releases_parallel_slot_v3(monkeypatch): + """ + A successful completion must release exactly the slot its pre-call + acquired, freeing capacity for the next request; without it every + completed request would keep occupying the gauge until TTL pruning. + """ + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=admitted_data, + call_type="", + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 1 + + await handler.async_log_success_event( + kwargs={ + "metadata": admitted_data["metadata"], + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + }, + response_obj=ModelResponse( + usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10) + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + +@pytest.mark.asyncio +async def test_read_only_gauge_check_counts_without_acquiring_v3(): + """ + read_only callers (e.g. the context-compaction pre-check) must observe + the in-flight count via the count script without registering a slot, and + a count-script failure must degrade to the local mirror instead of + raising. + """ + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + descriptors = [ + { + "key": "api_key", + "value": _api_key, + "rate_limit": {"max_parallel_requests": 5}, + } + ] + + captured_calls = [] + + async def fake_count(keys, args): + captured_calls.append((list(keys), list(args))) + return [3] + + handler.parallel_count_script = fake_count + + response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) + assert captured_calls == [ + ([counter_key], [PARALLEL_REQUEST_SLOT_TTL_SECONDS]) + ] + assert response["overall_code"] == "OK" + assert response["statuses"] == [ + { + "code": "OK", + "current_limit": 5, + "limit_remaining": 2, + "rate_limit_type": "max_parallel_requests", + "descriptor_key": "api_key", + } + ] + assert await local_cache.async_get_cache(key=counter_key) is None + + async def failing_count(keys, args): + raise ConnectionError("redis unavailable") + + handler.parallel_count_script = failing_count + await _seed_max_parallel_requests_slots( + local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"] + ) + response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) + assert response["overall_code"] == "OVER_LIMIT" + assert response["statuses"][0]["rate_limit_type"] == "max_parallel_requests" + + +@pytest.mark.asyncio +async def test_redis_release_script_updates_local_mirror_v3(): + """ + With Redis available, releases go through the release script with this + request's slot id per gauge key, and the returned in-flight counts are + mirrored into the local cache so the local first-pass check stays fresh. + """ + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + captured_calls = [] + + async def fake_release(keys, args): + captured_calls.append((list(keys), list(args))) + return [2] + + handler.parallel_release_script = fake_release + + await handler.async_log_failure_event( + kwargs={ + "metadata": { + MAX_PARALLEL_SLOT_ACQUIRED_KEY: { + "slot_id": "slot-redis-test", + "counter_keys": [counter_key], + } + }, + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + assert captured_calls == [([counter_key], ["slot-redis-test"])] + assert await local_cache.async_get_cache(key=counter_key) == 2 + + +@pytest.mark.asyncio +async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch): + """ + When the TPM reservation phase rejects a request AFTER the gauge slot was + acquired earlier in the same pre-call hook, the slot must be released + before the 429 is raised; otherwise every TPM rejection would leak a + slot until TTL pruning. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, max_parallel_requests=5, tpm_limit=100 + ) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None): + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "tokens", + "descriptor_key": "api_key", + } + ], + } + + monkeypatch.setattr(handler, "reserve_tpm_tokens", over_limit_reservation) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=counter_key) + ) == 0 + + +@pytest.mark.asyncio +async def test_in_memory_fallback_respects_mirrored_redis_count_v3(): + """ + When Redis scripting fails after having worked, the local cache holds the + integer in-flight count mirrored from the last successful script call. + The in-memory fallback must treat that count as real occupancy (and + release must decrement it, floored at 0), not start over from an empty + registry, which would double the admitted concurrency during a Redis + outage. + """ + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + async def failing_script(keys, args): + raise ConnectionError("redis unavailable") + + handler.parallel_acquire_script = failing_script + handler.parallel_release_script = failing_script + + await local_cache.async_set_cache(key=counter_key, value=5, local_only=True) + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + assert exc_info.value.status_code == 429 + + await local_cache.async_set_cache(key=counter_key, value=4, local_only=True) + admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=admitted_data, + call_type="", + ) + assert await local_cache.async_get_cache(key=counter_key) == 5 + + await handler.async_log_failure_event( + kwargs={ + "metadata": admitted_data["metadata"], + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + assert await local_cache.async_get_cache(key=counter_key) == 4 @pytest.mark.asyncio @@ -3338,7 +3914,9 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing limiter, cache, counter_key, user_api_key_dict = await _build_seeded_limiter() - assert await cache.async_get_cache(key=counter_key) == 1 + assert limiter._gauge_in_flight_from_cache_value( + await cache.async_get_cache(key=counter_key) + ) == 1 proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = limiter @@ -3354,7 +3932,15 @@ async def upstream(): gen = ProxyBaseLLMRequestProcessing.async_sse_data_generator( response=upstream(), user_api_key_dict=user_api_key_dict, - request_data={"model": "claude-test"}, + request_data={ + "model": "claude-test", + "metadata": { + MAX_PARALLEL_SLOT_ACQUIRED_KEY: { + "slot_id": _TEST_SLOT_ID, + "counter_keys": [counter_key], + } + }, + }, proxy_logging_obj=proxy_logging_obj, ) await gen.__anext__() @@ -3365,7 +3951,9 @@ async def upstream(): await gen.aclose() await _drain_release_task() - assert await cache.async_get_cache(key=counter_key) == 0 + assert limiter._gauge_in_flight_from_cache_value( + await cache.async_get_cache(key=counter_key) + ) == 0 @pytest.mark.parametrize("disconnect", ["cancel", "aclose"]) @@ -3399,7 +3987,15 @@ async def upstream(): gen = proxy_server.async_data_generator( response=upstream(), user_api_key_dict=user_api_key_dict, - request_data={"model": "gpt-test"}, + request_data={ + "model": "gpt-test", + "metadata": { + MAX_PARALLEL_SLOT_ACQUIRED_KEY: { + "slot_id": _TEST_SLOT_ID, + "counter_keys": [counter_key], + } + }, + }, ) await gen.__anext__() if disconnect == "cancel": @@ -3408,7 +4004,9 @@ async def upstream(): else: await gen.aclose() await _drain_release_task() - assert await cache.async_get_cache(key=counter_key) == 0 + assert limiter._gauge_in_flight_from_cache_value( + await cache.async_get_cache(key=counter_key) + ) == 0 finally: if saved_hook is not None: proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = ( @@ -3452,12 +4050,22 @@ async def upstream(): gen = proxy_server.async_data_generator( response=upstream(), user_api_key_dict=user_api_key_dict, - request_data={"model": "gpt-test"}, + request_data={ + "model": "gpt-test", + "metadata": { + MAX_PARALLEL_SLOT_ACQUIRED_KEY: { + "slot_id": _TEST_SLOT_ID, + "counter_keys": [counter_key], + } + }, + }, ) await gen.__anext__() await gen.aclose() await _drain_release_task() - assert await cache.async_get_cache(key=counter_key) == 0 + assert limiter._gauge_in_flight_from_cache_value( + await cache.async_get_cache(key=counter_key) + ) == 0 finally: if saved_hook is not None: proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = ( diff --git a/tests/test_litellm/proxy/proxy_server/test_background_health.py b/tests/test_litellm/proxy/proxy_server/test_background_health.py index ee8d8b227796..dca93e137acf 100644 --- a/tests/test_litellm/proxy/proxy_server/test_background_health.py +++ b/tests/test_litellm/proxy/proxy_server/test_background_health.py @@ -378,8 +378,12 @@ async def test_adaptive_router_flusher_loop_flushes_each_router(monkeypatch): fake_ar.queue.flush_state_to_db = AsyncMock() fake_ar.queue.flush_session_to_db = AsyncMock() + from litellm.types.router import TaggedPreRoutingStrategy + fake_router = MagicMock() - fake_router.adaptive_routers = {"alpha": fake_ar} + fake_router.adaptive_routers = { + "alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)] + } monkeypatch.setattr(proxy_server, "llm_router", fake_router) monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 45d354196803..bd8e92c3cc22 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -22,6 +22,7 @@ _scrub_db_overlay_remote_module_loads, _scrub_guardrail_inner, resolve_complexity_router_plugins, + resolve_routing_plugins, ) from .conftest import normalize @@ -185,6 +186,75 @@ def test_resolve_complexity_router_plugins_rejects_synchronous_run_method(tmp_pa ) +# --------------------------------------------------------------------------- +# resolve_routing_plugins +# --------------------------------------------------------------------------- + + +def test_resolve_routing_plugins_resolves_dotted_paths(tmp_path): + plugin_file = tmp_path / "rs_plugin.py" + plugin_file.write_text( + "class _Plugin:\n" + " async def run(self, context):\n" + " return context\n" + "\n" + "rs_plugin_instance = _Plugin()\n" + ) + + resolved = resolve_routing_plugins( + plugin_paths=["rs_plugin.rs_plugin_instance"], + config_file_path=str(tmp_path / "config.yaml"), + source_label="router_settings.plugins", + ) + + assert len(resolved) == 1 + assert type(resolved[0]).__name__ == "_Plugin" + + +def test_resolve_routing_plugins_passes_through_instances(tmp_path): + class _Plugin: + async def run(self, context): + return context + + instance = _Plugin() + resolved = resolve_routing_plugins( + plugin_paths=[instance], + config_file_path=None, + source_label="router_settings.plugins", + ) + assert resolved == [instance] + + +def test_resolve_routing_plugins_rejects_non_routing_plugin(tmp_path): + plugin_file = tmp_path / "bad_rs_plugin.py" + plugin_file.write_text("not_a_plugin = object()\n") + + with pytest.raises(ValueError, match="router_settings.plugins"): + resolve_routing_plugins( + plugin_paths=["bad_rs_plugin.not_a_plugin"], + config_file_path=str(tmp_path / "config.yaml"), + source_label="router_settings.plugins", + ) + + +def test_resolve_routing_plugins_rejects_synchronous_run(tmp_path): + plugin_file = tmp_path / "sync_rs_plugin.py" + plugin_file.write_text( + "class _SyncPlugin:\n" + " def run(self, context):\n" + " return context\n" + "\n" + "sync_plugin_instance = _SyncPlugin()\n" + ) + + with pytest.raises(ValueError, match="does not implement the RoutingPlugin interface"): + resolve_routing_plugins( + plugin_paths=["sync_rs_plugin.sync_plugin_instance"], + config_file_path=str(tmp_path / "config.yaml"), + source_label="router_settings.plugins", + ) + + # --------------------------------------------------------------------------- # ProxyConfig.__init__ # --------------------------------------------------------------------------- @@ -793,6 +863,62 @@ async def test_ProxyConfig_load_config_minimal_yaml(tmp_path, monkeypatch): } +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch): + """Regression: router_settings.plugins dotted-path strings must be resolved to + live RoutingPlugin instances on the created Router. Previously they were passed + through as raw strings and only blew up at request time when the pipeline tried + to `await "some.string".run(context)`.""" + plugin_file = tmp_path / "rs_plugin.py" + plugin_file.write_text( + "class _Plugin:\n" + " async def run(self, context):\n" + " return context\n" + "\n" + "rs_plugin_instance = _Plugin()\n" + ) + f = tmp_path / "c.yaml" + f.write_text( + "model_list: []\n" + "general_settings: {}\n" + "litellm_settings: {}\n" + "router_settings:\n" + " plugins:\n" + " - rs_plugin.rs_plugin_instance\n" + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + router, _model_list, _general_settings = await ProxyConfig().load_config( + router=None, config_file_path=str(f) + ) + + assert len(router.routing_plugins) == 1 + assert type(router.routing_plugins[0]).__name__ == "_Plugin" + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_rejects_bad_router_settings_plugin(tmp_path, monkeypatch): + plugin_file = tmp_path / "bad_rs_plugin.py" + plugin_file.write_text("not_a_plugin = object()\n") + f = tmp_path / "c.yaml" + f.write_text( + "model_list: []\n" + "general_settings: {}\n" + "litellm_settings: {}\n" + "router_settings:\n" + " plugins:\n" + " - bad_rs_plugin.not_a_plugin\n" + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + with pytest.raises(ValueError, match="does not implement the RoutingPlugin interface"): + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + @pytest.mark.asyncio async def test_ProxyConfig_load_config_wires_general_settings_url_validation(tmp_path, monkeypatch): """Regression for #26599: SSRF settings in general_settings must reach litellm globals.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_misc.py b/tests/test_litellm/proxy/proxy_server/test_routes_misc.py index 0c45e31afd2d..677ab8765bb2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_misc.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_misc.py @@ -94,7 +94,11 @@ def test_adaptive_router_state_returns_snapshots(client, auth_as, monkeypatch): snap = {"router_name": "ar-1", "queue_depth": 0, "posteriors": []} bandit = MagicMock() bandit.get_state_snapshot = AsyncMock(return_value=snap) - fake_router.adaptive_routers = {"ar-1": bandit} + from litellm.types.router import TaggedPreRoutingStrategy + + fake_router.adaptive_routers = { + "ar-1": [TaggedPreRoutingStrategy(tags=(), strategy=bandit)] + } monkeypatch.setattr(ps, "llm_router", fake_router) with auth_as(LitellmUserRoles.PROXY_ADMIN): diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index ac43418e41f5..a860a91e08b6 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -243,6 +243,91 @@ def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user): }, }, ) - assert ( - response.status_code != 400 - ), f"Clean Bedrock ingest_options should not be rejected: {response.json()}" + assert response.status_code != 400, ( + f"Clean Bedrock ingest_options should not be rejected: {response.json()}" + ) + + +def test_rag_query_returns_response_cost_header(client_internal_user): + """ + /v1/rag/query must surface the completion cost via the + x-litellm-response-cost response header, like /v1/chat/completions does. + """ + from litellm.types.utils import ModelResponse + + mock_response = ModelResponse( + id="chatcmpl-test", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "The codename is AZURE-FALCON-42."}, + "finish_reason": "stop", + } + ], + model="gpt-4o-mini", + usage={"prompt_tokens": 35, "completion_tokens": 14, "total_tokens": 49}, + ) + mock_response._hidden_params["response_cost"] = 3.45e-06 + + with patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ), patch("litellm.vector_store_registry", None), patch( + "litellm.proxy.proxy_server.prisma_client", None + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is the codename?"}], + "retrieval_config": { + "vector_store_id": "vs_test_123", + "custom_llm_provider": "openai", + }, + }, + ) + + assert response.status_code == 200, response.json() + assert response.headers.get("x-litellm-response-cost") == "3.45e-06" + + +def test_rag_query_stream_returns_event_stream(client_internal_user): + """ + A stream=true /v1/rag/query must return an SSE response. Returning the raw + stream wrapper makes FastAPI try to serialize it, which raises and turns + every streaming RAG query into a 500; the stream then never drains, so its + single billing event (which carries the folded sub-call costs) never fires. + """ + import litellm as litellm_module + + async def fake_aquery(**kwargs): + return await litellm_module.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What is the codename?"}], + mock_response="The codename is AZURE-FALCON-42.", + stream=True, + api_key="test-key", + ) + + with patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=fake_aquery), + ), patch("litellm.vector_store_registry", None), patch("litellm.proxy.proxy_server.prisma_client", None): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is the codename?"}], + "retrieval_config": { + "vector_store_id": "vs_test_123", + "custom_llm_provider": "openai", + }, + "stream": True, + }, + ) + + assert response.status_code == 200, response.text + assert response.headers.get("content-type", "").startswith("text/event-stream") + assert '"object":"chat.completion.chunk"' in response.text + assert "data: [DONE]" in response.text diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 0b304f2fec74..79a29811ff61 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2279,6 +2279,7 @@ async def _reserve_for_stream(counter_cache, key_cache, proxy_logging_obj, token def _drive_streaming_cancel(valid_token, iterator_hook): streaming_logging_obj = MagicMock() streaming_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock() return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=MagicMock(), user_api_key_dict=valid_token, @@ -2387,6 +2388,7 @@ async def one_chunk(user_api_key_dict, response, request_data): streaming_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=asyncio.CancelledError() ) + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock() generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=MagicMock(), diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index aa1911f80bcf..ebfbb46053d3 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -17,6 +17,7 @@ ProxyBaseLLMRequestProcessing, ProxyConfig, _await_llm_call_cancelling_on_disconnect, + _bill_partial_streamed_spend_on_disconnect, _buffer_first_chunk_honoring_disconnect, _cancel_llm_call_on_client_disconnect, _ClientDisconnectedBeforeFirstChunk, @@ -4871,3 +4872,217 @@ async def real_limiter_pre_call(**kwargs): }, call_type="acompletion", ) + + +class _RecordingSuccessLogger(CustomLogger): + def __init__(self): + super().__init__() + self.success_events = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_events.append({"kwargs": kwargs, "response_obj": response_obj}) + + +class TestStreamingClientDisconnectBilling: + """ + A client disconnect throws GeneratorExit into the proxy streaming + generator; neither the success nor failure logging callback fires from the + stream wrapper, so without disconnect-time finalization the chunks already + streamed (and any sub-call cost folded into the logging object) never + reach spend tracking. + """ + + async def _start_partial_stream(self): + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "tell me a story"}], + mock_response="The codename is AZURE-FALCON-42 and the story is long.", + stream=True, + api_key="test-key", + ) + stream_iter = response.__aiter__() + await stream_iter.__anext__() + await stream_iter.__anext__() + return response + + @pytest.mark.asyncio + async def test_disconnect_bills_partial_streamed_spend(self): + recorder = _RecordingSuccessLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recorder] + try: + response = await self._start_partial_stream() + logging_obj = response.logging_obj + logging_obj.model_call_details["additional_response_cost"] = 0.002 + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": logging_obj}, + response=response, + stream_completed=False, + client_disconnected=True, + ) + + for _ in range(50): + if recorder.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + finally: + litellm.callbacks = original_callbacks + + assert len(recorder.success_events) == 1 + standard_logging_object = recorder.success_events[0]["kwargs"]["standard_logging_object"] + assert standard_logging_object["total_tokens"] > 0 + assert standard_logging_object["response_cost"] >= 0.002 + + @pytest.mark.asyncio + async def test_completed_stream_does_not_double_bill_on_late_disconnect(self): + recorder = _RecordingSuccessLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recorder] + try: + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello there", + stream=True, + api_key="test-key", + ) + async for _ in response: + pass + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=response, + stream_completed=False, + client_disconnected=True, + ) + + for _ in range(50): + if recorder.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + finally: + litellm.callbacks = original_callbacks + + assert len(recorder.success_events) == 1 + + @pytest.mark.asyncio + async def test_disconnect_bills_partial_spend_for_router_stream(self): + """ + The router wraps streamed responses in FallbackStreamWrapper, whose + __anext__ bypasses the base class, so its own chunk list stays empty + unless it aliases the inner stream's chunks; without the alias the + disconnect path sees no chunks and bills nothing for router requests, + which is every proxy request. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + } + ] + ) + recorder = _RecordingSuccessLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recorder] + try: + response = await router.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "tell me a story"}], + mock_response="The codename is AZURE-FALCON-42 and the story is long.", + stream=True, + ) + stream_iter = response.__aiter__() + await stream_iter.__anext__() + await stream_iter.__anext__() + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=response, + stream_completed=False, + client_disconnected=True, + ) + + for _ in range(50): + if recorder.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + finally: + litellm.callbacks = original_callbacks + + assert len(recorder.success_events) == 1 + standard_logging_object = recorder.success_events[0]["kwargs"]["standard_logging_object"] + assert standard_logging_object["total_tokens"] > 0 + + @pytest.mark.asyncio + async def test_disconnect_billing_does_not_double_release_slot(self): + """ + The disconnect billing fires a success event whose limiter callback + already releases the max_parallel_requests slot. The shielded cleanup + must therefore NOT also release the slot explicitly; two releases of + the same acquisition race and double-decrement under the limiter's + in-memory fallback. + """ + import types + + original_callbacks = litellm.callbacks + litellm.callbacks = [_RecordingSuccessLogger()] + try: + response = await self._start_partial_stream() + proxy_logging_obj = types.SimpleNamespace( + _arelease_max_parallel_requests_on_disconnect=AsyncMock(), + ) + + billed = await _bill_partial_streamed_spend_on_disconnect( + {"litellm_logging_obj": response.logging_obj}, response + ) + assert billed is True + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=response, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=proxy_logging_obj, + ) + finally: + litellm.callbacks = original_callbacks + + proxy_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_not_called() + + @pytest.mark.asyncio + async def test_disconnect_without_billable_chunks_releases_slot(self): + """ + When there is nothing to bill (no chunks streamed), no success event + fires, so the slot would leak unless the cleanup releases it + explicitly. The explicit release must run exactly once in that case. + """ + import types + + response = await self._start_partial_stream() + # No chunks to assemble -> billing dispatches no success event. + empty_response = types.SimpleNamespace(chunks=[], messages=None) + proxy_logging_obj = types.SimpleNamespace( + _arelease_max_parallel_requests_on_disconnect=AsyncMock(), + ) + + await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( + request=None, + request_data={"litellm_logging_obj": response.logging_obj}, + response=empty_response, + stream_completed=False, + client_disconnected=True, + user_api_key_dict=MagicMock(), + proxy_logging_obj=proxy_logging_obj, + ) + + proxy_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 6b0c0dba40fc..5d2236fd9183 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -582,6 +582,79 @@ def test_skip_server_startup( ), f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() + @patch("uvicorn.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_limit_concurrency_passed_to_uvicorn( + self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run + ): + """--limit_concurrency must reach uvicorn.run so uvicorn sheds load with 503 + past the cap; omitted values stay absent and non-positive values are rejected.""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.side_effect = lambda *a, **k: { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, ["--local", "--limit_concurrency", "250"] + ) + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_uvicorn_run.assert_called_once() + assert mock_uvicorn_run.call_args.kwargs.get("limit_concurrency") == 250 + + mock_uvicorn_run.reset_mock() + result = runner.invoke(run_server, ["--local"]) + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_uvicorn_run.assert_called_once() + assert "limit_concurrency" not in mock_uvicorn_run.call_args.kwargs + + for invalid_value in ("0", "-1"): + mock_uvicorn_run.reset_mock() + result = runner.invoke( + run_server, + ["--local", "--limit_concurrency", invalid_value], + ) + assert result.exit_code == 2 + assert "Invalid value for '--limit_concurrency'" in result.output + mock_uvicorn_run.assert_not_called() + @pytest.mark.parametrize( "timeout_config,expected_timeout", [ diff --git a/tests/test_litellm/proxy/types_utils/test_get_instance_fn_runtime_gate.py b/tests/test_litellm/proxy/types_utils/test_get_instance_fn_runtime_gate.py index bf77ef816413..3d76ad54a9c5 100644 --- a/tests/test_litellm/proxy/types_utils/test_get_instance_fn_runtime_gate.py +++ b/tests/test_litellm/proxy/types_utils/test_get_instance_fn_runtime_gate.py @@ -64,6 +64,65 @@ def test_dotted_module_path_is_unaffected_by_gate(): assert result == "loaded" +def test_installed_package_resolved_when_local_file_absent(tmp_path, monkeypatch): + # Regression: with config_file_path set (startup load path) but no local + # module file next to it, get_instance_fn must fall back to importing the + # dotted name as an installed package. Previously it raised ImportError + # ("Could not find module file ..."), so plugins shipped as pip packages + # (e.g. router_settings/complexity_router plugins) could not be referenced. + pkg_dir = tmp_path / "site" + pkg_dir.mkdir() + (pkg_dir / "my_installed_plugin.py").write_text( + "class _P:\n" + " async def run(self, context):\n" + " return context\n" + "\n" + "instance = _P()\n" + ) + monkeypatch.syspath_prepend(str(pkg_dir)) + config_dir = tmp_path / "cfg" + config_dir.mkdir() + + result = get_instance_fn( + value="my_installed_plugin.instance", + config_file_path=str(config_dir / "config.yaml"), + ) + + assert type(result).__name__ == "_P" + + +def test_local_module_file_wins_over_installed_package(tmp_path, monkeypatch): + # A local module file next to the config must still take precedence over an + # installed package of the same dotted name -- the fallback only kicks in + # when no local file exists. + pkg_dir = tmp_path / "site" + pkg_dir.mkdir() + (pkg_dir / "shadowed_mod.py").write_text("value = 'from-installed'\n") + monkeypatch.syspath_prepend(str(pkg_dir)) + config_dir = tmp_path / "cfg" + config_dir.mkdir() + (config_dir / "shadowed_mod.py").write_text("value = 'from-local-file'\n") + + result = get_instance_fn( + value="shadowed_mod.value", + config_file_path=str(config_dir / "config.yaml"), + ) + + assert result == "from-local-file" + + +def test_missing_module_everywhere_raises_import_error(tmp_path): + # Neither a local file nor an installed package: the fallback import must + # surface a real ImportError rather than silently succeeding. + config_dir = tmp_path / "cfg" + config_dir.mkdir() + with pytest.raises(ImportError): + get_instance_fn( + value="definitely_not_a_real_module_xyz.instance", + config_file_path=str(config_dir / "config.yaml"), + ) + + def test_pass_through_route_threads_config_file_path(): # ``create_pass_through_route`` must forward ``config_file_path`` so # an operator with ``custom_handler: s3://...`` declared in diff --git a/tests/test_litellm/rag/__init__.py b/tests/test_litellm/rag/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py new file mode 100644 index 000000000000..584124ba06ac --- /dev/null +++ b/tests/test_litellm/rag/test_main.py @@ -0,0 +1,266 @@ +""" +Tests for the RAG query pipeline in litellm/rag/main.py. + +The RAG pipeline forwards its kwargs (including the parent litellm_logging_obj) +into @client-decorated sub-calls (vector store search, completion). Each logging +object allows exactly one async_success event, so if sub-calls are not marked as +internal, the vector store search consumes the slot first and the LLM +completion's usage/cost is never logged (spend tracking and budget enforcement +are bypassed). These tests pin the invariant that the single billing event for +aquery carries the completion response with real usage and cost. +""" + +import asyncio +from unittest.mock import patch + +import pytest + +import litellm +from litellm._internal_context import is_internal_call +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import CallTypes, ModelResponse + + +class RecordingLogger(CustomLogger): + def __init__(self): + super().__init__() + self.success_events = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_events.append({"kwargs": kwargs, "response_obj": response_obj}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_router", [False, True]) +async def test_aquery_single_billing_event_carries_completion_usage_and_cost(use_router): + """ + litellm.aquery must produce exactly one success event, and that event must + carry the LLM completion (a ModelResponse with non-zero usage and cost), + not the vector store search response. The proxy always passes a router, so + both the router and non-router completion branches are pinned. + """ + recording_logger = RecordingLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recording_logger] + + router_kwargs = {} + if use_router: + router_kwargs["router"] = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + } + ] + ) + + try: + response = await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What is the secret project codename?"}], + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + mock_response="The secret project codename is AZURE-FALCON-42.", + **router_kwargs, + ) + + assert isinstance(response, ModelResponse) + assert is_internal_call.get() is False + + for _ in range(50): + if recording_logger.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + finally: + litellm.callbacks = original_callbacks + + assert len(recording_logger.success_events) == 1 + event = recording_logger.success_events[0] + + response_obj = event["response_obj"] + assert isinstance(response_obj, ModelResponse) + assert response_obj.usage.total_tokens > 0 + + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["call_type"] == "aquery" + assert standard_logging_object["total_tokens"] > 0 + assert standard_logging_object["prompt_tokens"] > 0 + assert standard_logging_object["completion_tokens"] > 0 + assert standard_logging_object["response_cost"] > 0 + + +@pytest.mark.asyncio +async def test_aquery_response_hidden_params_carry_completion_cost(): + """ + The aquery response must expose the completion's response_cost via hidden + params, so the proxy can return the x-litellm-response-cost header. + """ + response = await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + mock_response="hi there", + ) + + assert isinstance(response, ModelResponse) + response_cost = response._hidden_params.get("response_cost") + assert response_cost is not None + assert response_cost > 0 + + +@pytest.mark.asyncio +async def test_aquery_billed_cost_includes_priced_vector_store_search(): + """ + When the vector store provider prices search calls (e.g. per-query cost), + that cost must be folded into the aquery billing instead of being dropped + with the suppressed sub-call event. + """ + recording_logger = RecordingLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recording_logger] + + try: + with patch("litellm.rag.main.vector_store_search_cost", return_value=(0.002, 0.0)): + response = await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + mock_response="hi there", + ) + + for _ in range(50): + if recording_logger.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + finally: + litellm.callbacks = original_callbacks + + assert isinstance(response, ModelResponse) + total_cost = response._hidden_params.get("response_cost") + assert total_cost is not None + assert total_cost > 0.002 + + assert len(recording_logger.success_events) == 1 + standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] == total_cost + + +@pytest.mark.asyncio +async def test_aquery_with_rerank_bills_once_and_folds_rerank_cost(): + """ + When rerank is enabled, its sub-call must run under the internal-call + context (no standalone billing event) and its cost must be folded into + the single aquery billing event. + """ + from litellm.types.rerank import RerankResponse + + recording_logger = RecordingLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recording_logger] + rerank_seen = {} + + async def fake_arerank(**kwargs): + rerank_seen["internal"] = is_internal_call.get() + rerank_result = RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={}) + rerank_result._hidden_params["response_cost"] = 0.001 + return rerank_result + + try: + with patch("litellm.arerank", side_effect=fake_arerank): + response = await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1}, + mock_response="hi there", + ) + + for _ in range(50): + if recording_logger.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + finally: + litellm.callbacks = original_callbacks + + assert rerank_seen["internal"] is True + assert is_internal_call.get() is False + + assert isinstance(response, ModelResponse) + total_cost = response._hidden_params.get("response_cost") + assert total_cost is not None + assert total_cost > 0.001 + + assert len(recording_logger.success_events) == 1 + standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"] + assert standard_logging_object["call_type"] == "aquery" + assert standard_logging_object["response_cost"] == total_cost + + +@pytest.mark.asyncio +async def test_aquery_streaming_bills_sub_call_costs_into_final_event(): + """ + On the streaming path the response cost is computed from the assembled + chunks after the pipeline returns, so there is no response object to fold + sub-call costs into. The pipeline must instead carry the accumulated + search and rerank cost through the logging object so the single streamed + billing event includes it; otherwise a caller passing stream=true incurs + priced vector search and rerank costs that never reach spend tracking. + """ + from litellm.types.rerank import RerankResponse + + recording_logger = RecordingLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recording_logger] + rerank_seen = {} + + async def fake_arerank(**kwargs): + rerank_seen["internal"] = is_internal_call.get() + rerank_result = RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={}) + rerank_result._hidden_params["response_cost"] = 0.001 + return rerank_result + + try: + with ( + patch("litellm.rag.main.vector_store_search_cost", return_value=(0.002, 0.0)), + patch("litellm.arerank", side_effect=fake_arerank), + ): + response = await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1}, + mock_response="hi there", + stream=True, + ) + async for _ in response: + pass + + for _ in range(50): + if recording_logger.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + finally: + litellm.callbacks = original_callbacks + + assert rerank_seen["internal"] is True + assert is_internal_call.get() is False + + assert len(recording_logger.success_events) == 1 + standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"] + assert standard_logging_object["call_type"] == "aquery" + assert standard_logging_object["response_cost"] >= 0.003 + + +def test_rag_call_types_are_registered(): + """ + query/aquery/ingest/aingest are @client-decorated entry points, so their + function names must resolve to CallTypes members (deployment hooks and + call-type driven logic silently no-op for unregistered call types). + """ + assert CallTypes("query") is CallTypes.query + assert CallTypes("aquery") is CallTypes.aquery + assert CallTypes("ingest") is CallTypes.ingest + assert CallTypes("aingest") is CallTypes.aingest diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 6fdbb0741aa6..a1347aa111c0 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -8,6 +8,7 @@ from fastapi import HTTPException import importlib +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) @@ -455,7 +456,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch Regression test for 872e5b98...: Ensure responses-side tool discovery enables list-tools SpendLogs logging flags. """ - mock_get_tools = AsyncMock(return_value=[]) + mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=[], outcomes={})) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", mock_get_tools, @@ -509,7 +510,7 @@ def test_get_parent_request_tags_from_nested_litellm_params(): @pytest.mark.asyncio async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): - mock_get_tools = AsyncMock(return_value=[]) + mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=[], outcomes={})) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", mock_get_tools, diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py b/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py index 2703d7b28c6b..9e3f6ec1a569 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py @@ -21,6 +21,11 @@ from litellm.types.router import LiteLLM_Params, RequestType +def _adaptive(r, name): + """Registries hold tag-scoped strategy lists; these tests use a single tagless entry.""" + return r.adaptive_routers[name][0].strategy + + def _params(**overrides): base = {"model": "auto_router/adaptive_router"} base.update(overrides) @@ -122,7 +127,7 @@ def test_init_adaptive_router_reads_cost_from_litellm_params(): ] ) assert "smart-cheap-router" in r.adaptive_routers - assert r.adaptive_routers["smart-cheap-router"].model_to_cost == { + assert _adaptive(r, "smart-cheap-router").model_to_cost == { "fast": 0.00000015, "smart": 0.0000050, } @@ -176,7 +181,7 @@ def _router_with_adaptive() -> Router: @pytest.mark.asyncio async def test_async_pre_routing_hook_dispatches_to_adaptive_router(): r = _router_with_adaptive() - ar = r.adaptive_routers["smart-cheap-router"] + ar = _adaptive(r, "smart-cheap-router") ar.pick_model = AsyncMock(return_value="smart") # type: ignore[assignment] response = await r.async_pre_routing_hook( @@ -195,7 +200,7 @@ async def test_async_pre_routing_hook_dispatches_to_adaptive_router(): @pytest.mark.asyncio async def test_async_pre_routing_hook_pick_model_not_passed_session_id(): r = _router_with_adaptive() - ar = r.adaptive_routers["smart-cheap-router"] + ar = _adaptive(r, "smart-cheap-router") ar.pick_model = AsyncMock(return_value="fast") # type: ignore[assignment] response = await r.async_pre_routing_hook( @@ -211,7 +216,7 @@ async def test_async_pre_routing_hook_pick_model_not_passed_session_id(): @pytest.mark.asyncio async def test_async_pre_routing_hook_returns_none_for_unrelated_model(): r = _router_with_adaptive() - ar = r.adaptive_routers["smart-cheap-router"] + ar = _adaptive(r, "smart-cheap-router") ar.pick_model = AsyncMock() # type: ignore[assignment] response = await r.async_pre_routing_hook( model="some-other-model", @@ -233,7 +238,7 @@ async def test_async_pre_routing_hook_stashes_chosen_model_in_metadata(): `x-litellm-adaptive-router-model` response header. """ r = _router_with_adaptive() - r.adaptive_routers["smart-cheap-router"].pick_model = AsyncMock( # type: ignore[assignment] + _adaptive(r, "smart-cheap-router").pick_model = AsyncMock( # type: ignore[assignment] return_value="smart" ) @@ -250,7 +255,7 @@ async def test_async_pre_routing_hook_stashes_chosen_model_in_metadata(): async def test_async_pre_routing_hook_creates_metadata_when_missing(): """If no metadata was passed in, the hook should create one to stash the chosen model.""" r = _router_with_adaptive() - r.adaptive_routers["smart-cheap-router"].pick_model = AsyncMock( # type: ignore[assignment] + _adaptive(r, "smart-cheap-router").pick_model = AsyncMock( # type: ignore[assignment] return_value="fast" ) @@ -300,8 +305,8 @@ def test_two_adaptive_routers_can_coexist_on_one_router(): ] ) assert set(r.adaptive_routers.keys()) == {"cheap-router", "premium-router"} - assert r.adaptive_routers["cheap-router"].config.available_models == ["fast"] - assert r.adaptive_routers["premium-router"].config.available_models == ["smart"] + assert _adaptive(r, "cheap-router").config.available_models == ["fast"] + assert _adaptive(r, "premium-router").config.available_models == ["smart"] @pytest.mark.asyncio @@ -339,8 +344,8 @@ async def test_async_pre_routing_hook_dispatches_to_correct_router_when_multiple }, ] ) - cheap = r.adaptive_routers["cheap-router"] - premium = r.adaptive_routers["premium-router"] + cheap = _adaptive(r, "cheap-router") + premium = _adaptive(r, "premium-router") cheap.pick_model = AsyncMock(return_value="fast") # type: ignore[assignment] premium.pick_model = AsyncMock(return_value="smart") # type: ignore[assignment] @@ -410,12 +415,12 @@ def test_finalize_adaptive_router_if_configured_initializes_and_is_idempotent(): # Router __init__ already called _finalize_adaptive_router_if_configured. assert "my-router" in r.adaptive_routers - original = r.adaptive_routers["my-router"] + original = _adaptive(r, "my-router") # Calling again must be idempotent: the existing AdaptiveRouter instance # is preserved, not rebuilt. r._finalize_adaptive_router_if_configured() - assert r.adaptive_routers["my-router"] is original + assert _adaptive(r, "my-router") is original def test_finalize_prunes_stale_adaptive_router_hooks_from_callbacks(): diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py b/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py index d6d89c8e8112..5662870a5cb9 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py @@ -13,6 +13,7 @@ AdaptiveRouterConfig, AdaptiveRouterPreferences, RequestType, + TaggedPreRoutingStrategy, ) @@ -33,6 +34,10 @@ def _make_router(name: str = "r1") -> AdaptiveRouter: ) +def _entry(name: str = "r1") -> list: + return [TaggedPreRoutingStrategy(tags=(), strategy=_make_router(name))] + + # ---- snapshot helper --------------------------------------------------- @@ -127,7 +132,7 @@ async def test_endpoint_rejects_non_admin_role(monkeypatch): from litellm.proxy import proxy_server fake_router = MagicMock() - fake_router.adaptive_routers = {"r1": _make_router()} + fake_router.adaptive_routers = {"r1": _entry()} monkeypatch.setattr(proxy_server, "llm_router", fake_router) non_admin = UserAPIKeyAuth( @@ -144,7 +149,7 @@ async def test_endpoint_returns_snapshot_list_for_admin(monkeypatch): from litellm.proxy import proxy_server fake_router = MagicMock() - fake_router.adaptive_routers = {"r1": _make_router("r1")} + fake_router.adaptive_routers = {"r1": _entry("r1")} monkeypatch.setattr(proxy_server, "llm_router", fake_router) admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) @@ -164,8 +169,8 @@ async def test_endpoint_returns_one_snapshot_per_router(monkeypatch): fake_router = MagicMock() fake_router.adaptive_routers = { - "r1": _make_router("r1"), - "r2": _make_router("r2"), + "r1": _entry("r1"), + "r2": _entry("r2"), } monkeypatch.setattr(proxy_server, "llm_router", fake_router) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 12b2c9abefb0..280a0fe072a8 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -30,6 +30,11 @@ ComplexityRouterConfig, ComplexityTier, ) +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + TaggedPreRoutingStrategy, +) @pytest.fixture @@ -953,7 +958,7 @@ def test_hybrid_initialization_waits_for_later_pool_deployments(self): ] ) - adaptive = router.adaptive_routers["hybrid"] + adaptive = router.adaptive_routers["hybrid"][0].strategy assert adaptive.model_to_cost == { "cheap": pytest.approx(0.00000015), "premium": pytest.approx(0.000005), @@ -962,6 +967,138 @@ def test_hybrid_initialization_waits_for_later_pool_deployments(self): assert adaptive.model_to_prefs["premium"].quality_tier == 3 +class TestComplexityRouterTagBasedRouting: + """Regression tests for https://github.com/BerriAI/litellm/issues/33655. + + Two complexity-router deployments can share a public model_name while + carrying different tags. Both must register, and the request's tags must + pick the matching config before classification (previously the second + deployment was rejected and every request used the first config).""" + + @staticmethod + def _tagged_config(routed_model: str, tags: list) -> dict: + return { + "model_name": "smart", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": routed_model, + "complexity_router_config": { + "tiers": { + "SIMPLE": [routed_model], + "MEDIUM": [routed_model], + "COMPLEX": [routed_model], + "REASONING": [routed_model], + } + }, + "tags": tags, + }, + } + + def _router(self) -> Router: + return Router( + model_list=[ + self._tagged_config("gpt-cn", ["cn"]), + self._tagged_config("gpt-us", ["us"]), + ] + ) + + def test_both_tagged_configs_register_under_same_model_name(self): + router = self._router() + registered = router.complexity_routers["smart"] + assert len(registered) == 2 + assert {entry.tags for entry in registered} == {("cn",), ("us",)} + + def test_duplicate_model_name_with_same_tags_still_rejected(self): + with pytest.raises(ValueError, match="already exists"): + Router( + model_list=[ + self._tagged_config("gpt-cn", ["cn"]), + self._tagged_config("gpt-cn-2", ["cn"]), + ] + ) + + @pytest.mark.asyncio + async def test_request_tags_select_matching_complexity_config(self): + router = self._router() + cn = await router.async_pre_routing_hook( + model="smart", + request_kwargs={"metadata": {"tags": ["cn"]}}, + messages=[{"role": "user", "content": "hi"}], + ) + us = await router.async_pre_routing_hook( + model="smart", + request_kwargs={"metadata": {"tags": ["us"]}}, + messages=[{"role": "user", "content": "hi"}], + ) + assert cn is not None and cn.model == "gpt-cn" + assert us is not None and us.model == "gpt-us" + + +class TestPreRoutingStrategyRegistry: + """Directly exercise the tag-scoped registry/selection helpers behind #33655.""" + + def _router(self) -> Router: + return Router(model_list=[{"model_name": "x", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) + + @staticmethod + def _deployment(tags: list) -> Deployment: + return Deployment( + model_name="smart", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini", tags=tags), + ) + + def test_deployment_tags_normalizes_to_tuple(self): + router = self._router() + assert router._deployment_tags(self._deployment(["cn", "row"])) == ("cn", "row") + untagged = Deployment(model_name="smart", litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini")) + assert router._deployment_tags(untagged) == () + + def test_register_scopes_by_tags_and_rejects_exact_duplicate(self): + router = self._router() + registry: dict = {} + router._register_pre_routing_strategy( + registry=registry, deployment=self._deployment(["cn"]), strategy="CN", strategy_label="Test" + ) + router._register_pre_routing_strategy( + registry=registry, deployment=self._deployment(["us"]), strategy="US", strategy_label="Test" + ) + assert [entry.tags for entry in registry["smart"]] == [("cn",), ("us",)] + assert router._has_registered_strategy(registry, "smart", ("cn",)) is True + assert router._has_registered_strategy(registry, "smart", ("row",)) is False + with pytest.raises(ValueError, match="already exists"): + router._register_pre_routing_strategy( + registry=registry, deployment=self._deployment(["cn"]), strategy="CN2", strategy_label="Test" + ) + + def test_select_prefers_request_tag_then_default_then_first(self): + router = self._router() + cn, us, fallback = object(), object(), object() + router.complexity_routers = { + "smart": [ + TaggedPreRoutingStrategy(tags=("cn",), strategy=cn), + TaggedPreRoutingStrategy(tags=("us",), strategy=us), + ] + } + assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["us"]}}) is us + assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["cn"]}}) is cn + assert router._select_pre_routing_strategy("missing", {"metadata": {"tags": ["cn"]}}) is None + + router.complexity_routers = { + "smart": [ + TaggedPreRoutingStrategy(tags=("cn",), strategy=cn), + TaggedPreRoutingStrategy(tags=("default",), strategy=fallback), + ] + } + assert router._select_pre_routing_strategy("smart", {}) is fallback + router.complexity_routers = { + "smart": [ + TaggedPreRoutingStrategy(tags=("cn",), strategy=cn), + TaggedPreRoutingStrategy(tags=("us",), strategy=us), + ] + } + assert router._select_pre_routing_strategy("smart", {}) is cn + + class TestAsyncPreRoutingHookMultiFormat: """Test async_pre_routing_hook with multiple input formats.""" @@ -2905,9 +3042,7 @@ async def run(self, context): assert result.model == "gpt-4o-nano" @pytest.mark.asyncio - async def test_no_user_message_prefers_default_model_over_medium_tier_without_plugins( - self, mock_router_instance - ): + async def test_no_user_message_prefers_default_model_over_medium_tier_without_plugins(self, mock_router_instance): """Regression: without plugins configured, the no-user-message path must keep its pre-existing default_model-first priority over the MEDIUM tier exactly as before -- closing the plugin-bypass gap must not silently flip model selection for the (much @@ -2984,3 +3119,260 @@ async def run(self, context): assert first.model == "gpt-4o-mini" assert second.model == "gpt-4o-mini" assert spy.call_count == 2 + + +class TestEscalationKeywords: + """Test user-triggered escalation: a keyword in the prompt bumps the resolved tier + one step higher so a user can force a stronger model when unhappy with results.""" + + @staticmethod + def _request_kwargs(session_id: str) -> Dict: + return {"metadata": {"session_id": session_id}} + + def test_default_escalation_keyword(self, complexity_router): + assert complexity_router.escalation_keywords == ["LITELLM ESCALATE"] + + def test_escalation_triggered_is_case_sensitive(self, complexity_router): + assert complexity_router._escalation_triggered("please LITELLM ESCALATE now") is True + assert complexity_router._escalation_triggered("please litellm escalate now") is False + assert complexity_router._escalation_triggered("how do I escalate this ticket") is False + + def test_escalate_tier_bumps_one_step(self, complexity_router): + assert complexity_router._escalate_tier(ComplexityTier.SIMPLE) == ComplexityTier.MEDIUM + assert complexity_router._escalate_tier(ComplexityTier.MEDIUM) == ComplexityTier.COMPLEX + assert complexity_router._escalate_tier(ComplexityTier.COMPLEX) == ComplexityTier.REASONING + + def test_escalate_tier_caps_at_highest_configured(self, complexity_router): + assert complexity_router._escalate_tier(ComplexityTier.REASONING) == ComplexityTier.REASONING + + def test_escalate_tier_skips_unconfigured_intermediate(self, mock_router_instance): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": "o1-preview"}}, + ) + assert router._escalate_tier(ComplexityTier.SIMPLE) == ComplexityTier.REASONING + + def test_tier_for_model_returns_most_severe(self, mock_router_instance): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "shared", "COMPLEX": "shared", "REASONING": "top"} + }, + ) + assert router._tier_for_model("shared") == ComplexityTier.COMPLEX + assert router._tier_for_model("top") == ComplexityTier.REASONING + assert router._tier_for_model("unknown") is None + + @pytest.mark.asyncio + async def test_escalation_bumps_classified_tier(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=basic_config, + ) + # Baseline: this prompt classifies SIMPLE. + baseline = await router.async_pre_routing_hook( + model="test-model", request_kwargs={}, messages=[{"role": "user", "content": "Hello there!"}] + ) + assert baseline.model == "gpt-4o-mini" + + escalated = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "LITELLM ESCALATE Hello there!"}], + ) + assert escalated.model == "gpt-4o" # SIMPLE bumped to MEDIUM + + @pytest.mark.asyncio + async def test_lowercase_keyword_does_not_escalate(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=basic_config, + ) + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "litellm escalate Hello there!"}], + ) + assert result.model == "gpt-4o-mini" # not escalated + + @pytest.mark.asyncio + async def test_custom_escalation_keyword(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "escalation_keywords": ["MAKE IT BETTER"]}, + ) + # The default keyword no longer triggers once a custom list is supplied. + default = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "LITELLM ESCALATE Hello there!"}], + ) + assert default.model == "gpt-4o-mini" + + custom = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "MAKE IT BETTER Hello there!"}], + ) + assert custom.model == "gpt-4o" + + @pytest.mark.asyncio + async def test_empty_keyword_list_disables_escalation(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "escalation_keywords": []}, + ) + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "LITELLM ESCALATE Hello there!"}], + ) + assert result.model == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_escalation_caps_at_highest_tier(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=basic_config, + ) + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[ + { + "role": "user", + "content": "LITELLM ESCALATE Let's think step by step and reason through this carefully.", + } + ], + ) + assert result.model == "o1-preview" # already REASONING, stays there + + @pytest.mark.asyncio + async def test_escalation_bumps_keyword_tier_override(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "keyword_tier_rules": [{"keywords": ["billing"], "tier": "SIMPLE"}], + }, + ) + baseline = await router.async_pre_routing_hook( + model="test-model", request_kwargs={}, messages=[{"role": "user", "content": "a billing question"}] + ) + assert baseline.model == "gpt-4o-mini" + + escalated = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "LITELLM ESCALATE a billing question"}], + ) + assert escalated.model == "gpt-4o" # override SIMPLE bumped to MEDIUM + + @pytest.mark.asyncio + async def test_escalation_overrides_session_pin_and_persists(self, mock_router_instance, basic_config): + """Mid-session escalation bumps relative to the pinned model (never below it) and + the bumped model persists for later turns.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "session_affinity": True}, + ) + request_kwargs = self._request_kwargs("session-1") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=[{"role": "user", "content": "Hello!"}] + ) + assert first.model == "gpt-4o-mini" # pinned SIMPLE + + with patch.object(router, "aclassify", wraps=router.aclassify) as spy_aclassify: + escalated = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "LITELLM ESCALATE"}], + ) + spy_aclassify.assert_not_called() + assert escalated.model == "gpt-4o" # bumped relative to the SIMPLE pin, not reclassified + + # The bump persists: a later ordinary turn stays on the escalated model. + later = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=[{"role": "user", "content": "thanks"}] + ) + assert later.model == "gpt-4o" + + # Escalating again climbs one more tier. + again = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "LITELLM ESCALATE still not good"}], + ) + assert again.model == "claude-sonnet-4-20250514" # MEDIUM bumped to COMPLEX + + def test_blank_escalation_keywords_are_stripped(self): + """Blank/whitespace-only phrases are dropped so `"" in message` can't escalate + every request; surrounding whitespace on real phrases is trimmed.""" + assert ComplexityRouterConfig( + tiers={"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + escalation_keywords=["", " "], + ).escalation_keywords == [] + assert ComplexityRouterConfig( + tiers={"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + escalation_keywords=[" LITELLM ESCALATE ", ""], + ).escalation_keywords == ["LITELLM ESCALATE"] + + @pytest.mark.asyncio + async def test_blank_escalation_keyword_does_not_escalate_everything( + self, mock_router_instance, basic_config + ): + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "escalation_keywords": [""]}, + ) + assert router.escalation_keywords == [] + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "Hello there!"}], + ) + assert result.model == "gpt-4o-mini" # not escalated + + def test_escalated_pin_stays_on_same_model_at_ceiling(self, mock_router_instance): + """At the highest configured tier escalation keeps the exact pinned model, even + when that tier's pool has peers `get_model_for_tier` could randomly pick instead.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]} + }, + ) + for pinned in ("o1-a", "o1-b", "o1-c"): + assert router._escalated_pin(pinned) == pinned + + @pytest.mark.asyncio + async def test_session_escalation_at_ceiling_keeps_multi_model_pin(self, mock_router_instance): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}, + "session_affinity": True, + }, + ) + cache_key = router._get_session_affinity_cache_key("session-top", {}) + await mock_router_instance.cache.async_set_cache(key=cache_key, value="o1-b") + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("session-top"), + messages=[{"role": "user", "content": "LITELLM ESCALATE do better"}], + ) + assert result.model == "o1-b" # unchanged: no random hop to o1-a / o1-c diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/test_litellm/test_gpt_realtime_mode.py new file mode 100644 index 000000000000..80cb3cc85f05 --- /dev/null +++ b/tests/test_litellm/test_gpt_realtime_mode.py @@ -0,0 +1,82 @@ +import json +import typing +from pathlib import Path + +import pytest + +import litellm +from litellm.types.utils import ModelInfoBase + +REALTIME_ONLY_GPT_MODELS = ( + "azure/gpt-realtime-2025-08-28", + "azure/gpt-realtime-1.5-2026-02-23", + "azure/gpt-realtime-mini-2025-10-06", + "gpt-realtime", + "gpt-realtime-1.5", + "gpt-realtime-2", + "gpt-realtime-2.1", + "gpt-realtime-2.1-mini", + "gpt-realtime-mini", + "gpt-realtime-2025-08-28", + "gpt-realtime-mini-2025-10-06", + "gpt-realtime-mini-2025-12-15", +) + +REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS = ( + "azure/eu/gpt-4o-mini-realtime-preview-2024-12-17", + "azure/eu/gpt-4o-realtime-preview-2024-10-01", + "azure/eu/gpt-4o-realtime-preview-2024-12-17", + "azure/gpt-4o-mini-realtime-preview-2024-12-17", + "azure/gpt-4o-realtime-preview-2024-10-01", + "azure/gpt-4o-realtime-preview-2024-12-17", + "azure/us/gpt-4o-mini-realtime-preview-2024-12-17", + "azure/us/gpt-4o-realtime-preview-2024-10-01", + "azure/us/gpt-4o-realtime-preview-2024-12-17", + "gpt-4o-mini-realtime-preview", + "gpt-4o-mini-realtime-preview-2024-12-17", + "gpt-4o-realtime-preview", + "gpt-4o-realtime-preview-2024-12-17", + "gpt-4o-realtime-preview-2025-06-03", +) + +ALL_REALTIME_ONLY_GPT_MODELS = REALTIME_ONLY_GPT_MODELS + REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS + + +def _load_cost_map() -> dict: + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + with open(json_path) as f: + return json.load(f) + + +def test_realtime_is_a_valid_mode_literal(): + hints = typing.get_type_hints(ModelInfoBase, include_extras=False) + assert "realtime" in typing.get_args(hints["mode"]) + + +@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS) +def test_realtime_only_gpt_models_are_mode_realtime(model): + """These models only serve /v1/realtime and are rejected by /v1/chat/completions + ("This is not a chat model ..."), so they must not be tagged mode=chat.""" + info = _load_cost_map()[model] + assert info["supported_endpoints"] == ["/v1/realtime"] + assert info["mode"] == "realtime" + + +@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS) +def test_realtime_only_gpt_4o_models_are_mode_realtime(model): + """gpt-4o(-mini)-realtime-preview are realtime-only and must not be mode=chat.""" + assert _load_cost_map()[model]["mode"] == "realtime" + + +def test_get_model_info_reports_realtime_mode(): + assert litellm.get_model_info("gpt-realtime-mini")["mode"] == "realtime" + + +def test_backup_matches_main_for_realtime_models(): + repo_root = Path(__file__).parents[2] + with open(repo_root / "model_prices_and_context_window.json") as f: + main_cost = json.load(f) + with open(repo_root / "litellm" / "model_prices_and_context_window_backup.json") as f: + backup_cost = json.load(f) + for model in ALL_REALTIME_ONLY_GPT_MODELS: + assert backup_cost.get(model) == main_cost.get(model) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 073ff17991e6..edd93cbebe0a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4717,6 +4717,55 @@ def test_reports_key_missing(self): assert "TENCENT_API_KEY" in result["missing_keys"] +class TestVertexEmbeddingEncodingFormat: + """vertex_ai/gemini embeddings must accept encoding_format="float" — it's + the OpenAI SDK default and float lists are exactly what the vertex API + returns. Other values keep the unsupported-param behavior (drop with + drop_params, raise otherwise). Issue #33173.""" + + def test_encoding_format_float_is_accepted_and_dropped(self): + optional_params = litellm.utils.get_optional_params_embeddings( + model="gemini-embedding-001", + encoding_format="float", + custom_llm_provider="vertex_ai", + ) + assert "encoding_format" not in optional_params + + def test_encoding_format_float_accepted_for_gemini_provider(self): + optional_params = litellm.utils.get_optional_params_embeddings( + model="gemini-embedding-001", + encoding_format="float", + custom_llm_provider="gemini", + ) + assert "encoding_format" not in optional_params + + def test_encoding_format_base64_still_rejected_without_drop_params(self): + with pytest.raises(Exception) as excinfo: + litellm.utils.get_optional_params_embeddings( + model="gemini-embedding-001", + encoding_format="base64", + custom_llm_provider="vertex_ai", + ) + assert "encoding_format" in str(excinfo.value) + + def test_encoding_format_base64_dropped_with_drop_params(self): + optional_params = litellm.utils.get_optional_params_embeddings( + model="gemini-embedding-001", + encoding_format="base64", + custom_llm_provider="vertex_ai", + drop_params=True, + ) + assert "encoding_format" not in optional_params + + def test_dimensions_still_mapped(self): + optional_params = litellm.utils.get_optional_params_embeddings( + model="gemini-embedding-001", + encoding_format="float", + dimensions=256, + custom_llm_provider="vertex_ai", + ) + assert optional_params.get("outputDimensionality") == 256 + @pytest.mark.parametrize( "model", diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index e1f902967700..a2e2ca21d00f 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -269,4 +269,22 @@ describe("ComplexityRouterConfig", () => { ); expect(screen.getAllByText("This tier is required")).toHaveLength(1); }); + + it("renders the escalation keywords section with current keywords when the handler is provided", () => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Escalation Keywords")); + expect(screen.getByText("Escalation Keywords")).toBeInTheDocument(); + expect(screen.getByText("LITELLM ESCALATE")).toBeInTheDocument(); + }); + + it("hides the escalation keywords section when no handler is provided", () => { + renderWithProviders(); + expect(screen.queryByText("Advanced: Escalation Keywords")).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 855a1b27df90..8008012a95c2 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -4,6 +4,7 @@ import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; +import EscalationKeywords from "./EscalationKeywords"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; @@ -61,6 +62,8 @@ interface ComplexityRouterConfigProps { onEmbeddingModelChange?: (model: string) => void; matchThreshold?: number; onMatchThresholdChange?: (threshold: number) => void; + escalationKeywords?: string[]; + onEscalationKeywordsChange?: (keywords: string[]) => void; showValidationErrors?: boolean; } @@ -101,6 +104,8 @@ const ComplexityRouterConfig: React.FC = ({ onEmbeddingModelChange = () => {}, matchThreshold = 0.5, onMatchThresholdChange = () => {}, + escalationKeywords = [], + onEscalationKeywordsChange, showValidationErrors = false, }) => { // Embedding models can't serve a chat-completion role, so they're excluded here. @@ -213,6 +218,19 @@ const ComplexityRouterConfig: React.FC = ({ ), children: , }, + ...(onEscalationKeywordsChange + ? [ + { + key: "escalation", + label: ( + + Advanced: Escalation Keywords + + ), + children: , + }, + ] + : []), ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange ? [ { diff --git a/ui/litellm-dashboard/src/components/add_model/EscalationKeywords.tsx b/ui/litellm-dashboard/src/components/add_model/EscalationKeywords.tsx new file mode 100644 index 000000000000..c232eb4c8018 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/EscalationKeywords.tsx @@ -0,0 +1,45 @@ +import { InfoCircleOutlined } from "@ant-design/icons"; +import { Select as AntdSelect, Tooltip, Typography } from "antd"; +import React from "react"; + +const { Text } = Typography; + +export const DEFAULT_ESCALATION_KEYWORDS = ["LITELLM ESCALATE"]; + +interface EscalationKeywordsProps { + keywords: string[]; + onChange: (keywords: string[]) => void; +} + +const EscalationKeywords: React.FC = ({ keywords, onChange }) => { + return ( +
+
+ + Escalation Keywords + + + + +
+ + Optional: when a user message contains one of these phrases, the request is bumped one tier higher than it would + otherwise route to. Matching is case-sensitive, so "LITELLM ESCALATE" only fires on the exact, shouted + form. Leave empty to disable. + + +
+ ); +}; + +export default EscalationKeywords; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 5122d54db9bb..6e7bc49afcea 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -14,6 +14,7 @@ import ComplexityRouterConfig, { DEFAULT_TIER_DISTANCE_PENALTY, } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; +import { DEFAULT_ESCALATION_KEYWORDS } from "./EscalationKeywords"; import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching"; import { buildComplexityRouterConfig, @@ -52,6 +53,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState(false); const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); + const [escalationKeywords, setEscalationKeywords] = useState(DEFAULT_ESCALATION_KEYWORDS); const [showValidationErrors, setShowValidationErrors] = useState(false); // Semantic router config (existing) @@ -141,6 +143,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc semanticMatchingEnabled, embeddingModel, matchThreshold, + escalationKeywords, adaptive, adaptiveWeights, tierDistancePenalty, @@ -316,6 +319,8 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc onEmbeddingModelChange={setEmbeddingModel} matchThreshold={matchThreshold} onMatchThresholdChange={setMatchThreshold} + escalationKeywords={escalationKeywords} + onEscalationKeywordsChange={setEscalationKeywords} showValidationErrors={showValidationErrors} /> diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 85a15ffad45c..0c9c19d1286f 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -21,6 +21,7 @@ const baseParams: BuildComplexityRouterConfigParams = { semanticMatchingEnabled: false, embeddingModel: undefined, matchThreshold: 0.5, + escalationKeywords: ["LITELLM ESCALATE"], adaptive: false, adaptiveWeights: { quality: 0.3, cost: 0.7 }, tierDistancePenalty: 0.5, @@ -28,9 +29,22 @@ const baseParams: BuildComplexityRouterConfigParams = { }; describe("buildComplexityRouterConfig", () => { - it("emits only tiers and classifier_type when nothing else is configured", () => { + it("emits tiers, classifier_type, and escalation_keywords when nothing else is configured", () => { const config = buildComplexityRouterConfig(baseParams); - expect(config).toEqual({ tiers, classifier_type: "heuristic" }); + expect(config).toEqual({ tiers, classifier_type: "heuristic", escalation_keywords: ["LITELLM ESCALATE"] }); + }); + + it("trims escalation keywords and drops blank entries", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + escalationKeywords: [" LITELLM ESCALATE ", "", " ", "MAKE IT BETTER"], + }); + expect(config.escalation_keywords).toEqual(["LITELLM ESCALATE", "MAKE IT BETTER"]); + }); + + it("emits an empty escalation_keywords list so clearing the field disables escalation", () => { + const config = buildComplexityRouterConfig({ ...baseParams, escalationKeywords: [] }); + expect(config.escalation_keywords).toEqual([]); }); it("passes through a tier configured with more than one model as a pool", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 3c3f21163b3c..0b92dc1b02da 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -16,6 +16,7 @@ export interface BuildComplexityRouterConfigParams { semanticMatchingEnabled: boolean; embeddingModel: string | undefined; matchThreshold: number; + escalationKeywords: string[]; adaptive: boolean; adaptiveWeights: AdaptiveRouterWeights; tierDistancePenalty: number; @@ -31,6 +32,7 @@ export interface ComplexityRouterConfigPayload { semantic_keyword_matching?: boolean; embedding_model?: string; match_threshold?: number; + escalation_keywords?: string[]; adaptive?: boolean; adaptive_weights?: AdaptiveRouterWeights; tier_distance_penalty?: number; @@ -69,11 +71,13 @@ export const buildComplexityRouterConfig = ({ semanticMatchingEnabled, embeddingModel, matchThreshold, + escalationKeywords, adaptive, adaptiveWeights, tierDistancePenalty, adaptiveEligible, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { + const cleanedEscalationKeywords = escalationKeywords.map((keyword) => keyword.trim()).filter(Boolean); // Trim keywords and drop empty ones; drop any rule left with no keywords. Clicking // "Add keyword rule" seeds a rule with an empty keywords list, so without this an // unfilled row (common in the heuristic flow, where getSemanticConfigError doesn't run) @@ -88,6 +92,7 @@ export const buildComplexityRouterConfig = ({ ...(classifierType === "llm" && classifierLlmConfig && { classifier_llm_config: classifierLlmConfig }), ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }), + escalation_keywords: cleanedEscalationKeywords, ...(semanticMatchingEnabled && { semantic_keyword_matching: true, embedding_model: embeddingModel, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9ad0b8d1101f..0d8f55164f97 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21690,7 +21690,7 @@ export interface components { * CallTypes * @enum {string} */ - CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; /** CallbackDelete */ CallbackDelete: { /** Callback Name */ @@ -21870,6 +21870,11 @@ export interface components { }; /** ChatCompletionCachedContent */ ChatCompletionCachedContent: { + /** + * Ttl + * @enum {string} + */ + ttl?: "5m" | "1h"; /** * Type * @constant