diff --git a/.gitignore b/.gitignore index 0c976a1a226..b812d45e349 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,9 @@ litellm/rust_bridge/_native*.so litellm/rust_bridge/_native*.pyd litellm-rust/target/ +# Python package build output +dist/ + bun.lockb **/.DS_Store .aider* diff --git a/dist/litellm-1.79.1.tar.gz b/dist/litellm-1.79.1.tar.gz deleted file mode 100644 index 5980922c1b5..00000000000 Binary files a/dist/litellm-1.79.1.tar.gz and /dev/null differ diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c5083a63012..167f448b2bb 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -28,12 +28,6 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" -[[package]] -name = "autocfg" -version = "1.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" - [[package]] name = "axum" version = "0.7.9" @@ -392,21 +386,6 @@ dependencies = [ "percent-encoding", ] -[[package]] -name = "futures" -version = "0.3.32" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" -dependencies = [ - "futures-channel", - "futures-core", - "futures-executor", - "futures-io", - "futures-sink", - "futures-task", - "futures-util", -] - [[package]] name = "futures-channel" version = "0.3.32" @@ -423,17 +402,6 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" -[[package]] -name = "futures-executor" -version = "0.3.32" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" -dependencies = [ - "futures-core", - "futures-task", - "futures-util", -] - [[package]] name = "futures-io" version = "0.3.32" @@ -469,7 +437,6 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ - "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -928,15 +895,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "indoc" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" -dependencies = [ - "rustversion", -] - [[package]] name = "ipnet" version = "2.12.0" @@ -1089,15 +1047,6 @@ version = "2.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - [[package]] name = "mime" version = "0.3.17" @@ -1225,29 +1174,26 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" dependencies = [ - "cfg-if", - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] name = "pyo3-async-runtimes" -version = "0.23.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e" +checksum = "b3ef68daa7316a3fac65e5e18b2203f010346de1c1c53456811a2624673ab046" dependencies = [ - "futures", + "futures-channel", + "futures-util", "once_cell", "pin-project-lite", "pyo3", @@ -1256,19 +1202,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" dependencies = [ - "once_cell", "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" dependencies = [ "libc", "pyo3-build-config", @@ -1276,9 +1221,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -1288,13 +1233,12 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", "syn", ] @@ -1974,9 +1918,9 @@ dependencies = [ [[package]] name = "target-lexicon" -version = "0.12.16" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "thiserror" @@ -2250,12 +2194,6 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - [[package]] name = "untrusted" version = "0.9.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index fa933856ba7..ab6e92eaa9b 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -16,8 +16,8 @@ rust-version = "1.86" litellm-core = { path = "crates/core" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } axum = "0.7" -pyo3 = "0.23.5" -pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] } +pyo3 = "0.29.0" +pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } serde = { version = "1.0", features = ["derive"] } diff --git a/litellm-rust/crates/ai-gateway/src/python/config.rs b/litellm-rust/crates/ai-gateway/src/python/config.rs index 6ec9595469d..54b7a53bafa 100644 --- a/litellm-rust/crates/ai-gateway/src/python/config.rs +++ b/litellm-rust/crates/ai-gateway/src/python/config.rs @@ -17,7 +17,7 @@ use crate::gil; /// Load the router's `model_list` from `config_path` via the Python reader. pub fn load_router_from_config(config_path: &str) -> CoreResult { gil::record_acquisition(); - Python::with_gil(|py| { + Python::attach(|py| { let model_list = py .import("litellm.proxy.read_model_list") .and_then(|module| module.getattr("read_model_list")) diff --git a/litellm-rust/crates/python-bridge/src/gil.rs b/litellm-rust/crates/python-bridge/src/gil.rs index dc1b591735c..e887c8ec1e3 100644 --- a/litellm-rust/crates/python-bridge/src/gil.rs +++ b/litellm-rust/crates/python-bridge/src/gil.rs @@ -2,7 +2,7 @@ //! //! A single chokepoint for releasing the GIL around blocking work. Every //! blocking call in the bridge goes through [`release_gil`] instead of calling -//! `Python::allow_threads` directly, so the release count stays accurate and we +//! `Python::detach` directly, so the release count stays accurate and we //! have one place to extend later (timing histograms, per-call labels, etc.). use std::sync::atomic::{AtomicU64, Ordering}; @@ -23,7 +23,7 @@ where T: Send, { GIL_RELEASES.fetch_add(1, Ordering::Relaxed); - py.allow_threads(f) + py.detach(f) } /// Total GIL releases performed by the bridge so far. diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 0f3ac8fdc6c..c34e8476b41 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -174,7 +174,7 @@ fn aocr( .await .map_err(|err| Python::with_gil(|py| core_error_to_pyerr(py, err)))?; - Python::with_gil(|py| json_to_py(py, value)) + Python::attach(|py| json_to_py(py, value)) }) } diff --git a/litellm/__init__.py b/litellm/__init__.py index 6e2a03b7c7c..2f6643c644c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -315,6 +315,11 @@ disable_token_counter: bool = False 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 608fdebc1d9..94c86e07ff5 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -296,18 +296,148 @@ class AnthropicCacheControlHook(CustomPromptManagement): 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 7222c9d0502..d4850d50778 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] +def create_uuid7() -> str: + """Generate an RFC 9562 conformant UUIDv7 string. - # 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 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) - # 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 + # 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)) - # 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()}" + # 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 36d17596873..3b3c6a6ce29 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1453,6 +1453,9 @@ class Logging(LiteLLMLoggingBaseClass): 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 e006662ec4d..256fee6b166 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 ebee9323766..703ccf13c27 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/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 0d02b4fa969..4fd49a35417 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -75,8 +75,9 @@ class AnthropicResponsesStreamWrapper: # ---- message_start ---- if event_type == "response.created": - self._sent_message_start = True - self._chunk_queue.append(self._make_message_start()) + if not self._sent_message_start: + self._sent_message_start = True + self._chunk_queue.append(self._make_message_start()) return # ---- content_block_start for a new output message item ---- diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index ed936f6233a..682adf5a8ff 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_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 = usage["prompt_tokens"] * model_info["input_cost_per_token"] + 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 6fd68921fb0..3584297b35f 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 1a24088396f..e5afc81b641 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": [ @@ -16272,7 +16272,7 @@ "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/glm-5p2": { - "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -16686,7 +16686,7 @@ "supports_vision": false }, "fireworks_ai/glm-5p2": { - "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -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 421f1dcfbea..d2f3efbc54e 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 @@ class MCPRequestHandler: 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 @@ class MCPRequestHandler: 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 3e3e549008d..74752809e86 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 00000000000..6f27c1c0472 --- /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 115ff2e492c..732381edb0e 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.llms.custom_httpx.http_handler import get_async_httpx_client 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 upstream 401/403 and its ``WWW-Authenticate`` header from the exception tree, or ``None``. - 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) - - 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 @@ class MCPServerManager: 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 @@ class MCPServerManager: 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 @@ class MCPServerManager: 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 7ca4923337c..94271c54f4b 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -19,12 +19,20 @@ import httpx 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 @@ if MCP_AVAILABLE: # 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 @@ if MCP_AVAILABLE: # 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 26322e9c58b..a8ab0937124 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -348,6 +348,7 @@ if MCP_AVAILABLE: CallToolResult, EmbeddedResource, ImageContent, + ListToolsResult, Prompt, TextContent, ) @@ -356,6 +357,14 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: ######################################################## @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 @@ if MCP_AVAILABLE: # 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: _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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: # 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 [] + return AggregateToolListing(tools=[], outcomes={}) - # 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) - - # 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 @@ if MCP_AVAILABLE: 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 fa57a2b3eb2..2f6b54a264a 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 fa354b8cccb..6ed283d898b 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 a610e44e69c..38900260c98 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -14,6 +14,9 @@ from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HE 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 4d07d4c043c..1a1b355cb17 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 02bb66388ca..1dc0ee3f947 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: @@ -851,9 +925,12 @@ class ProxyBaseLLMRequestProcessing: # If conversion fails, use original spend pass + model_name = ProxyBaseLLMRequestProcessing._get_deployment_model_name(litellm_logging_obj) + headers = { "x-litellm-call-id": call_id, "x-litellm-model-id": model_id, + "x-litellm-model-name": model_name, "x-litellm-cache-key": cache_key, "x-litellm-model-api-base": ( api_base.split("?")[0] if api_base else None @@ -1322,6 +1399,27 @@ class ProxyBaseLLMRequestProcessing: model_id = model_info.get("id", "") or "" return model_id + @staticmethod + def _get_deployment_model_name( + litellm_logging_obj: LiteLLMLoggingObj | None, + ) -> str | None: + """Extract the underlying deployment model string (e.g. ``azure/gpt-4o``). + + The router rewrites the response ``model`` field to the model-group alias + the client requested, so neither the response body nor the existing + headers expose the concrete deployment model. The router records it under + ``litellm_params`` metadata as ``deployment``, so read it back from there. + """ + litellm_params = getattr(litellm_logging_obj, "litellm_params", None) + if not isinstance(litellm_params, dict): + return None + for key in ("litellm_metadata", "metadata"): + metadata = litellm_params.get(key, {}) or {} + deployment = metadata.get("deployment") + if deployment: + return deployment + return None + @staticmethod def _response_cost_from_logging_obj( *, @@ -2575,6 +2673,8 @@ class ProxyBaseLLMRequestProcessing: 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 +2686,28 @@ class ProxyBaseLLMRequestProcessing: 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 +2796,13 @@ class ProxyBaseLLMRequestProcessing: 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 +2845,8 @@ class ProxyBaseLLMRequestProcessing: 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 00000000000..58ed3a8942f --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/__init__.py @@ -0,0 +1,44 @@ +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 00000000000..36a09a4ea25 --- /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/guardrails/guardrail_hooks/straiker/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py new file mode 100644 index 00000000000..ba4c712764e --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py @@ -0,0 +1,71 @@ +from typing import TYPE_CHECKING + +import litellm +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .straiker import StraikerGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + +_OPTIONAL_INIT_FIELDS = ( + "timeout", + "max_retries", + "initial_backoff", + "max_backoff", + "unreachable_fallback", + "fail_on_error", + "max_payload_bytes", + "custom_headers", + "metadata", + "verbose", +) + + +def _get_config_value(litellm_params: "LitellmParams", optional_params: object, attribute_name: str) -> object: + if optional_params is not None: + if isinstance(optional_params, dict): + value = optional_params.get(attribute_name) + else: + value = getattr(optional_params, attribute_name, None) + if value is not None: + return value + return getattr(litellm_params, attribute_name, None) + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + optional_params = getattr(litellm_params, "optional_params", None) + api_key = litellm_params.api_key + if not api_key: + raise ValueError("api_key is required for straiker") + + api_base = litellm_params.api_base or "https://api.prod.straiker.ai" + default_app = getattr(litellm_params, "default_app", None) or getattr(litellm_params, "source", None) + source = default_app if isinstance(default_app, str) and default_app else "LiteLLM Gateway" + kwargs: dict[str, object] = { + field: value + for field in _OPTIONAL_INIT_FIELDS + for value in [_get_config_value(litellm_params, optional_params, field)] + if value is not None + } + _callback = StraikerGuardrail( + api_key=api_key, + api_base=api_base if isinstance(api_base, str) else "https://api.prod.straiker.ai", + source=source, + guardrail_name=guardrail.get("guardrail_name", "straiker"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + **kwargs, + ) + + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.STRAIKER.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.STRAIKER.value: StraikerGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py new file mode 100644 index 00000000000..5c9f93fc2cd --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py @@ -0,0 +1,541 @@ +from __future__ import annotations + +import asyncio +import json +import random +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, NoReturn +from urllib.parse import urlsplit + +import httpx +from pydantic import ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm._version import version as litellm_version +from litellm.exceptions import ( + BadRequestError, + GuardrailRaisedException, + ModifyResponseException, + Timeout, +) +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + get_session_id_from_request_data, + log_guardrail_information, +) +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +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.straiker import ( + STRAIKER_WEBHOOK_SCHEMA_VERSION, + StraikerGuardrailConfigModel, + StraikerWebhookApplication, + StraikerWebhookContent, + StraikerWebhookContext, + StraikerWebhookEvent, + StraikerWebhookIdentity, + StraikerWebhookRequest, + StraikerWebhookResponse, + StraikerWebhookStream, + StraikerWebhookUsage, +) +from litellm.types.utils import GenericGuardrailAPIInputs, Usage + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + +GUARDRAIL_NAME = "straiker" +DEFAULT_BLOCK_MESSAGE = "Content violates policy" +DEFAULT_API_BASE = "https://api.prod.straiker.ai" +DEFAULT_MAX_PAYLOAD_BYTES = 524288 +WEBHOOK_PATH = "/api/v1/detect/webhook" +RETRY_STATUS = frozenset({408, 429, 500, 502, 503, 504}) +UNREACHABLE_STATUS = frozenset({502, 503, 504}) +_APPLICATION_METADATA_KEYS = frozenset({"agent_id", "app_name"}) +_OPAQUE_METADATA_SCALAR_TYPES = (str, int, float, bool) + + +@dataclass(frozen=True, slots=True) +class _WebhookFailure: + message: str + is_unreachable: bool + + +def _as_dict(value: object) -> dict: + return value if isinstance(value, dict) else {} + + +def _merged_metadata(request_data: dict) -> dict: + return { + **_as_dict(request_data.get("metadata")), + **_as_dict(request_data.get("litellm_metadata")), + } + + +def _as_optional_str(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _build_webhook_metadata(request_data: dict, default_metadata: dict[str, str]) -> dict[str, object] | None: + out: dict[str, object] = {} + for key, value in _as_dict(request_data.get("metadata")).items(): + if key in _APPLICATION_METADATA_KEYS or key.startswith("user_api"): + continue + if key == "session_id": + continue + if isinstance(value, _OPAQUE_METADATA_SCALAR_TYPES): + out[key] = value + out.update(default_metadata) + return out or None + + +def _extract_identity(request_data: dict) -> StraikerWebhookIdentity: + meta = _merged_metadata(request_data) + return StraikerWebhookIdentity( + litellm_key=_as_optional_str(meta.get("user_api_key_alias")) + or _as_optional_str(meta.get("user_api_key_hash")) + or _as_optional_str(meta.get("user_api_key_token")), + litellm_team=_as_optional_str(meta.get("user_api_key_team_alias")) + or _as_optional_str(meta.get("user_api_key_team_id")), + litellm_user_id=_as_optional_str(meta.get("user_api_key_user_id")), + litellm_user_email=_as_optional_str(meta.get("user_api_key_user_email")), + litellm_org_id=_as_optional_str(meta.get("user_api_key_org_id")), + end_user_id=_as_optional_str(meta.get("user_api_key_end_user_id")), + ) + + +def _resolve_provider(request_data: dict, model: str | None) -> str | None: + litellm_params = _as_dict(request_data.get("litellm_params")) + custom_llm_provider = request_data.get("custom_llm_provider") or litellm_params.get("custom_llm_provider") + if custom_llm_provider: + return custom_llm_provider + if not model: + return None + try: + _, provider, _, _ = get_llm_provider( + model=model, + api_base=request_data.get("api_base") or litellm_params.get("api_base"), + api_key=request_data.get("api_key") or litellm_params.get("api_key"), + ) + except BadRequestError: + return None + return provider or None + + +def _resolve_destination(request_data: dict) -> str | None: + litellm_params = _as_dict(request_data.get("litellm_params")) + api_base = request_data.get("api_base") or litellm_params.get("api_base") + if not isinstance(api_base, str): + return None + try: + return urlsplit(api_base).hostname + except ValueError: + return None + + +def _resolve_call_surface(logging_obj: LiteLLMLoggingObj | None, request_data: dict) -> str: + call_type = ( + (getattr(logging_obj, "call_type", None) if logging_obj is not None else None) + or request_data.get("call_type") + or request_data.get("litellm_call_type") + ) + return call_type if isinstance(call_type, str) and call_type else "unknown" + + +def _response_finish_reason(response: Any) -> str | None: + choices = getattr(response, "choices", None) + if not isinstance(choices, list): + return None + for choice in choices: + reason = getattr(choice, "finish_reason", None) + if isinstance(reason, str) and reason: + return reason + return None + + +def _build_usage(response: object) -> StraikerWebhookUsage | None: + usage = getattr(response, "usage", None) + if not isinstance(usage, Usage): + return None + input_tokens = usage.prompt_tokens + output_tokens = usage.completion_tokens + if input_tokens is None and output_tokens is None: + return None + return StraikerWebhookUsage(input_tokens=input_tokens, output_tokens=output_tokens) + + +def _is_streamed_request(request_data: dict) -> bool: + if request_data.get("stream") is True: + return True + body = _as_dict(_as_dict(request_data.get("proxy_server_request")).get("body")) + return body.get("stream") is True + + +class StraikerGuardrail(CustomGuardrail): + @staticmethod + def get_config_model() -> type[GuardrailConfigModel]: + return StraikerGuardrailConfigModel + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + return [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + def __init__( + self, + api_key: str, + api_base: str = DEFAULT_API_BASE, + source: str = "LiteLLM Gateway", + timeout: float = 5.0, + max_retries: int = 2, + initial_backoff: float = 0.1, + max_backoff: float = 2.0, + unreachable_fallback: Literal["fail_open", "fail_closed"] = "fail_closed", + fail_on_error: bool = True, + max_payload_bytes: int = DEFAULT_MAX_PAYLOAD_BYTES, + custom_headers: dict[str, str] | None = None, + metadata: dict[str, str] | None = None, + verbose: bool = False, + async_handler: httpx.AsyncClient | None = None, + **kwargs: object, + ) -> None: + if not api_key: + raise ValueError("api_key must be non-empty") + if unreachable_fallback not in ("fail_open", "fail_closed"): + raise ValueError(f"unreachable_fallback must be 'fail_open' or 'fail_closed'; got {unreachable_fallback!r}") + + self.api_key = api_key + self.api_base = api_base.rstrip("/") + self.source = source + self.timeout = float(timeout) + self.max_retries = max(0, int(max_retries)) + self.initial_backoff = max(0.0, float(initial_backoff)) + self.max_backoff = max(self.initial_backoff, float(max_backoff)) + self.unreachable_fallback = unreachable_fallback + self.fail_on_error = fail_on_error + self.max_payload_bytes = int(max_payload_bytes) + self.custom_headers = dict(custom_headers) if custom_headers else {} + self.default_metadata = dict(metadata) if metadata else {} + self.verbose = bool(verbose) + + self.streaming_end_of_stream_only = True + self.streaming_buffer_until_moderated = True + + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + ) + + kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) + super().__init__(**kwargs) + + def _webhook_url(self) -> str: + return f"{self.api_base}{WEBHOOK_PATH}" + + def _headers(self) -> dict[str, str]: + reserved = {"authorization", "content-type", "x-straiker-webhook-format"} + extra = {k: v for k, v in self.custom_headers.items() if k.lower() not in reserved} + return { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + "X-Straiker-Webhook-Format": "litellm", + **extra, + } + + def _build_application(self, request_data: dict) -> StraikerWebhookApplication: + meta = _merged_metadata(request_data) + agent_id = _as_optional_str(meta.get("agent_id")) + return StraikerWebhookApplication( + source=agent_id or self.source, + name=_as_optional_str(meta.get("app_name")), + ) + + def _build_context( + self, + request_data: dict, + model: str | None, + logging_obj: LiteLLMLoggingObj | None, + ) -> StraikerWebhookContext: + return StraikerWebhookContext( + call_surface=_resolve_call_surface(logging_obj, request_data), + model=model, + model_provider=_resolve_provider(request_data, model), + destination=_resolve_destination(request_data), + session_id=get_session_id_from_request_data(request_data), + litellm_call_id=getattr(logging_obj, "litellm_call_id", None) if logging_obj else None, + litellm_trace_id=getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None, + litellm_version=litellm_version, + ) + + def _build_envelope( + self, + *, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None, + ) -> StraikerWebhookRequest: + model = inputs.get("model") or request_data.get("model") + call_id = getattr(logging_obj, "litellm_call_id", None) if logging_obj else None + event_id = f"{call_id or 'litellm'}:{input_type}" + + content = StraikerWebhookContent( + texts=list(inputs.get("texts") or []), + images=list(inputs.get("images") or []), + structured_messages=inputs.get("structured_messages"), + tools=inputs.get("tools"), + tool_calls=inputs.get("tool_calls"), + ) + + if input_type == "request": + event = StraikerWebhookEvent(type="pre_call", id=event_id) + return StraikerWebhookRequest( + event=event, + request=content, + context=self._build_context(request_data, model, logging_obj), + identity=_extract_identity(request_data), + application=self._build_application(request_data), + metadata=_build_webhook_metadata(request_data, self.default_metadata), + ) + + response_obj = request_data.get("response") + content.finish_reason = _response_finish_reason(response_obj) + original_messages = request_data.get("messages") + request_content = StraikerWebhookContent( + structured_messages=original_messages if isinstance(original_messages, list) else None, + ) + phase: Literal["none", "assembled"] = "assembled" if _is_streamed_request(request_data) else "none" + event = StraikerWebhookEvent(type="post_call", id=event_id, stream=StraikerWebhookStream(phase=phase)) + return StraikerWebhookRequest( + event=event, + request=request_content, + response=content, + context=self._build_context(request_data, model, logging_obj), + identity=_extract_identity(request_data), + application=self._build_application(request_data), + usage=_build_usage(response_obj), + metadata=_build_webhook_metadata(request_data, self.default_metadata), + ) + + async def _post_webhook(self, payload: dict) -> tuple[StraikerWebhookResponse | None, _WebhookFailure | None]: + try: + body = json.dumps(payload).encode("utf-8") + except (TypeError, ValueError, OverflowError) as error: + return None, _WebhookFailure(f"request serialization failed: {error}", is_unreachable=False) + body_bytes = len(body) + if body_bytes > self.max_payload_bytes: + return None, _WebhookFailure( + f"payload {body_bytes}B exceeds max_payload_bytes {self.max_payload_bytes}", + is_unreachable=False, + ) + + url = self._webhook_url() + headers = self._headers() + attempts = self.max_retries + 1 + last_failure: _WebhookFailure | None = None + + if self.verbose: + verbose_proxy_logger.info( + json.dumps( + { + "event": "straiker.webhook_request", + "url": url, + "bytes": body_bytes, + "payload": payload, + }, + default=str, + ) + ) + + for attempt in range(attempts): + try: + resp = await self.async_handler.post(url, content=body, headers=headers, timeout=self.timeout) + if resp.status_code == 200: + try: + body = resp.json() + parsed = StraikerWebhookResponse.model_validate(body) + except (ValidationError, json.JSONDecodeError) as ve: + return None, _WebhookFailure(f"invalid response schema: {ve}", is_unreachable=False) + if self.verbose: + verbose_proxy_logger.info( + json.dumps( + { + "event": "straiker.webhook_response", + "status_code": resp.status_code, + "body": body, + }, + default=str, + ) + ) + return parsed, None + last_failure = _WebhookFailure( + f"HTTP {resp.status_code}: {resp.text[:200]}", + is_unreachable=resp.status_code in UNREACHABLE_STATUS, + ) + if resp.status_code not in RETRY_STATUS: + return None, last_failure + except (httpx.RequestError, asyncio.TimeoutError, Timeout) as e: + last_failure = _WebhookFailure(f"{type(e).__name__}: {e}", is_unreachable=True) + except (json.JSONDecodeError, TypeError, ValueError) as e: + return None, _WebhookFailure(f"{type(e).__name__}: {e}", is_unreachable=False) + + if attempt < attempts - 1: + backoff = min(self.initial_backoff * (2**attempt), self.max_backoff) + await asyncio.sleep(random.uniform(0, backoff)) + + return None, last_failure or _WebhookFailure("unknown error", is_unreachable=True) + + def _record( + self, + *, + request_data: dict, + logging_obj: LiteLLMLoggingObj | None, + parsed: StraikerWebhookResponse, + ) -> None: + if not self.verbose: + return + response_obj = request_data.get("response") + hidden = getattr(response_obj, "_hidden_params", None) + if isinstance(hidden, dict): + straiker_hidden = hidden.setdefault("straiker", {}) + if isinstance(straiker_hidden, dict): + straiker_hidden.update({"action": parsed.action, "turn_id": parsed.turn_id}) + + def _fail( + self, + *, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + error: str, + is_unreachable: bool, + ) -> GenericGuardrailAPIInputs: + fail_open = (is_unreachable and self.unreachable_fallback == "fail_open") or not self.fail_on_error + verbose_proxy_logger.error( + json.dumps( + { + "event": "straiker.error", + "input_type": input_type, + "error": error, + "fail_open": fail_open, + }, + default=str, + ) + ) + if fail_open: + return inputs + self._block( + request_data=request_data, + input_type=input_type, + message=f"Straiker detection unavailable: {error}", + ) + + def _block( + self, + *, + request_data: dict, + input_type: Literal["request", "response"], + message: str, + ) -> NoReturn: + if input_type == "request": + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name or GUARDRAIL_NAME, + message=message, + should_wrap_with_default_message=False, + ) + raise ModifyResponseException( + message=message, + model=request_data.get("model", "unknown") or "unknown", + request_data=request_data, + guardrail_name=self.guardrail_name or GUARDRAIL_NAME, + original_response=request_data.get("response"), + ) + + @staticmethod + def _intervened_inputs( + inputs: GenericGuardrailAPIInputs, + parsed: StraikerWebhookResponse, + ) -> GenericGuardrailAPIInputs: + return_inputs: GenericGuardrailAPIInputs = {} + return_inputs.update(inputs) + if parsed.texts is not None: + return_inputs["texts"] = parsed.texts + return return_inputs + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + try: + envelope = self._build_envelope( + inputs=inputs, + request_data=request_data, + input_type=input_type, + logging_obj=logging_obj, + ) + payload = envelope.model_dump(mode="json", exclude_none=True) + except (ValidationError, TypeError, ValueError) as error: + return self._fail( + inputs=inputs, + request_data=request_data, + input_type=input_type, + error=str(error), + is_unreachable=False, + ) + + parsed, failure = await self._post_webhook(payload) + if failure is not None: + return self._fail( + inputs=inputs, + request_data=request_data, + input_type=input_type, + error=failure.message, + is_unreachable=failure.is_unreachable, + ) + + if parsed is None: + return self._fail( + inputs=inputs, + request_data=request_data, + input_type=input_type, + error="empty response from Straiker", + is_unreachable=False, + ) + self._record(request_data=request_data, logging_obj=logging_obj, parsed=parsed) + + if parsed.schema_version is not None and parsed.schema_version != STRAIKER_WEBHOOK_SCHEMA_VERSION: + verbose_proxy_logger.warning( + json.dumps( + { + "event": "straiker.schema_drift", + "expected": STRAIKER_WEBHOOK_SCHEMA_VERSION, + "received": parsed.schema_version, + } + ) + ) + + if parsed.action == "BLOCKED": + self._block( + request_data=request_data, + input_type=input_type, + message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, + ) + if parsed.action == "GUARDRAIL_INTERVENED": + is_streamed_response = input_type == "response" and _is_streamed_request(request_data) + if parsed.texts is None or is_streamed_response: + self._block( + request_data=request_data, + input_type=input_type, + message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, + ) + return self._intervened_inputs(inputs, parsed) + return inputs diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d60c17c744f..22ea9fe176a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -7,6 +7,7 @@ This is currently in development and not yet ready for production. import asyncio import binascii import os +import uuid from datetime import datetime from typing import ( TYPE_CHECKING, @@ -185,6 +186,69 @@ end 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 @@ RATE_LIMIT_DESCRIPTORS_KEY = "_litellm_rate_limit_descriptors" # 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 @@ _LITELLM_STASH_KEYS: Tuple[str, ...] = ( 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ``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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 - ## CHECK IN-MEMORY CACHE - cache_values = await self.internal_usage_cache.async_batch_get_cache( - keys=keys_to_fetch, + 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] + + 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] - 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 + 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. - ## 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 + 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)) - # 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, + 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) - # 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] + 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, + 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=counter_key, - value=counter_value, - ttl=self.window_size, + value=new_value, + ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, 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, - ) - - 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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. + 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. """ - if not user_api_key_dict.api_key or user_api_key_dict.max_parallel_requests is None: + 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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. + 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 via TPM_RESERVATION_RELEASED_KEY: if both this hook and + 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/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 15d1876e5a2..9ddc7ce2caf 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -45,6 +45,7 @@ _EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session- # Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores # (covers UUIDs and most common session-id formats). _SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") +_ANTHROPIC_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]+$") def _sanitize_for_log(value: Any) -> str: @@ -426,6 +427,30 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str ) +def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None: + if not isinstance(metadata, dict): + return None + + user_id = metadata.get("user_id") + if isinstance(user_id, dict): + session_id = user_id.get("session_id") + if isinstance(session_id, str) and _ANTHROPIC_SESSION_ID_VALUE_RE.fullmatch(session_id): + return session_id + return None + if not isinstance(user_id, str): + return None + + session_marker = "_session_" + session_marker_index = user_id.rfind(session_marker) + if session_marker_index == -1: + return None + + session_id = user_id[session_marker_index + len(session_marker) :] + if not session_id or not _ANTHROPIC_SESSION_ID_VALUE_RE.fullmatch(session_id): + return None + return session_id + + def is_claude_code_user_agent(user_agent: str) -> bool: """Claude Code identifies itself as ``claude-cli/ ...``; the IDE extensions and the Agent SDK run through the same CLI and share that prefix.""" @@ -935,6 +960,15 @@ class LiteLLMProxyRequestSetup: data["litellm_session_id"] = chain_id data["litellm_trace_id"] = chain_id verbose_proxy_logger.debug(f"Extracted chain_id from header (trace-id/session-id): {chain_id}") + else: + body_metadata = data.get("metadata") + session_id = _get_anthropic_session_id_from_metadata(body_metadata) + if session_id: + metadata_from_headers["session_id"] = session_id + data["litellm_session_id"] = session_id + if isinstance(body_metadata, dict) and isinstance(body_metadata.get("user_id"), dict): + body_metadata["user_id"] = session_id + verbose_proxy_logger.debug("Extracted session_id from Anthropic metadata.user_id") if isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name].update(metadata_from_headers) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d920ee474cc..d50db8324ef 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -723,12 +723,13 @@ if MCP_AVAILABLE: """ 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 2aff663038b..acb2e50c79b 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.secret_managers.main import get_secret_str 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 @@ def create_pass_through_route( 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/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 6a673f6bebb..932216141fe 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -48,18 +48,18 @@ def _safe_response_text(httpx_response: httpx.Response) -> str: class PassThroughEndpointLogging: def __init__(self): - self.TRACKED_VERTEX_ROUTES = [ + self.TRACKED_VERTEX_METHOD_ROUTES = ( "generateContent", "streamGenerateContent", "predict", "rawPredict", "streamRawPredict", "search", - "batchPredictionJobs", "predictLongRunning", "embedContent", "batchEmbedContents", - ] + ) + self.TRACKED_VERTEX_RESOURCE_ROUTES = ("batchPredictionJobs",) # Anthropic self.TRACKED_ANTHROPIC_ROUTES = ["/messages", "/v1/messages/batches"] @@ -339,11 +339,10 @@ class PassThroughEndpointLogging: **kwargs, ) - def is_vertex_route(self, url_route: str): - for route in self.TRACKED_VERTEX_ROUTES: - if route in url_route: - return True - return False + def is_vertex_route(self, url_route: str) -> bool: + if any(f":{method}" in url_route for method in self.TRACKED_VERTEX_METHOD_ROUTES): + return True + return any(resource in url_route for resource in self.TRACKED_VERTEX_RESOURCE_ROUTES) def is_anthropic_route(self, url_route: str): for route in self.TRACKED_ANTHROPIC_ROUTES: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 9bed3657b20..dc5bde8cb0b 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -802,6 +802,19 @@ class ProxyInitializationHelpers: ), 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 dcde9a27ec0..8936f6e9ca9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1076,9 +1076,10 @@ async def proxy_startup_event(app: FastAPI): # 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 @@ class ProxyConfig: 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 f7f6adaa8a2..27ffc49901b 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -11,14 +11,16 @@ from typing import Any, Dict, Optional, Tuple 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 e2206541fc6..8d7aedce4d0 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 9f36e729330..7a52fdfdb87 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -19,6 +19,7 @@ from typing import ( Any, AsyncGenerator, Awaitable, + Callable, ClassVar, Dict, List, @@ -49,7 +50,7 @@ from litellm.proxy._types import ( from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse -from litellm.types.utils import CallTypes, CallTypesLiteral +from litellm.types.utils import CallTypes, CallTypesLiteral, ModelInfo try: from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( @@ -2583,35 +2584,30 @@ class ProxyLogging: 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. + 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. - 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. + 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): """ @@ -6101,6 +6097,7 @@ def create_model_info_response( include_metadata: bool = False, fallback_type: Optional[str] = None, llm_router: Optional["Router"] = None, + get_model_info: Callable[[str], ModelInfo] = litellm.get_model_info, ) -> ModelInfoResponse: """ Create a standardized OpenAI-compatible model object. @@ -6118,25 +6115,37 @@ def create_model_info_response( "owned_by": provider, } - # Surface context-window limits for OpenAI-compatible discovery clients. - # Only emitted when known, so wildcard routes and limitless backends stay clean. - # Limits are best-effort enrichment, so a single malformed deployment degrades - # to the base response rather than 500-ing the whole listing. + try: + model_cost_info: ModelInfo | None = get_model_info(model_id) + except Exception as e: + verbose_proxy_logger.debug( + "create_model_info_response: cost map lookup failed for %s: %s", + model_id, + e, + ) + model_cost_info = None + + max_input_tokens: int | None = None + max_output_tokens: int | None = None + if model_cost_info is not None: + cost_map_input = model_cost_info.get("max_input_tokens") + if cost_map_input is not None: + max_input_tokens = int(cost_map_input) + cost_map_output = model_cost_info.get("max_output_tokens") + if cost_map_output is not None: + max_output_tokens = int(cost_map_output) + if llm_router is not None: - try: - model_group_info = llm_router.get_model_group_info(model_id) - except Exception as e: - verbose_proxy_logger.debug( - "create_model_info_response: get_model_group_info failed for %s: %s", - model_id, - e, - ) - model_group_info = None - if model_group_info is not None: - if model_group_info.max_input_tokens is not None: - base["max_input_tokens"] = int(model_group_info.max_input_tokens) - if model_group_info.max_output_tokens is not None: - base["max_output_tokens"] = int(model_group_info.max_output_tokens) + configured_input, configured_output = llm_router.get_configured_token_limits(model_id) + if configured_input is not None: + max_input_tokens = configured_input + if configured_output is not None: + max_output_tokens = configured_output + + if max_input_tokens is not None: + base["max_input_tokens"] = max_input_tokens + if max_output_tokens is not None: + base["max_output_tokens"] = max_output_tokens if not include_metadata: return base diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 6b5f087f902..29891ccfd24 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -11,12 +11,14 @@ __all__ = ["ingest", "aingest", "query", "aquery"] 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 @@ from typing import ( 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 e03f0296109..392bb7bcab2 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -260,7 +260,7 @@ class LiteLLM_Proxy_MCP_Handler: # 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 @@ class LiteLLM_Proxy_MCP_Handler: 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 78e156801f8..b1a5405ebf1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -33,6 +33,7 @@ from typing import ( Optional, Set, Tuple, + TypeVar, Union, cast, ) @@ -86,7 +87,11 @@ from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler 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, @@ -175,6 +180,7 @@ from litellm.types.router import ( MockRouterTestingParams, ModelGroupInfo, OptionalPreCallChecks, + PreRoutingStrategy, RetryPolicy, RouterCacheEnum, RouterGeneralSettings, @@ -186,6 +192,7 @@ from litellm.types.router import ( RoutingPlugin, RoutingStrategy, SearchToolTypedDict, + TaggedPreRoutingStrategy, ) from litellm.types.services import ServiceTypes from litellm.types.utils import ( @@ -260,6 +267,9 @@ def _cost_value_as_float(value: Union[str, int, float, None]) -> Optional[float] return None +_PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") + + class RoutingArgs(enum.Enum): ttl = 60 # 1min (RPM/TPM expire key) @@ -487,10 +497,10 @@ class Router: 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 @@ -2037,6 +2047,9 @@ class Router: 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() @@ -7568,6 +7581,11 @@ class Router: 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. @@ -7603,11 +7621,12 @@ class Router: 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: """ @@ -7658,20 +7677,54 @@ class Router: 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 @@ -7694,23 +7747,31 @@ class Router: 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: """ @@ -7763,18 +7824,18 @@ class Router: 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) ) @@ -7826,11 +7887,12 @@ class Router: 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: """ @@ -8459,6 +8521,27 @@ class Router: raise Exception("Model Name invalid - {}".format(type(model))) return None + def get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]": + """ + Return (max_input_tokens, max_output_tokens) explicitly configured in a concrete + deployment's model_info for model_name, via O(1) index lookup. + + Returns (None, None) for wildcard-expanded or unknown names. Unlike + get_model_group_info, this never triggers pattern matching or deep copies, so it + is safe to call per listed model on the /v1/models hot path. + """ + deployment = self.get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return (None, None) + + model_info = deployment.model_info + max_input = model_info.get("max_input_tokens") + max_output = model_info.get("max_output_tokens") + return ( + int(max_input) if max_input is not None else None, + int(max_output) if max_output is not None else None, + ) + def get_deployment_credentials_with_provider(self, model_id: str) -> Optional[Dict[str, Any]]: """ Get API credentials and provider info from a model name in model_list. @@ -10810,6 +10893,35 @@ class Router: 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, @@ -10832,12 +10944,7 @@ class Router: 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 fa6f14e9b26..695d8b8aeaa 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -28,6 +28,7 @@ from litellm.types.utils import ModelResponse from .config import ( DEFAULT_CODE_KEYWORDS, + DEFAULT_ESCALATION_KEYWORDS, DEFAULT_REASONING_KEYWORDS, DEFAULT_SIMPLE_KEYWORDS, DEFAULT_TECHNICAL_KEYWORDS, @@ -173,6 +174,11 @@ class ComplexityRouter(CustomLogger): 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 @@ class ComplexityRouter(CustomLogger): } 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 @@ class ComplexityRouter(CustomLogger): 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] = 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, - ) + 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, + ) response = await self._classify_and_route( model=model, @@ -1004,13 +1069,17 @@ class ComplexityRouter(CustomLogger): 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 @@ class ComplexityRouter(CustomLogger): ) 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 1f984798970..17c2c287dde 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -162,6 +162,9 @@ DEFAULT_TECHNICAL_KEYWORDS: list[str] = [ # 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 @@ class ComplexityRouterConfig(BaseModel): 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 3dda4e3990c..5b611971154 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -53,6 +53,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import ( 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,8 +128,10 @@ class SupportedGuardrailIntegrations(Enum): RUBRIK = "rubrik" VIGIL_GUARD = "vigil_guard" REPELLOAI = "repelloai" + SINGULR = "singulr" HEADROOM = "headroom" COMPRESR = "compresr" + STRAIKER = "straiker" class Role(Enum): @@ -932,6 +937,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 daac1e4506f..9f689a2dd31 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 3524a7eb7f7..098e99fe198 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -11,6 +11,14 @@ LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY = "litellm_pass_through_custom_body" # 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 00000000000..62d3b8653ef --- /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/proxy/guardrails/guardrail_hooks/straiker.py b/litellm/types/proxy/guardrails/guardrail_hooks/straiker.py new file mode 100644 index 00000000000..b4375237917 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/straiker.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk +from litellm.types.utils import ChatCompletionMessageToolCall + +from .base import GuardrailConfigModel + +StraikerWebhookEventType = Literal["pre_call", "post_call"] +StraikerWebhookStreamPhase = Literal["none", "assembled"] +StraikerWebhookAction = Literal["NONE", "BLOCKED", "GUARDRAIL_INTERVENED"] + +STRAIKER_WEBHOOK_SCHEMA_VERSION = "1" + + +class StraikerWebhookStream(BaseModel): + phase: StraikerWebhookStreamPhase = "none" + index: int | None = None + + +class StraikerWebhookEvent(BaseModel): + type: StraikerWebhookEventType + id: str + stream: StraikerWebhookStream = Field(default_factory=StraikerWebhookStream) + + +class StraikerWebhookContent(BaseModel): + model_config = ConfigDict(extra="ignore") + + texts: list[str] = Field(default_factory=list) + images: list[str] = Field(default_factory=list) + structured_messages: list[AllMessageValues] | None = None + tools: list[dict[str, object]] | None = None + tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None + finish_reason: str | None = None + + +class StraikerWebhookUsage(BaseModel): + input_tokens: int | None = None + output_tokens: int | None = None + + +class StraikerWebhookContext(BaseModel): + call_surface: str + model: str | None = None + model_provider: str | None = None + destination: str | None = None + session_id: str | None = None + litellm_call_id: str | None = None + litellm_trace_id: str | None = None + litellm_version: str | None = None + + +class StraikerWebhookIdentity(BaseModel): + litellm_key: str | None = None + litellm_team: str | None = None + litellm_user_id: str | None = None + litellm_user_email: str | None = None + litellm_org_id: str | None = None + end_user_id: str | None = None + + +class StraikerWebhookApplication(BaseModel): + source: str + name: str | None = None + + +class StraikerWebhookRequest(BaseModel): + schema_version: str = STRAIKER_WEBHOOK_SCHEMA_VERSION + event: StraikerWebhookEvent + request: StraikerWebhookContent + response: StraikerWebhookContent | None = None + context: StraikerWebhookContext + identity: StraikerWebhookIdentity + application: StraikerWebhookApplication + usage: StraikerWebhookUsage | None = None + metadata: dict[str, object] | None = None + + +class StraikerWebhookResponse(BaseModel): + model_config = ConfigDict(extra="allow") + + action: StraikerWebhookAction = "NONE" + blocked_reason: str | None = None + texts: list[str] | None = None + schema_version: str | None = None + turn_id: str | None = Field(default=None, alias="turnId") + + +class StraikerGuardrailConfigModelOptionalParams(BaseModel): + timeout: float | None = Field( + default=5.0, + gt=0.0, + description="Per-attempt HTTP timeout in seconds.", + ) + max_retries: int | None = Field( + default=2, + ge=0, + description="Retries on transient HTTP (408/429/5xx) and network errors.", + ) + initial_backoff: float | None = Field( + default=0.1, + ge=0.0, + description="Initial retry backoff in seconds.", + ) + max_backoff: float | None = Field( + default=2.0, + ge=0.0, + description="Maximum retry backoff in seconds.", + ) + unreachable_fallback: Literal["fail_open", "fail_closed"] | None = Field( + default="fail_closed", + description="Behavior when Straiker is unreachable after retries.", + ) + fail_on_error: bool | None = Field( + default=True, + description=( + "Behavior on any guardrail error, not just unreachability. True (default) blocks " + "the request on error; False logs and allows the request to proceed." + ), + ) + max_payload_bytes: int | None = Field( + default=524288, + gt=0, + description="Maximum serialized webhook payload size sent to Straiker.", + ) + custom_headers: dict[str, str] | None = Field( + default=None, + description="Additional HTTP headers sent to Straiker, excluding Authorization and the webhook-format header.", + ) + metadata: dict[str, str] | None = Field( + default=None, + description=( + "Default metadata key/values added to the webhook metadata bag on every request. " + "On key conflict with request-derived metadata, these configured values win." + ), + ) + verbose: bool | None = Field( + default=False, + description="Log webhook request/response payloads and record action/turn_id in response hidden params.", + ) + + +class StraikerGuardrailConfigModel(GuardrailConfigModel[StraikerGuardrailConfigModelOptionalParams]): + api_key: str = Field( + min_length=1, + description="Straiker DefendAI environment API key (Bearer token). Env: STRAIKER_API_KEY.", + json_schema_extra={"secret": True}, + ) + + api_base: str | None = Field( + default="https://api.prod.straiker.ai", + description="Straiker API base URL. Use the regional variant for non-US tenants.", + ) + + default_app: str | None = Field( + default="LiteLLM Gateway", + description=( + "Default application registered in the Straiker Defend Console. " + "Overridden per-request by metadata.agent_id when present." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Straiker" diff --git a/litellm/types/router.py b/litellm/types/router.py index 69a8ca9f19e..28e4a8272e8 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -5,7 +5,18 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc 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 04f1ff68c5d..ec8a9336ca7 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 a714c9ac4c9..15f5e5d3563 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3196,6 +3196,12 @@ def get_optional_params_embeddings( 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 ffbc0dcd098..5cf99ba8bac 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": [ @@ -16272,7 +16272,7 @@ "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/glm-5p2": { - "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -16686,7 +16686,7 @@ "supports_vision": false }, "fireworks_ai/glm-5p2": { - "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -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/pyproject.toml b/pyproject.toml index 45ae3c179d5..108f28cb124 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ proxy = [ "boto3>=1.43.1,<2.0", "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", - "mcp>=1.26.0,<2.0", + "mcp>=1.28.1,<2.0", "litellm-proxy-extras==0.4.78", "litellm-enterprise==0.1.51", "RestrictedPython>=8.1,<9.0", @@ -91,7 +91,7 @@ extra_proxy = [ "google-cloud-iam>=2.19.1,<3.0", # Not in PyPI proxy extra. "resend>=2.23.0,<3.0", - "redisvl>=0.4.1,<1.0; python_version < '3.14'", + "redisvl>=0.4.1,<1.0", "a2a-sdk>=1.1.0,<2.0", ] utils = [ @@ -136,7 +136,7 @@ proxy-runtime = [ "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", "azure-storage-file-datalake>=12.20.0,<13.0", - "pypdf>=6.12.0,<7.0; python_version < '3.14'", + "pypdf>=6.12.0,<7.0", "llm-sandbox>=0.3.39,<1.0", "detect-secrets>=1.5.0,<2.0", ] @@ -181,7 +181,7 @@ dev = [ "pytest-rerunfailures==15.1", "pytest-cov==5.0.0", "parameterized==0.9.0", - "openapi-core==0.22.0; python_version < '3.14'", + "openapi-core==0.22.0", "pytest-timeout==2.4.0", "vcrpy==8.2.1", "pytest-recording==0.13.4", diff --git a/router_plugins.json b/router_plugins.json new file mode 100644 index 00000000000..ffcddf89fd1 --- /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 0e1eafb5196..67e9f4f78a7 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 2082f2c9de4..555ac0482e2 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 9f4a00fe06f..b5681ff76ad 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 2c6070c437a..d3b6d42bc24 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 18aff2edc98..00000000000 --- 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/claude_code/_gpt_cells.py b/tests/e2e/claude_code/_gpt_cells.py index 870e9cea918..70e2b5ed18f 100644 --- a/tests/e2e/claude_code/_gpt_cells.py +++ b/tests/e2e/claude_code/_gpt_cells.py @@ -16,12 +16,12 @@ cover "OpenAI plus the big three clouds": carries only the open-weight gpt-oss MaaS models -The openai and azure_openai columns run unconditionally, like every -other live column: the environments that run the suite carry -`OPENAI_API_KEY` and `AZURE_API_BASE` + `AZURE_API_KEY` pointing at a -resource with gpt-5.6 deployments. The bedrock_mantle column is -opt-in via `COMPAT_MANTLE_CELLS=1` because the AWS account is still -waiting on the Bedrock Mantle allowlist for the `openai.gpt-5.6-*` +The azure_openai column runs unconditionally when Azure gpt-5.6 +deployments exist. The openai column is opt-in via +`COMPAT_OPENAI_GPT_CELLS=1` because under the full stage suite those +cells routinely burn minutes on Claude CLI timeouts. The bedrock_mantle +column is opt-in via `COMPAT_MANTLE_CELLS=1` because the AWS account is +still waiting on the Bedrock Mantle allowlist for the `openai.gpt-5.6-*` models; until the flag is set each Mantle cell skips and its matrix cell publishes as `not_tested` instead of a credential-shaped red. The `vertex_ai_gpt` column needs no flag either way: its cells report @@ -35,6 +35,7 @@ import os import pytest MANTLE_CELLS_ENV = "COMPAT_MANTLE_CELLS" +OPENAI_GPT_CELLS_ENV = "COMPAT_OPENAI_GPT_CELLS" VERTEX_AI_GPT_NOT_APPLICABLE_REASON = ( "GCP Vertex AI does not offer OpenAI's closed-weight GPT-5.6 family " @@ -59,3 +60,19 @@ def skip_unless_mantle_cells_enabled() -> None: f"Bedrock Mantle GPT-5.6 cells are opt-in; set {MANTLE_CELLS_ENV}=1 " "once the AWS account is allowlisted for the openai.gpt-5.6-* models" ) + + +def skip_unless_openai_gpt_cells_enabled() -> None: + """Skip OpenAI GPT-5.6 columns unless `COMPAT_OPENAI_GPT_CELLS` opts them in. + + Under the full stage suite these cells routinely hit 120s Claude CLI + timeouts and rate-limit-shaped retries across Sol/Terra/Luna, burning + ~8+ minutes per cell without a stable green. Opt in when exercising + the OpenAI GPT translation path in isolation. + """ + if os.environ.get(OPENAI_GPT_CELLS_ENV, "").strip().lower() in {"1", "true", "yes"}: + return + pytest.skip( + f"OpenAI GPT-5.6 cells are opt-in; set {OPENAI_GPT_CELLS_ENV}=1 " + "to run them (stage suite timeouts under concurrent load)" + ) diff --git a/tests/e2e/claude_code/basic_messaging_non_streaming/test_openai.py b/tests/e2e/claude_code/basic_messaging_non_streaming/test_openai.py index b0d143fa5e0..323c2f11173 100644 --- a/tests/e2e/claude_code/basic_messaging_non_streaming/test_openai.py +++ b/tests/e2e/claude_code/basic_messaging_non_streaming/test_openai.py @@ -23,6 +23,7 @@ green if all three pass. from __future__ import annotations from claude_code._basic_messaging import run_basic_messaging_cell +from claude_code._gpt_cells import skip_unless_openai_gpt_cells_enabled OPENAI_MODELS = [ "gpt-5-6-sol-openai", @@ -34,6 +35,7 @@ OPENAI_MODELS = [ def test_basic_messaging_non_streaming_openai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a non-empty reply from each GPT-5.6 tier.""" + skip_unless_openai_gpt_cells_enabled() run_basic_messaging_cell( compat_result=compat_result, models=OPENAI_MODELS, diff --git a/tests/e2e/claude_code/basic_messaging_streaming/test_openai.py b/tests/e2e/claude_code/basic_messaging_streaming/test_openai.py index 402c763496b..a7945fb92c0 100644 --- a/tests/e2e/claude_code/basic_messaging_streaming/test_openai.py +++ b/tests/e2e/claude_code/basic_messaging_streaming/test_openai.py @@ -25,6 +25,7 @@ green if all three pass. from __future__ import annotations from claude_code._basic_messaging import run_basic_messaging_cell +from claude_code._gpt_cells import skip_unless_openai_gpt_cells_enabled OPENAI_MODELS = [ "gpt-5-6-sol-openai", @@ -36,6 +37,7 @@ OPENAI_MODELS = [ def test_basic_messaging_streaming_openai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a non-empty streamed reply from each GPT-5.6 tier.""" + skip_unless_openai_gpt_cells_enabled() run_basic_messaging_cell( compat_result=compat_result, models=OPENAI_MODELS, diff --git a/tests/e2e/claude_code/tool_use/test_openai.py b/tests/e2e/claude_code/tool_use/test_openai.py index dbe60a65281..ffb7e795c2b 100644 --- a/tests/e2e/claude_code/tool_use/test_openai.py +++ b/tests/e2e/claude_code/tool_use/test_openai.py @@ -29,6 +29,7 @@ from typing import Any, Mapping, Sequence import pytest from claude_code._env import require_proxy +from claude_code._gpt_cells import skip_unless_openai_gpt_cells_enabled from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, @@ -69,6 +70,7 @@ def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool: def test_tool_use_openai(compat_result): """Drive the `claude` CLI against the LiteLLM proxy and assert a tool call was emitted on the wire by each GPT-5.6 tier.""" + skip_unless_openai_gpt_cells_enabled() proxy = require_proxy(compat_result) outcomes = run_claude_models_parallel( diff --git a/tests/e2e/claude_code/tool_use_streaming/test_openai.py b/tests/e2e/claude_code/tool_use_streaming/test_openai.py index 895f88d994b..a5ce31b1fd6 100644 --- a/tests/e2e/claude_code/tool_use_streaming/test_openai.py +++ b/tests/e2e/claude_code/tool_use_streaming/test_openai.py @@ -30,6 +30,7 @@ from typing import Any, Mapping, Sequence import pytest from claude_code._env import require_proxy +from claude_code._gpt_cells import skip_unless_openai_gpt_cells_enabled from claude_code.cli_driver import ( ClaudeCLIError, failure_diagnostic, @@ -87,6 +88,7 @@ def _count_input_json_deltas(events: Sequence[Mapping[str, Any]]) -> int: def test_tool_use_streaming_openai(compat_result): + skip_unless_openai_gpt_cells_enabled() proxy = require_proxy(compat_result) outcomes = run_claude_models_parallel( diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 3aec104c861..88a9deecb7e 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -15,14 +15,14 @@ shared fixtures build on it. 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 fac266149f4..633351ca97c 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -7,14 +7,20 @@ - {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"} - {id: quota_management.budget.key.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: key, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes key spend after the window; a blocked key serves again"} +- {id: quota_management.budget.team.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: team, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes a team's spend after the window; every key on the team serves again"} +- {id: quota_management.budget.organization.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "An org budget resets after its window; keys under the org serve again"} +- {id: quota_management.budget.internal_user.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "An internal user's budget resets after its window; their personal and team-member keys serve again"} - {id: quota_management.budget.team_member.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "Member per-team budget reset keeps advancing window after window"} - {id: quota_management.budget.key_multi_window.blocks_then_resets, module: quota_management, tier: P1, behavior: budget, variant: key_multi_window, assertions: [blocks_then_resets], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_limits enforce within a short window and serve again in the next"} - {id: quota_management.budget.key_multi_window.resets_windows_independently, module: quota_management, tier: P2, behavior: budget, variant: key_multi_window, assertions: [resets_windows_independently], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "Each window of a multi-window budget resets on its own schedule"} diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index a117cbd570d..29d54b011be 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 22f7581818f..00000000000 --- 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 f8cda509c63..00000000000 --- 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 00000000000..e4f59f5c4d2 --- /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 2a87ef7259d..5258b751a8c 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 4795c3b9f54..bae858d50af 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 15cd789664e..8e6e596bcd3 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 6aaffdd208e..f99fa8d86b3 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 @@ schema: the session lifecycle, the canonical response event sequence with a 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 921010e5eae..e735d9c01b5 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 @@ references the proxy resolves at call time, so adding a provider is a new type 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 18da1305c13..4f2dc874a33 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 00000000000..77fef574706 --- /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 00000000000..a1dac3fdac4 --- /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 00000000000..eaa49af5b69 --- /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 82c276d0b64..39832d1a17f 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 7b37c3af98e..c2f5dcdd57c 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -33,12 +33,26 @@ _TEAM_READY_SLEEP_SECONDS = 0.4 class UserNewBody(BaseModel): max_budget: float + budget_duration: str | None = None 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] @@ -51,6 +65,7 @@ class CustomerNewBody(BaseModel): class OrgNewBody(BaseModel): organization_alias: str max_budget: float + budget_duration: str | None = None class OrgNewResponse(BaseModel): @@ -61,6 +76,14 @@ class OrgDeleteBody(BaseModel): organization_ids: list[str] +class OrgInfoParams(BaseModel): + organization_id: str + + +class OrgInfoResponse(BaseModel): + budget_id: str | None = None + + class TeamMember(BaseModel): role: str user_id: str @@ -69,6 +92,7 @@ class TeamMember(BaseModel): class TeamNewBody(BaseModel): team_alias: str max_budget: float | None = None + budget_duration: str | None = None organization_id: str | None = None budget_limits: list[BudgetWindow] | None = None @@ -244,12 +268,12 @@ class BudgetClient: # ---- internal user -------------------------------------------------- - def create_user(self, *, max_budget: float) -> str: + def create_user(self, *, max_budget: float, budget_duration: str | None = None) -> str: return unwrap( self.gateway.transport.post( "/user/new", headers=self.gateway.transport.master, - json=UserNewBody(max_budget=max_budget), + json=UserNewBody(max_budget=max_budget, budget_duration=budget_duration), response_type=UserNewResponse, ) ).user_id @@ -262,6 +286,19 @@ class BudgetClient: 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: @@ -275,16 +312,36 @@ class BudgetClient: # ---- organization --------------------------------------------------- - def create_org(self, *, max_budget: float, alias: str) -> str: + def create_org(self, *, max_budget: float, alias: str, budget_duration: str | None = None) -> str: return unwrap( self.gateway.transport.post( "/organization/new", headers=self.gateway.transport.master, - json=OrgNewBody(organization_alias=alias, max_budget=max_budget), + json=OrgNewBody( + organization_alias=alias, + max_budget=max_budget, + budget_duration=budget_duration, + ), response_type=OrgNewResponse, ) ).organization_id + def org_budget_id(self, org_id: str) -> str | None: + """The id of the budget row backing an org; its budget_reset_at is read via + budget_info (LIT-4570: /organization/new stores budget_duration without + scheduling budget_reset_at, so the reset job's first tick schedules it).""" + result = self.gateway.transport.get( + "/organization/info", + headers=self.gateway.transport.master, + params=OrgInfoParams(organization_id=org_id), + response_type=OrgInfoResponse, + ) + match result: + case Success(data=data): + return data.budget_id + case _: + return None + def delete_org(self, org_id: str) -> None: _ = self.gateway.transport.delete( "/organization/delete", @@ -300,6 +357,7 @@ class BudgetClient: *, alias: str, max_budget: float | None = None, + budget_duration: str | None = None, organization_id: str | None = None, budget_limits: list[BudgetWindow] | None = None, ) -> str: @@ -310,6 +368,7 @@ class BudgetClient: json=TeamNewBody( team_alias=alias, max_budget=max_budget, + budget_duration=budget_duration, organization_id=organization_id, budget_limits=budget_limits, ), diff --git a/tests/e2e/quota_management/budgets/conftest.py b/tests/e2e/quota_management/budgets/conftest.py index 236822f4309..4299d2ffd49 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 dbe1cfa4ea8..c4ad0c38f31 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 @@ Each entity is an E2ECase (lifecycle.E2ECase) driven by run_case: init() creates 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 @@ import pytest 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,18 +70,85 @@ class _BudgetCase: 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): + """A user's max_budget follows the person, not the key. The capped user holds + two personal keys (no team, no key budgets) plus a team-member key on an + uncapped team; once the first personal key is refused, the other two must be + refused as well - a second key is not a fresh allowance, and since #32005 the + user budget draws down team keys too. All refusals must be 429 budget_exceeded.""" + def init(self) -> None: user_id = self.client.create_user(max_budget=3e-6) self._undo.append(lambda: self.client.delete_user(user_id)) - # personal key (no team) -> the user budget governs self.key = self.client.generate_key(user_id=user_id) self._undo.append(lambda: self.client.delete_key(self.key)) + self._second_key = self.client.generate_key(user_id=user_id) + self._undo.append(lambda: self.client.delete_key(self._second_key)) + team_id = self.client.create_team(alias=f"e2e-budget-team-{unique_marker()}") + self._undo.append(lambda: self.client.delete_team(team_id)) + self.client.add_team_member(team_id, user_id) + self._team_key = self.client.generate_key(team_id=team_id, user_id=user_id) + self._undo.append(lambda: self.client.delete_key(self._team_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]}" + ) + for label, key in (("second personal key", self._second_key), ("team-member key", self._team_key)): + result = self.client.chat(key, "claude-haiku-4-5", f"spend {unique_marker()}", max_tokens=16) + assert is_budget_block(result) and result.status_code == 429, ( + f"the {label} of a user over budget must get the same 429 budget_exceeded, " + f"got {result.status_code}: {result.body[:200]}" + ) class EndUserBudgetCase(_BudgetCase): @@ -97,34 +165,67 @@ class EndUserBudgetCase(_BudgetCase): 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): + """Member A's per-team budget is tiny while the team and both members' user + budgets are roomy (100.0), so the only cap that can trip is A's: a block + proves member-level enforcement and must be a 429 budget_exceeded. Teammate + B, uncapped on the same team, must keep serving after A is cut off, proving + the member cap does not leak onto the team or its members.""" + def init(self) -> None: - # Member's per-team budget is tiny while the team has a large budget, so a - # block proves member-level (not team-level) enforcement. - team_id = self.client.create_team( + self._team_id = self.client.create_team( alias=f"e2e-budget-team-{unique_marker()}", max_budget=100.0 ) - self._undo.append(lambda: self.client.delete_team(team_id)) - user_id = self.client.create_user(max_budget=100.0) - self._undo.append(lambda: self.client.delete_user(user_id)) - self.client.add_team_member(team_id, user_id, max_budget_in_team=3e-6) - self.key = self.client.generate_key(team_id=team_id, user_id=user_id) + self._undo.append(lambda: self.client.delete_team(self._team_id)) + self._member_id = self.client.create_user(max_budget=100.0) + self._undo.append(lambda: self.client.delete_user(self._member_id)) + self.client.add_team_member(self._team_id, self._member_id, max_budget_in_team=3e-6) + self.key = self.client.generate_key(team_id=self._team_id, user_id=self._member_id) self._undo.append(lambda: self.client.delete_key(self.key)) + teammate_id = self.client.create_user(max_budget=100.0) + self._undo.append(lambda: self.client.delete_user(teammate_id)) + self.client.add_team_member(self._team_id, teammate_id) + self._teammate_key = self.client.generate_key(team_id=self._team_id, user_id=teammate_id) + self._undo.append(lambda: self.client.delete_key(self._teammate_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]}" + ) + teammate = self.client.chat( + self._teammate_key, + "claude-haiku-4-5", + f"spend {unique_marker()}", + max_tokens=16, + ) + require_successful_call(teammate) def _case_id(case_cls: Type[_BudgetCase]) -> str: @@ -138,6 +239,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_budget_reset_e2e.py b/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py index bdbee027f28..b57ad23ebf5 100644 --- a/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py @@ -1,12 +1,4 @@ -"""Live e2e: a key budget resets (zeroes spend) after its budget_duration. - -Short budget_duration (30s) + the fast-rescheduled reset job: a key blocked for -exceeding its max_budget starts succeeding again once the duration elapses and the -reset job zeroes key.spend. Closes the reset-zeroing gap in -BUDGET_TEST_COVERAGE_MATRIX.md (reset_budget_for_litellm_keys), which the unit -suite covers but no live test did - distinct from the per-window reset in -test_multi_window_budget_e2e.py. -""" +"""Live e2e: an entity blocked over its max_budget serves again after its budget_duration window.""" import time @@ -19,42 +11,107 @@ from lifecycle import ResourceManager pytestmark = pytest.mark.e2e +TINY_CAP = 3e-6 +WINDOW = "30s" +RESET_DEADLINE_SECONDS = 150 + def _call(client: BudgetClient, key: str): - return client.chat( - key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16 - ) + return client.chat(key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16) -@pytest.mark.covers("quota_management.budget.key.resets_after_window") -def test_key_budget_resets_after_duration( - client: BudgetClient, resources: ResourceManager -) -> None: - key = client.generate_key(max_budget=3e-6, budget_duration="30s") - resources.defer(lambda: client.delete_key(key)) - - # 1. exceed the budget -> litellm returns budget_exceeded - blocked = False - for _ in range(20): +def _drive_to_block(client: BudgetClient, key: str) -> None: + """Spend until the cap blocks a call, staying under one window so the block + is observed before the reset job can fire; fail hard if enforcement never trips.""" + for _ in range(12): result = _call(client, key) if is_budget_block(result): - blocked = True - break + return require_successful_call(result) time.sleep(2) - assert blocked, "key budget never enforced" + pytest.fail("budget never enforced before the window could reset") - # 2. once the 30s duration elapses + the reset job runs, key.spend zeroes and - # calls flow again. The window is wall-clock-aligned, so the reset lands up to - # a window later, then the rescheduler (~15-20s) zeroes the spend; allow - # generous headroom over that. A stuck rescheduler is caught by the wait-loop - # timeout, not this elapsed bound. - start = time.monotonic() - while time.monotonic() < start + 150: + +def _poll_until_serves_again(client: BudgetClient, key: str) -> None: + """Poll past the window until the blocked key serves again; every refusal must + stay a budget block, so a crashed reset path or provider error fails loudly.""" + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + while time.monotonic() < deadline: time.sleep(5) result = _call(client, key) if result.ok: - assert time.monotonic() - start < 120, "reset too slow for a 30s budget" return - assert is_budget_block(result), f"non-budget error: {result.body[:200]}" - pytest.fail("key budget never reset within 150s") + if not is_budget_block(result): + pytest.fail(f"non-budget error during reset wait: HTTP {result.status_code}: {result.body[:200]}") + pytest.fail(f"budget never reset within {RESET_DEADLINE_SECONDS}s") + + +class TestBudgetResetDiagonal: + @pytest.mark.covers("quota_management.budget.key.resets_after_window") + def test_bare_key_budget_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(max_budget=TINY_CAP, budget_duration=WINDOW) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.team.resets_after_window") + def test_team_budget_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team( + alias=f"e2e-team-reset-{unique_marker()}", max_budget=TINY_CAP, budget_duration=WINDOW + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.organization.resets_after_window") + def test_org_budget_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + org_id = client.create_org( + max_budget=TINY_CAP, alias=f"e2e-org-reset-{unique_marker()}", budget_duration=WINDOW + ) + resources.defer(lambda: client.delete_org(org_id)) + team_id = client.create_team(alias=f"e2e-org-team-{unique_marker()}", organization_id=org_id) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(key)) + + budget_id = client.org_budget_id(org_id) + assert budget_id, "org created without a budget row" + deadline = time.monotonic() + 30 + while not any(row.budget_reset_at for row in client.budget_info(budget_id)): + if time.monotonic() > deadline: + pytest.fail("org budget window never scheduled by the reset job") + time.sleep(2) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.internal_user.resets_after_window") + def test_personal_key_user_budget_resets_after_window( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + user_id = client.create_user(max_budget=TINY_CAP, budget_duration=WINDOW) + resources.defer(lambda: client.delete_user(user_id)) + key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.internal_user.resets_after_window") + def test_team_member_key_user_budget_resets_after_window( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + user_id = client.create_user(max_budget=TINY_CAP, budget_duration=WINDOW) + resources.defer(lambda: client.delete_user(user_id)) + team_id = client.create_team(alias=f"e2e-user-team-reset-{unique_marker()}") + resources.defer(lambda: client.delete_team(team_id)) + client.add_team_member(team_id, user_id, max_budget_in_team=100.0) + key = client.generate_key(team_id=team_id, user_id=user_id) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) diff --git a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py index 23f3e162761..45e2bf539d4 100644 --- a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py @@ -23,14 +23,16 @@ pytestmark = pytest.mark.e2e WINDOW_SECONDS = 30 # the tight window; calls succeed again only after it elapses # Prefer the OpenAI cheap model for this polling test: under the full stage suite # Claude chat latency + ALB target idle timeout (~60s) can surface as awselb 502 -# HTML mid-wait, which is not a budget signal. gpt-5.5 + 1 token stays well under -# that ceiling so the wait loop measures window reset, not provider/ALB timeout. +# HTML mid-wait, which is not a budget signal. gpt-5.5 stays well under that +# ceiling so the wait loop measures window reset, not provider/ALB timeout. +# max_tokens must be >1: gpt-5.5 refuses completions that hit the output limit +# mid-message when capped at 1 token. MODEL = CHEAP_OPENAI_MODEL def _call(client: BudgetClient, key: str): return client.chat( - key, MODEL, f"window {unique_marker()}", max_tokens=1 + key, MODEL, f"window {unique_marker()}", max_tokens=16 ) 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 00000000000..87855a9a1c1 --- /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 00000000000..4dc7a2df647 --- /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 4a5a73bb5e4..59dee5e65b3 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 062ef8d73da..6baebc4c28c 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 0e80764236b..434af15b182 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 046cdd80c2b..344d8ab5c13 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 ff89c3845e4..ef374de5e2a 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 00000000000..28fb0846481 --- /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 515bf1233aa..f1e6539439a 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -935,8 +935,8 @@ async def test_get_tools_from_mcp_servers(): 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 @@ -978,9 +978,9 @@ async def test_get_tools_from_mcp_servers(): 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" # @@ -1015,8 +1015,8 @@ async def test_get_tools_from_mcp_servers(): 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)}") @@ -2436,11 +2436,12 @@ async def test_filter_tools_by_allowed_tools_integration(): 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 ( @@ -2549,11 +2550,12 @@ async def test_filter_tools_by_disallowed_tools_integration(): 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 ( @@ -2650,11 +2652,12 @@ async def test_filter_tools_no_restrictions_integration(): 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 848a6c28a57..a969d21a681 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 4664cc86303..70c1f65b541 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 copy 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 @@ -1533,3 +1535,242 @@ class TestApplyToAnthropicMessagesRequest: 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 00000000000..a4250acf1dc --- /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/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 450f69fb87c..9b5197d9028 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -3,12 +3,11 @@ Tests for AnthropicResponsesStreamWrapper (litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py) """ +import asyncio import os import sys -sys.path.insert( - 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../..")) -) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, @@ -22,6 +21,60 @@ def _process_all(events: list) -> list: return list(wrapper._chunk_queue) +def _drain_async(events: list) -> list: + async def _gen(): + for event in events: + yield event + + async def _run() -> list: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=_gen(), model="m") + return [chunk async for chunk in wrapper] + + return asyncio.run(_run()) + + +class TestMessageStartEmittedExactlyOnce: + """The ``__anext__`` fallback emits ``message_start`` before consuming the + stream, so ``_process_event`` must not emit a second one when + ``response.created`` later arrives. Two ``message_start`` events (byte + identical, same id) break strict Anthropic SDK clients (e.g. Claude Code) + with 'Content block is not a thinking block' once thinking blocks follow.""" + + def test_response_created_does_not_duplicate_message_start(self): + chunks = _drain_async( + [ + {"type": "response.created"}, + {"type": "response.output_text.delta", "item_id": "m1", "delta": "hi"}, + ] + ) + message_starts = [c for c in chunks if c["type"] == "message_start"] + assert len(message_starts) == 1 + + def test_message_start_is_first_event(self): + chunks = _drain_async([{"type": "response.created"}]) + assert chunks[0]["type"] == "message_start" + + +class TestProcessEventResponseCreatedGuard: + """``_process_event`` must emit ``message_start`` exactly once even if + ``response.created`` arrives more than once. The guard mirrors the + ``__anext__`` fallback's ``_sent_message_start`` flag, so a direct caller + and the async fallback can never double-emit. This also exercises the + guard's emit-branch, which the async path never reaches because the + fallback sets the flag before the upstream stream is consumed.""" + + def test_first_response_created_emits_message_start(self): + chunks = _process_all([{"type": "response.created"}]) + assert len(chunks) == 1 + assert chunks[0]["type"] == "message_start" + assert chunks[0]["message"]["model"] == "m" + + def test_second_response_created_is_skipped(self): + chunks = _process_all([{"type": "response.created"}, {"type": "response.created"}]) + message_starts = [c for c in chunks if c["type"] == "message_start"] + assert len(message_starts) == 1 + + class TestProcessEventTextDeltaWithoutOutputItemAdded: """Streams that skip response.output_item.added (e.g. LMStudio) must still open a text block before any delta and never emit index -1.""" 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 3c410cf84df..6ab0f2c08ab 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 @@ class TestAnthropicThinkingSignatureSelfHeal: ) 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 @@ class TestAnthropicThinkingSignatureSelfHeal: 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 00000000000..3297750fa6e --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -0,0 +1,69 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm +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 +# Read the cached rate from the price map so this test tracks the shipped value +# (glm-5p2 is $0.14/1M) instead of hardcoding a number that breaks when it changes. +CACHE_READ_COST = litellm.get_model_info(model=MODEL, custom_llm_provider="fireworks_ai")["cache_read_input_token_cost"] +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 6f132aaae9c..9375f7481c8 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 @@ class TestMCPRequestHandler: 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 00000000000..cb27e992ecb --- /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 16c36af5156..fdc77d19d73 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 test_aggregate_list_tools_absorbs_one_unauthenticated_server(): ), 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 test_single_server_route_also_absorbs_upstream_auth_error(): ), 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 test_aggregate_with_single_accessible_server_still_absorbs(): 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 7e59904a39c..ae4f12fc1e1 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 test_get_tools_from_mcp_servers_continues_when_one_server_fails(): ) # 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 test_get_tools_from_mcp_servers_handles_all_servers_failing(): ) # 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 test_list_tools_single_server_unprefixed_names(): "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 test_list_tools_single_server_unprefixed_names(): ) # 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 test_list_tools_multiple_servers_prefixed_names(): "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 test_list_tools_multiple_servers_prefixed_names(): ) # 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 test_list_tools_filters_by_key_team_permissions(): "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 test_list_tools_filters_by_key_team_permissions(): ) # 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 test_list_tools_with_team_tool_permissions_inheritance(): "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 test_list_tools_with_team_tool_permissions_inheritance(): ) # 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 test_list_tools_with_no_tool_permissions_shows_all(): "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 test_list_tools_with_no_tool_permissions_shows_all(): ) # 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 test_list_tools_strips_prefix_when_matching_permissions(): "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 test_list_tools_strips_prefix_when_matching_permissions(): ) # 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 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab ): 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 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab 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 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab 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 test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow(): 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 test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow(): "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 80bf08a5eba..55b6bbbdbc2 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 @@ from unittest.mock import AsyncMock, MagicMock, patch 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 @@ class TestMCPServerManager: 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 @@ class TestMCPServerManager: 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 @@ class TestMCPToolsListAuthSurfacing: 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 @@ class TestMCPToolsListAuthSurfacing: 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 @@ class TestMCPToolsListAuthSurfacing: ) ) - 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 b8f0b205831..1da44029b5c 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 @@ from unittest.mock import AsyncMock, MagicMock, patch 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 @@ class TestCallToolRestApiVirtualTools: 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 @@ class TestCallToolRestApiVirtualTools: 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 99e05182361..d4ba66c4381 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 @@ class TestListToolsRestAPI: ) 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 @@ class TestListToolsRestAPI: 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 @@ class TestListToolsRestAPI: ) 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 @@ class TestToolResponseMcpInfoEnrichment: "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 cc4a7d5bfb4..2da645bf4e1 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 17ff700791f..b5d8727f7e6 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 typing import Optional 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 @@ class TestGetEndUserIdFromRequestBodyWithStandardHeaders: 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 00000000000..14d8e90e027 --- /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/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py new file mode 100644 index 00000000000..ca57118ee9d --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -0,0 +1,733 @@ +import json +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from litellm.exceptions import GuardrailRaisedException, ModifyResponseException +from litellm.proxy.guardrails.guardrail_hooks.straiker import initialize_guardrail +from litellm.proxy.guardrails.guardrail_hooks.straiker.straiker import ( + StraikerGuardrail, +) +from litellm.proxy.guardrails.guardrail_registry import ( + guardrail_class_registry, + guardrail_initializer_registry, +) +from litellm.types.proxy.guardrails.guardrail_hooks.straiker import ( + StraikerGuardrailConfigModel, + StraikerGuardrailConfigModelOptionalParams, +) +from litellm.types.utils import Choices, Message, ModelResponse, Usage + + +def _mock_response(action: str, turn_id: str = "turn-1", schema_version: str = "1", **extra) -> MagicMock: + resp = MagicMock(spec=httpx.Response) + resp.status_code = 200 + resp.json.return_value = { + "schema_version": schema_version, + "action": action, + "turn_id": turn_id, + **extra, + } + resp.text = "" + return resp + + +def _make_guardrail(**overrides) -> StraikerGuardrail: + defaults = dict( + api_key="test-key", + api_base="https://test.straiker.ai", + max_retries=0, + guardrail_name="straiker", + event_hook="pre_call", + async_handler=MagicMock(spec=httpx.AsyncClient), + ) + defaults.update(overrides) + g = StraikerGuardrail(**defaults) + g.async_handler.post = AsyncMock() + return g + + +def _logging_obj() -> MagicMock: + obj = MagicMock() + obj.litellm_call_id = "call-123" + obj.litellm_trace_id = "trace-456" + obj.call_type = "acompletion" + return obj + + +def _posted_payload(g: StraikerGuardrail) -> dict: + return json.loads(g.async_handler.post.call_args.kwargs["content"]) + + +def test_registry_membership(): + assert "straiker" in guardrail_initializer_registry + assert guardrail_class_registry["straiker"] is StraikerGuardrail + + +def test_config_model_wiring(): + assert StraikerGuardrailConfigModel.ui_friendly_name() == "Straiker" + assert StraikerGuardrail.get_config_model() is StraikerGuardrailConfigModel + fields = StraikerGuardrailConfigModel.model_fields + assert "api_key" in fields + assert "api_base" in fields + assert "default_app" in fields + assert "source" not in fields + assert "optional_params" in fields + assert "timeout" not in fields + assert "verbose" not in fields + + +def test_init_rejects_empty_api_key(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="") + + +def test_init_rejects_invalid_fallback(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="k", unreachable_fallback="nope") + + +def test_supported_hooks_limited_to_pre_and_post(): + from litellm.types.guardrails import GuardrailEventHooks + + assert StraikerGuardrail.get_supported_event_hooks() == [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + +def test_during_call_mode_rejected_at_init(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="k", event_hook="during_call") + + +def test_streaming_attrs_hardcoded_to_buffered(): + g = _make_guardrail() + assert g.streaming_buffer_until_moderated is True + assert g.streaming_end_of_stream_only is True + + +def test_streaming_flags_not_configurable(): + fields = StraikerGuardrailConfigModelOptionalParams.model_fields + assert "streaming_buffer_until_moderated" not in fields + assert "streaming_end_of_stream_only" not in fields + assert "streaming_sampling_rate" not in fields + + +def test_initializer_builds_working_callback(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams(guardrail="straiker", mode="pre_call", api_key="abc", api_base="https://x.straiker.ai") + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.api_base == "https://x.straiker.ai" + + +def test_initializer_maps_default_app_to_source(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams( + guardrail="straiker", + mode="pre_call", + api_key="abc", + default_app="My App", + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert callback.source == "My App" + + +def test_initializer_reads_optional_params_flattened_like_ui(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + timeout=9.5, + verbose=True, + unreachable_fallback="fail_open", + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 9.5 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + assert callback.api_base == "https://x.straiker.ai" + + +def test_initializer_reads_nested_optional_params(): + from types import SimpleNamespace + + from litellm.types.guardrails import LitellmParams + + params = LitellmParams.model_construct( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + optional_params=SimpleNamespace( + timeout=7.25, + verbose=True, + unreachable_fallback="fail_open", + ), + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 7.25 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + + +def test_initializer_reads_dict_optional_params(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams.model_construct( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + optional_params={"timeout": 7.25, "verbose": True, "unreachable_fallback": "fail_open"}, + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 7.25 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + + +@pytest.mark.asyncio +async def test_request_envelope_transport_and_shape(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + inputs = {"texts": ["hello world"], "model": "gpt-4o-mini"} + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello world"}], + "metadata": {"user_api_key_alias": "team-key", "agent_id": "chatbot-app", "app_name": "Chatbot"}, + } + + out = await g.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request", logging_obj=_logging_obj()) + + assert out is inputs + url = g.async_handler.post.call_args.args[0] + assert url == "https://test.straiker.ai/api/v1/detect/webhook" + headers = g.async_handler.post.call_args.kwargs["headers"] + assert headers["X-Straiker-Webhook-Format"] == "litellm" + assert headers["Authorization"] == "Bearer test-key" + + payload = _posted_payload(g) + assert payload["schema_version"] == "1" + assert payload["event"]["type"] == "pre_call" + assert payload["event"]["id"] == "call-123:request" + assert payload["request"]["texts"] == ["hello world"] + assert payload["context"]["litellm_call_id"] == "call-123" + assert payload["identity"]["litellm_key"] == "team-key" + assert payload["application"] == {"source": "chatbot-app", "name": "Chatbot"} + assert "session_id" not in payload["application"] + assert "user_name" not in payload["application"] + assert "user_role" not in payload["application"] + assert "response" not in payload + assert "metadata" not in payload + + +@pytest.mark.asyncio +async def test_webhook_metadata_session_id_and_opaque_passthrough(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "litellm_session_id": "sess-from-litellm", + "metadata": { + "agent_id": "chatbot-app", + "app_name": "Chatbot", + "user_api_key_alias": "team-key", + "custom_tag": "experiment-7", + "client_ip": "10.0.0.1", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + payload = _posted_payload(g) + assert payload["application"] == {"source": "chatbot-app", "name": "Chatbot"} + assert payload["identity"]["litellm_key"] == "team-key" + assert payload["context"]["session_id"] == "sess-from-litellm" + assert "session_id" not in payload["metadata"] + assert payload["metadata"] == { + "custom_tag": "experiment-7", + "client_ip": "10.0.0.1", + } + + +@pytest.mark.asyncio +async def test_webhook_metadata_never_forwards_proxy_internal_keys(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "custom_tag": "experiment-7", + "user_api_key": "sk-hashed-secret", + "user_api_end_user_max_budget": 12.5, + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == {"custom_tag": "experiment-7"} + + +@pytest.mark.asyncio +async def test_default_metadata_injected_and_config_wins_on_clash(): + g = _make_guardrail(metadata={"tenant": "acme", "custom_tag": "config-value"}) + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": {"custom_tag": "request-value", "client_ip": "10.0.0.1"}, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == { + "client_ip": "10.0.0.1", + "custom_tag": "config-value", + "tenant": "acme", + } + + +@pytest.mark.asyncio +async def test_default_metadata_present_without_request_metadata(): + g = _make_guardrail(metadata={"tenant": "acme"}) + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m"}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == {"tenant": "acme"} + + +@pytest.mark.asyncio +async def test_context_session_id_from_request_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "metadata": {"session_id": "sess-meta"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + payload = _posted_payload(g) + assert payload["context"]["session_id"] == "sess-meta" + assert "metadata" not in payload + + +@pytest.mark.asyncio +async def test_identity_key_and_team_coalesce_alias_over_id(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_alias": "prod-key", + "user_api_key_hash": "hash-abc", + "user_api_key_team_alias": "growth", + "user_api_key_team_id": "team-9", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["litellm_key"] == "prod-key" + assert identity["litellm_team"] == "growth" + assert "key" not in identity + assert "team" not in identity + + +@pytest.mark.asyncio +async def test_identity_key_and_team_fall_back_to_hash_and_id(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_hash": "hash-abc", + "user_api_key_team_id": "team-9", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["litellm_key"] == "hash-abc" + assert identity["litellm_team"] == "team-9" + + +@pytest.mark.asyncio +async def test_identity_end_user_from_resolved_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_end_user_id": "eu-meta", + "user_api_key_user_id": "default_user_id", + }, + "user": "eu-body", + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["end_user_id"] == "eu-meta" + assert identity["litellm_user_id"] == "default_user_id" + assert _posted_payload(g)["application"] == {"source": g.source} + + +@pytest.mark.asyncio +async def test_identity_end_user_absent_without_resolved_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "user": "eu-body", "metadata": {"user_api_key_user_id": "default_user_id"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert "end_user_id" not in _posted_payload(g)["identity"] + + +@pytest.mark.asyncio +async def test_application_source_from_agent_id(): + g = _make_guardrail(source="litellm") + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "metadata": {"agent_id": "analytics-app", "app_name": "Analytics"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["application"] == {"source": "analytics-app", "name": "Analytics"} + +@pytest.mark.asyncio +async def test_request_block_raises_guardrail_exception_with_reason(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("BLOCKED", blocked_reason="prompt injection") + with pytest.raises(GuardrailRaisedException) as exc: + await g.apply_guardrail( + inputs={"texts": ["attack"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert "prompt injection" in str(exc.value) + + +@pytest.mark.asyncio +async def test_guardrail_intervened_writes_back_modified_text_only(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + inputs = {"texts": ["my ssn is 123"], "images": ["img-a"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out["texts"] == ["[redacted]"] + assert out["images"] == ["img-a"] + + +@pytest.mark.asyncio +async def test_streamed_response_intervention_converts_to_block(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "p"}], + "stream": True, + "response": response, + } + with pytest.raises(ModifyResponseException): + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_non_streamed_response_intervention_redacts(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "p"}], + "response": response, + } + out = await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + assert out["texts"] == ["[redacted]"] + + +@pytest.mark.asyncio +async def test_guardrail_intervened_without_texts_blocks(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED") + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["my ssn is 123"]}, + request_data={"model": "m"}, + input_type="request", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_streamed_via_proxy_server_request_body_converts_to_block(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "proxy_server_request": {"body": {"stream": True}}, + "response": response, + } + with pytest.raises(ModifyResponseException): + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_response_envelope_and_block_replaces_response(): + g = _make_guardrail(verbose=True) + g.async_handler.post.return_value = _mock_response("BLOCKED") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "original prompt"}], + "stream": True, + "response": response, + } + with pytest.raises(ModifyResponseException) as exc: + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + assert exc.value.original_response is response + payload = _posted_payload(g) + assert payload["event"]["type"] == "post_call" + assert payload["event"]["stream"]["phase"] == "assembled" + assert payload["response"]["texts"] == ["secret"] + assert payload["response"]["finish_reason"] == "stop" + assert payload["request"]["structured_messages"] == [{"role": "user", "content": "original prompt"}] + + +@pytest.mark.asyncio +async def test_post_call_fail_closed_raises_modify_response_exception(): + g = _make_guardrail(unreachable_fallback="fail_closed") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = {"model": "gpt-4o-mini", "response": response} + with pytest.raises(ModifyResponseException) as exc: + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + assert exc.value.original_response is response + + +@pytest.mark.asyncio +async def test_usage_tokens_on_post_call(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="hi", role="assistant"))], + model="gpt-4o-mini", + usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18), + ) + await g.apply_guardrail( + inputs={"texts": ["hi"], "model": "gpt-4o-mini"}, + request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hey"}], "response": response}, + input_type="response", + logging_obj=_logging_obj(), + ) + usage = _posted_payload(g)["usage"] + assert usage == {"input_tokens": 11, "output_tokens": 7} + + +@pytest.mark.asyncio +async def test_usage_absent_on_pre_call(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={"model": "gpt-4o-mini"}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert "usage" not in _posted_payload(g) + + +@pytest.mark.asyncio +async def test_allow_returns_inputs_unchanged(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + inputs = {"texts": ["fine"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_unreachable_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_unreachable_fail_open_passes_through(): + g = _make_guardrail(unreachable_fallback="fail_open") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_fail_on_error_false_allows_on_bad_status(): + g = _make_guardrail(unreachable_fallback="fail_closed", fail_on_error=False) + bad = MagicMock(spec=httpx.Response) + bad.status_code = 400 + bad.text = "bad request" + g.async_handler.post.return_value = bad + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_non_retryable_status_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed", fail_on_error=True) + bad = MagicMock(spec=httpx.Response) + bad.status_code = 401 + bad.text = "unauthorized" + g.async_handler.post.return_value = bad + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_payload_size_guard_fails_closed(): + g = _make_guardrail(max_payload_bytes=10) + inputs = {"texts": ["x" * 5000]} + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + g.async_handler.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_payload_size_guard_blocks_even_with_fail_open(): + g = _make_guardrail(max_payload_bytes=10, unreachable_fallback="fail_open") + inputs = {"texts": ["x" * 5000]} + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + g.async_handler.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_invalid_response_schema_blocks_even_with_fail_open(): + g = _make_guardrail(unreachable_fallback="fail_open") + bad = MagicMock(spec=httpx.Response) + bad.status_code = 200 + bad.json.return_value = {"action": "NOT_A_VALID_ACTION"} + bad.text = "" + g.async_handler.post.return_value = bad + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_unreachable_http_status_fail_open_passes(): + g = _make_guardrail(unreachable_fallback="fail_open") + resp = MagicMock(spec=httpx.Response) + resp.status_code = 503 + resp.text = "service unavailable" + g.async_handler.post.return_value = resp + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_unreachable_http_status_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed") + resp = MagicMock(spec=httpx.Response) + resp.status_code = 503 + resp.text = "service unavailable" + g.async_handler.post.return_value = resp + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) 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 e7d2909263a..c76e1a60afd 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 @@ import litellm 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 test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ # 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 test_async_log_success_event_counts_non_chat_response_tokens( @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 - ) - - # 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-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 + + await parallel_request_handler.async_log_failure_event( + 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"] ) - # 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 + 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 = [] + + 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 test_async_streaming_data_generator_releases_counter_on_disconnect_v3( 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 test_async_streaming_data_generator_releases_counter_on_disconnect_v3( 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 test_async_data_generator_releases_counter_on_disconnect_v3(disconnect 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 test_async_data_generator_releases_counter_on_disconnect_v3(disconnect 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 test_async_data_generator_releases_counter_when_wrapped_v3(): 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/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 89d100cc3a4..fdf629c36bd 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -396,6 +396,122 @@ def test_is_langfuse_route(): assert handler.is_langfuse_route("") is False +def test_is_vertex_route_ignores_plain_predict_path_segment(): + """ + Regression for LIT-4527: a custom (non-Vertex) passthrough URL whose path + contains a plain `predict`/`search` segment must not be classified as a + Vertex route, otherwise its success logging is routed into the Vertex + handler, the Vertex-shaped transform fails, and no log row is recorded. + + Real Vertex custom methods are invoked with the GCP `resource:method` colon + syntax, so only the colon form should count as Vertex. + """ + handler = PassThroughEndpointLogging() + + assert ( + handler.is_vertex_route( + "https://upstream.example.com/ml/api/v1/time-series-forecast/predict" + ) + is False + ) + assert handler.is_vertex_route("https://upstream.example.com/api/v1/search") is False + assert ( + handler.is_vertex_route("https://upstream.example.com/predict/generateContent") + is False + ) + + assert ( + handler.is_vertex_route( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/l/publishers/google/models/m:predict" + ) + is True + ) + assert ( + handler.is_vertex_route( + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/l/publishers/anthropic/models/claude:rawPredict" + ) + is True + ) + assert ( + handler.is_vertex_route( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/l/publishers/google/models/m:generateContent" + ) + is True + ) + assert ( + handler.is_vertex_route( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/l/publishers/google/models/m:predictLongRunning" + ) + is True + ) + assert ( + handler.is_vertex_route("https://discoveryengine.googleapis.com/v1/x:search") + is True + ) + + assert ( + handler.is_vertex_route( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/l/batchPredictionJobs" + ) + is True + ) + + +@pytest.mark.asyncio +async def test_custom_passthrough_predict_path_logs_via_generic_handler(): + """ + Regression for LIT-4527: an upstream-successful custom passthrough request to + a `/predict` path (non-Vertex body) must still produce a log row. Before the + fix it was routed into VertexPassthroughLoggingHandler, which failed on the + non-Vertex body and dropped the log entirely. + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, + ) + + handler = PassThroughEndpointLogging() + handler._handle_logging = AsyncMock() + + mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {} + + mock_response = MagicMock(spec=httpx.Response) + mock_response.text = '{"forecast": [1, 2, 3]}' + + url_route = "https://upstream.example.com/ml/api/v1/time-series-forecast/predict" + passthrough_logging_payload = PassthroughStandardLoggingPayload( + url=url_route, + request_body={"series": [1, 2]}, + request_method="POST", + ) + + with patch( + "litellm.proxy.pass_through_endpoints.success_handler.VertexPassthroughLoggingHandler.vertex_passthrough_handler" + ) as mock_vertex_handler: + await handler.pass_through_async_success_handler( + httpx_response=mock_response, + response_body={"forecast": [1, 2, 3]}, + logging_obj=mock_logging_obj, + url_route=url_route, + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"series": [1, 2]}, + passthrough_logging_payload=passthrough_logging_payload, + ) + + mock_vertex_handler.assert_not_called() + handler._handle_logging.assert_awaited_once() + logged_object = handler._handle_logging.call_args.kwargs[ + "standard_logging_response_object" + ] + assert logged_object == {"response": '{"forecast": [1, 2, 3]}'} + + @pytest.mark.asyncio async def test_langfuse_passthrough_no_logging(): """ 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 ee8d8b22779..dca93e137ac 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 45d35419680..bd8e92c3cc2 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 @@ from litellm.proxy.proxy_server import ( _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 0c45e31afd2..677ab8765bb 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 656e1406f07..15a117bd6fc 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -242,3 +242,88 @@ class TestRagIngestSSRFBlocked: 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 0b304f2fec7..1db76aed61d 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2279,7 +2279,8 @@ 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 - return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock() + generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=MagicMock(), user_api_key_dict=valid_token, request_data=_request_body(), @@ -2287,6 +2288,7 @@ def _drive_streaming_cancel(valid_token, iterator_hook): serialize_chunk=lambda chunk: chunk, serialize_error=lambda exc: str(exc), ) + return generator, streaming_logging_obj @pytest.mark.asyncio @@ -2305,7 +2307,7 @@ async def test_streaming_cancel_before_any_chunk_reconciles_to_input_cost( yield "" # make this an async generator raise asyncio.CancelledError() - generator = _drive_streaming_cancel(valid_token, cancel_before_chunk) + generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_before_chunk) received = [] with pytest.raises(asyncio.CancelledError): async for chunk in generator: @@ -2318,6 +2320,7 @@ async def test_streaming_cancel_before_any_chunk_reconciles_to_input_cost( key="spend:key:key-cancel-no-chunk" ) == pytest.approx(0.5) assert reservation["finalized"] is True + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() @pytest.mark.asyncio @@ -2336,7 +2339,7 @@ async def test_streaming_cancel_after_chunk_keeps_reservation( yield "data: chunk\n\n" raise asyncio.CancelledError() - generator = _drive_streaming_cancel(valid_token, cancel_after_chunk) + generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_after_chunk) received = [] with pytest.raises(asyncio.CancelledError): async for chunk in generator: @@ -2348,6 +2351,7 @@ async def test_streaming_cancel_after_chunk_keeps_reservation( key="spend:key:key-cancel-after-chunk" ) == pytest.approx(2.0) assert reservation.get("finalized") is not True + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() @pytest.mark.asyncio @@ -2382,6 +2386,7 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_ streaming_logging_obj = MagicMock() streaming_logging_obj.async_post_call_streaming_iterator_hook = one_chunk + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock() # On the slow path the per-chunk hook is awaited before the chunk is yielded # to the client; cancel there. Nothing has reached the client yet. streaming_logging_obj.async_post_call_streaming_hook = AsyncMock( @@ -2411,6 +2416,7 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_ key="spend:key:key-cancel-slowpath" ) == pytest.approx(0.5) assert reservation["finalized"] is True + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() @pytest.mark.asyncio @@ -2427,7 +2433,7 @@ async def test_streaming_disconnect_after_consuming_chunk_keeps_reservation( yield "data: a\n\n" yield "data: b\n\n" - generator = _drive_streaming_cancel(valid_token, two_chunks) + generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, two_chunks) # Client consumes one chunk, then disconnects. aclose() raises GeneratorExit # at the suspended yield, after the chunk already reached the client. @@ -2440,6 +2446,7 @@ async def test_streaming_disconnect_after_consuming_chunk_keeps_reservation( key="spend:key:key-disconnect-after-chunk" ) == pytest.approx(2.0) assert reservation.get("finalized") is not True + streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index aa1911f80bc..ebfbb46053d 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 @@ from litellm.proxy.common_request_processing import ( 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 @@ class TestPreCallWithFallbacksOnLocalRateLimit: }, 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_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index d2b8b7ec23d..47879ee96ad 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2560,6 +2560,108 @@ def test_add_litellm_metadata_from_request_headers_generic_session_id_header(): assert data["litellm_trace_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" +def test_add_litellm_metadata_from_anthropic_user_id_sets_session_id(): + data = { + "metadata": { + "user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01" + } + } + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + assert data["litellm_session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + assert "litellm_trace_id" not in data + + +def test_add_litellm_metadata_from_anthropic_user_id_dict_sets_session_id(): + data = { + "metadata": { + "user_id": { + "device_id": "device", + "account_uuid": "account", + "session_id": "sess_4f8c1d2a-1234", + } + } + } + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["user_id"] == "sess_4f8c1d2a-1234" + assert data["metadata"]["session_id"] == "sess_4f8c1d2a-1234" + assert data["litellm_session_id"] == "sess_4f8c1d2a-1234" + assert "litellm_trace_id" not in data + + +def test_add_litellm_metadata_from_headers_session_id_beats_anthropic_user_id(): + data = { + "metadata": { + "user_id": "user_abc123_account__session_body-session-id", + } + } + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-litellm-session-id": "header-session-id"}, + data=data, + _metadata_variable_name="metadata", + ) + assert data["metadata"]["session_id"] == "header-session-id" + assert data["litellm_session_id"] == "header-session-id" + assert data["litellm_trace_id"] == "header-session-id" + + +def test_add_litellm_metadata_from_headers_session_id_beats_anthropic_user_id_dict(): + data = { + "metadata": { + "user_id": { + "session_id": "body-session-id", + } + } + } + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-litellm-session-id": "header-session-id"}, + data=data, + _metadata_variable_name="metadata", + ) + assert data["metadata"]["session_id"] == "header-session-id" + assert data["litellm_session_id"] == "header-session-id" + assert data["litellm_trace_id"] == "header-session-id" + + +@pytest.mark.parametrize( + "user_id", + [ + "user_abc123_account__session_", + "user_abc123_account_", + "user_abc123_account__session_invalid!", + ], +) +def test_add_litellm_metadata_from_anthropic_user_id_ignores_invalid_session_id(user_id: str): + data = {"metadata": {"user_id": user_id}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="metadata" + ) + assert data == {"metadata": {"user_id": user_id}} + + +@pytest.mark.parametrize( + "user_id", + [ + {}, + {"session_id": 123}, + {"session_id": "invalid session id"}, + {"session_id": ""}, + ], +) +def test_add_litellm_metadata_from_anthropic_user_id_dict_ignores_invalid_session_id( + user_id: object, +): + data = {"metadata": {"user_id": user_id}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="metadata" + ) + assert data == {"metadata": {"user_id": user_id}} + + def test_add_litellm_metadata_from_request_headers_explicit_header_beats_generic(): """Explicit x-litellm-trace-id wins over a generic x-*-session-id header.""" headers = { diff --git a/tests/test_litellm/proxy/test_model_id_header_propagation.py b/tests/test_litellm/proxy/test_model_id_header_propagation.py index e48168f89b3..f7f3eabae6d 100644 --- a/tests/test_litellm/proxy/test_model_id_header_propagation.py +++ b/tests/test_litellm/proxy/test_model_id_header_propagation.py @@ -200,6 +200,60 @@ def test_get_custom_headers_without_model_id(): assert headers["x-litellm-model-id"] in [None, ""] +class _FakeLoggingObj: + def __init__(self, litellm_params): + self.litellm_params = litellm_params + self.litellm_call_id = "test-call-id" + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_get_custom_headers_includes_deployment_model_name(metadata_key): + """ + x-litellm-model-name should expose the underlying deployment model string, + which the router records under litellm_params[metadata]["deployment"]. + """ + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.tpm_limit = 1000 + mock_user_api_key_dict.rpm_limit = 100 + + logging_obj = _FakeLoggingObj( + litellm_params={metadata_key: {"deployment": "azure/gpt-4o-2024-08-06"}} + ) + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + model_id="deployment-uuid", + request_data={}, + hidden_params={}, + litellm_logging_obj=logging_obj, + ) + + assert headers["x-litellm-model-name"] == "azure/gpt-4o-2024-08-06" + assert headers["x-litellm-model-id"] == "deployment-uuid" + + +def test_get_custom_headers_omits_model_name_when_deployment_missing(): + """ + Without a deployment model string, x-litellm-model-name must not be emitted + (rather than leaking an empty/None value). + """ + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.tpm_limit = 1000 + mock_user_api_key_dict.rpm_limit = 100 + + logging_obj = _FakeLoggingObj(litellm_params={"metadata": {}}) + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + model_id="deployment-uuid", + request_data={}, + hidden_params={}, + litellm_logging_obj=logging_obj, + ) + + assert "x-litellm-model-name" not in headers + + def test_get_custom_headers_with_empty_string_model_id(): """ Test that get_custom_headers handles empty string model_id correctly. diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 6b0c0dba40f..5d2236fd918 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -582,6 +582,79 @@ class TestProxyInitializationHelpers: ), 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/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index a909c510581..d2bdb1764a4 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -476,101 +476,118 @@ class TestPostCallFailureHookLiftsRecoveredPartialSpend: assert "response_cost" not in request_data +from typing import cast + from litellm.proxy.utils import create_model_info_response -from litellm.types.router import ModelGroupInfo +from litellm.types.utils import ModelInfo -def _router_returning(model_group_info): - router = MagicMock() - router.get_model_group_info = MagicMock(return_value=model_group_info) - return router +def _fake_model_info(**fields: int) -> ModelInfo: + return cast(ModelInfo, dict(fields)) -def test_create_model_info_response_includes_max_tokens_when_available(): - router = _router_returning( - ModelGroupInfo( - model_group="qwen-vllm", - providers=["hosted_vllm"], - max_input_tokens=32768, - max_output_tokens=8192, - ) +def _raise_unmapped(model_id: str) -> ModelInfo: + raise ValueError(f"This model isn't mapped yet: {model_id}") + + +def test_create_model_info_response_includes_max_tokens_from_lookup(): + response = create_model_info_response( + model_id="some-model", + provider="openai", + llm_router=None, + get_model_info=lambda _model: _fake_model_info( + max_input_tokens=128000, max_output_tokens=16384 + ), ) + assert response["id"] == "some-model" + assert response["object"] == "model" + assert response["max_input_tokens"] == 128000 + assert response["max_output_tokens"] == 16384 + + +def test_create_model_info_response_does_not_call_router_group_info(): + router = MagicMock() + router.get_configured_token_limits.return_value = (None, None) + response = create_model_info_response( - model_id="qwen-vllm", provider="openai", llm_router=router + model_id="some-model", + provider="openai", + llm_router=router, + get_model_info=lambda _model: _fake_model_info( + max_input_tokens=128000, max_output_tokens=16384 + ), ) - router.get_model_group_info.assert_called_once_with("qwen-vllm") - assert response["id"] == "qwen-vllm" - assert response["object"] == "model" - assert response["max_input_tokens"] == 32768 - assert response["max_output_tokens"] == 8192 + router.get_model_group_info.assert_not_called() + assert response["max_input_tokens"] == 128000 + + +def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map(): + router = MagicMock() + router.get_configured_token_limits.return_value = (32000, 8000) + + response = create_model_info_response( + model_id="my-custom-deployment", + provider="openai", + llm_router=router, + get_model_info=_raise_unmapped, + ) + + router.get_model_group_info.assert_not_called() + assert response["max_input_tokens"] == 32000 + assert response["max_output_tokens"] == 8000 + + +def test_create_model_info_response_deployment_limits_override_cost_map(): + router = MagicMock() + router.get_configured_token_limits.return_value = (200000, None) + + response = create_model_info_response( + model_id="gpt-4o", + provider="openai", + llm_router=router, + get_model_info=lambda _model: _fake_model_info( + max_input_tokens=128000, max_output_tokens=16384 + ), + ) + + assert response["max_input_tokens"] == 200000 + assert response["max_output_tokens"] == 16384 def test_create_model_info_response_emits_integer_token_counts(): - # ModelGroupInfo types the limits as float; OpenAI-compatible clients expect - # plain integers, so the response must not leak 128000.0. - router = _router_returning( - ModelGroupInfo( - model_group="gpt-4o", - providers=["openai"], - max_input_tokens=128000.0, - max_output_tokens=16384.0, - ) - ) - response = create_model_info_response( - model_id="gpt-4o", provider="openai", llm_router=router + model_id="some-model", + provider="openai", + llm_router=None, + get_model_info=lambda _model: _fake_model_info( + max_input_tokens=128000, max_output_tokens=16384 + ), ) - assert response["max_input_tokens"] == 128000 assert isinstance(response["max_input_tokens"], int) - assert response["max_output_tokens"] == 16384 assert isinstance(response["max_output_tokens"], int) def test_create_model_info_response_omits_unknown_individual_limit(): - router = _router_returning( - ModelGroupInfo( - model_group="partial", - providers=["openai"], - max_input_tokens=4096, - max_output_tokens=None, - ) - ) - response = create_model_info_response( - model_id="partial", provider="openai", llm_router=router + model_id="some-embedding", + provider="openai", + llm_router=None, + get_model_info=lambda _model: _fake_model_info(max_input_tokens=8191), ) - assert response["max_input_tokens"] == 4096 + assert response["max_input_tokens"] == 8191 assert "max_output_tokens" not in response -def test_create_model_info_response_omits_limits_when_both_none(): - router = _router_returning( - ModelGroupInfo( - model_group="no-limits", - providers=["openai"], - max_input_tokens=None, - max_output_tokens=None, - ) - ) - +def test_create_model_info_response_omits_limits_when_lookup_raises(): response = create_model_info_response( - model_id="no-limits", provider="openai", llm_router=router - ) - - assert "max_input_tokens" not in response - assert "max_output_tokens" not in response - - -def test_create_model_info_response_omits_limits_when_group_unknown(): - # Wildcard routes / access groups have no ModelGroupInfo. - router = _router_returning(None) - - response = create_model_info_response( - model_id="openai/*", provider="openai", llm_router=router + model_id="openai/*", + provider="openai", + llm_router=None, + get_model_info=_raise_unmapped, ) assert response["id"] == "openai/*" @@ -578,32 +595,33 @@ def test_create_model_info_response_omits_limits_when_group_unknown(): assert "max_output_tokens" not in response -def test_create_model_info_response_degrades_when_group_info_raises(): - # A malformed deployment must not turn the listing into a 500; the entry - # falls back to the base fields without limits. - router = MagicMock() - router.get_model_group_info = MagicMock(side_effect=ValueError("bad deployment")) - - response = create_model_info_response( - model_id="broken", provider="openai", llm_router=router - ) - - assert response["id"] == "broken" - assert "max_input_tokens" not in response - assert "max_output_tokens" not in response - - def test_create_model_info_response_no_router_keeps_base_fields(): response = create_model_info_response( - model_id="some-model", provider="openai", llm_router=None + model_id="totally-unknown-model-xyz", + provider="openai", + llm_router=None, + get_model_info=_raise_unmapped, ) assert response == { - "id": "some-model", + "id": "totally-unknown-model-xyz", "object": "model", "created": response["created"], "owned_by": "openai", } + + +def test_create_model_info_response_reads_real_cost_map(): + response = create_model_info_response( + model_id="gpt-4o", provider="openai", llm_router=None + ) + + assert isinstance(response["max_input_tokens"], int) + assert response["max_input_tokens"] > 0 + assert isinstance(response["max_output_tokens"], int) + assert response["max_output_tokens"] > 0 + + class TestPostCallFailureHookLLMExceptionAlerting: """The llm_exceptions alert is for infra / LLM-API failures, not user errors (https://github.com/BerriAI/litellm/issues/3395). Already-normalized 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 bf77ef81641..3d76ad54a9c 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/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index b8e4013c960..59268e1427b 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -103,18 +103,16 @@ def test_is_known_vector_store_index_error_path_no_registry(monkeypatch): def test_create_model_info_response_happy_path_no_metadata(): result = create_model_info_response(model_id="gpt-4o", provider="openai") - assert result == { - "id": "gpt-4o", - "object": "model", - "created": result["created"], - "owned_by": "openai", - } snapshot = { "id": result["id"], "object": result["object"], "owned_by": result["owned_by"], "created_is_int": isinstance(result["created"], int), "metadata_absent": "metadata" not in result, + "max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int) + and result["max_input_tokens"] > 0, + "max_output_tokens_positive_int": isinstance(result["max_output_tokens"], int) + and result["max_output_tokens"] > 0, } assert snapshot == { "id": "gpt-4o", @@ -122,6 +120,8 @@ def test_create_model_info_response_happy_path_no_metadata(): "owned_by": "openai", "created_is_int": True, "metadata_absent": True, + "max_input_tokens_positive_int": True, + "max_output_tokens_positive_int": True, } diff --git a/tests/test_litellm/rag/__init__.py b/tests/test_litellm/rag/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py new file mode 100644 index 00000000000..584124ba06a --- /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 6fdbb0741aa..a1347aa111c 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 @@ import pytest 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 604155e1221..a4e803f59ad 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 import Router 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 d6d89c8e811..5662870a5cb 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 @@ from litellm.types.router import ( 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 12b2c9abefb..280a0fe072a 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 @@ from litellm.router_strategy.complexity_router.config import ( ComplexityRouterConfig, ComplexityTier, ) +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + TaggedPreRoutingStrategy, +) @pytest.fixture @@ -953,7 +958,7 @@ class TestRouterComplexityDeploymentMethods: ] ) - 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 @@ class TestRouterComplexityDeploymentMethods: 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 @@ class TestRoutingPlugins: 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 @@ class TestRoutingPlugins: 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 00000000000..4413cbc12ef --- /dev/null +++ b/tests/test_litellm/test_gpt_realtime_mode.py @@ -0,0 +1,90 @@ +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(monkeypatch): + """get_model_info must resolve the retag against the bundled cost map, not the + hosted map fetched from main, which lags this repo until the next promotion.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + try: + assert litellm.get_model_info("gpt-realtime-mini")["mode"] == "realtime" + finally: + litellm.get_model_info.cache_clear() + + +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_router.py b/tests/test_litellm/test_router.py index c2c98c8869c..55c09e6cac4 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5430,3 +5430,50 @@ class TestRouterRequestTimeoutPropagation: ) == 60 ) + + +def test_get_configured_token_limits_reads_deployment_model_info(): + router = litellm.Router( + model_list=[ + { + "model_name": "my-custom-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"max_input_tokens": 32000, "max_output_tokens": 8000}, + } + ] + ) + + assert router.get_configured_token_limits("my-custom-model") == (32000, 8000) + + +def test_get_configured_token_limits_returns_none_for_unset_or_unknown(): + router = litellm.Router( + model_list=[ + { + "model_name": "no-limits-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + } + ] + ) + + assert router.get_configured_token_limits("no-limits-model") == (None, None) + assert router.get_configured_token_limits("not-a-real-model") == (None, None) + + +def test_get_configured_token_limits_skips_wildcard_pattern_matching(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + "model_info": {"max_input_tokens": 12345}, + } + ] + ) + + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert router.get_configured_token_limits( + "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" + ) == (None, None) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 073ff17991e..a1a9448cc58 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4317,7 +4317,7 @@ _FIREWORKS_MODELS = [ "accounts/fireworks/models/glm-5p2", 1.4e-06, 4.4e-06, - 2.6e-07, + 1.4e-07, 1048576, 131072, False, @@ -4717,6 +4717,55 @@ class TestValidateEnvironmentTencent: 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/e2e_tests/tests/modelsPage/credentials.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts new file mode 100644 index 00000000000..8b7824813a4 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/credentials.spec.ts @@ -0,0 +1,75 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Role, users } from "../../fixtures/users"; + +test.describe("Edit LLM credential", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + const masterKey = users[Role.ProxyAdmin].password; + const SEED_API_KEY = "sk-e2e-credential-ABCDEFGHIJKLMNOP"; + const SEED_API_BASE = "https://api.openai.com/v1"; + const NEW_API_BASE = "https://proxy.e2e.example.com/v1"; + + let credentialName: string; + + test.beforeEach(async ({ page }) => { + credentialName = `e2e-cred-${Date.now()}`; + const res = await page.request.post("/credentials", { + headers: { Authorization: `Bearer ${masterKey}` }, + data: { + credential_name: credentialName, + credential_values: { api_key: SEED_API_KEY, api_base: SEED_API_BASE }, + credential_info: { custom_llm_provider: "openai" }, + }, + }); + expect(res.ok(), `POST /credentials for ${credentialName}`).toBe(true); + }); + + test.afterEach(async ({ page }) => { + await page.request.delete(`/credentials/${credentialName}`, { + headers: { Authorization: `Bearer ${masterKey}` }, + }); + }); + + test("changing only the api base does not overwrite the stored api key with its masked value", async ({ page }) => { + await page.goto("/ui"); + await page.getByText("Models + Endpoints").click(); + await page.getByRole("tab", { name: "LLM Credentials" }).click(); + + const row = page.locator("tr", { hasText: credentialName }); + await expect(row).toBeVisible({ timeout: 15_000 }); + await row.getByRole("button").first().click(); + + const modal = page.locator(".ant-modal-content").filter({ hasText: "Edit Credential" }); + await expect(modal).toBeVisible({ timeout: 10_000 }); + + const apiKeyField = modal.locator("#api_key"); + const apiBaseField = modal.locator("#api_base"); + await expect(apiKeyField).toBeVisible({ timeout: 15_000 }); + + await expect(apiKeyField, "form pre-fills the api key with the backend's masked value").toHaveValue(/\*{2,}/); + await expect(apiKeyField).not.toHaveValue(SEED_API_KEY); + + await apiBaseField.fill(NEW_API_BASE); + + const patchPromise = page.waitForRequest( + (req) => req.method() === "PATCH" && req.url().includes(`/credentials/${credentialName}`), + ); + await modal.getByRole("button", { name: "Update Credential" }).click(); + const patchReq = await patchPromise; + const patchBody = JSON.parse(patchReq.postData() ?? "{}"); + + expect(patchBody.credential_values.api_base, "UI sends the edited api base").toBe(NEW_API_BASE); + expect("api_key" in patchBody.credential_values, "UI must not send the masked api key back on update").toBe(false); + + await expect(page.getByText("Credential updated successfully")).toBeVisible({ timeout: 10_000 }); + + const infoRes = await page.request.get(`/credentials/by_name/${credentialName}`, { + headers: { Authorization: `Bearer ${masterKey}` }, + }); + expect(infoRes.ok()).toBe(true); + const cred = await infoRes.json(); + expect(cred.credential_values.api_base, "edited api base persisted to the backend").toBe(NEW_API_BASE); + expect(cred.credential_values.api_key, "a stored api key is still present (returned masked)").toMatch(/\*{2,}/); + }); +}); diff --git a/ui/litellm-dashboard/public/assets/logos/straiker.svg b/ui/litellm-dashboard/public/assets/logos/straiker.svg new file mode 100644 index 00000000000..bdfe0405736 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/straiker.svg @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index 2ad5819b5f0..a40587cb3ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -300,4 +300,10 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, + straiker: { + provider: "Straiker", + guardrailNameSuggestion: "Straiker Guardrail", + mode: "pre_call", + defaultOn: false, + }, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index f81277f13c3..ba11d3d400d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -442,6 +442,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Security", "Policy", "Prompt Injection"], providerKey: "Repelloai", }, + { + id: "straiker", + name: "Straiker", + description: + "Defend AI Agentic Guardrails: Indirect/Direct Prompt Injection, Tool Misuse, Malicious MCP and Skills", + category: "partner", + logo: `${ASSET_PREFIX}straiker.svg`, + tags: ["Agentic", "Prompt Injection", "Tool Misuse", "MCP", "Skills"], + providerKey: "Straiker", + }, ]; export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...PARTNER_GUARDRAIL_CARDS]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index e8c4810ca69..a2873797096 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -166,6 +166,7 @@ export const guardrailLogoMap: Record = { Akto: `${asset_logos_folder}akto.svg`, "Qostodian Nexus": `${asset_logos_folder}qohash.jpg`, "RepelloAI Argus": `${asset_logos_folder}repelloai.png`, + Straiker: `${asset_logos_folder}straiker.svg`, }; export const getGuardrailLogoAndName = (guardrailValue: string): { logo: string; displayName: string } => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.test.tsx index 2530ce9fda7..e5740a89822 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.test.tsx @@ -1,7 +1,8 @@ import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { tagListCall } from "@/components/networking"; +import { tagDeleteCall, tagListCall } from "@/components/networking"; import TagManagement from "./index"; @@ -12,10 +13,23 @@ vi.mock("@/components/networking", () => ({ modelInfoCall: vi.fn(), })); +vi.mock("@/components/molecules/notifications_manager", () => ({ + __esModule: true, + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + vi.mock("./TagTable", () => ({ __esModule: true, - default: ({ isLoading }: { isLoading?: boolean }) => ( -
{isLoading ? "table-loading" : "table-loaded"}
+ default: ({ isLoading, onDelete }: { isLoading?: boolean; onDelete: (tagName: string) => void }) => ( +
+ {isLoading ? "table-loading" : "table-loaded"} + +
), })); @@ -30,6 +44,7 @@ vi.mock("./components/CreateTagModal", () => ({ })); const mockTagListCall = vi.mocked(tagListCall); +const mockTagDeleteCall = vi.mocked(tagDeleteCall); describe("TagManagement loading state", () => { beforeEach(() => { @@ -57,3 +72,41 @@ describe("TagManagement loading state", () => { expect(mockTagListCall).toHaveBeenCalledWith("sk-test"); }); }); + +describe("TagManagement delete flow", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTagListCall.mockResolvedValue({}); + }); + + it("should confirm deletion through the shared DeleteResourceModal and call tagDeleteCall with the tag name", async () => { + const user = userEvent.setup(); + mockTagDeleteCall.mockResolvedValue({}); + render(); + await screen.findByText("table-loaded"); + + expect(screen.queryByText("Tag Information")).not.toBeInTheDocument(); + + await user.click(screen.getByTestId("mock-delete-trigger")); + + expect(await screen.findByText("Tag Information")).toBeInTheDocument(); + expect(screen.getByText("test-tag")).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: /delete/i })); + + expect(mockTagDeleteCall).toHaveBeenCalledWith("sk-test", "test-tag"); + }); + + it("should not call tagDeleteCall when the deletion is cancelled", async () => { + const user = userEvent.setup(); + render(); + await screen.findByText("table-loaded"); + + await user.click(screen.getByTestId("mock-delete-trigger")); + await screen.findByText("Tag Information"); + + await user.click(screen.getByRole("button", { name: "Cancel" })); + + expect(mockTagDeleteCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx index 16b92bd4553..301d9264437 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx @@ -7,6 +7,7 @@ import { tagCreateCall, tagListCall, tagDeleteCall } from "@/components/networki import { Tag } from "@/components/tag_management/types"; import TagTable from "./TagTable"; import NotificationsManager from "@/components/molecules/notifications_manager"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import CreateTagModal from "./components/CreateTagModal"; interface ModelInfo { @@ -33,6 +34,7 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) => const [editTag, setEditTag] = useState(false); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [tagToDelete, setTagToDelete] = useState(null); + const [isDeleting, setIsDeleting] = useState(false); const [lastRefreshed, setLastRefreshed] = useState(""); const [availableModels, setAvailableModels] = useState([]); @@ -87,6 +89,7 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) => const confirmDelete = async () => { if (!accessToken || !tagToDelete) return; + setIsDeleting(true); try { await tagDeleteCall(accessToken, tagToDelete); NotificationsManager.success("Tag deleted successfully"); @@ -94,9 +97,11 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) => } catch (error) { console.error("Error deleting tag:", error); NotificationsManager.fromBackend("Error deleting tag: " + error); + } finally { + setIsDeleting(false); + setIsDeleteModalOpen(false); + setTagToDelete(null); } - setIsDeleteModalOpen(false); - setTagToDelete(null); }; useEffect(() => { @@ -189,40 +194,19 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) => /> {/* Delete Confirmation Modal */} - {isDeleteModalOpen && ( -
-
- -
-
-
-
-

Delete Tag

-
-

Are you sure you want to delete this tag?

-
-
-
-
-
- - -
-
-
-
- )} + { + setIsDeleteModalOpen(false); + setTagToDelete(null); + }} + onOk={confirmDelete} + confirmLoading={isDeleting} + /> )} 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 e1f90296770..a2e2ca21d00 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 855a1b27df9..8008012a95c 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 00000000000..c232eb4c801 --- /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 5122d54db9b..6e7bc49afce 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 85a15ffad45..0c9c19d1286 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 3c3f21163b3..0b92dc1b02d 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/components/model_add/credentials.tsx b/ui/litellm-dashboard/src/components/model_add/credentials.tsx index 53df9015c97..82320b7ff8d 100644 --- a/ui/litellm-dashboard/src/components/model_add/credentials.tsx +++ b/ui/litellm-dashboard/src/components/model_add/credentials.tsx @@ -27,6 +27,7 @@ import EditCredentialsModal from "./EditCredentialModal"; import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { isProxyAdminRole } from "@/utils/roles"; +import { stripMaskedSecrets } from "@/utils/maskedSecretUtils"; interface CredentialsPanelProps { uploadProps: UploadProps; } @@ -52,9 +53,11 @@ const CredentialsPanel: React.FC = ({ uploadProps }) => { return; } - const filter_credential_values = Object.entries(values) - .filter(([key]) => !restrictedFields.includes(key)) - .reduce((acc, [key, value]) => ({ ...acc, [key]: value }), {}); + const filter_credential_values = stripMaskedSecrets( + Object.entries(values) + .filter(([key]) => !restrictedFields.includes(key)) + .reduce((acc, [key, value]) => ({ ...acc, [key]: value }), {}), + ); // Transform form values into credential structure const newCredential = { credential_name: values.credential_name, diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 5b7f82c44b4..8aaabdc50a2 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -22,6 +22,7 @@ import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import { CheckIcon, CopyIcon } from "lucide-react"; import { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../utils/dataUtils"; +import { isMaskedSecret, stripMaskedSecrets } from "../utils/maskedSecretUtils"; import { formItemValidateJSON, truncateString } from "../utils/textUtils"; import AutoRouterConnectionTest from "./add_model/auto_router_connection_test"; import { AutoRouterTestTarget, buildAutoRouterTestTargets } from "./add_model/build_auto_router_test_targets"; @@ -58,18 +59,6 @@ interface ModelInfoViewProps { modelAccessGroups: string[] | null; } -// The /model/info response redacts secrets by masking them (e.g. "sk-1****2345"), -// not by removing them. The edit form must never echo a masked value back on save: -// the backend would encrypt the asterisks and overwrite the real secret. A run of -// 2+ mask chars only appears in masker output (real config — incl. wildcard model -// names like "openai/*" — carries at most a single "*"), so this reliably detects a -// redacted value without a provider-metadata lookup. API-key rotation goes through -// UpdateModelCredentialsModal instead, which sends only the new key. -const isMaskedSecret = (value: unknown): boolean => typeof value === "string" && /\*{2,}/.test(value); - -const stripMaskedSecrets = (params: Record): Record => - Object.fromEntries(Object.entries(params).filter(([, value]) => !isMaskedSecret(value))); - const normalizeTierModels = (value: unknown): string[] => { if (Array.isArray(value)) return value; if (typeof value === "string" && value) return [value]; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9ad0b8d1101..0d8f55164f9 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 diff --git a/ui/litellm-dashboard/src/utils/maskedSecretUtils.ts b/ui/litellm-dashboard/src/utils/maskedSecretUtils.ts new file mode 100644 index 00000000000..101316bbd82 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/maskedSecretUtils.ts @@ -0,0 +1,10 @@ +// The proxy redacts secrets in API responses by masking them (e.g. "sk-1****2345"), +// not by removing them. Edit forms must never echo a masked value back on save: the +// backend would encrypt the asterisks and overwrite the real secret. A run of 2+ mask +// chars only appears in masker output (real config -- incl. wildcard model names like +// "openai/*" -- carries at most a single "*"), so this reliably detects a redacted +// value without a provider-metadata lookup. +export const isMaskedSecret = (value: unknown): boolean => typeof value === "string" && /\*{2,}/.test(value); + +export const stripMaskedSecrets = (params: Record): Record => + Object.fromEntries(Object.entries(params).filter(([, value]) => !isMaskedSecret(value))); diff --git a/uv.lock b/uv.lock index 8dfc4bd5fcc..89524353d0c 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-13T21:03:01.672393Z" +exclude-newer = "2026-07-15T02:04:45.513604Z" exclude-newer-span = "P3D" [manifest] @@ -1047,7 +1047,7 @@ name = "coloredlogs" version = "15.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "humanfriendly", marker = "python_full_version < '3.14'" }, + { name = "humanfriendly" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/c7/eed8f27100517e8c0e6b923d5f0845d0cb99763da6fdee00478f91db7325/coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0", size = 278520, upload-time = "2021-06-11T10:22:45.202Z" } wheels = [ @@ -2966,7 +2966,7 @@ name = "humanfriendly" version = "10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pyreadline3", marker = "python_full_version < '3.14' and sys_platform == 'win32'" }, + { name = "pyreadline3", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/3f/2c29224acb2e2df4d2046e4c73ee2662023c58ff5b113c4c1adac0886c43/humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc", size = 360702, upload-time = "2021-09-17T21:40:43.31Z" } wheels = [ @@ -3293,10 +3293,10 @@ name = "jsonschema-path" version = "0.3.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pathable", marker = "python_full_version < '3.14'" }, - { name = "pyyaml", marker = "python_full_version < '3.14'" }, - { name = "referencing", marker = "python_full_version < '3.14'" }, - { name = "requests", marker = "python_full_version < '3.14'" }, + { name = "pathable" }, + { name = "pyyaml" }, + { name = "referencing" }, + { name = "requests" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6e/45/41ebc679c2a4fced6a722f624c18d658dee42612b83ea24c1caf7c0eb3a8/jsonschema_path-0.3.4.tar.gz", hash = "sha256:8365356039f16cc65fddffafda5f58766e34bebab7d6d105616ab52bc4297001", size = 11159, upload-time = "2025-01-24T14:33:16.547Z" } wheels = [ @@ -3779,7 +3779,7 @@ extra-proxy = [ { name = "google-cloud-iam" }, { name = "google-cloud-kms" }, { name = "prisma" }, - { name = "redisvl", marker = "python_full_version < '3.14'" }, + { name = "redisvl" }, { name = "resend" }, ] google = [ @@ -3841,7 +3841,7 @@ proxy-runtime = [ { name = "opentelemetry-instrumentation-fastapi" }, { name = "opentelemetry-sdk" }, { name = "prometheus-client" }, - { name = "pypdf", marker = "python_full_version < '3.14'" }, + { name = "pypdf" }, { name = "sentry-sdk" }, ] semantic-router = [ @@ -3897,7 +3897,7 @@ dev = [ { name = "fastapi-offline" }, { name = "flake8" }, { name = "langfuse" }, - { name = "openapi-core", marker = "python_full_version < '3.14'" }, + { name = "openapi-core" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-instrumentation-fastapi" }, @@ -3990,7 +3990,7 @@ requires-dist = [ { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, { name = "mangum", marker = "extra == 'proxy-runtime'", specifier = ">=0.17.0,<1.0" }, - { name = "mcp", marker = "extra == 'proxy'", specifier = ">=1.26.0,<2.0" }, + { name = "mcp", marker = "extra == 'proxy'", specifier = ">=1.28.1,<2.0" }, { name = "mlflow", marker = "extra == 'mlflow'", specifier = ">=3.11.1,<4.0" }, { name = "numpy", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.26.0" }, { name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" }, @@ -4008,13 +4008,13 @@ requires-dist = [ { name = "pydantic-settings", marker = "extra == 'proxy'", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, - { name = "pypdf", marker = "python_full_version < '3.14' and extra == 'proxy-runtime'", specifier = ">=6.12.0,<7.0" }, + { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.12.0,<7.0" }, { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0" }, { name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" }, { name = "pyyaml", marker = "extra == 'cli'", specifier = ">=6.0.3,<7.0" }, { name = "pyyaml", marker = "extra == 'proxy'", specifier = ">=6.0.3,<7.0" }, - { name = "redisvl", marker = "python_full_version < '3.14' and extra == 'extra-proxy'", specifier = ">=0.4.1,<1.0" }, + { name = "redisvl", marker = "extra == 'extra-proxy'", specifier = ">=0.4.1,<1.0" }, { name = "requests", marker = "extra == 'cli'", specifier = ">=2.32.0,<3.0" }, { name = "resend", marker = "extra == 'extra-proxy'", specifier = ">=2.23.0,<3.0" }, { name = "restrictedpython", marker = "extra == 'proxy'", specifier = ">=8.1,<9.0" }, @@ -4072,7 +4072,7 @@ dev = [ { name = "fastapi-offline", specifier = "==1.7.6" }, { name = "flake8", specifier = "==7.3.0" }, { name = "langfuse", specifier = "==2.59.7" }, - { name = "openapi-core", marker = "python_full_version < '3.14'", specifier = "==0.22.0" }, + { name = "openapi-core", specifier = "==0.22.0" }, { name = "opentelemetry-api", specifier = "==1.28.0" }, { name = "opentelemetry-exporter-otlp", specifier = "==1.28.0" }, { name = "opentelemetry-instrumentation-fastapi", specifier = "==0.49b0" }, @@ -4431,7 +4431,7 @@ wheels = [ [[package]] name = "mcp" -version = "1.26.0" +version = "1.28.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -4449,9 +4449,9 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fc/6d/62e76bbb8144d6ed86e202b5edd8a4cb631e7c8130f3f4893c3f90262b10/mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66", size = 608005, upload-time = "2026-01-24T19:40:32.468Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6e/77/9450b8f251a13affb6281997d0523c4615f8a8b35d0b21ff30db3a5aac9d/mcp-1.28.1.tar.gz", hash = "sha256:d51e36a5f5644faea4f85ea649bfffa6bc6c26770d42798ad6a3de3d2ba69683", size = 638501, upload-time = "2026-06-26T12:57:29.093Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/d9/eaa1f80170d2b7c5ba23f3b59f766f3a0bb41155fbc32a69adfa1adaaef9/mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca", size = 233615, upload-time = "2026-01-24T19:40:30.652Z" }, + { url = "https://files.pythonhosted.org/packages/e2/5e/d118fce19f87a2e7d8101c35c8ae0ec289098a4df0ff244cec23e415aca0/mcp-1.28.1-py3-none-any.whl", hash = "sha256:2726bca5e7193f61c5dde8b12500a6de2d9acf6d1a1c0be9e8c2e706437991df", size = 222620, upload-time = "2026-06-26T12:57:27.218Z" }, ] [[package]] @@ -4481,7 +4481,7 @@ version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fd/15/76f86faa0902836cc133939732f7611ace68cf54148487a99c539c272dc8/ml_dtypes-0.4.1.tar.gz", hash = "sha256:fad5f2de464fd09127e49b7fd1252b9006fb43d2edc1ff112d390c324af5ca7a", size = 692594, upload-time = "2024-09-13T19:07:11.624Z" } wheels = [ @@ -4980,14 +4980,14 @@ name = "openapi-core" version = "0.22.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "isodate", marker = "python_full_version < '3.14'" }, - { name = "jsonschema", marker = "python_full_version < '3.14'" }, - { name = "jsonschema-path", marker = "python_full_version < '3.14'" }, - { name = "more-itertools", marker = "python_full_version < '3.14'" }, - { name = "openapi-schema-validator", marker = "python_full_version < '3.14'" }, - { name = "openapi-spec-validator", marker = "python_full_version < '3.14'" }, - { name = "typing-extensions", marker = "python_full_version < '3.14'" }, - { name = "werkzeug", marker = "python_full_version < '3.14'" }, + { name = "isodate" }, + { name = "jsonschema" }, + { name = "jsonschema-path" }, + { name = "more-itertools" }, + { name = "openapi-schema-validator" }, + { name = "openapi-spec-validator" }, + { name = "typing-extensions" }, + { name = "werkzeug" }, ] sdist = { url = "https://files.pythonhosted.org/packages/fd/65/ee75f25b9459a02df6f713f8ffde5dacb57b8b4e45145cde4cab28b5abba/openapi_core-0.22.0.tar.gz", hash = "sha256:b30490dfa74e3aac2276105525590135212352f5dd7e5acf8f62f6a89ed6f2d0", size = 109242, upload-time = "2025-12-22T19:19:49.608Z" } wheels = [ @@ -4999,9 +4999,9 @@ name = "openapi-schema-validator" version = "0.6.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jsonschema", marker = "python_full_version < '3.14'" }, - { name = "jsonschema-specifications", marker = "python_full_version < '3.14'" }, - { name = "rfc3339-validator", marker = "python_full_version < '3.14'" }, + { name = "jsonschema" }, + { name = "jsonschema-specifications" }, + { name = "rfc3339-validator" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8b/f3/5507ad3325169347cd8ced61c232ff3df70e2b250c49f0fe140edb4973c6/openapi_schema_validator-0.6.3.tar.gz", hash = "sha256:f37bace4fc2a5d96692f4f8b31dc0f8d7400fd04f3a937798eaf880d425de6ee", size = 11550, upload-time = "2025-01-10T18:08:22.268Z" } wheels = [ @@ -5013,10 +5013,10 @@ name = "openapi-spec-validator" version = "0.7.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "jsonschema", marker = "python_full_version < '3.14'" }, - { name = "jsonschema-path", marker = "python_full_version < '3.14'" }, - { name = "lazy-object-proxy", marker = "python_full_version < '3.14'" }, - { name = "openapi-schema-validator", marker = "python_full_version < '3.14'" }, + { name = "jsonschema" }, + { name = "jsonschema-path" }, + { name = "lazy-object-proxy" }, + { name = "openapi-schema-validator" }, ] sdist = { url = "https://files.pythonhosted.org/packages/82/af/fe2d7618d6eae6fb3a82766a44ed87cd8d6d82b4564ed1c7cfb0f6378e91/openapi_spec_validator-0.7.2.tar.gz", hash = "sha256:cc029309b5c5dbc7859df0372d55e9d1ff43e96d678b9ba087f7c56fc586f734", size = 36855, upload-time = "2025-06-07T14:48:56.299Z" } wheels = [ @@ -7163,16 +7163,16 @@ name = "redisvl" version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "coloredlogs", marker = "python_full_version < '3.14'" }, - { name = "ml-dtypes", marker = "python_full_version < '3.14'" }, + { name = "coloredlogs" }, + { name = "ml-dtypes" }, { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" }, - { name = "pydantic", marker = "python_full_version < '3.14'" }, - { name = "python-ulid", marker = "python_full_version < '3.14'" }, - { name = "pyyaml", marker = "python_full_version < '3.14'" }, - { name = "redis", marker = "python_full_version < '3.14'" }, - { name = "tabulate", marker = "python_full_version < '3.14'" }, - { name = "tenacity", marker = "python_full_version < '3.14'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "pydantic" }, + { name = "python-ulid" }, + { name = "pyyaml" }, + { name = "redis" }, + { name = "tabulate" }, + { name = "tenacity" }, ] sdist = { url = "https://files.pythonhosted.org/packages/21/33/ab14865a0b2a31b1d003c29e7e8ea3a7a2f2c8ecb24e58e58d606e1f031b/redisvl-0.4.1.tar.gz", hash = "sha256:fd6a36426ba94792c0efca20915c31232d4ee3cc58eb23794a62c142696401e6", size = 77688, upload-time = "2025-02-21T22:51:41.389Z" } wheels = [ @@ -7406,7 +7406,7 @@ name = "rfc3339-validator" version = "0.1.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "six", marker = "python_full_version < '3.14'" }, + { name = "six" }, ] sdist = { url = "https://files.pythonhosted.org/packages/28/ea/a9387748e2d111c3c2b275ba970b735e04e15cdb1eb30693b6b5708c4dbd/rfc3339_validator-0.1.4.tar.gz", hash = "sha256:138a2abdf93304ad60530167e51d2dfb9549521a836871b88d7f4695d0022f6b", size = 5513, upload-time = "2021-05-12T16:37:54.178Z" } wheels = [