From 40f02b2eb520f2b8178cd1a3c56ca4c3549639d7 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 00:21:22 -0700 Subject: [PATCH 01/60] refactor(mcp): consolidate exception-tree walkers into one shared faults traversal --- .../mcp_server/faults/__init__.py | 2 + .../mcp_server/faults/traversal.py | 35 ++++++++++++++ .../mcp_server/mcp_server_manager.py | 42 ++++------------ .../mcp_server/semantic_tool_filter.py | 22 ++++----- ruff-strict-budget.json | 4 +- .../mcp_server/faults/test_traversal.py | 48 +++++++++++++++++++ .../test_mcp_oauth_passthrough_tools.py | 30 ++++++++++++ .../mcp_server/test_semantic_tool_filter.py | 28 +++++++++++ type-discipline-budget.json | 2 +- 9 files changed, 164 insertions(+), 49 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/faults/traversal.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py diff --git a/litellm/proxy/_experimental/mcp_server/faults/__init__.py b/litellm/proxy/_experimental/mcp_server/faults/__init__.py index da078f0e242..1b9ee77d795 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/__init__.py +++ b/litellm/proxy/_experimental/mcp_server/faults/__init__.py @@ -15,6 +15,7 @@ from litellm.proxy._experimental.mcp_server.faults.render_oauth import ( dcr_fault_detail, render_token_fault, ) +from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree from litellm.proxy._experimental.mcp_server.faults.types import ( CallerRejected, CredentialSource, @@ -34,5 +35,6 @@ __all__ = [ "classify_upstream_dcr_rejection", "classify_upstream_token_rejection", "dcr_fault_detail", + "iter_exception_tree", "render_token_fault", ] diff --git a/litellm/proxy/_experimental/mcp_server/faults/traversal.py b/litellm/proxy/_experimental/mcp_server/faults/traversal.py new file mode 100644 index 00000000000..78e94e22e70 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/faults/traversal.py @@ -0,0 +1,35 @@ +"""Shared exception-tree traversal for fault classification. + +Failures cross the MCP SDK's anyio task groups wrapped in ``ExceptionGroup``s and chained through +``raise ... from`` causes, so every classifier that needs an exception buried in the tree (an +upstream ``httpx.Response``, a context-window overflow) has to walk the same shapes. One traversal +with one deliberate order keeps blame assignment consistent across classifiers: explicit links are +searched before incidental ones, so an exception raised while handling the real failure can never +shadow the failure itself. +""" + +from __future__ import annotations + +from collections.abc import Iterator + + +def iter_exception_tree(exc: BaseException) -> Iterator[BaseException]: + """Yield ``exc`` and every exception reachable from it, explicit links first: each node's + ``raise ... from`` cause subtree, then ``ExceptionGroup`` members in raise order, then the + incidental ``__context__`` chain last. Cycle-safe via identity tracking, and iterative so a + deep chain cannot overflow the interpreter stack.""" + seen: set[int] = set() + stack = [exc] + while stack: + current = stack.pop() + if id(current) in seen: + continue + seen.add(id(current)) + yield current + 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__) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1d681b43b9e..f2d3f568635 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -51,6 +51,7 @@ 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.faults import iter_exception_tree from litellm.proxy._experimental.mcp_server.elicitation_handler import ( MCP_ELICITATION_AVAILABLE, ) @@ -440,44 +441,17 @@ def _extract_upstream_auth_failure( upstream MCP server. The MCP SDK wraps transport errors in anyio ``ExceptionGroup`` objects and - may chain through ``__cause__`` / ``__context__``. We inspect all of those - layers for an ``httpx.Response``-bearing exception (typically - ``httpx.HTTPStatusError``) and extract the status code and any upstream - ``WWW-Authenticate`` header. + may chain through ``__cause__`` / ``__context__``; ``iter_exception_tree`` + visits all of those layers, explicit links first. The first exception + bearing a real ``httpx.Response`` with a 401/403 wins, and its status code + and upstream ``WWW-Authenticate`` header are extracted. 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)) - + for current in iter_exception_tree(exc): 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__) - + if isinstance(response, httpx.Response) and response.status_code in (401, 403): + return response.status_code, response.headers.get("www-authenticate") return None diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index e12c6cdbd56..b22dd64e7fc 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional from litellm._logging import verbose_logger from litellm.exceptions import ContextWindowExceededError from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers +from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR if TYPE_CHECKING: @@ -33,18 +34,15 @@ class SemanticToolFilterContextWindowError(Exception): ) -def _is_context_window_error(error: Optional[BaseException], max_depth: int = 5) -> bool: - """Detect a context-window overflow anywhere in an exception's cause chain.""" - current = error - for _ in range(max_depth): - if current is None: - return False - if isinstance(current, ContextWindowExceededError): - return True - if ExceptionCheckers.is_error_str_context_window_exceeded(str(current)): - return True - current = current.__cause__ or current.__context__ - return False +def _is_context_window_error(error: Optional[BaseException]) -> bool: + """Detect a context-window overflow anywhere in an exception's tree.""" + if error is None: + return False + return any( + isinstance(current, ContextWindowExceededError) + or ExceptionCheckers.is_error_str_context_window_exceeded(str(current)) + for current in iter_exception_tree(error) + ) class SemanticMCPToolFilter: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index dcde6fd1641..448a0079674 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -60,7 +60,7 @@ "limit": 4 }, "BLE001": { - "limit": 2903 + "limit": 2902 }, "C401": { "limit": 11 @@ -363,6 +363,6 @@ "limit": 105 }, "UP045": { - "limit": 18462 + "limit": 18461 } } diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py new file mode 100644 index 00000000000..a12c02339e6 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py @@ -0,0 +1,48 @@ +"""Traversal contract for the shared exception-tree walk: the root is yielded first, explicit +links win (the ``raise ... from`` cause subtree, then ExceptionGroup members in raise order, +then the incidental ``__context__`` chain last), and adversarial shapes terminate.""" + +from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree + + +def test_yields_the_root_itself_first(): + exc = ValueError("root") + assert list(iter_exception_tree(exc)) == [exc] + + +def test_cause_subtree_is_exhausted_before_context(): + deep = KeyError("deep") + cause = RuntimeError("cause") + cause.__cause__ = deep + context = OSError("context") + root = ValueError("root") + root.__cause__ = cause + root.__context__ = context + assert list(iter_exception_tree(root)) == [root, cause, deep, context] + + +def test_group_members_yield_in_raise_order_between_cause_and_context(): + first = KeyError("first") + second = IndexError("second") + group = BaseExceptionGroup("group", [first, second]) + cause = RuntimeError("cause") + context = OSError("context") + group.__cause__ = cause + group.__context__ = context + assert list(iter_exception_tree(group)) == [group, cause, first, second, context] + + +def test_terminates_on_a_cause_cycle(): + a = ValueError("a") + b = RuntimeError("b") + a.__cause__ = b + b.__cause__ = a + assert list(iter_exception_tree(a)) == [a, b] + + +def test_node_reachable_as_both_cause_and_context_yields_once(): + inner = KeyError("inner") + root = ValueError("root") + root.__cause__ = inner + root.__context__ = inner + assert list(iter_exception_tree(root)) == [root, inner] 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..617b5aa17fc 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 @@ -50,6 +50,36 @@ def test_extract_upstream_auth_failure_returns_none_for_non_auth(): assert _extract_upstream_auth_failure(RuntimeError("boom")) is None +def _auth_status_error(status_code: int, www_authenticate: str) -> httpx.HTTPStatusError: + response = httpx.Response( + status_code=status_code, + headers={"www-authenticate": www_authenticate}, + request=httpx.Request("GET", "https://upstream/mcp"), + ) + return httpx.HTTPStatusError(str(status_code), request=response.request, response=response) + + +def test_extract_upstream_auth_failure_finds_401_behind_cause_chain(): + wrapper = RuntimeError("wrapped") + wrapper.__cause__ = _auth_status_error(401, "Bearer") + assert _extract_upstream_auth_failure(wrapper) == (401, "Bearer") + + +def test_extract_upstream_auth_failure_finds_401_behind_context_chain(): + wrapper = RuntimeError("wrapped") + wrapper.__context__ = _auth_status_error(401, "Bearer") + assert _extract_upstream_auth_failure(wrapper) == (401, "Bearer") + + +def test_extract_upstream_auth_failure_prefers_causal_chain_over_context(): + """A 403 raised incidentally while handling the real 401 (surviving only as ``__context__``) + must not shadow the 401 on the explicit ``raise ... from`` chain.""" + wrapper = RuntimeError("wrapped") + wrapper.__cause__ = _auth_status_error(401, "Bearer realm=real") + wrapper.__context__ = _auth_status_error(403, "Bearer realm=incidental") + assert _extract_upstream_auth_failure(wrapper) == (401, "Bearer realm=real") + + @pytest.mark.asyncio async def test_fetch_tools_from_passthrough_raises_on_upstream_401(): manager = MCPServerManager() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 9bc0a525326..28ee7435702 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -1664,3 +1664,31 @@ def test_is_context_window_error_detection_variants(): assert _is_context_window_error(ValueError("Invalid 'input[0]': maximum input length is 8192 tokens.")) assert not _is_context_window_error(ValueError("A generic API error occurred.")) assert not _is_context_window_error(None) + + +def test_is_context_window_error_sees_through_trees_the_chain_walk_missed(): + """Overflow shapes the old single-path depth-5 chain walk could not reach: hidden in + ``__context__`` behind a non-matching ``__cause__``, buried inside an anyio-style + ``ExceptionGroup``, and chained deeper than five links.""" + import litellm + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + _is_context_window_error, + ) + + def _cwe() -> litellm.ContextWindowExceededError: + return litellm.ContextWindowExceededError(message="overflow", model="m", llm_provider="openai") + + shadowed = ValueError("wrapper") + shadowed.__cause__ = TypeError("unrelated failure") + shadowed.__context__ = _cwe() + assert _is_context_window_error(shadowed) + + grouped = BaseExceptionGroup("task group", [RuntimeError("sibling"), _cwe()]) + assert _is_context_window_error(grouped) + + deep: BaseException = _cwe() + for depth in range(6): + wrapper = ValueError(f"layer {depth}") + wrapper.__cause__ = deep + deep = wrapper + assert _is_context_window_error(deep) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 87b4c96e323..83291fac388 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23409 + "limit": 23408 }, "LIT002": { "limit": 27511 From ae952ce971ff52c91a17abb5e89bd1062383820a Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 18:35:58 -0700 Subject: [PATCH 02/60] feat(mcp): support MCP servers on the Anthropic /v1/messages API MCP tool calling worked on /v1/chat/completions and /v1/responses but not on /v1/messages. Those are the only two surfaces with an MCP gateway entry point, so a litellm_proxy MCP reference reached Anthropic verbatim inside tools and the API rejected the request with "Input tag 'mcp' found using 'type' does not match any of the expected tags". The playground never surfaced this because it dropped the reference before sending, and disabled the MCP selector for the endpoint. Add the third entry point in anthropic_messages_handler, ahead of the provider branch so it covers the native path and both bridges from one place. The gateway expands the reference against the caller's own credentials and access control, which is the whole point of routing it through litellm rather than handing the url to the provider. /v1/messages needs Anthropic's own tool shape, so transform_mcp_tool_to_anthropic_tool joins the OpenAI chat and Responses transforms alongside it. The tool loop speaks tool_use and tool_result rather than OpenAI tool_calls, and reuses the existing FakeAnthropicMessagesStreamIterator to re-stream the result, the same pattern the websearch interception already uses on this route. Argument extraction moves into the shared extractor: an Anthropic tool_use block carries its arguments under `input`, and reading only `arguments` failed silently, executing the tool with every argument dropped. On the frontend the request builder declared selectedMCPTools and never read it, so no tools key was ever sent. Wire it through a shared block builder and add the endpoint to MCP_SUPPORTED_ENDPOINTS, which is what greys the selector out. Resolves LIT-4517 Resolves LIT-4518 --- litellm/experimental_mcp_client/tools.py | 13 ++ .../messages/handler.py | 35 ++++ .../messages/mcp_handler.py | 174 ++++++++++++++++++ .../mcp/litellm_proxy_mcp_handler.py | 5 + .../experimental_mcp_client/test_tools.py | 49 +++++ .../messages/test_mcp_handler.py | 122 ++++++++++++ .../mcp/test_litellm_proxy_mcp_handler.py | 42 +++++ .../playground/components/chat_ui/ChatUI.tsx | 12 +- .../llm_calls/anthropic_messages.tsx | 14 +- .../components/llm_calls/mcp_tool_blocks.ts | 79 ++++++++ 10 files changed, 542 insertions(+), 3 deletions(-) create mode 100644 litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py create mode 100644 ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.ts diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index c65b266bd02..1bd65847616 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -9,6 +9,7 @@ from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition +from litellm.types.llms.anthropic import AnthropicInputSchema, AnthropicMessagesTool from litellm.types.utils import ChatCompletionMessageToolCall @@ -75,6 +76,18 @@ def transform_mcp_tool_to_openai_responses_api_tool( ) +def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessagesTool: + """Convert an MCP tool to an Anthropic Messages API tool.""" + normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema) + + return AnthropicMessagesTool( + name=mcp_tool.name, + description=mcp_tool.description or "", + input_schema=AnthropicInputSchema(**normalized_parameters), + type="custom", + ) + + async def load_mcp_tools( session: ClientSession, format: Literal["mcp", "openai"] = "mcp" ) -> Union[List[MCPTool], List[ChatCompletionToolParam]]: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index dd983f0c344..499f5bc486c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -477,6 +477,41 @@ def anthropic_messages_handler( mock_response=litellm_params.mock_response, ) + # Expand litellm_proxy MCP references through the MCP gateway before dispatch, so every + # downstream path (native passthrough and both bridges) gets real tools rather than a + # reference the provider cannot resolve. Popped from kwargs so it never reaches the provider. + skip_mcp_handler = kwargs.pop("_skip_mcp_handler", False) + if not skip_mcp_handler and tools: + from litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler import ( + anthropic_messages_with_mcp, + ) + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + + if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): + return anthropic_messages_with_mcp( + max_tokens=max_tokens, + messages=messages, + model=model, + metadata=metadata, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + container=container, + api_key=api_key, + api_base=api_base, + client=client, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + anthropic_messages_provider_config: Optional[BaseAnthropicMessagesConfig] = None if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py new file mode 100644 index 00000000000..392b9e2e02d --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py @@ -0,0 +1,174 @@ +""" +MCP gateway support for the Anthropic `/v1/messages` API. + +Mirrors ``litellm.responses.mcp.chat_completions_handler`` but speaks the +Anthropic Messages shapes: tools carry an ``input_schema``, the model asks for a +tool through a ``tool_use`` content block, and results are fed back as +``tool_result`` blocks in a user message. +""" + +from typing import Any, AsyncIterator, Mapping, Sequence, Union + +from litellm._logging import verbose_logger +from litellm.types.llms.anthropic import ( + AnthropicMessagesTool, + AnthropicMessagesToolResultParam, + AnthropicMessagesUserMessageParam, +) +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, +) + +MAX_MCP_TOOL_USE_ITERATIONS = 10 + + +def _get_response_content(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, Any]]: + content = response.get("content") + if not isinstance(content, list): + return () + return tuple(block for block in content if isinstance(block, dict)) + + +def _extract_tool_use_blocks(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, Any]]: + """Return the ``tool_use`` content blocks the model emitted.""" + return tuple(block for block in _get_response_content(response) if block.get("type") == "tool_use") + + +def _get_stop_reason(response: AnthropicMessagesResponse) -> Union[str, None]: + stop_reason = response.get("stop_reason") + return stop_reason if isinstance(stop_reason, str) else None + + +def _build_tool_result_message(tool_results: Sequence[Mapping[str, Any]]) -> AnthropicMessagesUserMessageParam: + """Turn executed tool results into the user message Anthropic expects.""" + return AnthropicMessagesUserMessageParam( + role="user", + content=tuple( + AnthropicMessagesToolResultParam( + type="tool_result", + tool_use_id=str(result.get("tool_call_id") or ""), + content=str(result.get("result") or ""), + ) + for result in tool_results + ), + ) + + +def _resolve_user_api_key_auth( + kwargs: Mapping[str, Any], +) -> Any: # any-ok: UserAPIKeyAuth is proxy-only, importing it here would create a cycle + """`/v1/messages` is a LITELLM_METADATA_ROUTE, so the auth object rides in litellm_metadata.""" + litellm_metadata = kwargs.get("litellm_metadata") or {} + metadata = kwargs.get("metadata") or {} + return ( + kwargs.get("user_api_key_auth") + or litellm_metadata.get("user_api_key_auth") + or metadata.get("user_api_key_auth") + ) + + +async def anthropic_messages_with_mcp( + max_tokens: int, + messages: Sequence[Mapping[str, Any]], + model: str, + tools: Union[Sequence[Mapping[str, Any]], None] = None, + **kwargs: Any, # kwargs-ok: forwarded verbatim to litellm.anthropic_messages, which owns the param contract +) -> Union[AnthropicMessagesResponse, AsyncIterator[Any]]: + """ + Expand litellm_proxy MCP references for `/v1/messages` and run the tool loop. + + The MCP gateway owns the expansion so the reference resolves against the + caller's own credentials and access control, rather than being handed to the + upstream provider as a url it cannot reach. + """ + import litellm + from litellm.experimental_mcp_client.tools import ( + transform_mcp_tool_to_anthropic_tool, + ) + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + + mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + + if not mcp_references: + return await litellm.anthropic_messages( + max_tokens=max_tokens, + messages=list(messages), + model=model, + tools=list(tools) if tools else None, + _skip_mcp_handler=True, + **kwargs, + ) + + user_api_key_auth = _resolve_user_api_key_auth(kwargs) + + ( + deduplicated_mcp_tools, + tool_server_map, + ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( + user_api_key_auth, + mcp_references, + litellm_trace_id=kwargs.get("litellm_trace_id"), + ) + + anthropic_tools: Sequence[AnthropicMessagesTool] = tuple( + transform_mcp_tool_to_anthropic_tool(mcp_tool) for mcp_tool in deduplicated_mcp_tools + ) + all_tools = [*anthropic_tools, *(other_tools or ())] + + should_auto_execute = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( + mcp_tools_with_litellm_proxy=mcp_references + ) + stream = bool(kwargs.pop("stream", False)) + + base_call_args: Mapping[str, Any] = { + "max_tokens": max_tokens, + "model": model, + "tools": all_tools or None, + "_skip_mcp_handler": True, + **kwargs, + } + + if not should_auto_execute: + return await litellm.anthropic_messages(messages=list(messages), stream=stream, **base_call_args) + + working_messages: Sequence[Mapping[str, Any]] = tuple(messages) + response: AnthropicMessagesResponse = await litellm.anthropic_messages( + messages=list(working_messages), stream=False, **base_call_args + ) + + for _ in range(MAX_MCP_TOOL_USE_ITERATIONS): + if _get_stop_reason(response) != "tool_use": + break + + tool_use_blocks = _extract_tool_use_blocks(response) + if not tool_use_blocks: + break + + tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=tool_server_map, + tool_calls=list(tool_use_blocks), + user_api_key_auth=user_api_key_auth, + litellm_trace_id=kwargs.get("litellm_trace_id"), + ) + + working_messages = ( + *working_messages, + {"role": "assistant", "content": list(_get_response_content(response))}, + _build_tool_result_message(tool_results), + ) + response = await litellm.anthropic_messages(messages=list(working_messages), stream=False, **base_call_args) + else: + verbose_logger.warning( + f"MCP tool loop hit its {MAX_MCP_TOOL_USE_ITERATIONS} iteration cap for model {model}; " + "returning the last response" + ) + + if stream: + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + + return FakeAnthropicMessagesStreamIterator(response) + return response diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index e03f0296109..d2c9f220690 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -541,7 +541,10 @@ class LiteLLM_Proxy_MCP_Handler: tool_arguments = function_block.get("arguments") else: tool_name = tool_call.get("name") + # Anthropic tool_use blocks carry the arguments under `input` tool_arguments = tool_call.get("arguments") + if tool_arguments is None: + tool_arguments = tool_call.get("input") else: tool_call_id = getattr(tool_call, "call_id", None) or getattr(tool_call, "id", None) @@ -552,6 +555,8 @@ class LiteLLM_Proxy_MCP_Handler: else: tool_name = getattr(tool_call, "name", None) tool_arguments = getattr(tool_call, "arguments", None) + if tool_arguments is None: + tool_arguments = getattr(tool_call, "input", None) return tool_name, tool_arguments, tool_call_id diff --git a/tests/test_litellm/experimental_mcp_client/test_tools.py b/tests/test_litellm/experimental_mcp_client/test_tools.py index 786bbf7dcc9..625ab56951f 100644 --- a/tests/test_litellm/experimental_mcp_client/test_tools.py +++ b/tests/test_litellm/experimental_mcp_client/test_tools.py @@ -18,6 +18,7 @@ from mcp.types import ( from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.tools import ( + transform_mcp_tool_to_anthropic_tool, _get_function_arguments, _normalize_mcp_input_schema, call_mcp_tool, @@ -250,3 +251,51 @@ def test_transform_mcp_tool_to_openai_responses_api_tool(): assert "query" in openai_tool["parameters"]["properties"] assert openai_tool["parameters"]["required"] == ["query"] assert openai_tool["parameters"]["additionalProperties"] == False + + +def test_transform_mcp_tool_to_anthropic_tool(): + """ + Regression test (LIT-4517): MCP tools must reach /v1/messages in Anthropic's + own tool shape. + + Given: An MCP tool + When: It is transformed for the Anthropic Messages API + Then: It carries name/description/input_schema, the shape that endpoint + accepts, rather than an OpenAI function block + + /v1/messages rejects an OpenAI-shaped tool outright ("Input tag 'function' + does not match any of the expected tags"), so reusing either OpenAI + transform here loses every MCP tool. + """ + tool = MCPTool( + name="read_wiki_structure", + description="Get a list of documentation topics", + inputSchema={ + "type": "object", + "properties": {"repoName": {"type": "string"}}, + "required": ["repoName"], + }, + ) + + anthropic_tool = transform_mcp_tool_to_anthropic_tool(tool) + + assert anthropic_tool["name"] == "read_wiki_structure" + assert anthropic_tool["description"] == "Get a list of documentation topics" + assert anthropic_tool["type"] == "custom" + assert anthropic_tool["input_schema"]["type"] == "object" + assert "repoName" in anthropic_tool["input_schema"]["properties"] + assert anthropic_tool["input_schema"]["required"] == ["repoName"] + assert "function" not in anthropic_tool, "Anthropic tools must not carry an OpenAI function block" + assert "parameters" not in anthropic_tool, "Anthropic names the schema input_schema, not parameters" + + +def test_transform_mcp_tool_to_anthropic_tool_normalizes_empty_schema(): + """A tool with no declared arguments must still present a valid object schema.""" + anthropic_tool = transform_mcp_tool_to_anthropic_tool( + MCPTool(name="noargs", description=None, inputSchema={}) + ) + + assert anthropic_tool["name"] == "noargs" + assert anthropic_tool["description"] == "" + assert anthropic_tool["input_schema"]["type"] == "object" + assert anthropic_tool["input_schema"]["properties"] == {} diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py new file mode 100644 index 00000000000..3faa6b1e4e2 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py @@ -0,0 +1,122 @@ +import os +import sys +from unittest.mock import AsyncMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( + anthropic_messages_handler, +) +from litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler import ( + _build_tool_result_message, + _extract_tool_use_blocks, +) + +MCP_REFERENCE = { + "type": "mcp", + "server_label": "litellm", + "server_url": "litellm_proxy/mcp/deepwiki", + "require_approval": "never", +} + + +def test_anthropic_messages_handler_routes_litellm_proxy_mcp_to_the_gateway(): + """ + Regression test (LIT-4517): /v1/messages must expand a litellm_proxy MCP + reference through the MCP gateway. + + Given: A /v1/messages request whose tools carry a litellm_proxy MCP reference + When: The handler dispatches + Then: It hands off to the MCP gateway instead of the provider + + Without this hook the reference is forwarded to Anthropic verbatim and the API + rejects the request ("Input tag 'mcp' found using 'type' does not match any of + the expected tags"), because only /v1/chat/completions and /v1/responses ever + had a gateway entry point. This pins the wiring, not the helper: deleting the + dispatch makes the whole feature unreachable while every unit test still passes. + """ + with patch( + "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", + new=AsyncMock(return_value={"routed": True}), + ) as routed: + result = anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "hi"}], + model="claude-sonnet-4-5", + tools=[MCP_REFERENCE], + custom_llm_provider="anthropic", + ) + + assert routed.called, "A litellm_proxy MCP reference must be dispatched to the MCP gateway" + assert routed.call_args.kwargs["tools"] == [MCP_REFERENCE] + assert routed.call_args.kwargs["model"] == "claude-sonnet-4-5" + assert result is not None + + +def test_anthropic_messages_handler_skips_the_gateway_on_recursion(): + """The gateway's own follow-up call must not re-enter the gateway.""" + with patch( + "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", + new=AsyncMock(return_value={"routed": True}), + ) as routed: + with pytest.raises(Exception): + anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "hi"}], + model="claude-sonnet-4-5", + tools=[MCP_REFERENCE], + custom_llm_provider="anthropic", + _skip_mcp_handler=True, + ) + + assert not routed.called, "_skip_mcp_handler must stop the gateway from recursing" + + +def test_anthropic_messages_handler_leaves_native_tools_alone(): + """A plain Anthropic tool is not an MCP reference and must not reach the gateway.""" + with patch( + "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", + new=AsyncMock(return_value={"routed": True}), + ) as routed: + with pytest.raises(Exception): + anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "hi"}], + model="claude-sonnet-4-5", + tools=[{"name": "get_weather", "input_schema": {"type": "object"}}], + custom_llm_provider="anthropic", + ) + + assert not routed.called, "Only litellm_proxy MCP references belong to the gateway" + + +def test_extract_tool_use_blocks_ignores_text_blocks(): + """Only tool_use blocks drive the loop; text blocks are the model's prose.""" + response = { + "content": [ + {"type": "text", "text": "let me look that up"}, + {"type": "tool_use", "id": "toolu_1", "name": "read_wiki_structure", "input": {"repoName": "a/b"}}, + ] + } + + blocks = _extract_tool_use_blocks(response) + + assert len(blocks) == 1 + assert blocks[0]["name"] == "read_wiki_structure" + + +def test_build_tool_result_message_uses_anthropic_tool_result_blocks(): + """ + Results must go back as tool_result blocks in a user message. + + Anthropic pairs each result to its request by tool_use_id; the OpenAI shape + (a role="tool" message keyed by tool_call_id) is rejected here. + """ + message = _build_tool_result_message([{"tool_call_id": "toolu_1", "result": "9 sections", "name": "read_wiki"}]) + + assert message["role"] == "user" + assert list(message["content"]) == [ + {"type": "tool_result", "tool_use_id": "toolu_1", "content": "9 sections"} + ] 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..1c23b1a8b98 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 @@ -605,3 +605,45 @@ def test_completion_with_function_tools_works_without_fastapi_installed(): timeout=120, ) assert result.returncode == 0, result.stderr + + +def test_extract_tool_call_details_reads_anthropic_tool_use_input(): + """ + Regression test (LIT-4517): an Anthropic tool_use block carries its arguments + under `input`, not `arguments`. + + Given: A tool_use content block as /v1/messages returns it + When: The shared extractor reads it + Then: The arguments come back, so the MCP tool is called with them + + Reading only `arguments` fails silently rather than loudly: _parse_tool_arguments + turns the resulting None into {}, so the tool still executes, just with every + argument dropped. + """ + tool_use_block = { + "type": "tool_use", + "id": "toolu_01ABC", + "name": "read_wiki_structure", + "input": {"repoName": "BerriAI/litellm"}, + } + + name, arguments, call_id = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_use_block) + + assert name == "read_wiki_structure" + assert call_id == "toolu_01ABC" + assert arguments == {"repoName": "BerriAI/litellm"} + assert LiteLLM_Proxy_MCP_Handler._parse_tool_arguments(arguments) == {"repoName": "BerriAI/litellm"} + + +def test_extract_tool_call_details_still_prefers_openai_arguments(): + """The OpenAI chat shape must keep winning; `input` is only the fallback.""" + openai_tool_call = { + "id": "call_123", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + + name, arguments, call_id = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(openai_tool_call) + + assert name == "get_weather" + assert call_id == "call_123" + assert arguments == '{"city": "Paris"}' diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index d8f927b9a63..d2cf27e0c8b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -98,7 +98,12 @@ interface ChatUIProps { fixedModel?: string; } -const MCP_SUPPORTED_ENDPOINTS = new Set([EndpointType.CHAT, EndpointType.RESPONSES, EndpointType.MCP]); +const MCP_SUPPORTED_ENDPOINTS = new Set([ + EndpointType.CHAT, + EndpointType.RESPONSES, + EndpointType.MCP, + EndpointType.ANTHROPIC_MESSAGES, +]); const CUSTOM_MODEL_DEBOUNCE_WAIT_MS = 500; @@ -870,8 +875,11 @@ const ChatUI: React.FC = ({ selectedVectorStores.length > 0 ? selectedVectorStores : undefined, selectedGuardrails.length > 0 ? selectedGuardrails : undefined, selectedPolicies.length > 0 ? selectedPolicies : undefined, - selectedMCPServers, // Pass the selected tools array + selectedMCPServers, customProxyBaseUrl || undefined, + mcpServers, + mcpServerToolRestrictions, + mcpToolsets, ); } else if (endpointType === EndpointType.EMBEDDINGS) { await makeOpenAIEmbeddingsRequest( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx index ed2b4280b79..4319315396a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx @@ -1,6 +1,8 @@ import Anthropic from "@anthropic-ai/sdk"; import { MessageType } from "@/components/chat_ui/types"; import { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; +import { buildMcpToolBlocks } from "@/components/llm_calls/mcp_tool_blocks"; +import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; @@ -18,8 +20,11 @@ export async function makeAnthropicMessagesRequest( vector_store_ids?: string[], guardrails?: string[], policies?: string[], - selectedMCPTools?: string[], + selectedMCPServers?: string[], customBaseUrl?: string, + mcpServers?: MCPServer[], + mcpServerToolRestrictions?: Record, + mcpToolsets?: MCPToolset[], ) { if (!accessToken) { throw new Error("Virtual Key is required"); @@ -58,6 +63,13 @@ export async function makeAnthropicMessagesRequest( litellm_trace_id: traceId, }; + const tools = buildMcpToolBlocks({ + selectedMCPServers, + mcpServers, + mcpToolsets, + mcpServerToolRestrictions, + }); + if (tools.length > 0) requestBody.tools = tools; if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids; if (guardrails) requestBody.guardrails = guardrails; if (policies) requestBody.policies = policies; diff --git a/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.ts b/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.ts new file mode 100644 index 00000000000..42fa94d8208 --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.ts @@ -0,0 +1,79 @@ +import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; + +export const ALL_MCP_SERVERS_SENTINEL = "__all__"; +const TOOLSET_PREFIX = "toolset:"; + +export interface McpToolBlock { + type: "mcp"; + server_label: string; + server_url: string; + require_approval: "never"; + allowed_tools?: string[]; +} + +export interface BuildMcpToolBlocksArgs { + selectedMCPServers?: string[]; + mcpServers?: MCPServer[]; + mcpToolsets?: MCPToolset[]; + mcpServerToolRestrictions?: Record; +} + +/** + * Build the litellm_proxy MCP reference blocks for a playground request. + * + * Every endpoint that supports MCP sends the same reference shape; the gateway + * expands it server side and each endpoint's own transformation decides the + * final tool shape. Keeping one builder here stops the endpoints from drifting + * apart on routing name, label uniqueness, or escaping. + * + * server_name is used for both routing and labelling because it is the unique + * registered identifier; aliases can collide across servers, and a duplicated + * server_label causes silent tool-routing failures. + */ +export function buildMcpToolBlocks({ + selectedMCPServers, + mcpServers, + mcpToolsets, + mcpServerToolRestrictions, +}: BuildMcpToolBlocksArgs): McpToolBlock[] { + if (!selectedMCPServers || selectedMCPServers.length === 0) { + return []; + } + + if (selectedMCPServers.includes(ALL_MCP_SERVERS_SENTINEL)) { + return [ + { + type: "mcp", + server_label: "litellm", + server_url: "litellm_proxy/mcp", + require_approval: "never", + }, + ]; + } + + return selectedMCPServers.map((serverId) => { + if (serverId.startsWith(TOOLSET_PREFIX)) { + const toolsetId = serverId.slice(TOOLSET_PREFIX.length); + const toolset = mcpToolsets?.find((t) => t.toolset_id === toolsetId); + const toolsetName = toolset?.toolset_name || toolsetId; + return { + type: "mcp", + server_label: toolsetName, + server_url: `litellm_proxy/mcp/${encodeURIComponent(toolsetName)}`, + require_approval: "never", + }; + } + + const server = mcpServers?.find((s) => s.server_id === serverId); + const routeName = server?.server_name || serverId; + const allowedTools = mcpServerToolRestrictions?.[serverId] || []; + + return { + type: "mcp", + server_label: routeName, + server_url: `litellm_proxy/mcp/${encodeURIComponent(routeName)}`, + require_approval: "never", + ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), + }; + }); +} From cd3ac05a1fb6eedbbc078b38024f66693d3ef779 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 19:00:41 -0700 Subject: [PATCH 03/60] fix(mcp): forward the caller's MCP credentials from every gateway surface The /v1/messages handler resolved only the auth object and the trace id, so tool listing and tool execution ran without the caller's MCP auth headers. That fails quietly rather than loudly: the tool still executes, just with no credentials, so every server behind interactive OAuth, a bearer token or per-user env vars returns nothing while the model reports it has no access. Only a no-auth server looks healthy, which is exactly what the first proof used. Threading the missing arguments would have left the real problem in place. Each gateway surface rebuilds the same context by hand (responses/main.py twice, chat_completions_handler, mcp_streaming_iterator), which is why a new surface drops fields; this adds a fifth that dropped six of eight. Resolve it once into a frozen MCPRequestContext and have the handlers take that, so a field cannot be forgotten at a call site. chat_completions_handler now uses it too, and the resolver reads user_api_key_auth from both metadata keys because LITELLM_METADATA_ROUTES carry it in litellm_metadata while chat uses metadata. Also stop the loop when every tool call was skipped. tool_results is empty then, and the tool_result message built from it has empty content, which Anthropic rejects; the caller saw a 400 from mid-loop instead of the model's own answer. Tests pin both: dropping the headers from either listing or execution fails, and so does removing the empty-results guard. --- .../messages/mcp_handler.py | 38 +++--- .../responses/mcp/chat_completions_handler.py | 23 ++-- litellm/responses/mcp/request_context.py | 73 +++++++++++ .../messages/test_mcp_handler.py | 121 ++++++++++++++++++ 4 files changed, 222 insertions(+), 33 deletions(-) create mode 100644 litellm/responses/mcp/request_context.py diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py index 392b9e2e02d..813d4a62089 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py @@ -10,6 +10,7 @@ tool through a ``tool_use`` content block, and results are fed back as from typing import Any, AsyncIterator, Mapping, Sequence, Union from litellm._logging import verbose_logger +from litellm.responses.mcp.request_context import MCPRequestContext from litellm.types.llms.anthropic import ( AnthropicMessagesTool, AnthropicMessagesToolResultParam, @@ -54,19 +55,6 @@ def _build_tool_result_message(tool_results: Sequence[Mapping[str, Any]]) -> Ant ) -def _resolve_user_api_key_auth( - kwargs: Mapping[str, Any], -) -> Any: # any-ok: UserAPIKeyAuth is proxy-only, importing it here would create a cycle - """`/v1/messages` is a LITELLM_METADATA_ROUTE, so the auth object rides in litellm_metadata.""" - litellm_metadata = kwargs.get("litellm_metadata") or {} - metadata = kwargs.get("metadata") or {} - return ( - kwargs.get("user_api_key_auth") - or litellm_metadata.get("user_api_key_auth") - or metadata.get("user_api_key_auth") - ) - - async def anthropic_messages_with_mcp( max_tokens: int, messages: Sequence[Mapping[str, Any]], @@ -101,15 +89,18 @@ async def anthropic_messages_with_mcp( **kwargs, ) - user_api_key_auth = _resolve_user_api_key_auth(kwargs) + context = MCPRequestContext.resolve(kwargs=dict(kwargs), tools=tools) ( deduplicated_mcp_tools, tool_server_map, ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( - user_api_key_auth, + context.user_api_key_auth, mcp_references, - litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_trace_id=context.litellm_trace_id, + mcp_auth_header=context.mcp_auth_header, + mcp_server_auth_headers=context.mcp_server_auth_headers, + request_tags=list(context.request_tags) if context.request_tags else None, ) anthropic_tools: Sequence[AnthropicMessagesTool] = tuple( @@ -149,10 +140,21 @@ async def anthropic_messages_with_mcp( tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( tool_server_map=tool_server_map, tool_calls=list(tool_use_blocks), - user_api_key_auth=user_api_key_auth, - litellm_trace_id=kwargs.get("litellm_trace_id"), + user_api_key_auth=context.user_api_key_auth, + mcp_auth_header=context.mcp_auth_header, + mcp_server_auth_headers=context.mcp_server_auth_headers, + oauth2_headers=context.oauth2_headers, + raw_headers=context.raw_headers, + litellm_call_id=context.litellm_call_id, + litellm_trace_id=context.litellm_trace_id, + request_tags=list(context.request_tags) if context.request_tags else None, ) + # Every tool call was skipped, so there is nothing to feed back; a + # tool_result message with empty content is rejected by Anthropic. + if not tool_results: + break + working_messages = ( *working_messages, {"role": "assistant", "content": list(_get_response_content(response))}, diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index f2ccfd430ae..5c3e0cf0902 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -12,7 +12,7 @@ from typing import ( from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) -from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.responses.mcp.request_context import MCPRequestContext from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper @@ -114,20 +114,13 @@ async def acompletion_with_mcp( **kwargs, ) - # Extract user_api_key_auth from metadata or kwargs - user_api_key_auth = kwargs.get("user_api_key_auth") or ((kwargs.get("metadata", {}) or {}).get("user_api_key_auth")) - request_tags = LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(kwargs) - - # Extract MCP auth headers before fetching tools (needed for dynamic auth) - ( - mcp_auth_header, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( - secret_fields=kwargs.get("secret_fields"), - tools=tools, - ) + context = MCPRequestContext.resolve(kwargs=kwargs, tools=tools) + user_api_key_auth = context.user_api_key_auth + request_tags = list(context.request_tags) if context.request_tags else None + mcp_auth_header = context.mcp_auth_header + mcp_server_auth_headers = context.mcp_server_auth_headers + oauth2_headers = context.oauth2_headers + raw_headers = context.raw_headers # Process MCP tools (pass auth headers for dynamic auth) ( diff --git a/litellm/responses/mcp/request_context.py b/litellm/responses/mcp/request_context.py new file mode 100644 index 00000000000..fa03e677b39 --- /dev/null +++ b/litellm/responses/mcp/request_context.py @@ -0,0 +1,73 @@ +""" +The per-request context an MCP gateway handler needs. + +Listing and executing MCP tools both need the caller's identity, their MCP auth +headers, and the request's trace/tag identifiers. Every gateway surface resolves +the same set from its own kwargs, so resolving it in one place keeps a new +surface from silently dropping a field: omitting the auth headers, for instance, +still executes the tool, just with no credentials. +""" + +from dataclasses import dataclass +from typing import Any, Iterable, Mapping, Sequence, Union + + +@dataclass(frozen=True, slots=True) +class MCPRequestContext: + """Everything a gateway handler must forward to MCP tool listing and execution.""" + + user_api_key_auth: Any # any-ok: UserAPIKeyAuth is proxy-only; importing it here would create a cycle + mcp_auth_header: Union[str, None] = None + mcp_server_auth_headers: Union[Mapping[str, Mapping[str, str]], None] = None + oauth2_headers: Union[Mapping[str, str], None] = None + raw_headers: Union[Mapping[str, str], None] = None + request_tags: Union[Sequence[str], None] = None + litellm_trace_id: Union[str, None] = None + litellm_call_id: Union[str, None] = None + + @classmethod + def resolve( + cls, + kwargs: Mapping[str, Any], + tools: Union[Iterable[Any], None], + ) -> "MCPRequestContext": + """ + Build the context from a gateway handler's kwargs. + + ``user_api_key_auth`` is read from both metadata keys because routes differ: + LITELLM_METADATA_ROUTES (``/v1/messages``, ``/responses``) carry it in + ``litellm_metadata`` while ``/chat/completions`` uses ``metadata``. + """ + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + from litellm.responses.utils import ResponsesAPIRequestUtils + + litellm_metadata = kwargs.get("litellm_metadata") or {} + metadata = kwargs.get("metadata") or {} + user_api_key_auth = ( + kwargs.get("user_api_key_auth") + or litellm_metadata.get("user_api_key_auth") + or metadata.get("user_api_key_auth") + ) + + ( + mcp_auth_header, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( + secret_fields=kwargs.get("secret_fields"), + tools=tools, + ) + + return cls( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + request_tags=LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(dict(kwargs)), + litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_call_id=kwargs.get("litellm_call_id"), + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py index 3faa6b1e4e2..060c3e459d0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py @@ -120,3 +120,124 @@ def test_build_tool_result_message_uses_anthropic_tool_result_blocks(): assert list(message["content"]) == [ {"type": "tool_result", "tool_use_id": "toolu_1", "content": "9 sections"} ] + + +@pytest.mark.asyncio +async def test_anthropic_messages_with_mcp_forwards_the_callers_mcp_credentials(): + """ + Regression test (LIT-4517): the caller's MCP auth must reach both tool listing + and tool execution on /v1/messages. + + Given: A request carrying MCP auth headers and request tags + When: The gateway lists and then executes an MCP tool + Then: Both calls receive the caller's credentials, tags and trace ids + + Dropping them does not fail loudly; the tool still executes, just with no + credentials, so every auth-requiring MCP server (interactive OAuth, bearer + token, per-user env) silently returns nothing while the model claims it has + no access. Only a no-auth server would look healthy. + """ + from litellm.llms.anthropic.experimental_pass_through.messages import mcp_handler + from litellm.responses.mcp.request_context import MCPRequestContext + + context = MCPRequestContext( + user_api_key_auth="auth-object", + mcp_auth_header="legacy-header", + mcp_server_auth_headers={"deepwiki": {"authorization": "Bearer per-server"}}, + oauth2_headers={"authorization": "Bearer oauth"}, + raw_headers={"x-trace": "abc"}, + request_tags=["team-a"], + litellm_trace_id="trace-123", + litellm_call_id="call-456", + ) + + process = AsyncMock(return_value=([], {})) + execute = AsyncMock(return_value=[{"tool_call_id": "toolu_1", "result": "ok", "name": "t"}]) + responses = [ + {"stop_reason": "tool_use", "content": [{"type": "tool_use", "id": "toolu_1", "name": "t", "input": {}}]}, + {"stop_reason": "end_turn", "content": [{"type": "text", "text": "done"}]}, + ] + + with patch.object(MCPRequestContext, "resolve", return_value=context), patch.object( + mcp_handler.LiteLLM_Proxy_MCP_Handler + if hasattr(mcp_handler, "LiteLLM_Proxy_MCP_Handler") + else __import__( + "litellm.responses.mcp.litellm_proxy_mcp_handler", fromlist=["LiteLLM_Proxy_MCP_Handler"] + ).LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + new=process, + ), patch( + "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._execute_tool_calls", + new=execute, + ), patch( + "litellm.anthropic_messages", new=AsyncMock(side_effect=responses) + ): + await mcp_handler.anthropic_messages_with_mcp( + max_tokens=100, + messages=[{"role": "user", "content": "hi"}], + model="claude-sonnet-4-5", + tools=[MCP_REFERENCE], + ) + + listing = process.call_args.kwargs + assert listing["mcp_auth_header"] == "legacy-header", "tool listing must use the caller's MCP auth" + assert listing["mcp_server_auth_headers"] == {"deepwiki": {"authorization": "Bearer per-server"}} + assert listing["request_tags"] == ["team-a"] + assert listing["litellm_trace_id"] == "trace-123" + + execution = execute.call_args.kwargs + assert execution["user_api_key_auth"] == "auth-object" + assert execution["mcp_auth_header"] == "legacy-header", "tool execution must use the caller's MCP auth" + assert execution["mcp_server_auth_headers"] == {"deepwiki": {"authorization": "Bearer per-server"}} + assert execution["oauth2_headers"] == {"authorization": "Bearer oauth"} + assert execution["raw_headers"] == {"x-trace": "abc"} + assert execution["litellm_call_id"] == "call-456" + assert execution["litellm_trace_id"] == "trace-123" + assert execution["request_tags"] == ["team-a"] + + +@pytest.mark.asyncio +async def test_anthropic_messages_with_mcp_stops_when_every_tool_call_is_skipped(): + """ + Regression test (LIT-4517): a tool_use turn whose calls all get skipped must + end the loop, not send an empty tool_result message. + + Given: The model asks for a tool but the executor skips it (unresolvable name) + When: The gateway loop handles the empty result set + Then: It returns the last response instead of calling the model again + + _build_tool_result_message([]) produces a user message with empty content, and + Anthropic rejects that, so the caller would get an unhandled 400 from the middle + of the loop rather than the model's own answer. + """ + from litellm.llms.anthropic.experimental_pass_through.messages import mcp_handler + from litellm.responses.mcp.request_context import MCPRequestContext + + tool_use_response = { + "stop_reason": "tool_use", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "gone", "input": {}}], + } + anthropic_messages_mock = AsyncMock(return_value=tool_use_response) + + with patch.object( + MCPRequestContext, "resolve", return_value=MCPRequestContext(user_api_key_auth="auth") + ), patch( + "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform", + new=AsyncMock(return_value=([], {})), + ), patch( + "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._execute_tool_calls", + new=AsyncMock(return_value=[]), + ), patch( + "litellm.anthropic_messages", new=anthropic_messages_mock + ): + result = await mcp_handler.anthropic_messages_with_mcp( + max_tokens=100, + messages=[{"role": "user", "content": "hi"}], + model="claude-sonnet-4-5", + tools=[MCP_REFERENCE], + ) + + assert anthropic_messages_mock.await_count == 1, ( + "With no tool results there is nothing to send back, so the loop must not call the model again" + ) + assert result == tool_use_response From 56cda9f674d815a5f2686e29df9fb0b105a836f3 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 10:33:28 -0700 Subject: [PATCH 04/60] fix(mcp): sanitize Anthropic tool schemas and stop encoding gateway names Two review findings, both a chat-vs-messages divergence. transform_mcp_tool_to_anthropic_tool sent the MCP inputSchema to Anthropic almost as-is, while the chat path (_map_tool_helper) coerces the type to object, inlines legacy definitions with unpack_legacy_defs, and allow-lists keys to AnthropicInputSchema. So a tool whose schema carried $schema, legacy definitions or oneOf worked on /chat/completions and 400d on /v1/messages; a clean-schema server hid it. Both paths now run the same sanitize_input_schema_for_anthropic, extracted next to unpack_legacy_defs so they cannot drift again, and the chat path is refactored onto it rather than keeping its own copy. buildMcpToolBlocks percent-encoded the server and toolset names inside litellm_proxy/mcp/... urls, but the gateway resolves the name with a raw server_url.split("/")[-1] and never url-decodes, so a name with a space failed lookup. The already-working chat path does not encode; the shared builder now matches it. Tests pin both: reverting the transform to the unfiltered schema fails, and re-adding encodeURIComponent fails the builder test. --- litellm/experimental_mcp_client/tools.py | 8 ++- .../prompt_templates/common_utils.py | 26 ++++++++ litellm/llms/anthropic/chat/transformation.py | 29 ++------- .../experimental_mcp_client/test_tools.py | 42 +++++++++++++ .../llm_calls/mcp_tool_blocks.test.ts | 63 +++++++++++++++++++ .../components/llm_calls/mcp_tool_blocks.ts | 8 ++- 6 files changed, 147 insertions(+), 29 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.test.ts diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 1bd65847616..500d226752b 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -9,7 +9,7 @@ from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition -from litellm.types.llms.anthropic import AnthropicInputSchema, AnthropicMessagesTool +from litellm.types.llms.anthropic import AnthropicMessagesTool from litellm.types.utils import ChatCompletionMessageToolCall @@ -78,12 +78,14 @@ def transform_mcp_tool_to_openai_responses_api_tool( def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessagesTool: """Convert an MCP tool to an Anthropic Messages API tool.""" - normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema) + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + sanitize_input_schema_for_anthropic, + ) return AnthropicMessagesTool( name=mcp_tool.name, description=mcp_tool.description or "", - input_schema=AnthropicInputSchema(**normalized_parameters), + input_schema=sanitize_input_schema_for_anthropic(mcp_tool.inputSchema), type="custom", ) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 538d5f650ef..c43089950ee 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -42,6 +42,7 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: # newer pattern to avoid importing pydantic objects on __init__.py + from litellm.types.llms.anthropic import AnthropicInputSchema from litellm.types.llms.openai import ChatCompletionImageObject DEFAULT_USER_CONTINUE_MESSAGE = ChatCompletionUserMessage(content="Please continue.", role="user") @@ -1046,6 +1047,31 @@ def unpack_legacy_defs( return schema +def sanitize_input_schema_for_anthropic(input_schema: dict) -> "AnthropicInputSchema": + """Coerce an arbitrary tool input_schema into the shape Anthropic accepts. + + Anthropic requires ``type == "object"``, only recognises ``$defs`` (legacy + ``definitions`` / OpenAPI ``components.schemas`` refs must be inlined first), + and rejects keys outside ``AnthropicInputSchema``. Both the chat + (``AnthropicConfig._map_tool_helper``) and Anthropic Messages MCP paths run + a schema through here so an external MCP schema cannot succeed on one route + and 400 on the other. + """ + from litellm.types.llms.anthropic import AnthropicInputSchema + + normalized = dict(input_schema) if input_schema else {} + if normalized.get("type") != "object": + normalized["type"] = "object" + if "properties" not in normalized: + normalized["properties"] = {} + + normalized = unpack_legacy_defs(normalized, copy=True) + + allowed_keys = set(AnthropicInputSchema.__annotations__.keys()) + filtered = {key: value for key, value in normalized.items() if key in allowed_keys} + return AnthropicInputSchema(**filtered) + + def _get_image_mime_type_from_url(url: str) -> Optional[str]: """ Get mime type for common image URLs diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 0ec1f3eae13..5a0f274e3ca 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -29,7 +29,9 @@ from litellm.constants import ( RESPONSE_FORMAT_TOOL_NAME, ) from litellm.litellm_core_utils.core_helpers import map_finish_reason -from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_legacy_defs +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + sanitize_input_schema_for_anthropic, +) from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.anthropic import ( @@ -634,7 +636,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): mcp_server: Optional[AnthropicMcpServerTool] = None if tool["type"] == "function" or tool["type"] == "custom": - _input_schema: dict = tool["function"].get( + _input_schema = tool["function"].get( "parameters", { "type": "object", @@ -642,28 +644,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): }, ) - # Anthropic requires input_schema.type to be "object". Normalize - # schemas from external sources (MCP servers, OpenAI callers) that - # may omit the type field or use a non-object type. - if _input_schema.get("type") != "object": - litellm.verbose_logger.debug( - "_map_tool_helper: coercing input_schema type from %r to " - "'object' for Anthropic compatibility (tool: %s)", - _input_schema.get("type"), - tool["function"].get("name"), - ) - _input_schema = dict(_input_schema) # avoid mutating caller's dict - _input_schema["type"] = "object" - if "properties" not in _input_schema: - _input_schema["properties"] = {} - - # Inline legacy / OpenAPI $refs before the allow-list filter strips - # their backing def blocks (https://github.com/BerriAI/litellm/issues/26692). - _input_schema = unpack_legacy_defs(_input_schema, copy=True) - - _allowed_properties = set(AnthropicInputSchema.__annotations__.keys()) - input_schema_filtered = {k: v for k, v in _input_schema.items() if k in _allowed_properties} - input_anthropic_schema: AnthropicInputSchema = AnthropicInputSchema(**input_schema_filtered) + input_anthropic_schema = sanitize_input_schema_for_anthropic(_input_schema) _tool = AnthropicMessagesTool( name=tool["function"]["name"], diff --git a/tests/test_litellm/experimental_mcp_client/test_tools.py b/tests/test_litellm/experimental_mcp_client/test_tools.py index 625ab56951f..804e99b6f4e 100644 --- a/tests/test_litellm/experimental_mcp_client/test_tools.py +++ b/tests/test_litellm/experimental_mcp_client/test_tools.py @@ -299,3 +299,45 @@ def test_transform_mcp_tool_to_anthropic_tool_normalizes_empty_schema(): assert anthropic_tool["description"] == "" assert anthropic_tool["input_schema"]["type"] == "object" assert anthropic_tool["input_schema"]["properties"] == {} + + +def test_transform_mcp_tool_to_anthropic_tool_strips_keys_anthropic_rejects(): + """ + Regression test (LIT-4517): an MCP schema with keys Anthropic does not accept + must be sanitized, so the same tool cannot succeed on /chat/completions and 400 + on /v1/messages. + + Given: An MCP tool whose inputSchema carries $schema, legacy definitions and oneOf + When: It is transformed for the Anthropic Messages API + Then: Only keys in AnthropicInputSchema survive, matching the chat path + + The chat path runs the schema through the same sanitizer, so before this the two + routes diverged: a clean-schema server (deepwiki) worked on both, but a server + with a richer schema would be rejected only on messages. + """ + from litellm.types.llms.anthropic import AnthropicInputSchema + + tool = MCPTool( + name="rich", + description="tool with a dirty schema", + inputSchema={ + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + "$schema": "http://json-schema.org/draft-07/schema#", + "definitions": {"D": {"type": "string"}}, + "oneOf": [{"required": ["q"]}], + }, + ) + + anthropic_tool = transform_mcp_tool_to_anthropic_tool(tool) + schema_keys = set(anthropic_tool["input_schema"].keys()) + + assert schema_keys <= set(AnthropicInputSchema.__annotations__.keys()), ( + f"schema must only carry keys Anthropic accepts, got {schema_keys}" + ) + assert "$schema" not in schema_keys + assert "definitions" not in schema_keys + assert "oneOf" not in schema_keys + assert anthropic_tool["input_schema"]["properties"] == {"q": {"type": "string"}} + assert anthropic_tool["input_schema"]["required"] == ["q"] diff --git a/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.test.ts b/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.test.ts new file mode 100644 index 00000000000..62dd4d2631c --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.test.ts @@ -0,0 +1,63 @@ +import { describe, it, expect } from "vitest"; +import { buildMcpToolBlocks } from "./mcp_tool_blocks"; +import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; + +const server = (over: Partial): MCPServer => + ({ + server_id: "id-1", + server_name: "deepwiki", + alias: "wiki", + url: "", + transport: "http", + auth_type: "none", + ...over, + }) as any; + +describe("buildMcpToolBlocks", () => { + it("returns no blocks when nothing is selected", () => { + expect(buildMcpToolBlocks({ selectedMCPServers: [] })).toEqual([]); + expect(buildMcpToolBlocks({ selectedMCPServers: undefined })).toEqual([]); + }); + + it("routes by server_name, not alias, so colliding aliases cannot cross-route", () => { + const [block] = buildMcpToolBlocks({ + selectedMCPServers: ["id-1"], + mcpServers: [server({})], + }); + expect(block.server_url).toBe("litellm_proxy/mcp/deepwiki"); + expect(block.server_label).toBe("deepwiki"); + }); + + it("does not percent-encode the name; the gateway splits the raw path and never decodes", () => { + const [block] = buildMcpToolBlocks({ + selectedMCPServers: ["id-1"], + mcpServers: [server({ server_name: "my server" }) as any], + }); + expect(block.server_url).toBe("litellm_proxy/mcp/my server"); + expect(block.server_url).not.toContain("%20"); + }); + + it("passes per-server tool restrictions through as allowed_tools", () => { + const [block] = buildMcpToolBlocks({ + selectedMCPServers: ["id-1"], + mcpServers: [server({})], + mcpServerToolRestrictions: { "id-1": ["read_wiki_structure"] }, + }); + expect(block.allowed_tools).toEqual(["read_wiki_structure"]); + }); + + it("collapses the all-servers sentinel to a single proxy-wide block", () => { + expect(buildMcpToolBlocks({ selectedMCPServers: ["__all__", "id-1"] })).toEqual([ + { type: "mcp", server_label: "litellm", server_url: "litellm_proxy/mcp", require_approval: "never" }, + ]); + }); + + it("routes a toolset by its name", () => { + const toolset = { toolset_id: "ts-1", toolset_name: "docs" } as MCPToolset; + const [block] = buildMcpToolBlocks({ + selectedMCPServers: ["toolset:ts-1"], + mcpToolsets: [toolset], + }); + expect(block.server_url).toBe("litellm_proxy/mcp/docs"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.ts b/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.ts index 42fa94d8208..401d9fd9c84 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.ts +++ b/ui/litellm-dashboard/src/components/llm_calls/mcp_tool_blocks.ts @@ -29,6 +29,10 @@ export interface BuildMcpToolBlocksArgs { * server_name is used for both routing and labelling because it is the unique * registered identifier; aliases can collide across servers, and a duplicated * server_label causes silent tool-routing failures. + * + * The name is not percent-encoded: the gateway resolves it with a raw + * `server_url.split("/")[-1]` and never url-decodes, so an encoded name would + * fail server lookup rather than round-trip. */ export function buildMcpToolBlocks({ selectedMCPServers, @@ -59,7 +63,7 @@ export function buildMcpToolBlocks({ return { type: "mcp", server_label: toolsetName, - server_url: `litellm_proxy/mcp/${encodeURIComponent(toolsetName)}`, + server_url: `litellm_proxy/mcp/${toolsetName}`, require_approval: "never", }; } @@ -71,7 +75,7 @@ export function buildMcpToolBlocks({ return { type: "mcp", server_label: routeName, - server_url: `litellm_proxy/mcp/${encodeURIComponent(routeName)}`, + server_url: `litellm_proxy/mcp/${routeName}`, require_approval: "never", ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), }; From 31f293a9fc60a5ef7ff8c40ebfe4eab3fc5d11f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:49:02 -0400 Subject: [PATCH 05/60] feat(bedrock): forward bedrock_tags to CreateModelInvocationJob for batch jobs --- .../llms/bedrock/batches/transformation.py | 18 ++++ litellm/types/llms/bedrock.py | 7 +- .../bedrock/batches/test_transformation.py | 87 +++++++++++++++++++ 3 files changed, 111 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 4fcf7cf91cb..8648d6586e8 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -4,6 +4,7 @@ import time from typing import Any, Dict, List, Literal, Optional, Union, cast from httpx import Headers, Response +from pydantic import TypeAdapter, ValidationError from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, @@ -19,6 +20,7 @@ from litellm.types.llms.bedrock import ( BedrockOutputDataConfig, BedrockS3InputDataConfig, BedrockS3OutputDataConfig, + BedrockTag, ) from litellm.types.llms.openai import ( AllMessageValues, @@ -38,6 +40,18 @@ _S3_BATCH_FILE_UUID_SUFFIX_PATTERN = re.compile( r"-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.jsonl$" ) +_BEDROCK_TAGS_ADAPTER: TypeAdapter[list[BedrockTag]] = TypeAdapter(list[BedrockTag]) + + +def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]: + try: + return _BEDROCK_TAGS_ADAPTER.validate_python(raw_tags, strict=True) + except ValidationError as e: + raise ValueError( + "Invalid 'bedrock_tags' value. Expected a list of {'key': , 'value': } dicts, " + f"e.g. [{{'key': 'team', 'value': 'genai'}}]. Got: {raw_tags!r}" + ) from e + class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ @@ -201,6 +215,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): "roleArn": role_arn, } + bedrock_tags = litellm_params.get("bedrock_tags") or optional_params.get("bedrock_tags") + if bedrock_tags is not None: + bedrock_request["tags"] = _validate_bedrock_tags(bedrock_tags) + # Add optional parameters if provided completion_window = create_batch_data.get("completion_window") if completion_window: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index bdf6b8fefed..d9f8229dbed 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -985,6 +985,11 @@ class BedrockOutputDataConfig(TypedDict): s3OutputDataConfig: BedrockS3OutputDataConfig +class BedrockTag(TypedDict): + key: str + value: str + + class BedrockCreateBatchRequest(TypedDict, total=False): """ Request structure for creating a Bedrock batch inference job. @@ -999,7 +1004,7 @@ class BedrockCreateBatchRequest(TypedDict, total=False): outputDataConfig: BedrockOutputDataConfig timeoutDurationInHours: Optional[int] clientRequestToken: Optional[str] - tags: Optional[List[dict]] + tags: Optional[List[BedrockTag]] BedrockBatchJobStatus = Literal["Submitted", "InProgress", "Completed", "Failed", "Stopping", "Stopped"] diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index d1ad5943ae6..b38d271e210 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -258,6 +258,93 @@ def test_create_request_no_timeout_for_non_24h_window(config): assert "timeoutDurationInHours" not in mock_sign.call_args.kwargs["data"] +def test_create_request_forwards_bedrock_tags_from_litellm_params(config): + tags = [ + {"key": "application", "value": "genai-proxy"}, + {"key": "team", "value": "ml-platform"}, + ] + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_tags": tags, + }, + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == tags + + +def test_create_request_forwards_bedrock_tags_from_optional_params(config): + tags = [{"key": "env", "value": "prod"}] + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={"bedrock_tags": tags}, + litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r"}, + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == tags + + +def test_create_request_omits_tags_when_bedrock_tags_absent(config): + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={}, + litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r"}, + ) + assert "tags" not in mock_sign.call_args.kwargs["data"] + + +@pytest.mark.parametrize( + "bad_tags", + [ + ["application=genai-proxy"], + [{"key": "application"}], + [{"value": "genai-proxy"}], + [{"key": "application", "value": 42}], + {"key": "application", "value": "genai-proxy"}, + "application=genai-proxy", + ], +) +def test_create_request_rejects_malformed_bedrock_tags(config, bad_tags): + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + with pytest.raises(ValueError, match="Invalid 'bedrock_tags' value"): + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_tags": bad_tags, + }, + ) + mock_sign.assert_not_called() + + # --------------------------------------------------------------------------- # # transform_create_batch_response - status mapping + LiteLLMBatch shape # --------------------------------------------------------------------------- # From cf23df94313ae308484a41772e1e8aab23ded4d6 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 11:34:08 -0700 Subject: [PATCH 06/60] fix(mcp): require every reference to opt in before auto-executing tools _should_auto_execute_tools returned True as soon as any MCP reference set require_approval="never", so a request that mixed a "never" reference with an "always" or "manual" one auto-executed every tool call the model produced, including the approval-gated ones. A prompt could name the approval-required tool and have it run with no approval. Make the gate fail closed: auto-execute only when every reference opts in with "never". A single approval-required reference (including the object form or an unset value) returns the model's tool calls to the caller instead of running them, so an approval-gated tool can never be auto-invoked. This is the shared decision behind /chat/completions, /responses, the streaming iterator and the new /v1/messages path, so all four fail closed from one change. The common case, every reference "never", is unchanged. The alternative, executing the "never" calls and returning only the approval-required ones, needs partial execution that the Anthropic tool loop cannot express without fabricating tool_result blocks for the calls it withheld, so the whole-request fail-closed gate is the safe minimum. A future change can add per-call partial execution if a caller needs it. Test covers the mixed and manual cases; reverting to "any never" fails it. --- .../mcp/litellm_proxy_mcp_handler.py | 28 ++++++++++++------- .../mcp_tests/test_aresponses_api_with_mcp.py | 9 ++++++ 2 files changed, 27 insertions(+), 10 deletions(-) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index d2c9f220690..a94cd2413d8 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -478,17 +478,25 @@ class LiteLLM_Proxy_MCP_Handler: ) -> bool: """Check if we should auto-execute tool calls. - Only auto-execute tools if user passed a MCP tool with require_approval set to "never". - - + Auto-execution requires EVERY MCP reference to opt in with + ``require_approval="never"``. A single reference that requires approval + ("always", "manual", the object form, or an unset value) disables + auto-execution for the whole request. This fails closed: when an + approval-required reference shares a request with a "never" one, the + model's tool calls are returned to the caller instead of being run, so + an approval-gated tool can never be invoked without approval. Returns + False for an empty list. """ - for tool in mcp_tools_with_litellm_proxy: - if isinstance(tool, dict): - if tool.get("require_approval") == "never": - return True - elif getattr(tool, "require_approval", None) == "never": - return True - return False + references = list(mcp_tools_with_litellm_proxy or []) + if not references: + return False + for tool in references: + approval = ( + tool.get("require_approval") if isinstance(tool, dict) else getattr(tool, "require_approval", None) + ) + if approval != "never": + return False + return True @staticmethod def _extract_tool_calls_from_response(response: ResponsesAPIResponse) -> List[Any]: diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index 9cd45f3d6fc..32295310005 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -86,6 +86,15 @@ async def test_mcp_helper_methods(): LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(mcp_tools_always) == False ) + # A single approval-required reference must disable auto-execution for the + # whole request; otherwise a "never" reference alongside an "always" one + # would let the approval-gated tool run without approval. + mcp_tools_mixed = [{"require_approval": "never"}, {"require_approval": "always"}] + assert LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(mcp_tools_mixed) == False + mcp_tools_manual = [{"require_approval": "never"}, {"require_approval": "manual"}] + assert LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(mcp_tools_manual) == False + assert LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools([]) == False + print("✓ MCP helper methods test passed!") From 2e492f5cb7ee5764132e7123b787329a3e244c24 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 18 Jul 2026 16:27:59 -0700 Subject: [PATCH 07/60] fix(ui): hide guardrail group headers when only one group has entries The team settings guardrails dropdown always rendered the Global and Other headers, so a proxy with no global guardrails showed an empty Global heading above the list. --- .../src/components/team/TeamInfo.test.tsx | 99 ++++++++++++++++++- .../src/components/team/TeamInfo.tsx | 64 ++++++------ ui/litellm-dashboard/tests/test-utils.tsx | 4 +- 3 files changed, 132 insertions(+), 35 deletions(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index a57676d07b8..dcc72ccac9c 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1,8 +1,8 @@ import * as networking from "@/components/networking"; -import { screen, waitFor } from "@testing-library/react"; +import { screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import TeamInfoView from "./TeamInfo"; vi.mock("@/components/networking", () => ({ @@ -1024,4 +1024,99 @@ describe("TeamInfoView", () => { expect(payload).not.toHaveProperty("model_aliases"); }); }); + + describe("guardrails dropdown grouping", () => { + const guardrail = (name: string, defaultOn: boolean) => ({ + guardrail_name: name, + litellm_params: { default_on: defaultOn }, + }); + + const openGuardrailsDropdown = async (user: ReturnType) => { + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await waitFor(() => { + expect(screen.getByLabelText(/^Guardrails/)).toBeInTheDocument(); + }); + + const dropdownsBefore = new Set(document.querySelectorAll(".ant-select-dropdown")); + + await user.click(screen.getByLabelText(/^Guardrails/)); + + return waitFor( + () => { + const opened = Array.from(document.querySelectorAll(".ant-select-dropdown")).find( + (el) => !dropdownsBefore.has(el), + ); + expect(opened).toBeDefined(); + return opened as HTMLElement; + }, + { timeout: 5000 }, + ); + }; + + beforeEach(() => { + testQueryClient.clear(); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + }); + + it("should not render the Global or Other group headers when no global guardrails exist", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.getGuardrailsList).mockResolvedValue({ + guardrails: [guardrail("dwacxzcz", false), guardrail("dwadsa", false)], + }); + + const dropdown = await openGuardrailsDropdown(user); + + await waitFor(() => { + expect(within(dropdown).getByTitle("dwacxzcz")).toBeInTheDocument(); + }); + expect(within(dropdown).getByTitle("dwadsa")).toBeInTheDocument(); + expect(within(dropdown).queryByText("Global")).not.toBeInTheDocument(); + expect(within(dropdown).queryByText("Other")).not.toBeInTheDocument(); + }); + + it("should not render the Global or Other group headers when every guardrail is global", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.getGuardrailsList).mockResolvedValue({ + guardrails: [guardrail("always-on", true)], + }); + + const dropdown = await openGuardrailsDropdown(user); + + await waitFor(() => { + expect(within(dropdown).getByTitle("always-on")).toBeInTheDocument(); + }); + expect(within(dropdown).queryByText("Global")).not.toBeInTheDocument(); + expect(within(dropdown).queryByText("Other")).not.toBeInTheDocument(); + }); + + it("should render both group headers when global and non-global guardrails exist", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.getGuardrailsList).mockResolvedValue({ + guardrails: [guardrail("always-on", true), guardrail("opt-in", false)], + }); + + const dropdown = await openGuardrailsDropdown(user); + + await waitFor(() => { + expect(within(dropdown).getByText("Global")).toBeInTheDocument(); + }); + expect(within(dropdown).getByText("Other")).toBeInTheDocument(); + expect(within(dropdown).getByTitle("always-on")).toBeInTheDocument(); + expect(within(dropdown).getByTitle("opt-in")).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 74f201b40c6..eaf2faa08ee 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -14,7 +14,7 @@ import { teamMemberUpdateCall, teamUpdateCall, } from "@/components/networking"; -import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; +import { useGuardrails, GuardrailListItem } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils"; import { isProxyAdminRole } from "@/utils/roles"; @@ -679,6 +679,16 @@ const TeamInfoView: React.FC = ({ ? nonGlobalOptIns : [...Array.from(globalGuardrailNames).filter((n) => !optedOutGlobals.has(n)), ...nonGlobalOptIns]; + const allGuardrails: GuardrailListItem[] = guardrailsData?.guardrails ?? []; + const globalGuardrails = allGuardrails.filter((g) => g.litellm_params?.default_on); + const otherGuardrails = allGuardrails.filter((g) => !g.litellm_params?.default_on); + + const renderGuardrailOption = (g: GuardrailListItem, disabled: boolean) => ( + + {g.guardrail_name} + + ); + const preventTagMouseDown = (e: React.MouseEvent) => { e.preventDefault(); e.stopPropagation(); @@ -1264,36 +1274,28 @@ const TeamInfoView: React.FC = ({ optionLabelProp="label" tagRender={renderGuardrailTag} > - - - Global - - } - > - {(guardrailsData?.guardrails ?? []) - .filter((g) => g.litellm_params?.default_on) - .map((g) => ( - - {g.guardrail_name} - - ))} - - - {(guardrailsData?.guardrails ?? []) - .filter((g) => !g.litellm_params?.default_on) - .map((g) => ( - - {g.guardrail_name} - - ))} - + {globalGuardrails.length > 0 && otherGuardrails.length > 0 ? ( + <> + + + Global + + } + > + {globalGuardrails.map((g) => renderGuardrailOption(g, Boolean(killSwitchOn)))} + + + {otherGuardrails.map((g) => renderGuardrailOption(g, false))} + + + ) : ( + [ + ...globalGuardrails.map((g) => renderGuardrailOption(g, Boolean(killSwitchOn))), + ...otherGuardrails.map((g) => renderGuardrailOption(g, false)), + ] + )} diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx index ed1f248648e..ba07af0f376 100644 --- a/ui/litellm-dashboard/tests/test-utils.tsx +++ b/ui/litellm-dashboard/tests/test-utils.tsx @@ -3,7 +3,7 @@ import { render, RenderOptions } from "@testing-library/react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; // Create a client for testing -const queryClient = new QueryClient({ +export const testQueryClient = new QueryClient({ defaultOptions: { queries: { retry: false, @@ -20,7 +20,7 @@ const queryClient = new QueryClient({ }); const Providers: React.FC = ({ children }) => { - return {children}; + return {children}; }; export const renderWithProviders = (ui: React.ReactElement, options?: RenderOptions) => From 0268d0151636b7367d23ad5dcfe0e5448d7553f7 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 18 Jul 2026 16:44:27 -0700 Subject: [PATCH 08/60] fix(anthropic): only inject cache_control when the request carries none --- .../anthropic_cache_control_hook.py | 55 ++++++-- .../messages/handler.py | 24 +++- litellm/main.py | 1 + .../test_anthropic_cache_control_hook.py | 129 ++++++++++++++++++ 4 files changed, 192 insertions(+), 17 deletions(-) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 94c86e07ff5..54f0732fdf1 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -322,7 +322,9 @@ class AnthropicCacheControlHook(CustomPromptManagement): 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. + a common pattern, so injecting alongside them can exceed the cap. Tools + carry the mark either at the top level (Anthropic shape) or nested under + ``function`` (OpenAI shape); the Anthropic chat transform accepts both. """ if any(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages): return True @@ -330,7 +332,14 @@ class AnthropicCacheControlHook(CustomPromptManagement): 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 any( + isinstance(tool, dict) + and ( + tool.get("cache_control") is not None + or (isinstance(tool.get("function"), dict) and tool["function"].get("cache_control") is not None) + ) + for tool in tools + ) return False @staticmethod @@ -391,14 +400,24 @@ class AnthropicCacheControlHook(CustomPromptManagement): model: str, custom_llm_provider: str | None, tools: list | None = None, + is_first_pass: bool = True, ) -> None: - """For /chat/completions: add default injection points to the request params. + """For /chat/completions: resolve the injection points the request should carry. - 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. + Configured injection points win over the automatic defaults, but stand + down entirely when the client already marked its own cache_control + breakpoints (messages or tools): injecting alongside them clashes with + the client's caching strategy and can exceed the provider's four-block + limit. Only the first pass over a request may make that judgment; + ``acompletion`` re-enters ``completion`` after injection has already + run, and a later pass would mistake litellm's own injected marks for + client ones and drop the non-message points reserved for provider + transforms. Seeding the param lets the existing prompt-management gate + and the AnthropicCacheControlHook run unchanged. """ if non_default_params.get("cache_control_injection_points"): + if is_first_pass and AnthropicCacheControlHook._request_has_cache_control(messages, None, tools): + non_default_params.pop("cache_control_injection_points") return points = AnthropicCacheControlHook.get_default_injection_points( messages=messages, @@ -418,21 +437,35 @@ class AnthropicCacheControlHook(CustomPromptManagement): model: str | None = None, custom_llm_provider: str | None = None, tools: list[dict] | None = None, + is_first_pass: bool = True, ) -> 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. + Configured points stand down entirely when the client already marked + its own cache_control breakpoints anywhere in the request, judged only + on the first pass: the async entry re-dispatches into the sync handler + after injection has run, and a later pass would mistake litellm's own + injected marks for client ones and drop the written-back non-message + points. 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. """ + typed_messages = cast(list[AllMessageValues], messages) # cast-ok: Anthropic-shaped dicts from v1/messages 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) ) + if ( + configured + and is_first_pass + and AnthropicCacheControlHook._request_has_cache_control(typed_messages, system, tools) + ): + return messages, system 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 + messages=typed_messages, system=system, tools=tools, model=model, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 703ccf13c27..92151b45b75 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -351,9 +351,11 @@ async def anthropic_messages( custom_llm_provider=custom_llm_provider, # messages were already empty-text-block sanitized at the top of this # function and are NOT reassigned before this dispatch, so the handler - # can skip its (otherwise redundant) second full-messages scan. Passed - # explicitly (not via **kwargs) so it only affects this direct - # dispatch -- interceptor / sync entry points still sanitize. + # can skip its (otherwise redundant) second full-messages scan. It also + # tells the handler that cache_control injection already judged the + # pristine client input here. Passed explicitly (not via **kwargs) so + # it only affects this direct dispatch -- interceptor / sync entry + # points still sanitize. _litellm_messages_presanitized=True, **kwargs, ) @@ -419,8 +421,12 @@ def anthropic_messages_handler( # protection as the async wrapper. The async wrapper already sanitized and # does not reassign messages before dispatch, so it sets # ``_litellm_messages_presanitized`` to skip this redundant second - # full-messages scan. Pop it so it never leaks into provider params. - if not kwargs.pop("_litellm_messages_presanitized", False): + # full-messages scan. The same flag marks this call as a second pass for + # cache_control injection: the async wrapper already judged the pristine + # client input, and re-judging after injection would misread litellm's own + # marks as client ones. Pop it so it never leaks into provider params. + presanitized = kwargs.pop("_litellm_messages_presanitized", False) + if not presanitized: messages = strip_empty_text_blocks_from_anthropic_messages(messages) messages = sanitize_tool_use_ids_in_anthropic_messages(messages) @@ -429,7 +435,13 @@ def anthropic_messages_handler( ) messages, system = AnthropicCacheControlHook.maybe_inject_cache_control( - messages, system, kwargs, model=model, custom_llm_provider=custom_llm_provider, tools=tools + messages, + system, + kwargs, + model=model, + custom_llm_provider=custom_llm_provider, + tools=tools, + is_first_pass=not presanitized, ) metadata = validate_anthropic_api_metadata(metadata) diff --git a/litellm/main.py b/litellm/main.py index 3584297b35f..ab07da7528b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5080,6 +5080,7 @@ def completion( # type: ignore model=model, custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs tools=tools, + is_first_pass=not acompletion, ) if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( 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 70c1f65b541..85e77f2f493 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1622,6 +1622,13 @@ class TestEnableAnthropicPromptCaching: monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) assert [p["index"] for p in self._points(tools=tools)] == [None, -1] + def test_stands_down_when_tool_function_carries_cache_control(self, monkeypatch): + """OpenAI-shaped tools nest cache_control under ``function``; the Anthropic + chat transform honors that location, so the stand-down must see it too.""" + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + tools = [{"type": "function", "function": {"name": "t", "parameters": {}, "cache_control": {"type": "ephemeral"}}}] + assert self._points(tools=tools) == [] + 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) @@ -1726,6 +1733,128 @@ class TestEnableAnthropicPromptCaching: assert result_msgs == messages +class TestConfiguredInjectionPointsStandDown: + """Configured cache_control_injection_points must stand down entirely when the + client already set its own cache_control anywhere in the request (LIT-4582); + injecting alongside client breakpoints clashes with the client's caching + strategy and can push the request past Anthropic's four-block limit.""" + + CONFIGURED = [{"location": "message", "role": "system"}] + + CLEAN_MESSAGES: List[AllMessageValues] = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + ] + + MARKED_MESSAGES: List[AllMessageValues] = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]}, + ] + + V1_MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + + def _seed(self, params, messages, tools=None, is_first_pass=True): + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=params, + messages=messages, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=tools, + is_first_pass=is_first_pass, + ) + + def _inject(self, messages, kwargs, system="sys", tools=None, is_first_pass=True): + return AnthropicCacheControlHook.maybe_inject_cache_control( + messages, + system, + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=tools, + is_first_pass=is_first_pass, + ) + + def test_configured_points_dropped_when_messages_carry_cache_control(self): + params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} + self._seed(params, copy.deepcopy(self.MARKED_MESSAGES)) + assert "cache_control_injection_points" not in params + + @pytest.mark.parametrize( + "tool", + [ + {"type": "function", "function": {"name": "t", "parameters": {}}, "cache_control": {"type": "ephemeral"}}, + {"type": "function", "function": {"name": "t", "parameters": {}, "cache_control": {"type": "ephemeral"}}}, + ], + ids=["top_level", "nested_in_function"], + ) + def test_configured_points_dropped_when_tools_carry_cache_control(self, tool): + params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} + self._seed(params, copy.deepcopy(self.CLEAN_MESSAGES), tools=[tool]) + assert "cache_control_injection_points" not in params + + def test_configured_points_kept_when_request_is_unmarked(self): + configured = copy.deepcopy(self.CONFIGURED) + params = {"cache_control_injection_points": configured} + self._seed(params, copy.deepcopy(self.CLEAN_MESSAGES)) + assert params["cache_control_injection_points"] is configured + + def test_second_pass_keeps_points_despite_injected_marks(self): + """acompletion() re-enters completion() after injection ran, with only the + non-message points written back; the second pass must not misread litellm's + own marks as client ones and drop that remainder.""" + remainder = [{"location": "tool_config"}] + params = {"cache_control_injection_points": remainder} + self._seed(params, copy.deepcopy(self.MARKED_MESSAGES), is_first_pass=False) + assert params["cache_control_injection_points"] is remainder + + def test_v1_messages_stand_down_when_content_block_marked(self): + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]} + ] + kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} + result_msgs, result_sys = self._inject(copy.deepcopy(messages), kwargs) + assert result_msgs == messages + assert result_sys == "sys" + assert "cache_control_injection_points" not in kwargs + + def test_v1_messages_stand_down_when_system_block_marked(self): + """A configured point targeting a message must not fire when the client + marked the system prompt; the old behavior injected into the message + because only the exact targeted position was guarded.""" + system = [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}] + kwargs = {"cache_control_injection_points": [{"location": "message", "role": "user"}]} + result_msgs, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs, system=system) + assert result_msgs == self.V1_MESSAGES + assert result_sys == system + assert "cache_control_injection_points" not in kwargs + + def test_v1_messages_stand_down_when_tools_marked(self): + tools = [{"name": "t", "input_schema": {}, "cache_control": {"type": "ephemeral"}}] + kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} + result_msgs, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs, tools=tools) + assert result_msgs == self.V1_MESSAGES + assert result_sys == "sys" + assert "cache_control_injection_points" not in kwargs + + def test_v1_messages_configured_points_apply_when_unmarked(self): + kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} + _, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) + assert result_sys == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}] + + def test_v1_messages_second_pass_writes_back_remainder(self): + """The async entry re-dispatches into the sync handler after injecting; the + surviving tool_config remainder must survive that second pass even though + the messages now carry litellm's own marks.""" + marked = [ + {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]} + ] + remainder = [{"location": "tool_config"}] + kwargs = {"cache_control_injection_points": remainder} + result_msgs, _ = self._inject(copy.deepcopy(marked), kwargs, is_first_pass=False) + assert result_msgs == marked + assert kwargs["cache_control_injection_points"] == remainder + + 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 From ecff0c0a7de44c890b2c59a0bf6815b2789ddbc5 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 18 Jul 2026 16:53:44 -0700 Subject: [PATCH 09/60] fix(anthropic): state the async entry's first-pass judgment explicitly --- litellm/main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/main.py b/litellm/main.py index ab07da7528b..f96064c0591 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -522,6 +522,7 @@ async def acompletion( model=model, custom_llm_provider=cast(Optional[str], custom_llm_provider), # cast-ok: read from untyped kwargs tools=tools, + is_first_pass=True, ) if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( From e0648571ed9dc6d31337e975f0a55c36a77d985d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 18 Jul 2026 17:11:03 -0700 Subject: [PATCH 10/60] fix(anthropic): carry the stand-down judgment inside written-back injection points --- .../anthropic_cache_control_hook.py | 67 +++++++++++++------ .../messages/handler.py | 24 ++----- litellm/main.py | 2 - .../anthropic_cache_control_hook.py | 4 +- .../test_anthropic_cache_control_hook.py | 52 +++++++------- 5 files changed, 84 insertions(+), 65 deletions(-) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 54f0732fdf1..faedf8ae1a3 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -91,7 +91,9 @@ class AnthropicCacheControlHook(CustomPromptManagement): # Pass through non-message injection points for provider-specific handling if remaining_points: - non_default_params["cache_control_injection_points"] = remaining_points + non_default_params["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged( + remaining_points + ) return model, processed_messages, non_default_params @@ -310,6 +312,35 @@ class AnthropicCacheControlHook(CustomPromptManagement): return ChatCompletionCachedContent(type="ephemeral", ttl=ttl) return ChatCompletionCachedContent(type="ephemeral") + @staticmethod + def _stamped_as_judged(points: list[CacheControlInjectionPoint]) -> list[dict[str, object]]: + """Mark written-back points as having passed the client cache_control judgment. + + Builds copies because config-owned point dicts are shared across + requests; mutating them would leak the stamp into future requests. + """ + return [{**point, "_litellm_judged": True} for point in points] + + @staticmethod + def _should_stand_down( + points: list[CacheControlInjectionPoint], + messages: list[AllMessageValues], + system: str | list | None, + tools: list | None, + ) -> bool: + """Whether configured injection points must yield to client-set cache_control. + + Points that a prior pass over this request already judged and wrote + back carry the internal judged stamp; any re-entry (acompletion + re-entering completion, the async-to-sync /v1/messages dispatch, + interceptor sub-calls reusing the request kwargs) must not re-judge + them, because by then the messages carry litellm's own injected marks + and the judgment would misread those as client breakpoints. + """ + if all(point.get("_litellm_judged") for point in points): + return False + return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools) + @staticmethod def _request_has_cache_control( messages: list[AllMessageValues], @@ -400,7 +431,6 @@ class AnthropicCacheControlHook(CustomPromptManagement): model: str, custom_llm_provider: str | None, tools: list | None = None, - is_first_pass: bool = True, ) -> None: """For /chat/completions: resolve the injection points the request should carry. @@ -408,15 +438,16 @@ class AnthropicCacheControlHook(CustomPromptManagement): down entirely when the client already marked its own cache_control breakpoints (messages or tools): injecting alongside them clashes with the client's caching strategy and can exceed the provider's four-block - limit. Only the first pass over a request may make that judgment; - ``acompletion`` re-enters ``completion`` after injection has already - run, and a later pass would mistake litellm's own injected marks for - client ones and drop the non-message points reserved for provider - transforms. Seeding the param lets the existing prompt-management gate - and the AnthropicCacheControlHook run unchanged. + limit. The judgment happens once per request; points a prior pass + wrote back carry the judged stamp and are never re-judged (see + ``_should_stand_down``). Seeding the param lets the existing + prompt-management gate and the AnthropicCacheControlHook run + unchanged. """ if non_default_params.get("cache_control_injection_points"): - if is_first_pass and AnthropicCacheControlHook._request_has_cache_control(messages, None, tools): + if AnthropicCacheControlHook._should_stand_down( + non_default_params["cache_control_injection_points"], messages, None, tools + ): non_default_params.pop("cache_control_injection_points") return points = AnthropicCacheControlHook.get_default_injection_points( @@ -437,16 +468,14 @@ class AnthropicCacheControlHook(CustomPromptManagement): model: str | None = None, custom_llm_provider: str | None = None, tools: list[dict] | None = None, - is_first_pass: bool = True, ) -> Tuple[List[Dict], str | list | None]: """Extract cache_control_injection_points from kwargs and apply if present. Configured points stand down entirely when the client already marked - its own cache_control breakpoints anywhere in the request, judged only - on the first pass: the async entry re-dispatches into the sync handler - after injection has run, and a later pass would mistake litellm's own - injected marks for client ones and drop the written-back non-message - points. When none are configured but + its own cache_control breakpoints anywhere in the request. The + judgment happens once per request; points a prior pass wrote back + carry the judged stamp and are never re-judged (see + ``_should_stand_down``). 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 @@ -456,11 +485,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): 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) ) - if ( - configured - and is_first_pass - and AnthropicCacheControlHook._request_has_cache_control(typed_messages, system, tools) - ): + if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools): return messages, system injection_points: list[CacheControlInjectionPoint] = configured or [] if not injection_points and model is not None: @@ -480,7 +505,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): injection_points=injection_points, ) if remaining: - kwargs["cache_control_injection_points"] = remaining + kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(remaining) return messages, system @property diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 92151b45b75..703ccf13c27 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -351,11 +351,9 @@ async def anthropic_messages( custom_llm_provider=custom_llm_provider, # messages were already empty-text-block sanitized at the top of this # function and are NOT reassigned before this dispatch, so the handler - # can skip its (otherwise redundant) second full-messages scan. It also - # tells the handler that cache_control injection already judged the - # pristine client input here. Passed explicitly (not via **kwargs) so - # it only affects this direct dispatch -- interceptor / sync entry - # points still sanitize. + # can skip its (otherwise redundant) second full-messages scan. Passed + # explicitly (not via **kwargs) so it only affects this direct + # dispatch -- interceptor / sync entry points still sanitize. _litellm_messages_presanitized=True, **kwargs, ) @@ -421,12 +419,8 @@ def anthropic_messages_handler( # protection as the async wrapper. The async wrapper already sanitized and # does not reassign messages before dispatch, so it sets # ``_litellm_messages_presanitized`` to skip this redundant second - # full-messages scan. The same flag marks this call as a second pass for - # cache_control injection: the async wrapper already judged the pristine - # client input, and re-judging after injection would misread litellm's own - # marks as client ones. Pop it so it never leaks into provider params. - presanitized = kwargs.pop("_litellm_messages_presanitized", False) - if not presanitized: + # full-messages scan. Pop it so it never leaks into provider params. + if not kwargs.pop("_litellm_messages_presanitized", False): messages = strip_empty_text_blocks_from_anthropic_messages(messages) messages = sanitize_tool_use_ids_in_anthropic_messages(messages) @@ -435,13 +429,7 @@ def anthropic_messages_handler( ) messages, system = AnthropicCacheControlHook.maybe_inject_cache_control( - messages, - system, - kwargs, - model=model, - custom_llm_provider=custom_llm_provider, - tools=tools, - is_first_pass=not presanitized, + messages, system, kwargs, model=model, custom_llm_provider=custom_llm_provider, tools=tools ) metadata = validate_anthropic_api_metadata(metadata) diff --git a/litellm/main.py b/litellm/main.py index f96064c0591..3584297b35f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -522,7 +522,6 @@ async def acompletion( model=model, custom_llm_provider=cast(Optional[str], custom_llm_provider), # cast-ok: read from untyped kwargs tools=tools, - is_first_pass=True, ) if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( @@ -5081,7 +5080,6 @@ def completion( # type: ignore model=model, custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs tools=tools, - is_first_pass=not acompletion, ) if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( diff --git a/litellm/types/integrations/anthropic_cache_control_hook.py b/litellm/types/integrations/anthropic_cache_control_hook.py index 601978bb04f..efb189088b6 100644 --- a/litellm/types/integrations/anthropic_cache_control_hook.py +++ b/litellm/types/integrations/anthropic_cache_control_hook.py @@ -1,6 +1,6 @@ from typing import Literal, Optional, Union -from typing_extensions import TypedDict +from typing_extensions import NotRequired, TypedDict from litellm.types.llms.openai import ChatCompletionCachedContent @@ -12,6 +12,7 @@ class CacheControlMessageInjectionPoint(TypedDict): role: Optional[Literal["user", "system", "assistant"]] # Optional: target by role (user, system, assistant) index: Optional[Union[int, str]] # Optional: target by specific index control: Optional[ChatCompletionCachedContent] + _litellm_judged: NotRequired[bool] # Internal: written back by litellm once the client cache_control judgment ran class CacheControlToolConfigInjectionPoint(TypedDict): @@ -19,6 +20,7 @@ class CacheControlToolConfigInjectionPoint(TypedDict): location: Literal["tool_config"] control: Optional[ChatCompletionCachedContent] + _litellm_judged: NotRequired[bool] # Internal: written back by litellm once the client cache_control judgment ran CacheControlInjectionPoint = Union[ 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 85e77f2f493..d94f0d5f47e 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1265,8 +1265,11 @@ def test_cache_control_hook_reserves_slot_for_tool_config_point(): ) assert _count_cache_control(processed) == 3 - # The tool_config point is passed through for the provider transform. - assert non_default_params["cache_control_injection_points"] == [{"location": "tool_config"}] + # The tool_config point is passed through for the provider transform, + # stamped so re-entries never re-judge it against litellm's own marks. + assert non_default_params["cache_control_injection_points"] == [ + {"location": "tool_config", "_litellm_judged": True} + ] @pytest.mark.asyncio @@ -1753,17 +1756,16 @@ class TestConfiguredInjectionPointsStandDown: V1_MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] - def _seed(self, params, messages, tools=None, is_first_pass=True): + def _seed(self, params, messages, tools=None): AnthropicCacheControlHook.maybe_seed_default_injection_points( non_default_params=params, messages=messages, model="claude-sonnet-4-5", custom_llm_provider="anthropic", tools=tools, - is_first_pass=is_first_pass, ) - def _inject(self, messages, kwargs, system="sys", tools=None, is_first_pass=True): + def _inject(self, messages, kwargs, system="sys", tools=None): return AnthropicCacheControlHook.maybe_inject_cache_control( messages, system, @@ -1771,7 +1773,6 @@ class TestConfiguredInjectionPointsStandDown: model="claude-sonnet-4-5", custom_llm_provider="anthropic", tools=tools, - is_first_pass=is_first_pass, ) def test_configured_points_dropped_when_messages_carry_cache_control(self): @@ -1798,13 +1799,13 @@ class TestConfiguredInjectionPointsStandDown: self._seed(params, copy.deepcopy(self.CLEAN_MESSAGES)) assert params["cache_control_injection_points"] is configured - def test_second_pass_keeps_points_despite_injected_marks(self): + def test_judged_remainder_survives_reentry_despite_injected_marks(self): """acompletion() re-enters completion() after injection ran, with only the - non-message points written back; the second pass must not misread litellm's - own marks as client ones and drop that remainder.""" - remainder = [{"location": "tool_config"}] + stamped non-message points written back; the re-entry must not misread + litellm's own marks as client ones and drop that remainder.""" + remainder = [{"location": "tool_config", "_litellm_judged": True}] params = {"cache_control_injection_points": remainder} - self._seed(params, copy.deepcopy(self.MARKED_MESSAGES), is_first_pass=False) + self._seed(params, copy.deepcopy(self.MARKED_MESSAGES)) assert params["cache_control_injection_points"] is remainder def test_v1_messages_stand_down_when_content_block_marked(self): @@ -1841,18 +1842,23 @@ class TestConfiguredInjectionPointsStandDown: _, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) assert result_sys == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}] - def test_v1_messages_second_pass_writes_back_remainder(self): - """The async entry re-dispatches into the sync handler after injecting; the - surviving tool_config remainder must survive that second pass even though - the messages now carry litellm's own marks.""" - marked = [ - {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]} - ] - remainder = [{"location": "tool_config"}] - kwargs = {"cache_control_injection_points": remainder} - result_msgs, _ = self._inject(copy.deepcopy(marked), kwargs, is_first_pass=False) - assert result_msgs == marked - assert kwargs["cache_control_injection_points"] == remainder + def test_v1_messages_reentry_flow_preserves_tool_config_remainder(self): + """The advisor interceptor re-enters anthropic_messages() with the outer + request's kwargs and post-injection messages. The first pass applies the + message point and writes back a stamped tool_config remainder; the + re-entry must keep that remainder even though the messages and system + now carry litellm's own marks.""" + points = [{"location": "message", "role": "system"}, {"location": "tool_config"}] + kwargs = {"cache_control_injection_points": copy.deepcopy(points)} + msgs1, sys1 = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) + assert sys1[0]["cache_control"] == {"type": "ephemeral"} + expected_remainder = [{"location": "tool_config", "_litellm_judged": True}] + assert kwargs["cache_control_injection_points"] == expected_remainder + + msgs2, sys2 = self._inject(msgs1, kwargs, system=sys1) + assert kwargs["cache_control_injection_points"] == expected_remainder + assert msgs2 == msgs1 + assert sys2 == sys1 class TestAnthropicPromptCachingEnvVars: From 6d7a80ac755bf2e3c14e7b76af910ca28cea13c2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 23:59:18 -0700 Subject: [PATCH 11/60] feat(mcp): aggregate gateway DCR discovery front door behind mcp_gateway_dcr --- .../mcp_server/auth/user_api_key_auth_mcp.py | 156 ++++-- .../mcp_server/discoverable_endpoints.py | 98 ++++ .../_experimental/mcp_server/oauth_utils.py | 23 + .../auth/test_user_api_key_auth_mcp.py | 133 ++++++ .../mcp_server/test_discoverable_endpoints.py | 449 +++++------------- 5 files changed, 489 insertions(+), 370 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index d2f3efbc54e..418e63468b4 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 @@ -10,6 +10,10 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + get_request_base_url, + is_mcp_gateway_dcr_enabled, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, BridgeEnvelopeInvalid, @@ -120,6 +124,96 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +def _is_aggregate_gateway_dcr_challenge_scope( + route: str, + mcp_servers: list[str] | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + exc: Exception, +) -> bool: + """True when an unauthenticated request to the aggregate ``/mcp`` endpoint + should receive the RFC 9728 401 challenge that advertises the gateway as + the authorization server (``mcp_gateway_dcr`` front door). + + Fires only for a genuine 401 on the aggregate scope: any named target + (path or ``x-mcp-servers``) belongs to the per-server challenge paths, and + client-supplied MCP auth headers mean the caller is not a cold-start DCR + client. Fails closed to the original admission error otherwise.""" + if not is_mcp_gateway_dcr_enabled(): + return False + if not _is_litellm_auth_admission_error(exc): + return False + if mcp_servers: + return False + if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers): + return False + return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0 + + +def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> HTTPException: + """The RFC 9728 challenge for the aggregate endpoint: points the client at + the gateway's own protected-resource metadata so a DCR client discovers + the gateway as its authorization server and starts the sign-in flow. + + ``invalid_token`` adds the RFC 6750 error code for a request that DID + present a bearer that failed admission (expired or revoked), telling + spec-compliant clients to re-authorize rather than retry; a request with + no credentials at all gets the bare challenge per RFC 6750 section 3.1.""" + error_attr = 'error="invalid_token", ' if invalid_token else "" + resource_metadata_url = f"{get_request_base_url(request)}/.well-known/oauth-protected-resource/mcp" + return HTTPException( + status_code=401, + detail={ + "error": "authentication_required", + "message": "Authenticate with the gateway to use the MCP endpoint.", + }, + headers={"WWW-Authenticate": f'Bearer {error_attr}resource_metadata="{resource_metadata_url}"'}, + ) + + +def _admission_failure_fallback( + request: Request, + request_route: str, + mcp_servers: list[str] | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + exc: Exception, + bearer_presented: bool, +) -> UserAPIKeyAuth: + """Map a failed LiteLLM admission to its anonymous fallback or challenge. + + Two fallbacks exist, both gated on a genuine 401 with no client-supplied + MCP auth headers. The pass-through cold start (RFC 9728 / MCP + Authorization spec discovery return) admits anonymously so the route's + 401 emitter can produce the per-server challenge. The aggregate + gateway-DCR scope converts the failure into the gateway's own + resource_metadata challenge, with the RFC 6750 ``invalid_token`` error + code when the caller DID present a bearer (an expired gateway session + must re-authorize, not retry a dead token). Anything else re-raises the + original admission error unchanged.""" + mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) + if ( + mcp_servers_from_path is not None + and not _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers) + and _is_litellm_auth_admission_error(exc) + and _is_mcp_passthrough_cold_start( + mcp_servers_from_path, + client_ip=IPAddressUtils.get_mcp_client_ip(request), + ) + ): + verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") + return UserAPIKeyAuth() + if _is_aggregate_gateway_dcr_challenge_scope( + route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=exc, + ): + raise _aggregate_gateway_dcr_challenge(request, invalid_token=bearer_presented) from exc + raise exc + + class MCPRequestHandler: """ Class to handle MCP request processing, including: @@ -271,56 +365,32 @@ class MCPRequestHandler: elif oauth2_headers: # Authorization on a non-delegated server: the bearer must be a real # LiteLLM credential, so a failed validation is a genuine 401/403 and - # propagates. The sole anonymous fallback is the auth_type=none - # pass-through cold-start (RFC 9728 discovery return), gated on a 401 - # so a recognized-but-forbidden key still fails closed. - client_ip = IPAddressUtils.get_mcp_client_ip(request) + # propagates unless a fallback in _admission_failure_fallback applies. try: validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as e: - # ProxyException.code is normalized to str (possibly "None"), so - # compare both int and str forms rather than coercing. - status = e.status_code if isinstance(e, HTTPException) else e.code - is_unauthenticated = status in (401, "401") - mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) - if ( - is_unauthenticated - and mcp_servers_from_path is not None - and not _has_client_supplied_mcp_auth( - mcp_auth_header, - mcp_server_auth_headers, - ) - and _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) - ): - verbose_logger.debug( - "MCP pass-through return: forwarding Authorization as upstream OAuth token for delegated auth" - ) - validated_user_api_key_auth = UserAPIKeyAuth() - else: - raise + validated_user_api_key_auth = _admission_failure_fallback( + request=request, + request_route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=e, + bearer_presented=True, + ) else: try: validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as exc: - # Cold-start MCP OAuth discovery: RFC 9728 / MCP Authorization spec - # require unauthenticated requests to protected resources to receive - # 401 + WWW-Authenticate. Defer to _raise_preemptive_401_for_unauthenticated_servers - # for pass-through servers instead of surfacing a generic admission error. - mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) - client_ip = IPAddressUtils.get_mcp_client_ip(request) - if ( - mcp_servers_from_path is not None - and not _has_client_supplied_mcp_auth( - mcp_auth_header, - mcp_server_auth_headers, - ) - and _is_litellm_auth_admission_error(exc) - and _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) - ): - verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") - validated_user_api_key_auth = UserAPIKeyAuth() - else: - raise + validated_user_api_key_auth = _admission_failure_fallback( + request=request, + request_route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=exc, + bearer_presented=False, + ) return ( validated_user_api_key_auth, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 1af64749304..075ece42b04 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -42,6 +42,7 @@ from litellm.proxy._experimental.mcp_server.faults import ( from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, + is_mcp_gateway_dcr_enabled, validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -1708,6 +1709,12 @@ async def _build_oauth_protected_resource_response( global_mcp_server_manager, ) + # With the gateway-level DCR front door enabled, unnamed discovery + # describes the gateway itself as the authorization server for the + # aggregate /mcp resource instead of narrowing to one server. + if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): + return _build_aggregate_protected_resource_response(request) + request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) @@ -1838,6 +1845,92 @@ def _jwt_auth_issuers() -> list: return issuers +def _build_aggregate_protected_resource_response(request: Request) -> dict: + """RFC 9728 metadata for the aggregate /mcp resource: the gateway itself is + the authorization server. No per-server names or scopes leak here; access + is resolved after sign-in from the authenticated user's grants. + + The advertised authorization server is ``{base}/mcp`` (not the bare + origin) so RFC 8414 path-insertion resolves its metadata at + ``/.well-known/oauth-authorization-server/mcp``, a route this module + owns. The bare-origin well-known is registered first by the BYOK OAuth + feature and describes the BYOK flow, so it must not be the aggregate + discovery entry point (same pattern as the per-server documents, which + advertise ``{base}/{server_name}``).""" + request_base_url = get_request_base_url(request) + return { + "authorization_servers": [f"{request_base_url}/mcp"], + "resource": f"{request_base_url}/mcp", + "scopes_supported": [], + } + + +def _build_aggregate_authorization_server_response(request: Request) -> dict: + """RFC 8414 metadata for the gateway as the aggregate authorization server. + + The issuer is ``{base}/mcp`` and must stay equal to the value the + aggregate protected-resource document advertises: spec clients verify the + issuer in the metadata matches the one that derived the well-known URL. + Advertises the root /authorize, /token, and /register endpoints and + ``token_endpoint_auth_methods_supported: ["none", ...]`` because DCR + clients (Claude Desktop, MCP Inspector) register as public clients; PKCE + S256 is mandatory in the gateway's authorize flow.""" + request_base_url = get_request_base_url(request) + return { + "issuer": f"{request_base_url}/mcp", + "authorization_endpoint": f"{request_base_url}/authorize", + "token_endpoint": f"{request_base_url}/token", + "registration_endpoint": f"{request_base_url}/register", + "response_types_supported": ["code"], + "scopes_supported": [], + "grant_types_supported": ["authorization_code", "refresh_token"], + "code_challenge_methods_supported": ["S256"], + "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], + } + + +def _raise_404_unless_gateway_dcr_enabled() -> None: + """The aggregate well-known routes exist only under the gateway-level DCR + front door; flag-off they 404 exactly like the previously-absent routes so + discovery behavior is byte-identical for existing deployments.""" + if is_mcp_gateway_dcr_enabled(): + return + raise HTTPException(status_code=404, detail="Not Found") + + +# RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client +# pointed at {base}/mcp inserts the well-known segment before the resource +# path, so this exact route must exist for aggregate discovery to work at all. +# Declared before the parameterized well-known routes below: Starlette matches +# in registration order, and /.well-known/oauth-authorization-server/{name} +# would otherwise capture the "/mcp" suffix as a server name. +@router.get( + f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" +) +async def oauth_protected_resource_aggregate(request: Request): + """ + OAuth protected resource discovery for the aggregate /mcp endpoint + (gateway-level DCR front door; 404 when the flag is off). + """ + _raise_404_unless_gateway_dcr_enabled() + return _build_aggregate_protected_resource_response(request) + + +@router.get( + f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" +) +async def oauth_authorization_server_aggregate(request: Request): + """ + OAuth authorization server discovery for the aggregate /mcp endpoint, the + RFC 8414 path-inserted form for a client that treats {base}/mcp as its + authorization base URL (gateway-level DCR front door; 404 when the flag + is off, indistinguishable from an unknown server name on the + parameterized route below). + """ + _raise_404_unless_gateway_dcr_enabled() + return _build_aggregate_authorization_server_response(request) + + # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) @router.get( @@ -1897,6 +1990,11 @@ def _build_oauth_authorization_server_response( global_mcp_server_manager, ) + # With the gateway-level DCR front door enabled, unnamed discovery keeps + # advertising the gateway's own /authorize, /token, and /register. + if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): + return _build_aggregate_authorization_server_response(request) + request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 6edb22dd858..74b56cf424b 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -70,6 +70,29 @@ def _origin_label(scheme: str, netloc: str) -> str: return f"{scheme}://{netloc}" if netloc else f"{scheme}://" +MCP_GATEWAY_DCR_SETTING = "mcp_gateway_dcr" + + +def is_mcp_gateway_dcr_enabled() -> bool: + """True when ``general_settings.mcp_gateway_dcr`` opts this deployment into + the gateway-level DCR front door for the aggregate ``/mcp`` endpoint: root + OAuth discovery advertises the gateway itself as the authorization server + (instead of resolving the single configured oauth2 server), and the + anonymous aggregate 401 carries the RFC 9728 ``resource_metadata`` + challenge so DCR clients (Claude Desktop, MCP Inspector) can start the + sign-in flow. Off by default; flag-off behavior is unchanged.""" + from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load + + if not isinstance(general_settings, dict): + return False + raw = general_settings.get(MCP_GATEWAY_DCR_SETTING) + if isinstance(raw, bool): + return raw + if isinstance(raw, str): + return raw.strip().lower() == "true" + return False + + def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() 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 9375f7481c8..e1fd3bca7bb 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 @@ -6131,3 +6131,136 @@ class TestMCPDcrBridgeDelegateAdmission: route="/mcp/bridge_delegate_server", ) assert exc_info.value.status_code == 500 + + +@pytest.mark.asyncio +class TestAggregateGatewayDcrChallenge: + """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must + carry the RFC 9728 resource_metadata challenge pointing at the gateway's + own protected-resource metadata, and must NOT fire for named-server + targets, explicit litellm keys, non-401 failures, or with the flag off.""" + + _FLAG_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled" + _AUTH_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" + _EXPECTED_RESOURCE_METADATA = 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/mcp"' + + def _scope(self, path="/mcp", extra_headers=()): + return { + "type": "http", + "method": "POST", + "path": path, + "headers": [(b"host", b"testserver"), *extra_headers], + } + + def _auth_401(self): + async def _raise(api_key, request): + raise ProxyException( + message="Authentication Error: Invalid API key", + type="auth_error", + param="api_key", + code=401, + ) + + return _raise + + async def test_challenge_on_anonymous_aggregate_mcp(self): + """Anonymous request to the aggregate /mcp with the flag on: 401 plus + the bare bearer challenge (no error attribute, RFC 6750 section 3.1).""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + assert exc_info.value.status_code == 401 + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert www_authenticate == f"Bearer {self._EXPECTED_RESOURCE_METADATA}" + + async def test_challenge_invalid_token_on_failed_bearer(self): + """A bearer that fails LiteLLM admission at aggregate scope (an expired + gateway session, a revoked key) re-challenges with error=invalid_token + so a spec client re-authorizes instead of retrying the dead token.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"authorization", b"Bearer expired-session-token"),)) + ) + assert exc_info.value.status_code == 401 + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' + + async def test_no_challenge_when_flag_off(self): + """Flag off: the original admission error propagates untouched, both + with and without a bearer.""" + for extra_headers in ((), ((b"authorization", b"Bearer some-token"),)): + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=False), + ): + with pytest.raises(ProxyException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(extra_headers=extra_headers)) + assert str(exc_info.value.code) == "401" + + async def test_no_challenge_for_explicit_litellm_key(self): + """An explicit x-litellm-api-key declares a litellm-key client; a typo + there must surface the real auth error, never a DCR challenge that + would send SDKs into a sign-in flow.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-litellm-api-key", b"sk-typo"),)) + ) + + async def test_no_challenge_for_named_servers_header(self): + """x-mcp-servers names explicit targets; the per-server challenge paths + own those, so the aggregate challenge must not fire.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-mcp-servers", b"github"),)) + ) + + async def test_no_challenge_for_path_named_server(self): + """/mcp/{server} targets one server; the aggregate challenge must not + fire even when that server does not resolve.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request(self._scope(path="/mcp/github")) + + async def test_no_challenge_for_client_supplied_mcp_auth(self): + """Per-server x-mcp-{alias}-authorization headers mean the caller is + not a cold-start DCR client; keep the original error.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-mcp-github-authorization", b"Bearer upstream"),)) + ) + + async def test_no_challenge_for_non_401_failure(self): + """Only genuine 401s convert to a challenge; a 500 stays a 500.""" + + async def _raise_500(api_key, request): + raise ProxyException(message="boom", type="server_error", param=None, code=500) + + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=_raise_500), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + assert str(exc_info.value.code) == "500" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 6f2f24df8fa..30a814c8ea4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7132,371 +7132,166 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} +def _patch_gateway_dcr_flag(enabled: bool): + return patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.is_mcp_gateway_dcr_enabled", + return_value=enabled, + ) + + @pytest.mark.asyncio -async def test_persist_dcr_client_for_config_server_uses_side_store(): - """A config.yaml-declared OAuth2 DCR server has no LiteLLM_MCPServerTable row, so - update_mcp_server returns None. The minted client must then persist to the server-scoped - OAuth-client store keyed by server_id (never a shadow server row), overlay onto the in-memory - server so refresh can authenticate this process, and never call update_server(None) (which - previously raised AttributeError on .approval_status, was swallowed, and reported a 200 that - persisted nothing).""" +async def test_gateway_dcr_root_discovery_describes_gateway_not_single_server(): + """Flag on: root discovery must keep describing the gateway as the + authorization server for the aggregate /mcp resource even when exactly one + OAuth2 server exists (flag off, resolution narrows to that server; that + behavior is pinned by test_discovery_root_includes_server_name_prefix).""" + from fastapi import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _persist_dcr_client_registration, + _build_oauth_authorization_server_response, + _build_oauth_protected_resource_response, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - config_server = MCPServer( - server_id="config_faros", - name="config_faros", - server_name="config_faros", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - client_secret=None, - authorization_url="https://provider.example/oauth/authorize", - token_url="https://provider.example/oauth/token", - registration_url="https://provider.example/oauth/register", - ) + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - mock_upsert = AsyncMock() - mock_update_server = AsyncMock() + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} - with ( - patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch( - "litellm.proxy._experimental.mcp_server.db.update_mcp_server", - new=AsyncMock(return_value=None), - ), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server", - new=AsyncMock(return_value=None), - ), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=AsyncMock(return_value=None), - ), - patch( - "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", - new=mock_upsert, - ), - patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), - ): - result = await _persist_dcr_client_registration( - mcp_server=config_server, - registration_response={ - "client_id": "minted-client", - "client_secret": "minted-secret", - "token_endpoint_auth_method": "client_secret_basic", - }, - current_redirect_uri="https://proxy.litellm.example/callback", - ) + try: + with _patch_gateway_dcr_flag(True): + authorization_response = _build_oauth_authorization_server_response( + request=mock_request, + mcp_server_name=None, + ) + resource_response = await _build_oauth_protected_resource_response( + request=mock_request, + mcp_server_name=None, + use_standard_pattern=True, + ) - assert result == "persisted" + assert authorization_response["issuer"] == "https://llm.example.com/mcp" + assert authorization_response["authorization_endpoint"] == "https://llm.example.com/authorize" + assert authorization_response["token_endpoint"] == "https://llm.example.com/token" + assert authorization_response["registration_endpoint"] == "https://llm.example.com/register" + assert "none" in authorization_response["token_endpoint_auth_methods_supported"] + assert authorization_response["code_challenge_methods_supported"] == ["S256"] + assert authorization_response["scopes_supported"] == [] - mock_upsert.assert_called_once() - assert mock_upsert.call_args.kwargs["server_id"] == "config_faros" - stored = mock_upsert.call_args.kwargs["credentials"] - assert stored["client_id"] == "minted-client" - assert stored["client_secret"] == "minted-secret" - assert stored["token_endpoint_auth_method"] == "client_secret_basic" - assert stored["redirect_uris"] == ["https://proxy.litellm.example/callback"] - - assert config_server.client_id == "minted-client" - assert config_server.client_secret == "minted-secret" - assert config_server.token_endpoint_auth_method == "client_secret_basic" - - mock_update_server.assert_not_called() + assert resource_response["resource"] == "https://llm.example.com/mcp" + assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] + assert resource_response["scopes_supported"] == [] + finally: + global_mcp_server_manager.registry.clear() @pytest.mark.asyncio -async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch): - """On restart a config server's in-memory object has no client_id; hydration overlays the - persisted DCR client from the server-scoped store, decrypting the encrypted-at-rest blob, so the - refresh_token grant can authenticate as the registered client instead of re-authenticating.""" - import litellm.proxy.common_utils.encrypt_decrypt_utils as enc - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - from litellm.proxy._experimental.mcp_server.db import encrypt_credentials +async def test_gateway_dcr_named_discovery_unaffected_by_flag(): + """Flag on must not change named-server discovery: a named oauth2 server + still resolves to its own per-server document.""" + from fastapi import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - hydrate_config_server_dcr_client, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - server = MCPServer( - server_id="config_faros", - name="config_faros", - server_name="config_faros", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - ) - - monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-hydrate-key") - stored_blob = safe_dumps( - encrypt_credentials( - credentials={ - "client_id": "stored-client", - "client_secret": "stored-secret", - "token_endpoint_auth_method": "client_secret_basic", - "redirect_uris": ["https://proxy.litellm.example/callback"], - }, - encryption_key="salt-hydrate-key", - ) - ) - assert "stored-client" not in stored_blob and "stored-secret" not in stored_blob - - with ( - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=AsyncMock(return_value=stored_blob), - ), - ): - applied = await hydrate_config_server_dcr_client(server) - - assert applied is True - assert server.client_id == "stored-client" - assert server.client_secret == "stored-secret" - assert server.token_endpoint_auth_method == "client_secret_basic" - - -@pytest.mark.asyncio -async def test_reuse_config_server_reads_store_with_real_crypto(monkeypatch): - """A config-declared server (rowless) keeps its DCR client in the store, so the reuse read - resolves it from the store and decrypts the encrypted-at-rest client, mirroring the write path so - a re-authorize reuses the client instead of re-minting one.""" - import litellm.proxy.common_utils.encrypt_decrypt_utils as enc - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - from litellm.proxy._experimental.mcp_server.db import encrypt_credentials - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _reuse_persisted_dcr_client_if_available, + _build_oauth_authorization_server_response, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - server = MCPServer( - server_id="config_faros", - name="config_faros", - server_name="config_faros", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - ) + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-reuse-key") - blob = safe_dumps( - encrypt_credentials( - credentials={"client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, - encryption_key="salt-reuse-key", - ) - ) - assert "stored-client" not in blob - store_lookup = AsyncMock(return_value=blob) - with ( - patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=store_lookup, - ), - ): - result = await _reuse_persisted_dcr_client_if_available(server, current_redirect_uri="https://x/callback") + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} - assert result is True - assert server.client_id == "stored-client" - store_lookup.assert_awaited_once() + try: + with _patch_gateway_dcr_flag(True): + response = _build_oauth_authorization_server_response( + request=mock_request, + mcp_server_name="test_oauth", + ) + assert "/test_oauth/authorize" in response["authorization_endpoint"] + assert response["scopes_supported"] == ["read", "write"] + finally: + global_mcp_server_manager.registry.clear() -@pytest.mark.asyncio -async def test_temp_server_is_not_persisted_to_store(): - """A rowless server that is NOT config-declared (a throwaway /server/oauth/session server) must - not leave a permanent store row on persist, and the read must never consult the store for it. Its - minted client is overlaid in memory for the session only.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _persist_dcr_client_registration, - _reuse_persisted_dcr_client_if_available, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer +def test_aggregate_wellknown_routes_404_when_flag_off(): + """Flag off, the aggregate well-known routes answer 404 exactly like the + previously-absent routes: discovery behavior is byte-identical for + existing deployments.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient - temp = MCPServer( - server_id="temp-uuid", - name="temp", - server_name="temp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - authorization_url="https://p.example/authorize", - token_url="https://p.example/token", - registration_url="https://p.example/register", - ) + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router - upsert = AsyncMock() - store_read = AsyncMock(return_value=None) - with ( - patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=AsyncMock(return_value=None)), - patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), - patch("litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", new=upsert), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=store_read, - ), - patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), - ): - result = await _persist_dcr_client_registration( - temp, {"client_id": "temp-client", "client_secret": "s"}, "https://x/callback" - ) - reused = await _reuse_persisted_dcr_client_if_available( - MCPServer( - server_id="temp-uuid", - name="temp", - server_name="temp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - ), - current_redirect_uri="https://x/callback", - ) + app = FastAPI() + app.include_router(router) + client = TestClient(app) - assert result == "persisted" - assert temp.client_id == "temp-client" - upsert.assert_not_called() - store_read.assert_not_called() - assert reused is False + with _patch_gateway_dcr_flag(False): + assert client.get("/.well-known/oauth-protected-resource/mcp").status_code == 404 + assert client.get("/.well-known/oauth-authorization-server/mcp").status_code == 404 -@pytest.mark.asyncio -async def test_hydrate_does_not_overwrite_explicit_config_client_id(): - """An explicit client_id set in config.yaml wins: hydration must not overwrite it with a stale - persisted store client, and must not even read the store when config already supplied a client.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - hydrate_config_server_dcr_client, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer +def test_aggregate_wellknown_routes_serve_gateway_metadata_when_flag_on(): + """Flag on, both path-appended aggregate routes serve the gateway + documents. Exercises real routing, so this also pins registration order: + /.well-known/oauth-authorization-server/{name} would otherwise capture + the /mcp suffix as a server name and 404.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient - server = MCPServer( - server_id="config_static", - name="config_static", - server_name="config_static", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="explicit-from-config", - ) - store_read = AsyncMock( - return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []} - ) - with ( - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=store_read, - ), - ): - applied = await hydrate_config_server_dcr_client(server) - - assert applied is False - assert server.client_id == "explicit-from-config" - store_read.assert_not_called() - - -@pytest.mark.asyncio -async def test_reuse_does_not_inherit_store_client_when_a_row_exists(): - """Security: a server that HAS a LiteLLM_MCPServerTable row reads its DCR client only from that - row, never from the server-scoped store. server_id is caller-settable on create, so a submitted - server whose id collides with a config-declared server must not be able to load that config - server's client from the store and send it to its own token endpoint. A row that exists but has - no client_id yields no reusable client and must not fall back to the store.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _reuse_persisted_dcr_client_if_available, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - submitted = MCPServer( - server_id="collides_with_config", - name="submitted", - server_name="submitted", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - ) - - row_without_client = MagicMock() - row_without_client.credentials = None - row_without_client.server_id = "collides_with_config" - store_lookup = AsyncMock( - return_value={"client_id": "config-secret-client", "client_secret": "leak", "redirect_uris": []} - ) - with ( - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server", - new=AsyncMock(return_value=row_without_client), - ), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=store_lookup, - ), - ): - result = await _reuse_persisted_dcr_client_if_available(submitted, current_redirect_uri="https://x/callback") - - assert result is False - assert submitted.client_id is None - store_lookup.assert_not_called() - - -@pytest.mark.asyncio -async def test_load_servers_from_config_hydrates_dcr_clients(): - """load_servers_from_config must invoke DCR-client hydration so config servers pick up their - persisted client on startup; deleting the call site leaves a restarted server with no client_id - and forces re-authentication on every token expiry.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - hydrate_spy = AsyncMock() - with patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy): - await global_mcp_server_manager.load_servers_from_config({}) + global_mcp_server_manager.registry.clear() + app = FastAPI() + app.include_router(router) + client = TestClient(app) - hydrate_spy.assert_awaited_once() + with _patch_gateway_dcr_flag(True): + prm = client.get("/.well-known/oauth-protected-resource/mcp") + asm = client.get("/.well-known/oauth-authorization-server/mcp") + + assert prm.status_code == 200 + assert prm.json()["resource"] == "http://testserver/mcp" + assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] + + assert asm.status_code == 200 + assert asm.json()["issuer"] == "http://testserver/mcp" + assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" -@pytest.mark.asyncio -async def test_reload_servers_from_database_hydrates_dcr_clients(): - """load_servers_from_config runs before the DB connects at startup, so its hydration no-ops; - reload_servers_from_database runs after the DB connects and must hydrate config servers' persisted - DCR clients too, or a fresh pod has no client_id for a config server and forces re-authentication - on the first token refresh.""" - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, +def test_is_mcp_gateway_dcr_enabled_reads_general_settings(): + """The flag reader accepts YAML booleans and env-interpolated strings, and + fails closed on anything else.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + is_mcp_gateway_dcr_enabled, ) + from litellm.proxy.proxy_server import general_settings - prisma = MagicMock() - prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) - - hydrate_spy = AsyncMock() - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=prisma, - ), - patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy), + for raw, expected in ( + (True, True), + (False, False), + ("true", True), + ("True", True), + ("false", False), + ("yes", False), + (1, False), + (None, False), ): - await global_mcp_server_manager.reload_servers_from_database() + with patch.dict(general_settings, {"mcp_gateway_dcr": raw}): + assert is_mcp_gateway_dcr_enabled() is expected, f"raw={raw!r}" - hydrate_spy.assert_awaited_once() + with patch.dict(general_settings, {}, clear=True): + assert is_mcp_gateway_dcr_enabled() is False From 14b1647cd66e1da9216939a128b17a3268eee765 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 17:27:22 -0700 Subject: [PATCH 12/60] refactor(mcp): make the aggregate DCR front door always-on, remove the mcp_gateway_dcr flag The flag guarded no breaking change: the aggregate discovery lives at new /mcp-suffixed routes, the challenge only fires at aggregate scope, and the authorize/token/register/admission arms self-gate on the llm_dcrc_/llm_session_ prefixes. Bare-origin and per-server discovery are left exactly as they were, and a server literally named mcp keeps its own discovery via disambiguation, so turning it on for everyone changes nothing about existing flows. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 5 +- .../mcp_server/discoverable_endpoints.py | 55 +++-- .../_experimental/mcp_server/oauth_utils.py | 23 -- .../auth/test_user_api_key_auth_mcp.py | 24 +- .../mcp_server/test_discoverable_endpoints.py | 212 ++++++------------ 5 files changed, 104 insertions(+), 215 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 418e63468b4..815be2229fc 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 @@ -12,7 +12,6 @@ import litellm from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_request_base_url, - is_mcp_gateway_dcr_enabled, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, @@ -133,14 +132,12 @@ def _is_aggregate_gateway_dcr_challenge_scope( ) -> bool: """True when an unauthenticated request to the aggregate ``/mcp`` endpoint should receive the RFC 9728 401 challenge that advertises the gateway as - the authorization server (``mcp_gateway_dcr`` front door). + the authorization server. Fires only for a genuine 401 on the aggregate scope: any named target (path or ``x-mcp-servers``) belongs to the per-server challenge paths, and client-supplied MCP auth headers mean the caller is not a cold-start DCR client. Fails closed to the original admission error otherwise.""" - if not is_mcp_gateway_dcr_enabled(): - return False if not _is_litellm_auth_admission_error(exc): return False if mcp_servers: diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 075ece42b04..08827accf3c 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -42,7 +42,6 @@ from litellm.proxy._experimental.mcp_server.faults import ( from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, - is_mcp_gateway_dcr_enabled, validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -1709,12 +1708,6 @@ async def _build_oauth_protected_resource_response( global_mcp_server_manager, ) - # With the gateway-level DCR front door enabled, unnamed discovery - # describes the gateway itself as the authorization server for the - # aggregate /mcp resource instead of narrowing to one server. - if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): - return _build_aggregate_protected_resource_response(request) - request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) @@ -1889,13 +1882,20 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: } -def _raise_404_unless_gateway_dcr_enabled() -> None: - """The aggregate well-known routes exist only under the gateway-level DCR - front door; flag-off they 404 exactly like the previously-absent routes so - discovery behavior is byte-identical for existing deployments.""" - if is_mcp_gateway_dcr_enabled(): - return - raise HTTPException(status_code=404, detail="Not Found") +def _mcp_named_server_exists(request: Request) -> bool: + """True when a server literally named ``mcp`` is configured and visible to this caller. + + Its per-server authorization-server document is served at + ``/.well-known/oauth-authorization-server/mcp``, a single segment that collides with the + aggregate path. When such a server exists the real server wins the route, so that + deployment keeps its per-server discovery regardless of whether the aggregate front door + is on.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load + global_mcp_server_manager, + ) + + client_ip = IPAddressUtils.get_mcp_client_ip(request) + return global_mcp_server_manager.get_mcp_server_by_name("mcp", client_ip=client_ip) is not None # RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client @@ -1909,10 +1909,12 @@ def _raise_404_unless_gateway_dcr_enabled() -> None: ) async def oauth_protected_resource_aggregate(request: Request): """ - OAuth protected resource discovery for the aggregate /mcp endpoint - (gateway-level DCR front door; 404 when the flag is off). + OAuth protected resource discovery for the aggregate /mcp endpoint. + + The single-segment ``/mcp`` path does not collide with any per-server PRM pattern + (those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously + describes the aggregate resource. """ - _raise_404_unless_gateway_dcr_enabled() return _build_aggregate_protected_resource_response(request) @@ -1921,13 +1923,15 @@ async def oauth_protected_resource_aggregate(request: Request): ) async def oauth_authorization_server_aggregate(request: Request): """ - OAuth authorization server discovery for the aggregate /mcp endpoint, the - RFC 8414 path-inserted form for a client that treats {base}/mcp as its - authorization base URL (gateway-level DCR front door; 404 when the flag - is off, indistinguishable from an unknown server name on the - parameterized route below). + OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414 + path-inserted form for a client that treats {base}/mcp as its authorization base URL. + + This single-segment path collides with the parameterized ``/{mcp_server_name}`` route + below, so a server literally named ``mcp`` wins it and keeps its per-server discovery; + only when no such server exists is the aggregate document served. """ - _raise_404_unless_gateway_dcr_enabled() + if _mcp_named_server_exists(request): + return _build_oauth_authorization_server_response(request=request, mcp_server_name="mcp") return _build_aggregate_authorization_server_response(request) @@ -1990,11 +1994,6 @@ def _build_oauth_authorization_server_response( global_mcp_server_manager, ) - # With the gateway-level DCR front door enabled, unnamed discovery keeps - # advertising the gateway's own /authorize, /token, and /register. - if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): - return _build_aggregate_authorization_server_response(request) - request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 74b56cf424b..6edb22dd858 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -70,29 +70,6 @@ def _origin_label(scheme: str, netloc: str) -> str: return f"{scheme}://{netloc}" if netloc else f"{scheme}://" -MCP_GATEWAY_DCR_SETTING = "mcp_gateway_dcr" - - -def is_mcp_gateway_dcr_enabled() -> bool: - """True when ``general_settings.mcp_gateway_dcr`` opts this deployment into - the gateway-level DCR front door for the aggregate ``/mcp`` endpoint: root - OAuth discovery advertises the gateway itself as the authorization server - (instead of resolving the single configured oauth2 server), and the - anonymous aggregate 401 carries the RFC 9728 ``resource_metadata`` - challenge so DCR clients (Claude Desktop, MCP Inspector) can start the - sign-in flow. Off by default; flag-off behavior is unchanged.""" - from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load - - if not isinstance(general_settings, dict): - return False - raw = general_settings.get(MCP_GATEWAY_DCR_SETTING) - if isinstance(raw, bool): - return raw - if isinstance(raw, str): - return raw.strip().lower() == "true" - return False - - def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() 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 e1fd3bca7bb..568081e0673 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 @@ -6138,9 +6138,8 @@ class TestAggregateGatewayDcrChallenge: """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must carry the RFC 9728 resource_metadata challenge pointing at the gateway's own protected-resource metadata, and must NOT fire for named-server - targets, explicit litellm keys, non-401 failures, or with the flag off.""" + targets, explicit litellm keys, or non-401 failures.""" - _FLAG_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled" _AUTH_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" _EXPECTED_RESOURCE_METADATA = 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/mcp"' @@ -6164,11 +6163,10 @@ class TestAggregateGatewayDcrChallenge: return _raise async def test_challenge_on_anonymous_aggregate_mcp(self): - """Anonymous request to the aggregate /mcp with the flag on: 401 plus + """Anonymous request to the aggregate /mcp: 401 plus the bare bearer challenge (no error attribute, RFC 6750 section 3.1).""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(HTTPException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) @@ -6182,7 +6180,6 @@ class TestAggregateGatewayDcrChallenge: so a spec client re-authorizes instead of retrying the dead token.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(HTTPException) as exc_info: await MCPRequestHandler.process_mcp_request( @@ -6192,25 +6189,12 @@ class TestAggregateGatewayDcrChallenge: www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' - async def test_no_challenge_when_flag_off(self): - """Flag off: the original admission error propagates untouched, both - with and without a bearer.""" - for extra_headers in ((), ((b"authorization", b"Bearer some-token"),)): - with ( - patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=False), - ): - with pytest.raises(ProxyException) as exc_info: - await MCPRequestHandler.process_mcp_request(self._scope(extra_headers=extra_headers)) - assert str(exc_info.value.code) == "401" - async def test_no_challenge_for_explicit_litellm_key(self): """An explicit x-litellm-api-key declares a litellm-key client; a typo there must surface the real auth error, never a DCR challenge that would send SDKs into a sign-in flow.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6222,7 +6206,6 @@ class TestAggregateGatewayDcrChallenge: own those, so the aggregate challenge must not fire.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6234,7 +6217,6 @@ class TestAggregateGatewayDcrChallenge: fire even when that server does not resolve.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request(self._scope(path="/mcp/github")) @@ -6244,7 +6226,6 @@ class TestAggregateGatewayDcrChallenge: not a cold-start DCR client; keep the original error.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6259,7 +6240,6 @@ class TestAggregateGatewayDcrChallenge: with ( patch(self._AUTH_PATCH_TARGET, side_effect=_raise_500), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 30a814c8ea4..af4b2caca72 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7132,19 +7132,74 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} -def _patch_gateway_dcr_flag(enabled: bool): - return patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.is_mcp_gateway_dcr_enabled", - return_value=enabled, +def test_aggregate_wellknown_routes_serve_gateway_metadata(): + """Both path-appended aggregate routes serve the gateway documents. Exercises real + routing, so this also pins registration order: the parameterized + /.well-known/oauth-authorization-server/{name} route would otherwise capture the /mcp + suffix as a server name.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, ) + global_mcp_server_manager.registry.clear() + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + prm = client.get("/.well-known/oauth-protected-resource/mcp") + asm = client.get("/.well-known/oauth-authorization-server/mcp") + + assert prm.status_code == 200 + assert prm.json()["resource"] == "http://testserver/mcp" + assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] + + assert asm.status_code == 200 + assert asm.json()["issuer"] == "http://testserver/mcp" + assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" + assert "none" in asm.json()["token_endpoint_auth_methods_supported"] + + +def test_as_aggregate_route_prefers_a_real_server_named_mcp(): + """A server literally named ``mcp`` wins the single-segment + /.well-known/oauth-authorization-server/mcp route (it collides with the parameterized + /{server_name} route) and keeps its per-server discovery; the aggregate document is + served only when no such server exists.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + server_named_mcp = _create_oauth2_server(server_id="mcp_srv", name="mcp", server_name="mcp", alias="mcp") + global_mcp_server_manager.registry[server_named_mcp.server_id] = server_named_mcp + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + try: + asm = client.get("/.well-known/oauth-authorization-server/mcp") + assert asm.status_code == 200 + # the real server's own document (issuer is the bare origin, endpoint is /mcp/authorize), + # not the aggregate one (whose issuer would be {base}/mcp) + assert asm.json()["issuer"] == "http://testserver" + assert "/mcp/authorize" in asm.json()["authorization_endpoint"] + finally: + global_mcp_server_manager.registry.clear() + @pytest.mark.asyncio -async def test_gateway_dcr_root_discovery_describes_gateway_not_single_server(): - """Flag on: root discovery must keep describing the gateway as the - authorization server for the aggregate /mcp resource even when exactly one - OAuth2 server exists (flag off, resolution narrows to that server; that - behavior is pinned by test_discovery_root_includes_server_name_prefix).""" +async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): + """The always-on aggregate front door must not change bare-origin discovery: with one + oauth2 server configured, the no-suffix /.well-known/oauth-{authorization-server, + protected-resource} still resolves THAT server, so an existing single-server deployment's + discovery is unchanged. The aggregate document lives only at the /mcp-suffixed routes.""" from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -7164,134 +7219,15 @@ async def test_gateway_dcr_root_discovery_describes_gateway_not_single_server(): mock_request.headers = {} try: - with _patch_gateway_dcr_flag(True): - authorization_response = _build_oauth_authorization_server_response( - request=mock_request, - mcp_server_name=None, - ) - resource_response = await _build_oauth_protected_resource_response( - request=mock_request, - mcp_server_name=None, - use_standard_pattern=True, - ) - - assert authorization_response["issuer"] == "https://llm.example.com/mcp" - assert authorization_response["authorization_endpoint"] == "https://llm.example.com/authorize" - assert authorization_response["token_endpoint"] == "https://llm.example.com/token" - assert authorization_response["registration_endpoint"] == "https://llm.example.com/register" - assert "none" in authorization_response["token_endpoint_auth_methods_supported"] - assert authorization_response["code_challenge_methods_supported"] == ["S256"] - assert authorization_response["scopes_supported"] == [] - - assert resource_response["resource"] == "https://llm.example.com/mcp" - assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] - assert resource_response["scopes_supported"] == [] + authorization_response = _build_oauth_authorization_server_response( + request=mock_request, mcp_server_name=None + ) + resource_response = await _build_oauth_protected_resource_response( + request=mock_request, mcp_server_name=None, use_standard_pattern=True + ) + # per-server, not aggregate: the single server's name is in the endpoints + assert "/test_oauth/authorize" in authorization_response["authorization_endpoint"] + assert authorization_response["issuer"] == "https://llm.example.com" + assert resource_response["authorization_servers"] == ["https://llm.example.com/test_oauth"] finally: global_mcp_server_manager.registry.clear() - - -@pytest.mark.asyncio -async def test_gateway_dcr_named_discovery_unaffected_by_flag(): - """Flag on must not change named-server discovery: a named oauth2 server - still resolves to its own per-server document.""" - from fastapi import Request - - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _build_oauth_authorization_server_response, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - - global_mcp_server_manager.registry.clear() - oauth2_server = _create_oauth2_server() - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://llm.example.com/" - mock_request.headers = {} - - try: - with _patch_gateway_dcr_flag(True): - response = _build_oauth_authorization_server_response( - request=mock_request, - mcp_server_name="test_oauth", - ) - assert "/test_oauth/authorize" in response["authorization_endpoint"] - assert response["scopes_supported"] == ["read", "write"] - finally: - global_mcp_server_manager.registry.clear() - - -def test_aggregate_wellknown_routes_404_when_flag_off(): - """Flag off, the aggregate well-known routes answer 404 exactly like the - previously-absent routes: discovery behavior is byte-identical for - existing deployments.""" - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router - - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - with _patch_gateway_dcr_flag(False): - assert client.get("/.well-known/oauth-protected-resource/mcp").status_code == 404 - assert client.get("/.well-known/oauth-authorization-server/mcp").status_code == 404 - - -def test_aggregate_wellknown_routes_serve_gateway_metadata_when_flag_on(): - """Flag on, both path-appended aggregate routes serve the gateway - documents. Exercises real routing, so this also pins registration order: - /.well-known/oauth-authorization-server/{name} would otherwise capture - the /mcp suffix as a server name and 404.""" - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - - global_mcp_server_manager.registry.clear() - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - with _patch_gateway_dcr_flag(True): - prm = client.get("/.well-known/oauth-protected-resource/mcp") - asm = client.get("/.well-known/oauth-authorization-server/mcp") - - assert prm.status_code == 200 - assert prm.json()["resource"] == "http://testserver/mcp" - assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] - - assert asm.status_code == 200 - assert asm.json()["issuer"] == "http://testserver/mcp" - assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" - - -def test_is_mcp_gateway_dcr_enabled_reads_general_settings(): - """The flag reader accepts YAML booleans and env-interpolated strings, and - fails closed on anything else.""" - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - is_mcp_gateway_dcr_enabled, - ) - from litellm.proxy.proxy_server import general_settings - - for raw, expected in ( - (True, True), - (False, False), - ("true", True), - ("True", True), - ("false", False), - ("yes", False), - (1, False), - (None, False), - ): - with patch.dict(general_settings, {"mcp_gateway_dcr": raw}): - assert is_mcp_gateway_dcr_enabled() is expected, f"raw={raw!r}" - - with patch.dict(general_settings, {}, clear=True): - assert is_mcp_gateway_dcr_enabled() is False From 5e1050709dec103f021cd6246050fc7d10668012 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 22:37:35 -0700 Subject: [PATCH 13/60] fix(mcp): reserve mcp for the aggregate AS and root-path the discovery challenges Two RFC 9728 / 8414 discovery fixes on the aggregate front door, both raised by Bugbot on this PR The aggregate authorization-server document at /.well-known/oauth-authorization-server/mcp used to defer to a per-server row literally named "mcp", serving issuer {base} while the aggregate protected-resource document advertises {base}/mcp as its authorization server. A spec client following that chain fails the RFC 8414 issuer check and cannot sign in. The single segment /mcp is now reserved for the aggregate so the issuer stays {base}/mcp and matches the protected-resource document; a server named "mcp" keeps its standard two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp The 401 challenges built the resource_metadata URL as {base}/.well-known/oauth-protected-resource/mcp with no SERVER_ROOT_PATH segment, but the routes are registered with the path-inserted root segment, so a proxy mounted under a sub-path pointed DCR clients at a URL that 404s. Both the aggregate challenge and the pre-existing per-server pass-through challenge now derive the path from one well_known_root_suffix helper that the route registrations also use, so the advertised URL cannot drift from the served route --- .../mcp_server/auth/user_api_key_auth_mcp.py | 5 +- .../mcp_server/discoverable_endpoints.py | 54 +++++-------------- .../_experimental/mcp_server/oauth_utils.py | 12 +++++ .../proxy/_experimental/mcp_server/server.py | 6 ++- .../auth/test_user_api_key_auth_mcp.py | 17 ++++++ .../mcp_server/test_discoverable_endpoints.py | 45 ++++++++++++---- 6 files changed, 87 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 815be2229fc..f1fcc95c532 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 @@ -12,6 +12,7 @@ import litellm from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_request_base_url, + well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, @@ -157,7 +158,9 @@ def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> H spec-compliant clients to re-authorize rather than retry; a request with no credentials at all gets the bare challenge per RFC 6750 section 3.1.""" error_attr = 'error="invalid_token", ' if invalid_token else "" - resource_metadata_url = f"{get_request_base_url(request)}/.well-known/oauth-protected-resource/mcp" + resource_metadata_url = ( + f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp" + ) return HTTPException( status_code=401, detail={ diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 08827accf3c..9ea452b9aa8 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -43,6 +43,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, validate_trusted_redirect_uri, + well_known_root_suffix, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -50,7 +51,6 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body -from litellm.proxy.utils import get_server_root_path from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -1882,31 +1882,13 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: } -def _mcp_named_server_exists(request: Request) -> bool: - """True when a server literally named ``mcp`` is configured and visible to this caller. - - Its per-server authorization-server document is served at - ``/.well-known/oauth-authorization-server/mcp``, a single segment that collides with the - aggregate path. When such a server exists the real server wins the route, so that - deployment keeps its per-server discovery regardless of whether the aggregate front door - is on.""" - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load - global_mcp_server_manager, - ) - - client_ip = IPAddressUtils.get_mcp_client_ip(request) - return global_mcp_server_manager.get_mcp_server_by_name("mcp", client_ip=client_ip) is not None - - # RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client # pointed at {base}/mcp inserts the well-known segment before the resource # path, so this exact route must exist for aggregate discovery to work at all. # Declared before the parameterized well-known routes below: Starlette matches # in registration order, and /.well-known/oauth-authorization-server/{name} # would otherwise capture the "/mcp" suffix as a server name. -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp") async def oauth_protected_resource_aggregate(request: Request): """ OAuth protected resource discovery for the aggregate /mcp endpoint. @@ -1918,28 +1900,26 @@ async def oauth_protected_resource_aggregate(request: Request): return _build_aggregate_protected_resource_response(request) -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/mcp") async def oauth_authorization_server_aggregate(request: Request): """ OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414 path-inserted form for a client that treats {base}/mcp as its authorization base URL. - This single-segment path collides with the parameterized ``/{mcp_server_name}`` route - below, so a server literally named ``mcp`` wins it and keeps its per-server discovery; - only when no such server exists is the aggregate document served. + The single-segment /mcp is reserved for the aggregate so the discovery chain stays + consistent: the aggregate protected-resource document advertises {base}/mcp as its + authorization server, so the document served here must have issuer {base}/mcp. A server + literally named ``mcp`` therefore does not take this route; it keeps its standard + two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the + per-server row win here instead would serve an issuer of {base} against a resource that + advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door. """ - if _mcp_named_server_exists(request): - return _build_oauth_authorization_server_response(request=request, mcp_server_name="mcp") return _build_aggregate_authorization_server_response(request) # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp/{{mcp_server_name}}") async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_name: str): """ OAuth protected resource discovery endpoint using standard MCP URL pattern. @@ -1959,9 +1939,7 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam # LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp # Kept for backward compatibility with existing deployments -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/{{mcp_server_name}}/mcp") @router.get("/.well-known/oauth-protected-resource") async def oauth_protected_resource_mcp(request: Request, mcp_server_name: Optional[str] = None): """ @@ -2031,9 +2009,7 @@ def _build_oauth_authorization_server_response( # Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name} -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/mcp/{{mcp_server_name}}") async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_name: str): """ OAuth authorization server discovery endpoint using standard MCP URL pattern. @@ -2048,9 +2024,7 @@ async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_n # LiteLLM legacy pattern and root endpoint -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/{{mcp_server_name}}") @router.get("/.well-known/oauth-authorization-server") async def oauth_authorization_server_mcp(request: Request, mcp_server_name: Optional[str] = None): """ diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 6edb22dd858..ccee3fc8ac0 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -132,6 +132,18 @@ def get_request_base_url(request: Request) -> str: return urlunparse((scheme, _strip_default_port(scheme, netloc), parsed.path, "", "", "")) +def well_known_root_suffix() -> str: + """The ``SERVER_ROOT_PATH`` segment inserted into a ``.well-known`` path (RFC 8414 / 9728 + path insertion), empty for a root-mounted proxy or an explicit ``/``. + + The discovery route registrations and the 401 challenges that advertise those routes both + derive their path from this one function, so the ``resource_metadata`` URL a client is told + to fetch cannot drift from the route that actually serves it. + """ + root = os.getenv("SERVER_ROOT_PATH", "") + return "" if root == "/" else root + + def validate_loopback_redirect_uri(redirect_uri: str) -> None: """Require a loopback ``redirect_uri`` (OAuth 2.1 §4.1.2.1 + RFC 8252 §7.3 native-app pattern). MCP clients are native apps that listen on diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a8ab0937124..a9840bdc02f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -48,6 +48,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, + well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPToolResultError, @@ -3525,9 +3526,10 @@ if MCP_AVAILABLE: base_url = get_request_base_url(request) _path = scope.get("_original_path") or scope.get("path", "") or "" + suffix = well_known_root_suffix() if _path.startswith(f"/{server_name}/mcp"): - return f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" - return f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" + return f"{base_url}/.well-known/oauth-protected-resource{suffix}/{server_name}/mcp" + return f"{base_url}/.well-known/oauth-protected-resource{suffix}/mcp/{server_name}" def _get_passthrough_www_authenticate( scope: Scope, 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 568081e0673..7b05b8c9dd0 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 @@ -6189,6 +6189,23 @@ class TestAggregateGatewayDcrChallenge: www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' + async def test_challenge_inserts_server_root_path(self): + """With SERVER_ROOT_PATH set the resource_metadata URL must carry the same path-inserted + root segment the aggregate PRM route is registered with (both derive it from + well_known_root_suffix), so a DCR client behind a sub-path is pointed at a route that + exists instead of a 404. Regression: the challenge used to hard-code /mcp and omit the + root path the route inserts.""" + import os + + with ( + patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}), + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/litellm/mcp"' in www_authenticate + async def test_no_challenge_for_explicit_litellm_key(self): """An explicit x-litellm-api-key declares a litellm-key client; a typo there must surface the real auth error, never a DCR challenge that diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index af4b2caca72..7e9ff4692b5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7163,11 +7163,13 @@ def test_aggregate_wellknown_routes_serve_gateway_metadata(): assert "none" in asm.json()["token_endpoint_auth_methods_supported"] -def test_as_aggregate_route_prefers_a_real_server_named_mcp(): - """A server literally named ``mcp`` wins the single-segment - /.well-known/oauth-authorization-server/mcp route (it collides with the parameterized - /{server_name} route) and keeps its per-server discovery; the aggregate document is - served only when no such server exists.""" +def test_as_aggregate_route_reserves_mcp_for_the_aggregate(): + """The single-segment /.well-known/oauth-authorization-server/mcp is reserved for the + aggregate even when a server is literally named ``mcp``. The aggregate protected-resource + document advertises {base}/mcp as its authorization server, so the document served here + must carry issuer {base}/mcp for the RFC 8414 issuer check to pass. Letting the per-server + row win (issuer {base}) breaks that chain, so the aggregate wins and the mcp-named server + keeps its standard two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp.""" from fastapi import FastAPI from fastapi.testclient import TestClient @@ -7186,14 +7188,39 @@ def test_as_aggregate_route_prefers_a_real_server_named_mcp(): try: asm = client.get("/.well-known/oauth-authorization-server/mcp") assert asm.status_code == 200 - # the real server's own document (issuer is the bare origin, endpoint is /mcp/authorize), - # not the aggregate one (whose issuer would be {base}/mcp) - assert asm.json()["issuer"] == "http://testserver" - assert "/mcp/authorize" in asm.json()["authorization_endpoint"] + # the aggregate document, whose issuer matches what the aggregate PRM advertises + assert asm.json()["issuer"] == "http://testserver/mcp" + + prm = client.get("/.well-known/oauth-protected-resource/mcp") + assert prm.status_code == 200 + assert prm.json()["authorization_servers"] == [asm.json()["issuer"]] + + # the mcp-named server keeps its own document on the standard two-segment route + per_server = client.get("/.well-known/oauth-authorization-server/mcp/mcp") + assert per_server.status_code == 200 + assert "/mcp/authorize" in per_server.json()["authorization_endpoint"] finally: global_mcp_server_manager.registry.clear() +def test_well_known_root_suffix_reflects_server_root_path(): + """The single path segment both the discovery routes and the 401 challenges insert for RFC + 8414/9728 path insertion: empty for a root-mounted proxy or an explicit ``/``, the configured + path otherwise. Sharing this one function is what keeps the advertised resource_metadata URL + equal to the route that serves it.""" + import os + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import well_known_root_suffix + + with patch.dict(os.environ, {"SERVER_ROOT_PATH": ""}): + assert well_known_root_suffix() == "" + with patch.dict(os.environ, {"SERVER_ROOT_PATH": "/"}): + assert well_known_root_suffix() == "" + with patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}): + assert well_known_root_suffix() == "/litellm" + + @pytest.mark.asyncio async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): """The always-on aggregate front door must not change bare-origin discovery: with one From 70bc9523ba94615167b3728efeff6331306f4937 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 00:14:10 -0700 Subject: [PATCH 14/60] test(mcp): isolate MCP discovery tests from a leaked SERVER_ROOT_PATH tests/test_litellm/proxy/test_custom_proxy.py sets SERVER_ROOT_PATH at import time (its app mounts under a custom path) and never restores it, so in a shared shard the value leaks into the process. The discovery routes and the 401 challenges now read SERVER_ROOT_PATH to path-insert it where they previously ignored it, so a leaked value rewrites every resource_metadata URL and the exact-URL assertions in the delegate, pass-through, and aggregate challenge tests fail depending on shard order An autouse fixture clears SERVER_ROOT_PATH for the MCP discovery tests so they deterministically exercise the default root-mounted deployment; the tests that assert a sub-path deployment set the value explicitly within their own body. No assertion changed; the leak was invisible before only because the code ignored the variable --- .../_experimental/mcp_server/conftest.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/conftest.py diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py new file mode 100644 index 00000000000..b477bf3f406 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -0,0 +1,22 @@ +import os + +import pytest + + +@pytest.fixture(autouse=True) +def _hermetic_server_root_path(): + """Isolate MCP discovery tests from a leaked ``SERVER_ROOT_PATH``. + + ``tests/test_litellm/proxy/test_custom_proxy.py`` sets ``SERVER_ROOT_PATH`` at import time + (its app mounts under a custom path) and never restores it, so in a shared shard the value + leaks into this process. The discovery routes and the 401 challenges read it, so a leaked + value would silently rewrite every ``resource_metadata`` URL and make these tests depend on + shard ordering. Clearing it here pins the default (root-mounted) deployment; a test that + exercises a sub-path deployment sets the value explicitly within its own body. + """ + saved = os.environ.pop("SERVER_ROOT_PATH", None) + try: + yield + finally: + if saved is not None: + os.environ["SERVER_ROOT_PATH"] = saved From 8a57067d4a0e0f17605eb20911a4ba8b84598d35 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 17:04:00 -0700 Subject: [PATCH 15/60] refactor(mcp): scope the root-path helper to the aggregate front door only The SERVER_ROOT_PATH fix for the per-server pass-through challenge belongs with its sibling in exceptions.py (both fabricate a per-server resource_metadata URL and both omit the root segment), and both are pre-existing paths unrelated to the aggregate discovery this PR adds. Reverting the server.py change keeps this PR to the aggregate front door and avoids leaving the two per-server challenge builders inconsistent; the per-server root-path fix lands as its own change covering both sites. --- .../proxy/_experimental/mcp_server/server.py | 6 +- .../mcp_server/test_discoverable_endpoints.py | 370 ++++++++++++++++++ 2 files changed, 372 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a9840bdc02f..a8ab0937124 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -48,7 +48,6 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, - well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPToolResultError, @@ -3526,10 +3525,9 @@ if MCP_AVAILABLE: base_url = get_request_base_url(request) _path = scope.get("_original_path") or scope.get("path", "") or "" - suffix = well_known_root_suffix() if _path.startswith(f"/{server_name}/mcp"): - return f"{base_url}/.well-known/oauth-protected-resource{suffix}/{server_name}/mcp" - return f"{base_url}/.well-known/oauth-protected-resource{suffix}/mcp/{server_name}" + return f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" + return f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" def _get_passthrough_www_authenticate( scope: Scope, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 7e9ff4692b5..eb8b4a89721 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7132,6 +7132,376 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} +@pytest.mark.asyncio +async def test_persist_dcr_client_for_config_server_uses_side_store(): + """A config.yaml-declared OAuth2 DCR server has no LiteLLM_MCPServerTable row, so + update_mcp_server returns None. The minted client must then persist to the server-scoped + OAuth-client store keyed by server_id (never a shadow server row), overlay onto the in-memory + server so refresh can authenticate this process, and never call update_server(None) (which + previously raised AttributeError on .approval_status, was swallowed, and reported a 200 that + persisted nothing).""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + config_server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + mock_upsert = AsyncMock() + mock_update_server = AsyncMock() + + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.update_mcp_server", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", + new=mock_upsert, + ), + patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), + ): + result = await _persist_dcr_client_registration( + mcp_server=config_server, + registration_response={ + "client_id": "minted-client", + "client_secret": "minted-secret", + "token_endpoint_auth_method": "client_secret_basic", + }, + current_redirect_uri="https://proxy.litellm.example/callback", + ) + + assert result == "persisted" + + mock_upsert.assert_called_once() + assert mock_upsert.call_args.kwargs["server_id"] == "config_faros" + stored = mock_upsert.call_args.kwargs["credentials"] + assert stored["client_id"] == "minted-client" + assert stored["client_secret"] == "minted-secret" + assert stored["token_endpoint_auth_method"] == "client_secret_basic" + assert stored["redirect_uris"] == ["https://proxy.litellm.example/callback"] + + assert config_server.client_id == "minted-client" + assert config_server.client_secret == "minted-secret" + assert config_server.token_endpoint_auth_method == "client_secret_basic" + + mock_update_server.assert_not_called() + + +@pytest.mark.asyncio +async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch): + """On restart a config server's in-memory object has no client_id; hydration overlays the + persisted DCR client from the server-scoped store, decrypting the encrypted-at-rest blob, so the + refresh_token grant can authenticate as the registered client instead of re-authenticating.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + hydrate_config_server_dcr_client, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-hydrate-key") + stored_blob = safe_dumps( + encrypt_credentials( + credentials={ + "client_id": "stored-client", + "client_secret": "stored-secret", + "token_endpoint_auth_method": "client_secret_basic", + "redirect_uris": ["https://proxy.litellm.example/callback"], + }, + encryption_key="salt-hydrate-key", + ) + ) + assert "stored-client" not in stored_blob and "stored-secret" not in stored_blob + + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=stored_blob), + ), + ): + applied = await hydrate_config_server_dcr_client(server) + + assert applied is True + assert server.client_id == "stored-client" + assert server.client_secret == "stored-secret" + assert server.token_endpoint_auth_method == "client_secret_basic" + + +@pytest.mark.asyncio +async def test_reuse_config_server_reads_store_with_real_crypto(monkeypatch): + """A config-declared server (rowless) keeps its DCR client in the store, so the reuse read + resolves it from the store and decrypts the encrypted-at-rest client, mirroring the write path so + a re-authorize reuses the client instead of re-minting one.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-reuse-key") + blob = safe_dumps( + encrypt_credentials( + credentials={"client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, + encryption_key="salt-reuse-key", + ) + ) + assert "stored-client" not in blob + store_lookup = AsyncMock(return_value=blob) + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_lookup, + ), + ): + result = await _reuse_persisted_dcr_client_if_available(server, current_redirect_uri="https://x/callback") + + assert result is True + assert server.client_id == "stored-client" + store_lookup.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_temp_server_is_not_persisted_to_store(): + """A rowless server that is NOT config-declared (a throwaway /server/oauth/session server) must + not leave a permanent store row on persist, and the read must never consult the store for it. Its + minted client is overlaid in memory for the session only.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + temp = MCPServer( + server_id="temp-uuid", + name="temp", + server_name="temp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + authorization_url="https://p.example/authorize", + token_url="https://p.example/token", + registration_url="https://p.example/register", + ) + + upsert = AsyncMock() + store_read = AsyncMock(return_value=None) + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", new=upsert), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_read, + ), + patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + ): + result = await _persist_dcr_client_registration( + temp, {"client_id": "temp-client", "client_secret": "s"}, "https://x/callback" + ) + reused = await _reuse_persisted_dcr_client_if_available( + MCPServer( + server_id="temp-uuid", + name="temp", + server_name="temp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ), + current_redirect_uri="https://x/callback", + ) + + assert result == "persisted" + assert temp.client_id == "temp-client" + upsert.assert_not_called() + store_read.assert_not_called() + assert reused is False + + +@pytest.mark.asyncio +async def test_hydrate_does_not_overwrite_explicit_config_client_id(): + """An explicit client_id set in config.yaml wins: hydration must not overwrite it with a stale + persisted store client, and must not even read the store when config already supplied a client.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + hydrate_config_server_dcr_client, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_static", + name="config_static", + server_name="config_static", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="explicit-from-config", + ) + store_read = AsyncMock( + return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []} + ) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_read, + ), + ): + applied = await hydrate_config_server_dcr_client(server) + + assert applied is False + assert server.client_id == "explicit-from-config" + store_read.assert_not_called() + + +@pytest.mark.asyncio +async def test_reuse_does_not_inherit_store_client_when_a_row_exists(): + """Security: a server that HAS a LiteLLM_MCPServerTable row reads its DCR client only from that + row, never from the server-scoped store. server_id is caller-settable on create, so a submitted + server whose id collides with a config-declared server must not be able to load that config + server's client from the store and send it to its own token endpoint. A row that exists but has + no client_id yields no reusable client and must not fall back to the store.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + submitted = MCPServer( + server_id="collides_with_config", + name="submitted", + server_name="submitted", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + row_without_client = MagicMock() + row_without_client.credentials = None + row_without_client.server_id = "collides_with_config" + store_lookup = AsyncMock( + return_value={"client_id": "config-secret-client", "client_secret": "leak", "redirect_uris": []} + ) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=AsyncMock(return_value=row_without_client), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_lookup, + ), + ): + result = await _reuse_persisted_dcr_client_if_available(submitted, current_redirect_uri="https://x/callback") + + assert result is False + assert submitted.client_id is None + store_lookup.assert_not_called() + + +@pytest.mark.asyncio +async def test_load_servers_from_config_hydrates_dcr_clients(): + """load_servers_from_config must invoke DCR-client hydration so config servers pick up their + persisted client on startup; deleting the call site leaves a restarted server with no client_id + and forces re-authentication on every token expiry.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + hydrate_spy = AsyncMock() + with patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy): + await global_mcp_server_manager.load_servers_from_config({}) + + hydrate_spy.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_reload_servers_from_database_hydrates_dcr_clients(): + """load_servers_from_config runs before the DB connects at startup, so its hydration no-ops; + reload_servers_from_database runs after the DB connects and must hydrate config servers' persisted + DCR clients too, or a fresh pod has no client_id for a config server and forces re-authentication + on the first token refresh.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + + hydrate_spy = AsyncMock() + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma, + ), + patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy), + ): + await global_mcp_server_manager.reload_servers_from_database() + + hydrate_spy.assert_awaited_once() + + def test_aggregate_wellknown_routes_serve_gateway_metadata(): """Both path-appended aggregate routes serve the gateway documents. Exercises real routing, so this also pins registration order: the parameterized From ab02127b50a26ce9cd2a2ff33eed6a0ab7ff3123 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 18 Jul 2026 18:56:24 -0700 Subject: [PATCH 16/60] fix(proxy): treat malformed cost-map token limits as absent on /v1/models create_model_info_response cast cost-map max_input_tokens / max_output_tokens with unguarded int(). The surrounding try/except covers only the get_model_info lookup, so a deployment whose model_info carries a non-numeric limit (e.g. "128,000" or an empty string) raised inside the per-model listing loop and failed the entire GET /v1/models and /models response with a 500, taking healthy deployments down with it. A deployment's model_info is registered into litellm.model_cost verbatim, so the malformed value reaches the cost map and not just the router index. Router.get_configured_token_limits already coerced this safely for the deployment path; the cost-map path was missed, so the two together still regressed. Both now share coerce_token_limit in litellm_core_utils, which returns None for a malformed value so the listing omits that one limit instead of failing, matching the graceful degradation the endpoint had before the cost-map switch. --- litellm/litellm_core_utils/core_helpers.py | 29 +++++++++ litellm/proxy/utils.py | 9 +-- litellm/router.py | 13 +--- tests/test_litellm/proxy/test_proxy_utils.py | 64 +++++++++++++++++++- 4 files changed, 98 insertions(+), 17 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 002a46771e3..88dddb59cc7 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -57,6 +57,35 @@ def safe_divide( return numerator / denominator +def coerce_token_limit(value: object) -> int | None: + """ + Coerce a max_input_tokens / max_output_tokens value to an int, treating a + malformed value as absent. + + A deployment's model_info is registered into litellm.model_cost verbatim, so a + config value like "128,000" or "" reaches the /v1/models listing uncoerced from + both the router index and the cost map. Returning None omits that one limit + instead of failing the whole listing. + + Args: + value: The raw configured or cost-map value + + Returns: + The value as an int, or None if it is missing or not a usable number. + Bools are rejected because True/False is never a meaningful token limit. + """ + if isinstance(value, bool): + return None + if isinstance(value, int): + return value + if isinstance(value, (str, float)): + try: + return int(value) + except (TypeError, ValueError, OverflowError): + return None + return None + + _FINISH_REASON_MAP: dict[str, OpenAIChatCompletionFinishReason] = { # Anthropic "stop_sequence": "stop", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7a52fdfdb87..844d3c2ed26 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -104,6 +104,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert +from litellm.litellm_core_utils.core_helpers import coerce_token_limit from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -6128,12 +6129,8 @@ def create_model_info_response( 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) + max_input_tokens = coerce_token_limit(model_cost_info.get("max_input_tokens")) + max_output_tokens = coerce_token_limit(model_cost_info.get("max_output_tokens")) if llm_router is not None: configured_input, configured_output = llm_router.get_configured_token_limits(model_id) diff --git a/litellm/router.py b/litellm/router.py index 9e44edb1fb9..ae3f7ba11c2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -68,6 +68,7 @@ from litellm.litellm_core_utils.request_timeout_resolver import ( ) from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, + coerce_token_limit, get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.coroutine_checker import coroutine_checker @@ -8544,18 +8545,10 @@ class Router: if deployment is None: return (None, None) - def _as_int(value: object) -> "int | None": - if value is None or isinstance(value, bool): - return None - try: - return int(value) - except (TypeError, ValueError): - return None - model_info = deployment.model_info return ( - _as_int(model_info.get("max_input_tokens")), - _as_int(model_info.get("max_output_tokens")), + coerce_token_limit(model_info.get("max_input_tokens")), + coerce_token_limit(model_info.get("max_output_tokens")), ) def get_deployment_credentials_with_provider( diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9486646ea4a..5ace46fc775 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -478,11 +478,12 @@ class TestPostCallFailureHookLiftsRecoveredPartialSpend: from typing import cast +import litellm from litellm.proxy.utils import create_model_info_response from litellm.types.utils import ModelInfo -def _fake_model_info(**fields: int) -> ModelInfo: +def _fake_model_info(**fields: object) -> ModelInfo: return cast(ModelInfo, dict(fields)) @@ -581,6 +582,67 @@ def test_create_model_info_response_survives_malformed_configured_limits(): assert "max_output_tokens" not in response +@pytest.mark.parametrize("bad_value", ["128,000", "", "unlimited", [128000], {"max": 128000}, True]) +def test_create_model_info_response_survives_malformed_cost_map_limits(bad_value): + 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=bad_value, max_output_tokens=bad_value + ), + ) + + assert response["id"] == "some-model" + assert "max_input_tokens" not in response + assert "max_output_tokens" not in response + + +def test_create_model_info_response_keeps_valid_cost_map_limit_beside_malformed_one(): + 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="128,000", max_output_tokens=16384 + ), + ) + + assert "max_input_tokens" not in response + assert response["max_output_tokens"] == 16384 + + +def test_create_model_info_response_survives_malformed_limits_registered_by_router(): + """A deployment's model_info is registered into litellm.model_cost verbatim, so a + malformed configured limit reaches the listing through the real cost-map lookup and + not just the router index. Guarding only the index path still 500s the whole listing.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "openai/some-unmapped-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"max_input_tokens": "128,000"}, + } + ] + ) + + response = create_model_info_response( + model_id="openai/some-unmapped-model", + provider="openai", + llm_router=router, + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["id"] == "openai/some-unmapped-model" + assert "max_input_tokens" not in response + + def test_create_model_info_response_emits_integer_token_counts(): response = create_model_info_response( model_id="some-model", From f2e340cf2bfce78928fd6377d8afd2f9cb2e6b84 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:12:00 -0700 Subject: [PATCH 17/60] feat(rust): port BaseAWSLLM auth (credential resolution + SigV4) to litellm-core as a base provider (#33888) * feat(rust): add feature-gated Bedrock AWS auth Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(rust): move Bedrock auth into core Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(rust): fall through caller identity lookup errors Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(rust): add live Bedrock proof and CI coverage Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(rust): share in-memory cache with Bedrock auth Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(rust): preserve web identity credential expiry Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- .github/workflows/test-rust.yml | 6 + litellm-rust/Cargo.lock | 922 +++++++++++++++++- litellm-rust/crates/core/Cargo.toml | 19 + .../core/src/caching/in_memory_cache.rs | 258 +++++ litellm-rust/crates/core/src/caching/mod.rs | 1 + litellm-rust/crates/core/src/lib.rs | 1 + .../core/src/providers/bedrock/aws_base.rs | 724 ++++++++++++++ .../core/src/providers/bedrock/constants.rs | 14 + .../crates/core/src/providers/bedrock/mod.rs | 6 + litellm-rust/crates/core/src/providers/mod.rs | 2 + .../llms/bedrock/test_base_aws_llm.py | 25 + 11 files changed, 1930 insertions(+), 48 deletions(-) create mode 100644 litellm-rust/crates/core/src/caching/in_memory_cache.rs create mode 100644 litellm-rust/crates/core/src/caching/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/bedrock/aws_base.rs create mode 100644 litellm-rust/crates/core/src/providers/bedrock/constants.rs create mode 100644 litellm-rust/crates/core/src/providers/bedrock/mod.rs diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 13e1dc4ad5e..21e1bcb90c6 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -61,5 +61,11 @@ jobs: - name: Run Clippy run: cargo clippy --workspace --all-targets --locked -- -D warnings + - name: Run Clippy with Bedrock auth + run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings + - name: Run Rust tests run: cargo test --workspace --locked + + - name: Run core tests with Bedrock auth + run: cargo test -p litellm-core --features bedrock-auth --locked diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7daff1dfc91..402d16715a3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "arc-swap" +version = "1.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" +dependencies = [ + "rustversion", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -19,6 +28,358 @@ 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 = "aws-config" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47712fde1909402600ccfbb26e47d482d2e58bb9e9e603d9f17e67cc435a6319" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-sdk-sts", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 1.4.2", + "time", + "tokio", + "tracing", + "url", +] + +[[package]] +name = "aws-credential-types" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e93964ffdaf57857f544be3666a5f57570bb699e934700f11b49708f61bb556e" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "zeroize", +] + +[[package]] +name = "aws-lc-rs" +version = "1.17.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00bdb5da18dac48ca2cc7cd4a98e533e8635a58e2361d13a1a4ee3888e0d72f1" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.43.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43103168cc76fe62678a375e722fc9cb3a0146159ac5828bc4f0dfd755c2224c" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", + "pkg-config", +] + +[[package]] +name = "aws-runtime" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7816e98ee912159f45d307e5ee6bfea4a335a55aee15f7f3e32f81a6f3000f1d" +dependencies = [ + "aws-credential-types", + "aws-sigv4", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "bytes-utils", + "fastrand", + "http 1.4.2", + "http-body 1.0.1", + "percent-encoding", + "pin-project-lite", + "tracing", + "uuid", +] + +[[package]] +name = "aws-sdk-sts" +version = "1.108.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c72b08911d8128dd360fe1b22a9fec0fa8b552dde8ec828dcf20ef5ec974e9f" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-query", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-smithy-xml", + "aws-types", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sigv4" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "723c2234ad7511ceef63eab016b7ba6ff7c55590fefb96fa8467af014a07309f" +dependencies = [ + "aws-credential-types", + "aws-smithy-http", + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "form_urlencoded", + "hex", + "hmac", + "http 0.2.12", + "http 1.4.2", + "percent-encoding", + "sha2 0.11.0", + "time", + "tracing", +] + +[[package]] +name = "aws-smithy-async" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f02e407fb3b54891734224b9ffac8a71fdd35f542500fa1af95754a6b2beb316" +dependencies = [ + "futures-util", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "aws-smithy-http" +version = "0.64.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "bytes-utils", + "futures-core", + "futures-util", + "http 1.4.2", + "http-body 1.0.1", + "http-body-util", + "percent-encoding", + "pin-project-lite", + "pin-utils", + "tracing", +] + +[[package]] +name = "aws-smithy-http-client" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "635d23afda0a6ab48d666c4d447c4873e8d1e83518a2be2093122397e50b838e" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "h2 0.3.27", + "h2 0.4.15", + "http 0.2.12", + "http 1.4.2", + "http-body 0.4.6", + "hyper 0.14.32", + "hyper 1.10.1", + "hyper-rustls 0.24.2", + "hyper-rustls 0.27.9", + "hyper-util", + "pin-project-lite", + "rustls 0.21.12", + "rustls 0.23.41", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls 0.26.4", + "tower", + "tracing", +] + +[[package]] +name = "aws-smithy-json" +version = "0.63.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dc65a121adb4b33729919fcfa14fa36fb33c1555a8f06bb0e2188dbfdc1d9ef" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", +] + +[[package]] +name = "aws-smithy-observability" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e86338c869539a581bf161247762a6e87f92c5c075060057b5ed6d06632ed0c" +dependencies = [ + "aws-smithy-runtime-api", +] + +[[package]] +name = "aws-smithy-query" +version = "0.61.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd22a6ba36e3f113cb8d5b3d1fe0ed31c76ee608ef63322d753bb8d2c9479e77" +dependencies = [ + "aws-smithy-types", + "urlencoding", +] + +[[package]] +name = "aws-smithy-runtime" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bea94a9ff8464016338c851e24b472d7131c388c88898a502e781815b2ee6045" +dependencies = [ + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-http-client", + "aws-smithy-observability", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "http-body 0.4.6", + "http-body 1.0.1", + "http-body-util", + "pin-project-lite", + "pin-utils", + "tokio", + "tracing", +] + +[[package]] +name = "aws-smithy-runtime-api" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ed1ebe6e0a95ea84570225f5a8208dec4b8f77e61a9b0d6f51773fcb4612f0" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api-macros", + "aws-smithy-types", + "bytes", + "http 0.2.12", + "http 1.4.2", + "pin-project-lite", + "tokio", + "tracing", + "zeroize", +] + +[[package]] +name = "aws-smithy-runtime-api-macros" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "221eaa237ddf1ca79b60d1372aad77e47f9c0ea5b3ce5099da8c61d027dc77b3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "aws-smithy-schema" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "http 1.4.2", +] + +[[package]] +name = "aws-smithy-types" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6dc683efb34b9e755675b37fedbe0103141e5b6df7bdc9eb6967756a8c167d8" +dependencies = [ + "base64-simd", + "bytes", + "bytes-utils", + "futures-core", + "http 0.2.12", + "http 1.4.2", + "http-body 0.4.6", + "http-body 1.0.1", + "http-body-util", + "itoa", + "num-integer", + "pin-project-lite", + "pin-utils", + "ryu", + "serde", + "time", + "tokio", + "tokio-util", +] + +[[package]] +name = "aws-smithy-xml" +version = "0.61.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea3f68eec3607f02acd24067969ce2abc6ba16aa7d5ce59ca450ed2fb5f78957" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "xmlparser", +] + +[[package]] +name = "aws-types" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e957a6c6dbce82b7a91f44231c09273159703769f447cbe85e854dfe9cf67f86" +dependencies = [ + "aws-credential-types", + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "rustc_version", + "tracing", +] + [[package]] name = "axum" version = "0.7.9" @@ -30,10 +391,10 @@ dependencies = [ "base64", "bytes", "futures-util", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.10.1", "hyper-util", "itoa", "matchit", @@ -65,8 +426,8 @@ dependencies = [ "async-trait", "bytes", "futures-util", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "http-body-util", "mime", "pin-project-lite", @@ -83,6 +444,16 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64-simd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "339abbe78e73178762e23bea9dfd08e697eb3f3301cd4be981c0f78ba5859195" +dependencies = [ + "outref", + "vsimd", +] + [[package]] name = "bitflags" version = "2.13.0" @@ -98,6 +469,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -116,6 +496,16 @@ version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +[[package]] +name = "bytes-utils" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dafe3a8757b027e2be6e4e5601ed563c55989fcf1546e933c66c8eb3a058d35" +dependencies = [ + "bytes", + "either", +] + [[package]] name = "cc" version = "1.2.65" @@ -123,6 +513,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -138,6 +530,27 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "core-foundation" version = "0.10.1" @@ -163,6 +576,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -173,20 +595,56 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + [[package]] name = "data-encoding" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" + [[package]] name = "digest" version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", + "ctutils", ] [[package]] @@ -200,12 +658,30 @@ dependencies = [ "syn", ] +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -227,6 +703,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures-channel" version = "0.3.32" @@ -320,11 +802,41 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 5.3.0", "wasip2", "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + +[[package]] +name = "h2" +version = "0.3.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" +dependencies = [ + "bytes", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http 0.2.12", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "h2" version = "0.4.15" @@ -336,7 +848,7 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http", + "http 1.4.2", "indexmap", "slab", "tokio", @@ -356,6 +868,32 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", +] + +[[package]] +name = "http" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", +] + [[package]] name = "http" version = "1.4.2" @@ -366,6 +904,17 @@ dependencies = [ "itoa", ] +[[package]] +name = "http-body" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" +dependencies = [ + "bytes", + "http 0.2.12", + "pin-project-lite", +] + [[package]] name = "http-body" version = "1.0.1" @@ -373,7 +922,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http", + "http 1.4.2", ] [[package]] @@ -384,8 +933,8 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", "futures-core", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "pin-project-lite", ] @@ -401,6 +950,39 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + +[[package]] +name = "hyper" +version = "0.14.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" +dependencies = [ + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.3.27", + "http 0.2.12", + "http-body 0.4.6", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "socket2 0.5.10", + "tokio", + "tower-service", + "tracing", + "want", +] + [[package]] name = "hyper" version = "1.10.1" @@ -411,9 +993,9 @@ dependencies = [ "bytes", "futures-channel", "futures-core", - "h2", - "http", - "http-body", + "h2 0.4.15", + "http 1.4.2", + "http-body 1.0.1", "httparse", "httpdate", "itoa", @@ -423,18 +1005,34 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-rustls" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" +dependencies = [ + "futures-util", + "http 0.2.12", + "hyper 0.14.32", + "log", + "rustls 0.21.12", + "tokio", + "tokio-rustls 0.24.1", +] + [[package]] name = "hyper-rustls" version = "0.27.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ - "http", - "hyper", + "http 1.4.2", + "hyper 1.10.1", "hyper-util", - "rustls", + "rustls 0.23.41", + "rustls-native-certs", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tower-service", "webpki-roots", ] @@ -449,14 +1047,14 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "http", - "http-body", - "hyper", + "http 1.4.2", + "http-body 1.0.1", + "hyper 1.10.1", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -587,6 +1185,16 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + [[package]] name = "js-sys" version = "0.3.103" @@ -617,7 +1225,7 @@ dependencies = [ "reqwest", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "subtle", "tokio", "tokio-tungstenite", @@ -628,9 +1236,17 @@ dependencies = [ name = "litellm-core" version = "0.1.0" dependencies = [ + "aws-config", + "aws-credential-types", + "aws-sdk-sts", + "aws-sigv4", + "aws-smithy-runtime-api", + "aws-types", "rand 0.8.6", + "reqwest", "serde", "serde_json", + "sha2 0.10.9", "thiserror 2.0.18", "tokio", ] @@ -694,6 +1310,30 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -706,6 +1346,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "outref" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -718,6 +1364,18 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + [[package]] name = "portable-atomic" version = "1.13.1" @@ -733,6 +1391,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -834,8 +1498,8 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls", - "socket2", + "rustls 0.23.41", + "socket2 0.6.4", "thiserror 2.0.18", "tokio", "tracing", @@ -854,7 +1518,7 @@ dependencies = [ "rand 0.9.4", "ring", "rustc-hash", - "rustls", + "rustls 0.23.41", "rustls-pki-types", "slab", "thiserror 2.0.18", @@ -872,7 +1536,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2", + "socket2 0.6.4", "tracing", "windows-sys 0.60.2", ] @@ -892,6 +1556,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rand" version = "0.8.6" @@ -951,6 +1621,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "regex-lite" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" + [[package]] name = "reqwest" version = "0.12.28" @@ -962,26 +1638,26 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "h2", - "http", - "http-body", + "h2 0.4.15", + "http 1.4.2", + "http-body 1.0.1", "http-body-util", - "hyper", - "hyper-rustls", + "hyper 1.10.1", + "hyper-rustls 0.27.9", "hyper-util", "js-sys", "log", "percent-encoding", "pin-project-lite", "quinn", - "rustls", + "rustls 0.23.41", "rustls-pki-types", "serde", "serde_json", "serde_urlencoded", "sync_wrapper", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tokio-util", "tower", "tower-http", @@ -1014,16 +1690,38 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + +[[package]] +name = "rustls" +version = "0.21.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" +dependencies = [ + "log", + "ring", + "rustls-webpki 0.101.7", + "sct", +] + [[package]] name = "rustls" version = "0.23.41" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" dependencies = [ + "aws-lc-rs", "once_cell", "ring", "rustls-pki-types", - "rustls-webpki", + "rustls-webpki 0.103.13", "subtle", "zeroize", ] @@ -1050,12 +1748,23 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-webpki" +version = "0.101.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "rustls-webpki" version = "0.103.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -1082,6 +1791,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "sct" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -1105,6 +1824,12 @@ dependencies = [ "libc", ] +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + [[package]] name = "serde" version = "1.0.228" @@ -1178,8 +1903,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", ] [[package]] @@ -1189,8 +1914,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -1211,6 +1947,16 @@ version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "socket2" version = "0.6.4" @@ -1310,6 +2056,36 @@ dependencies = [ "syn", ] +[[package]] +name = "time" +version = "0.3.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18dfaaeddcb932337b5e7866ee7d0ce9b76d2fd092997146f187ec09b4558a50" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" + +[[package]] +name = "time-macros" +version = "0.2.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c431b87111666e491a90baa837f914fb45cd5dc3c268591b0220ff5057f2085f" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -1345,7 +2121,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", - "socket2", + "socket2 0.6.4", "tokio-macros", "windows-sys 0.61.2", ] @@ -1361,13 +2137,23 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-rustls" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" +dependencies = [ + "rustls 0.21.12", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls", + "rustls 0.23.41", "tokio", ] @@ -1379,11 +2165,11 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" dependencies = [ "futures-util", "log", - "rustls", + "rustls 0.23.41", "rustls-native-certs", "rustls-pki-types", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tungstenite", ] @@ -1425,8 +2211,8 @@ dependencies = [ "bitflags", "bytes", "futures-util", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "pin-project-lite", "tower", "tower-layer", @@ -1454,9 +2240,21 @@ checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "log", "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tracing-core" version = "0.1.36" @@ -1481,11 +2279,11 @@ dependencies = [ "byteorder", "bytes", "data-encoding", - "http", + "http 1.4.2", "httparse", "log", "rand 0.8.6", - "rustls", + "rustls 0.23.41", "rustls-pki-types", "sha1", "thiserror 1.0.69", @@ -1522,6 +2320,12 @@ dependencies = [ "serde", ] +[[package]] +name = "urlencoding" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" + [[package]] name = "utf-8" version = "0.7.6" @@ -1534,12 +2338,28 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "uuid" +version = "1.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "version_check" version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "vsimd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" + [[package]] name = "want" version = "0.3.1" @@ -1835,6 +2655,12 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "xmlparser" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" + [[package]] name = "yoke" version = "0.8.3" diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 9bd4634cc2a..65c6db7412c 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -10,6 +10,25 @@ rand.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true +sha2.workspace = true +aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } +aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true } +aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } +aws-sigv4 = { version = "1.5.1", optional = true } +aws-types = { version = "1.4.0", optional = true } +aws-smithy-runtime-api = { version = "1.13.0", optional = true } + +[features] +default = [] +bedrock-auth = [ + "dep:aws-config", + "dep:aws-credential-types", + "dep:aws-sdk-sts", + "dep:aws-sigv4", + "dep:aws-types", + "dep:aws-smithy-runtime-api", +] [dev-dependencies] +reqwest.workspace = true tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/core/src/caching/in_memory_cache.rs b/litellm-rust/crates/core/src/caching/in_memory_cache.rs new file mode 100644 index 00000000000..0ceeedb8b71 --- /dev/null +++ b/litellm-rust/crates/core/src/caching/in_memory_cache.rs @@ -0,0 +1,258 @@ +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashMap}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; +const DEFAULT_TTL: Duration = Duration::from_secs(600); + +pub struct InMemoryCache { + pub cache_dict: HashMap, + pub ttl_dict: HashMap, + pub expiration_heap: BinaryHeap>, + pub max_size_in_memory: usize, + pub default_ttl: Duration, + now: Box Duration + Send + Sync>, +} + +impl Default for InMemoryCache { + fn default() -> Self { + Self::new(None, None) + } +} + +impl InMemoryCache { + pub fn new(max_size_in_memory: Option, default_ttl: Option) -> Self { + Self::with_clock(max_size_in_memory, default_ttl, || { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + }) + } + + pub fn with_clock( + max_size_in_memory: Option, + default_ttl: Option, + now: impl Fn() -> Duration + Send + Sync + 'static, + ) -> Self { + Self { + cache_dict: HashMap::new(), + ttl_dict: HashMap::new(), + expiration_heap: BinaryHeap::new(), + max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + now: Box::new(now), + } + } + + pub fn evict_cache(&mut self) { + if self.max_size_in_memory == 0 { + return; + } + + let current_time = (self.now)(); + while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() { + if self.ttl_dict.get(&key).copied() != Some(expiration_time) { + self.expiration_heap.pop(); + } else if expiration_time <= current_time { + self.expiration_heap.pop(); + self.remove_key(&key); + } else { + break; + } + } + + while self.cache_dict.len() >= self.max_size_in_memory { + let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else { + break; + }; + if self.ttl_dict.get(&key).copied() == Some(expiration_time) { + self.remove_key(&key); + } + } + } + + pub fn allow_ttl_override(&self, key: &str) -> bool { + match self.ttl_dict.get(key).copied() { + None => true, + Some(expiration_time) => expiration_time < (self.now)(), + } + } + + pub fn set_cache(&mut self, key: impl Into, value: V, ttl: Option) { + if self.max_size_in_memory == 0 { + return; + } + + self.evict_cache(); + let key = key.into(); + self.cache_dict.insert(key.clone(), value); + if self.allow_ttl_override(&key) { + let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl); + self.ttl_dict.insert(key.clone(), expiration_time); + self.expiration_heap.push(Reverse((expiration_time, key))); + } + } + + // Generic values intentionally omit Python's per-item size check. + pub fn get_cache(&mut self, key: &str) -> Option { + if self.cache_dict.contains_key(key) { + if self.is_key_expired(key) { + self.remove_key(key); + return None; + } + return self.cache_dict.get(key).cloned(); + } + None + } + + pub fn get_ttl(&self, key: &str) -> Option { + self.ttl_dict.get(key).copied() + } + + pub fn delete_cache(&mut self, key: &str) { + self.remove_key(key); + } + + pub fn flush_cache(&mut self) { + self.cache_dict.clear(); + self.ttl_dict.clear(); + self.expiration_heap.clear(); + } + + fn is_key_expired(&self, key: &str) -> bool { + self.ttl_dict + .get(key) + .is_some_and(|expiration_time| *expiration_time < (self.now)()) + } + + fn remove_key(&mut self, key: &str) { + self.cache_dict.remove(key); + self.ttl_dict.remove(key); + } +} + +#[cfg(test)] +mod tests { + use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }; + + use super::InMemoryCache; + use std::time::Duration; + + fn cache(now: Arc, max_size: usize, default_ttl: Duration) -> InMemoryCache { + InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || { + Duration::from_secs(now.load(Ordering::Relaxed)) + }) + } + + #[test] + fn ttl_expiry_is_deterministic() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); + cache.set_cache("key", "value".to_string(), None); + assert_eq!(cache.get_cache("key"), Some("value".to_string())); + now.store(161, Ordering::Relaxed); + assert_eq!(cache.get_cache("key"), None); + assert_eq!(cache.get_ttl("key"), None); + } + + #[test] + fn default_and_per_set_ttl_are_applied() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); + cache.set_cache("default", "value".to_string(), None); + cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20))); + assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160))); + assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120))); + } + + #[test] + fn unexpired_entries_do_not_allow_ttl_override() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); + cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20))); + cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80))); + assert_eq!(cache.get_cache("key"), Some("second".to_string())); + assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120))); + now.store(121, Ordering::Relaxed); + cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80))); + assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201))); + } + + #[test] + fn max_size_evicts_earliest_expiration() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now, 2, Duration::from_secs(60)); + cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10))); + cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20))); + cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30))); + assert_eq!(cache.get_cache("early"), None); + assert!(cache.get_cache("late").is_some()); + assert!(cache.get_cache("new").is_some()); + } + + #[test] + fn expired_entries_are_evicted_before_live_entries() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now.clone(), 3, Duration::from_secs(60)); + cache.set_cache( + "expired-one", + "value".to_string(), + Some(Duration::from_secs(10)), + ); + cache.set_cache( + "expired-two", + "value".to_string(), + Some(Duration::from_secs(20)), + ); + cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100))); + now.store(121, Ordering::Relaxed); + cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100))); + assert_eq!(cache.get_cache("expired-one"), None); + assert_eq!(cache.get_cache("expired-two"), None); + assert!(cache.get_cache("live").is_some()); + assert!(cache.get_cache("new").is_some()); + } + + #[test] + fn stale_heap_entries_are_skipped() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now, 1, Duration::from_secs(60)); + cache.set_cache( + "removed", + "value".to_string(), + Some(Duration::from_secs(10)), + ); + cache.delete_cache("removed"); + cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20))); + cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30))); + assert_eq!(cache.get_cache("removed"), None); + assert_eq!(cache.get_cache("kept"), None); + assert!(cache.get_cache("new").is_some()); + } + + #[test] + fn delete_and_flush_remove_values_and_ttls() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now, 10, Duration::from_secs(60)); + cache.set_cache("one", "value".to_string(), None); + cache.set_cache("two", "value".to_string(), None); + cache.delete_cache("one"); + assert_eq!(cache.get_cache("one"), None); + cache.flush_cache(); + assert!(cache.cache_dict.is_empty()); + assert!(cache.ttl_dict.is_empty()); + assert!(cache.expiration_heap.is_empty()); + } + + #[test] + fn zero_max_size_does_not_cache() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now, 0, Duration::from_secs(60)); + cache.set_cache("key", "value".to_string(), None); + assert_eq!(cache.get_cache("key"), None); + assert!(cache.cache_dict.is_empty()); + } +} diff --git a/litellm-rust/crates/core/src/caching/mod.rs b/litellm-rust/crates/core/src/caching/mod.rs new file mode 100644 index 00000000000..5fb8a0e5174 --- /dev/null +++ b/litellm-rust/crates/core/src/caching/mod.rs @@ -0,0 +1 @@ +pub mod in_memory_cache; diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 117596f53d5..3989fb441bc 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,3 +1,4 @@ +pub mod caching; pub mod call_lifecycle; pub mod constants; pub mod error; diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs new file mode 100644 index 00000000000..82d1e8fdf91 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -0,0 +1,724 @@ +use std::collections::BTreeMap; +use std::sync::{Mutex, OnceLock}; +use std::time::Duration; +use std::time::{SystemTime, UNIX_EPOCH}; + +use crate::caching::in_memory_cache::InMemoryCache; +use crate::error::{CoreError, CoreResult}; +use aws_credential_types::provider::ProvideCredentials; +use aws_credential_types::Credentials; +use aws_sigv4::http_request::{ + sign, SignableBody, SignableRequest, SigningParams, SigningSettings, +}; +use aws_sigv4::sign::v4; +use aws_smithy_runtime_api::client::identity::Identity; +use sha2::{Digest, Sha256}; + +use super::constants::{ + AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN, + AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_STS_ENDPOINT, + AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, BEDROCK_SERVICE, + DEFAULT_SESSION_NAME_PREFIX, +}; + +const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); +const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); + +static IAM_CREDENTIALS_CACHE: OnceLock>> = OnceLock::new(); + +fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { + match flow { + AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL), + AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL), + AwsAuthFlow::WebIdentity { .. } + | AwsAuthFlow::AssumeRole { .. } + | AwsAuthFlow::Profile { .. } + | AwsAuthFlow::SessionToken { .. } => None, + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct AwsAuthConfig { + pub access_key_id: Option, + pub secret_access_key: Option, + pub session_token: Option, + pub region_name: Option, + pub session_name: Option, + pub profile_name: Option, + pub role_name: Option, + pub web_identity_token: Option, + pub sts_endpoint: Option, + pub external_id: Option, +} + +impl AwsAuthConfig { + fn with_environment(self, env_lookup: &dyn Fn(&str) -> Option) -> Self { + Self { + access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)), + secret_access_key: self + .secret_access_key + .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)), + session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)), + region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)), + session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)), + profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)), + role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)), + web_identity_token: self + .web_identity_token + .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)), + sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)), + external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum AwsAuthFlow { + WebIdentity { + token: String, + role: String, + session_name: String, + }, + AssumeRole { + role: String, + session_name: Option, + }, + Profile { + name: String, + }, + SessionToken { + access_key_id: String, + secret_access_key: String, + session_token: String, + }, + StaticKeys { + access_key_id: String, + secret_access_key: String, + region_name: String, + }, + DefaultChain, +} + +fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String { + let mut hasher = Sha256::new(); + hasher.update(format!("{config:?}:{flow:?}")); + format!("{:x}", hasher.finalize()) +} + +fn get_cached_credentials(key: &str) -> Option { + let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default())); + let mut entries = cache.lock().ok()?; + entries.get_cache(key) +} + +fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) { + let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default())); + if let Ok(mut entries) = cache.lock() { + entries.set_cache(key, credentials, Some(ttl)); + } +} + +fn role_identity(arn: &str) -> Option<(&str, &str, &str)> { + let mut parts = arn.splitn(6, ':'); + let ("arn", partition, _, _, account, resource) = ( + parts.next()?, + parts.next()?, + parts.next()?, + parts.next()?, + parts.next()?, + parts.next()?, + ) else { + return None; + }; + let role = if let Some(role) = resource.strip_prefix("role/") { + role.rsplit('/').next()? + } else { + resource.strip_prefix("assumed-role/")?.split('/').next()? + }; + Some((partition, account, role)) +} + +fn same_role_arns(target: &str, caller: &str) -> bool { + role_identity(target) == role_identity(caller) +} + +pub fn classify_auth( + config: AwsAuthConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> AwsAuthFlow { + let config = config.with_environment(env_lookup); + if let (Some(token), Some(role), Some(session_name)) = ( + config.web_identity_token.clone(), + config.role_name.clone(), + config.session_name.clone(), + ) { + return AwsAuthFlow::WebIdentity { + token, + role, + session_name, + }; + } + if let Some(role) = config.role_name.clone() { + return AwsAuthFlow::AssumeRole { + role, + session_name: config.session_name.clone(), + }; + } + if let Some(name) = config.profile_name { + return AwsAuthFlow::Profile { name }; + } + if let (Some(access_key_id), Some(secret_access_key), Some(session_token)) = ( + config.access_key_id.clone(), + config.secret_access_key.clone(), + config.session_token, + ) { + return AwsAuthFlow::SessionToken { + access_key_id, + secret_access_key, + session_token, + }; + } + if let (Some(access_key_id), Some(secret_access_key), Some(region_name)) = ( + config.access_key_id, + config.secret_access_key, + config.region_name, + ) { + return AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + region_name, + }; + } + AwsAuthFlow::DefaultChain +} + +pub async fn resolve_credentials( + config: AwsAuthConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let resolved = config.clone().with_environment(env_lookup); + let flow = classify_auth(config, env_lookup); + match flow { + AwsAuthFlow::SessionToken { + access_key_id, + secret_access_key, + session_token, + } => Ok(Credentials::new( + access_key_id, + secret_access_key, + Some(session_token), + None, + "litellm-static-session", + )), + AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + region_name, + } => { + let flow = AwsAuthFlow::StaticKeys { + access_key_id: access_key_id.clone(), + secret_access_key: secret_access_key.clone(), + region_name, + }; + let key = cache_key(&resolved, &flow); + if let Some(credentials) = get_cached_credentials(&key) { + return Ok(credentials); + } + let credentials = Credentials::new( + access_key_id, + secret_access_key, + None, + None, + "litellm-static", + ); + set_cached_credentials( + key, + credentials.clone(), + credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL), + ); + Ok(credentials) + } + AwsAuthFlow::Profile { name } => { + let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() + .profile_name(name) + .build(); + provider.provide_credentials().await.map_err(|error| { + CoreError::Auth(format!("AWS profile credentials failed: {error}")) + }) + } + AwsAuthFlow::AssumeRole { role, session_name } => { + if is_already_running_as_role(&role, &resolved).await? { + let ambient_flow = AwsAuthFlow::DefaultChain; + let key = cache_key(&resolved, &ambient_flow); + if let Some(credentials) = get_cached_credentials(&key) { + return Ok(credentials); + } + let provider = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await; + let credentials = provider.provide_credentials().await.map_err(|error| { + CoreError::Auth(format!("AWS default credentials failed: {error}")) + })?; + set_cached_credentials( + key, + credentials.clone(), + credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL), + ); + return Ok(credentials); + } + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = resolved.region_name.clone() { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = resolved.sts_endpoint.clone() { + loader = loader.endpoint_url(endpoint); + } + if let (Some(access_key_id), Some(secret_access_key)) = + (resolved.access_key_id, resolved.secret_access_key) + { + loader = loader.credentials_provider(Credentials::new( + access_key_id, + secret_access_key, + resolved.session_token, + None, + "litellm-role-source", + )); + } + let sdk_config = loader.load().await; + let builder = aws_config::sts::AssumeRoleProvider::builder(role); + let builder = match session_name { + Some(name) => builder.session_name(name), + None => builder.session_name(default_session_name()), + }; + let builder = match resolved.external_id { + Some(id) => builder.external_id(id), + None => builder, + }; + let provider = builder.configure(&sdk_config).build().await; + provider + .provide_credentials() + .await + .map_err(|error| CoreError::Auth(format!("AWS role credentials failed: {error}"))) + } + AwsAuthFlow::WebIdentity { + token, + role, + session_name, + } => { + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = resolved.region_name { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = resolved.sts_endpoint { + loader = loader.endpoint_url(endpoint); + } + let sdk_config = loader.load().await; + let client = aws_sdk_sts::Client::new(&sdk_config); + let response = client + .assume_role_with_web_identity() + .role_arn(role) + .role_session_name(session_name) + .web_identity_token(token) + .send() + .await + .map_err(|error| { + CoreError::Auth(format!("AWS web identity credentials failed: {error}")) + })?; + let credentials = response.credentials().ok_or_else(|| { + CoreError::Auth("AWS web identity response had no credentials".to_string()) + })?; + let expiration = SystemTime::try_from(*credentials.expiration()).map_err(|error| { + CoreError::Auth(format!("AWS web identity expiration was invalid: {error}")) + })?; + Ok(Credentials::new( + credentials.access_key_id(), + credentials.secret_access_key(), + Some(credentials.session_token().to_string()), + Some(expiration), + "litellm-web-identity", + )) + } + AwsAuthFlow::DefaultChain => { + let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain); + if let Some(credentials) = get_cached_credentials(&key) { + return Ok(credentials); + } + let provider = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await; + let credentials = provider.provide_credentials().await.map_err(|error| { + CoreError::Auth(format!("AWS default credentials failed: {error}")) + })?; + set_cached_credentials( + key, + credentials.clone(), + credential_cache_ttl(&AwsAuthFlow::DefaultChain).unwrap_or(AMBIENT_CREDENTIALS_TTL), + ); + Ok(credentials) + } + } +} + +async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> CoreResult { + if role_identity(role).is_none() { + return Ok(false); + } + if let (Ok(current_role), Ok(token_file)) = ( + std::env::var(AWS_ROLE_ARN), + std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), + ) { + if !token_file.is_empty() { + return Ok(same_role_arns(role, ¤t_role)); + } + } + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = config.region_name.clone() { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = config.sts_endpoint.clone() { + loader = loader.endpoint_url(endpoint); + } + let sdk_config = loader.load().await; + let response = match aws_sdk_sts::Client::new(&sdk_config) + .get_caller_identity() + .send() + .await + { + Ok(response) => response, + Err(_) => return Ok(false), + }; + Ok(response + .arn() + .is_some_and(|caller| same_role_arns(role, caller))) +} + +fn default_session_name() -> String { + let seconds = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()); + format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}") +} + +pub fn sign_bedrock_post( + url: &str, + body: &[u8], + headers: &BTreeMap, + region: &str, + credentials: &Credentials, + signing_time: SystemTime, +) -> CoreResult> { + let identity: Identity = credentials.clone().into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(region) + .name(BEDROCK_SERVICE) + .time(signing_time) + .settings(SigningSettings::default()) + .build() + .map(SigningParams::from) + .map_err(|error| CoreError::Auth(format!("AWS signing parameters failed: {error}")))?; + let header_refs = headers + .iter() + .map(|(name, value)| (name.as_str(), value.as_str())); + let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) + .map_err(|error| CoreError::Auth(format!("AWS signable request failed: {error}")))?; + let (instructions, _) = sign(request, ¶ms) + .map_err(|error| CoreError::Auth(format!("AWS request signing failed: {error}")))? + .into_parts(); + Ok(instructions + .headers() + .map(|(name, value)| { + let normalized_name = match name { + "authorization" => "Authorization", + "x-amz-date" => "X-Amz-Date", + "x-amz-security-token" => "X-Amz-Security-Token", + _ => name, + }; + (normalized_name.to_string(), value.to_string()) + }) + .collect()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn no_env(_: &str) -> Option { + None + } + + fn parity_inputs() -> (String, Vec, BTreeMap) { + ( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" + .to_string(), + br#"{"input":"hello"}"#.to_vec(), + BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]), + ) + } + + #[test] + fn classification_preserves_python_precedence() { + let config = AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + session_token: Some("token".into()), + region_name: Some("us-east-1".into()), + session_name: Some("session".into()), + profile_name: Some("profile".into()), + role_name: Some("role".into()), + web_identity_token: Some("oidc".into()), + ..Default::default() + }; + assert!(matches!( + classify_auth(config, &no_env), + AwsAuthFlow::WebIdentity { .. } + )); + } + + #[test] + fn classification_covers_fallthroughs() { + let env = |key: &str| match key { + AWS_PROFILE_NAME => Some("profile".into()), + _ => None, + }; + assert!(matches!( + classify_auth(AwsAuthConfig::default(), &env), + AwsAuthFlow::Profile { .. } + )); + assert!(matches!( + classify_auth( + AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + session_token: Some("token".into()), + ..Default::default() + }, + &no_env + ), + AwsAuthFlow::SessionToken { .. } + )); + assert!(matches!( + classify_auth( + AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + region_name: Some("us-east-1".into()), + ..Default::default() + }, + &no_env + ), + AwsAuthFlow::StaticKeys { .. } + )); + assert_eq!( + classify_auth(AwsAuthConfig::default(), &no_env), + AwsAuthFlow::DefaultChain + ); + } + + #[tokio::test] + async fn static_credentials_do_not_use_network() { + let credentials = resolve_credentials( + AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + region_name: Some("us-east-1".into()), + ..Default::default() + }, + &no_env, + ) + .await + .expect("static credentials"); + assert_eq!(credentials.access_key_id(), "ak"); + assert_eq!(credentials.session_token(), None); + } + + #[test] + fn cache_policy_matches_python_flows() { + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::StaticKeys { + access_key_id: "ak".into(), + secret_access_key: "sk".into(), + region_name: "us-east-1".into(), + }), + Some(STATIC_CREDENTIALS_TTL) + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::DefaultChain), + Some(AMBIENT_CREDENTIALS_TTL) + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::SessionToken { + access_key_id: "ak".into(), + secret_access_key: "sk".into(), + session_token: "token".into(), + }), + None + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::Profile { + name: "profile".into() + }), + None + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::AssumeRole { + role: "arn:aws:iam::123456789012:role/demo".into(), + session_name: None, + }), + None + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::WebIdentity { + token: "token".into(), + role: "arn:aws:iam::123456789012:role/demo".into(), + session_name: "session".into(), + }), + None + ); + } + + #[test] + fn cache_round_trip_preserves_credentials() { + let key = format!("cache-test-{}", std::process::id()); + let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test"); + set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL); + assert_eq!( + get_cached_credentials(&key).map(|value| value.access_key_id().to_string()), + Some("cache-ak".to_string()) + ); + } + + #[test] + fn same_role_comparison_matches_partition_account_and_role() { + assert!(same_role_arns( + "arn:aws:iam::123456789012:role/path/demo", + "arn:aws:sts::123456789012:assumed-role/demo/session" + )); + assert!(!same_role_arns( + "arn:aws:iam::123456789012:role/demo", + "arn:aws:iam::999999999999:role/demo" + )); + assert!(!same_role_arns( + "arn:aws:iam::123456789012:role/demo", + "arn:aws-cn:iam::123456789012:role/demo" + )); + assert!(!same_role_arns( + "arn:aws:iam::123456789012:user/demo", + "arn:aws:iam::123456789012:role/demo" + )); + } + + #[test] + fn signing_matches_botocore_golden_vector() { + let (url, body, headers) = parity_inputs(); + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + Some("session-token".to_string()), + None, + "test", + ); + let signed = sign_bedrock_post( + &url, + &body, + &headers, + "us-east-1", + &credentials, + UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + ) + .expect("golden signature"); + assert_eq!( + signed.get("X-Amz-Date").map(String::as_str), + Some("20240102T030405Z") + ); + assert_eq!( + signed.get("X-Amz-Security-Token").map(String::as_str), + Some("session-token") + ); + assert_eq!( + signed.get("Authorization").map(String::as_str), + Some("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464") + ); + } + + #[test] + fn signing_without_session_token_omits_security_header() { + let (url, body, headers) = parity_inputs(); + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + None, + None, + "test", + ); + let signed = sign_bedrock_post( + &url, + &body, + &headers, + "us-east-1", + &credentials, + UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + ) + .expect("signature"); + assert!(!signed.contains_key("X-Amz-Security-Token")); + } + + #[ignore] + #[tokio::test] + async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box> { + let access_key_id = std::env::var("AWS_BEDROCK_TEST_ACCESS_KEY_ID")?; + let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?; + let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec(); + let headers = + BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]); + let credentials = resolve_credentials( + AwsAuthConfig { + access_key_id: Some(access_key_id), + secret_access_key: Some(secret_access_key), + region_name: Some("us-west-2".to_string()), + ..Default::default() + }, + &no_env, + ) + .await?; + let client = reqwest::Client::new(); + let mut failures = Vec::new(); + + for region in ["us-west-2", "us-east-1"] { + let url = format!( + "https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke" + ); + let signed_headers = sign_bedrock_post( + &url, + &body, + &headers, + region, + &credentials, + SystemTime::now(), + )?; + let mut request = client.post(&url).body(body.clone()); + for (name, value) in &headers { + request = request.header(name, value); + } + for (name, value) in signed_headers { + request = request.header(name, value); + } + let response = request.send().await?; + let status = response.status(); + let response_body = response.text().await?; + let snippet: String = response_body.chars().take(240).collect(); + println!("region={region} status={status} response={snippet}"); + if status == reqwest::StatusCode::OK { + return Ok(()); + } + failures.push(format!("{region}: {status} {snippet}")); + } + + panic!( + "no Bedrock region returned HTTP 200: {}", + failures.join("; ") + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/bedrock/constants.rs b/litellm-rust/crates/core/src/providers/bedrock/constants.rs new file mode 100644 index 00000000000..a08ae9de146 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/constants.rs @@ -0,0 +1,14 @@ +pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID"; +pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY"; +pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN"; +pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME"; +pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME"; +pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME"; +pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME"; +pub const AWS_WEB_IDENTITY_TOKEN: &str = "AWS_WEB_IDENTITY_TOKEN"; +pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN"; +pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; +pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; +pub const BEDROCK_SERVICE: &str = "bedrock"; +pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session"; diff --git a/litellm-rust/crates/core/src/providers/bedrock/mod.rs b/litellm-rust/crates/core/src/providers/bedrock/mod.rs new file mode 100644 index 00000000000..8027260a78b --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/mod.rs @@ -0,0 +1,6 @@ +//! User-directed exception: this base provider owns AWS auth I/O for parity +//! with Python's `BaseAWSLLM`; the broader core purity guidance is reconciled +//! separately. + +pub mod aws_base; +mod constants; diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index dc9dc515e7d..805600d6dbe 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -1,5 +1,7 @@ pub mod anthropic; pub mod azure_ai; +#[cfg(feature = "bedrock-auth")] +pub mod bedrock; pub mod mistral; pub mod openai; pub mod vertex_ai; diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 470448251c9..aaf523eacd5 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -15,6 +15,7 @@ from typing import Any, Dict, Optional from unittest.mock import MagicMock, patch from botocore.awsrequest import AWSPreparedRequest, AWSRequest +from botocore.auth import SigV4Auth from botocore.credentials import Credentials import litellm @@ -768,6 +769,30 @@ def test_get_request_headers_with_sigv4(): assert result == mock_request.prepare.return_value +def test_sigv4_matches_rust_golden_vector(): + request = AWSRequest( + method="POST", + url="https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke", + data=b'{"input":"hello"}', + headers={"Content-Type": "application/json"}, + ) + credentials = Credentials( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + "session-token", + ) + with patch("botocore.auth.get_current_datetime", return_value=datetime(2024, 1, 2, 3, 4, 5)): + SigV4Auth(credentials, "bedrock", "us-east-1").add_auth(request) + assert request.headers["X-Amz-Date"] == "20240102T030405Z" + assert request.headers["X-Amz-Security-Token"] == "session-token" + assert ( + request.headers["Authorization"] + == "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, " + "SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, " + "Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" + ) + + def test_get_request_headers_with_api_key_bearer_token(): """ Test that get_request_headers uses the api_key parameter as a bearer token when provided From 7390f29b237bfd407926c7881afd1df7109fb106 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 00:18:29 -0700 Subject: [PATCH 18/60] feat(mcp): identity-only session tokens for the gateway DCR front door --- .../session_credentials.py | 190 ++++++++++ .../outbound_credentials/session_token.py | 356 ++++++++++++++++++ .../test_session_credentials.py | 135 +++++++ .../test_session_token.py | 206 ++++++++++ 4 files changed, 887 insertions(+) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py new file mode 100644 index 00000000000..08d5cc8b1f1 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -0,0 +1,190 @@ +"""Producer and consumer helpers for the gateway-level DCR session token. + +The aggregate ``/mcp`` front door (``mcp_gateway_dcr``) issues the identity-only session +tokens defined in :mod:`.session_token`. The gateway token endpoint mints them (producer) +after SSO sign-in, and at the MCP admission edge the gateway derives the session signing +key from the proxy ``master_key``, opens the bearer, and admits the request under the +recovered litellm user (consumer), reloading the live user record and policy before +anything runs. This module is the pure surface for both sides; the token-endpoint and +admission wiring live in their respective call sites. + +The signing key is derived with the same memory-hard scrypt construction as +:func:`~.bridge_credentials.envelope_keys_from_master_key` but under a distinct domain +label, so session tokens and bridge envelopes never share key material: a token of one +family is unverifiable in the other by key separation, on top of the distinct issuers, +prefixes, and claim shapes. +""" + +import hashlib +from datetime import datetime +from functools import lru_cache +from typing import Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + OpenedSessionToken, + SessionExpired, + SessionKeys, + SessionPrincipal, + is_session_refresh_token, + is_session_token, + open_session_refresh_token, + open_session_token, +) + +_SESSION_SIGNING_KEY_DOMAIN = b"litellm-mcp-gateway:session-signing:" + +# scrypt work factors (RFC 7914), identical to the envelope KDF: memory-hard so a captured +# session token is not a cheap offline oracle for the master key. +_SCRYPT_N = 2**15 +_SCRYPT_R = 8 +_SCRYPT_P = 1 +_SCRYPT_MAXMEM = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 +_DERIVED_KEY_BYTES = 32 + + +@lru_cache(maxsize=8) +def session_keys_from_master_key(master_key: str) -> SessionKeys: + """Derive the session signing key from the proxy master key. + + A memory-hard scrypt KDF (RFC 7914) over a session-specific domain-label salt yields a + 256-bit subkey from the one secret, so the producer (mint) and consumer (open) agree on + the key without persisting any. The domain label differs from both envelope labels in + :mod:`.bridge_credentials`, so compromise or misuse of one token family never crosses + into the other. The result is cached (the master key is fixed for a process); rotating + ``master_key`` invalidates every outstanding session, which is the intended behavior + for a signing-key change. + """ + signing = hashlib.scrypt( + master_key.encode(), + salt=_SESSION_SIGNING_KEY_DOMAIN, + n=_SCRYPT_N, + r=_SCRYPT_R, + p=_SCRYPT_P, + maxmem=_SCRYPT_MAXMEM, + dklen=_DERIVED_KEY_BYTES, + ).hex() + return SessionKeys(signing_key=SecretStr(signing)) + + +class NotSessionBearer(BaseModel): + """The bearer is not session-shaped; admission continues on its normal path.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["not_session_bearer"] = "not_session_bearer" + + +class SessionBearerAdmitted(BaseModel): + """A valid session access token: the principal to admit under after a live reload.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["admitted"] = "admitted" + principal: SessionPrincipal + + +class SessionBearerInvalid(BaseModel): + """The bearer is session-shaped but must not admit (expired, tampered, wrong key, or a + refresh token presented at the tool-call edge); admission fails closed with the + ``invalid_token`` challenge rather than falling through to another arm. ``expired`` + distinguishes a routine expiry (debug-log worthy) from a tampered or foreign token.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["invalid"] = "invalid" + expired: bool = False + + +SessionBearerResult: TypeAlias = NotSessionBearer | SessionBearerAdmitted | SessionBearerInvalid + + +def _strip_bearer(value: str) -> str: + parts = value.split(None, 1) + if len(parts) == 2 and parts[0].lower() == "bearer": + return parts[1] + return value + + +def is_session_bearer_shaped(authorization_value: str) -> bool: + """Cheap, keyless test that an ``Authorization`` value carries a session token of either + kind (optional ``Bearer`` scheme stripped). The admission edge engages the session arm + for an access token (to admit) and for a refresh token (to reject it explicitly, since + a refresh credential is never usable at the tool-call edge); anything else falls + through to normal admission.""" + candidate = _strip_bearer(authorization_value) + return is_session_token(candidate) or is_session_refresh_token(candidate) + + +def resolve_session_bearer( + authorization_value: str, + keys: SessionKeys, + now: datetime, +) -> SessionBearerResult: + """Classify an ``Authorization`` value presented at the aggregate MCP edge. + + Strips an optional ``Bearer`` scheme, then returns ``NotSessionBearer`` for a + non-session bearer (normal admission continues), ``SessionBearerAdmitted`` with the + recovered principal for a valid access token, and ``SessionBearerInvalid`` for a + session-shaped bearer that must not admit. Never raises: total over hostile input via + :func:`~.session_token.open_session_token`. + + A refresh token is ``SessionBearerInvalid`` here: it is a valid gateway credential but + only ever presented back to the token endpoint, so admission must fail it closed rather + than let it fall through to another arm. + """ + candidate = _strip_bearer(authorization_value) + if is_session_refresh_token(candidate): + return SessionBearerInvalid() + if not is_session_token(candidate): + return NotSessionBearer() + opened = open_session_token(candidate, keys, now) + if isinstance(opened, OpenedSessionToken): + return SessionBearerAdmitted(principal=opened.principal) + return SessionBearerInvalid(expired=isinstance(opened, SessionExpired)) + + +class SessionRefreshOpened(BaseModel): + """A valid session refresh token presented to the token endpoint: the principal to + re-validate and renew under.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["opened"] = "opened" + principal: SessionPrincipal + + +class SessionRefreshInvalid(BaseModel): + """The presented refresh grant is not a valid session refresh token for this client + (not refresh-shaped, will not open, or bound to a different ``client_id``); the token + endpoint fails the refresh closed.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["invalid"] = "invalid" + + +SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid + + +def open_session_refresh_bearer( + refresh_value: str, + keys: SessionKeys, + now: datetime, + expected_client_id: str, +) -> SessionRefreshResult: + """Open a session refresh token presented on a ``refresh_token`` grant. + + The token-endpoint mirror of :func:`resolve_session_bearer`: strips an optional + ``Bearer`` scheme, then returns ``SessionRefreshOpened`` with the recovered principal, + or ``SessionRefreshInvalid`` for anything that is not a valid session refresh token + issued to ``expected_client_id``. Never raises. The client binding (RFC 6749 section 6) + stops a refresh token stolen from one DCR client from being renewed through another; + ``client_id`` is not a secret (the caller presents it), so a plain equality check is + sufficient and, unlike ``hmac.compare_digest`` on ``str``, does not raise on non-ASCII. + """ + candidate = _strip_bearer(refresh_value) + if not is_session_refresh_token(candidate): + return SessionRefreshInvalid() + opened = open_session_refresh_token(candidate, keys, now) + if not isinstance(opened, OpenedSessionToken): + return SessionRefreshInvalid() + if opened.principal.client_id != expected_client_id: + return SessionRefreshInvalid() + return SessionRefreshOpened(principal=opened.principal) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py new file mode 100644 index 00000000000..78b1f7e4916 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -0,0 +1,356 @@ +"""Identity-only session tokens for the gateway-level (aggregate ``/mcp``) DCR front door. + +A DCR client that signs in through LiteLLM SSO holds ONE bearer that carries ONLY a +litellm identity; unlike the :mod:`.envelope` bridge bearer it seals no upstream +credential, because the custody model vaults every upstream token server-side in +``LiteLLM_MCPUserCredentials`` and egress resolves them by user at call time. The token +is therefore a stable REFERENCE, not an authorization: admission reloads the live user +record and policy on every request, so deactivating the user (or their team) kills +outstanding sessions immediately without a revocation store. + +Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT, +the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp`` +plus ``kind``, ``user_id``, and ``client_id``; ``client_id`` binds the refresh token +to the DCR client it was issued to (RFC 6749 section 6) and is carried on the access +token for parity and audit. There is no encrypted payload: nothing in a session token +is secret beyond the signature, and reprs never print the signed value because minted +tokens are ``SecretStr``. + +This module is pure and unwired: it imports nothing from endpoint or edge code, reads +no proxy globals, and takes all key material and the clock as explicit parameters. +Failures are values: :func:`open_session_token` and :func:`open_session_refresh_token` +are total over hostile, attacker-controlled input and return a +``SessionTokenOpenError`` variant rather than raising. PyJWT's ``iat``/``nbf``/``exp`` +validators are disabled for the same reasons documented in :mod:`.envelope` (they +raise on hostile claim types and compare against the wall clock instead of the +injected ``now``); the strict pydantic claims model is the sole, total type gate. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta +from typing import Literal, TypeAlias + +import jwt +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError + +SESSION_TOKEN_PREFIX = "llm_session_" +"""Marker prefix on every serialized session ACCESS token so the admission edge can cheaply +tell a gateway session from a litellm key, JWT, or bridge envelope before doing any +cryptography. Distinct from the ``llm_env_``/``llm_refresh_`` envelope prefixes.""" + +SESSION_REFRESH_PREFIX = "llm_srefresh_" +"""Marker prefix on every serialized session REFRESH token. A distinct prefix keeps the two +credentials routable without crypto and, together with the signed ``kind`` claim, stops one +from being presented where the other is expected: the refresh token is only ever presented +back to the token endpoint, never at the MCP edge.""" + +SESSION_ISSUER = "litellm-mcp-gateway" +"""``iss`` claim stamped into every session token and required back on open. Distinct from +the envelope issuer so a token of one family can never validate in the other even under a +hypothetical shared signing key.""" + +SESSION_TTL_SECONDS = 3600 +"""Session ACCESS token lifetime (1h), matching the access-envelope and BYOK session bearer +windows: a client-held credential never outlives a bounded window, and each refresh +re-validates the live user before re-minting.""" + +SESSION_REFRESH_TTL_SECONDS = 1209600 +"""Session REFRESH token lifetime (14 days), matching the refresh-envelope bound. Each +renewal re-validates the sealed user against the live record (deactivation gates it) and +rotates the refresh token, so the practical bound is idle time, not a fixed session.""" + +MAX_SESSION_TOKEN_BYTES = 4096 +"""Size cap on the serialized token (prefix + JWT, in bytes) and on any candidate accepted +by the openers. Session claims are small; the only variable-length field is ``client_id`` +(a sealed DCR client record), and 4096 leaves ample headroom under common 8-16KB header +limits while bounding hostile input before JWT parsing.""" + +_SESSION_JWT_ALGORITHM = "HS256" + +SessionTokenKind = Literal["session", "session_refresh"] +"""Which credential a session token is. Stamped into the signed claims and required to match +on open, so a signature-valid token of one kind cannot be replayed as the other even if its +wire prefix is swapped (the prefix is not part of the signed payload; this claim is).""" + + +class SessionPrincipal(BaseModel): + """The litellm user a session token identifies and the DCR client it was issued to. + + ``user_id`` is the SSO-established litellm user subject, never a credential: admission + reloads the live user record by it, so current role, team, and revocation state are + enforced at use time rather than frozen at mint time. ``client_id`` is the (stateless, + gateway-sealed) DCR client identifier the token was issued to; the token endpoint + requires it to match on the refresh grant. + """ + + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + client_id: str = Field(min_length=1) + + +class SessionKeys(BaseModel): + """Injected key material: the HS256 signing key. + + ``signing_key`` must be at least 32 bytes: HS256's HMAC-SHA256 has a 256-bit security + level, RFC 7518 requires a key of at least that size, and a shorter key makes PyJWT + emit ``InsecureKeyLengthWarning``. + """ + + model_config = ConfigDict(frozen=True) + signing_key: SecretStr = Field(min_length=32) + + +class MintedSessionToken(BaseModel): + """A minted session token: the client-held bearer value and when it expires.""" + + model_config = ConfigDict(frozen=True) + token: SecretStr + expires_at: datetime + + +class OpenedSessionToken(BaseModel): + """A validated session token of either kind: the principal it was minted for.""" + + model_config = ConfigDict(frozen=True) + principal: SessionPrincipal + + +class SessionTokenTooLarge(BaseModel): + """The serialized token exceeded ``MAX_SESSION_TOKEN_BYTES``; carries sizes only. Only + reachable through an oversized ``client_id``, which registration should have bounded.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_token_too_large"] = "session_token_too_large" + size_bytes: int + max_bytes: int + + +SessionTokenMintError: TypeAlias = SessionTokenTooLarge + + +class NotASessionToken(BaseModel): + """The candidate does not carry the expected session prefix.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["not_a_session_token"] = "not_a_session_token" + + +class SessionBadSignature(BaseModel): + """The JWT signature does not verify under the provided signing key.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_bad_signature"] = "session_bad_signature" + + +class SessionExpired(BaseModel): + """The token's ``exp`` is not in the future relative to the provided ``now``.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_expired"] = "session_expired" + + +class SessionMalformed(BaseModel): + """The token is not a well-formed session token: undecodable JWT, wrong issuer, wrong + ``kind``, or missing/mistyped/extra claims.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_malformed"] = "session_malformed" + + +SessionTokenOpenError: TypeAlias = NotASessionToken | SessionBadSignature | SessionExpired | SessionMalformed + + +class _SessionClaims(BaseModel): + """Decoded-claims boundary that pins the exact shape the mints emit. + + ``user_id``/``client_id`` mirror the ``min_length`` constraints of + :class:`SessionPrincipal` so any claim set that validates here also constructs a + principal, keeping the openers raise-free: a correctly signed JWT with an empty + identity claim fails here and maps to ``SessionMalformed``. ``strict`` rejects coerced + types (``exp: "123"``) and ``extra="forbid"`` rejects any claim the gateway never + mints; PyJWT's own registered-claim validators are disabled at decode (see module + docstring), so this model is the sole, total type gate for every claim. + """ + + model_config = ConfigDict(frozen=True, strict=True, extra="forbid") + iss: str + iat: int + exp: int + kind: SessionTokenKind + user_id: str = Field(min_length=1) + client_id: str = Field(min_length=1) + + +def is_session_token(candidate: str) -> bool: + """Cheap prefix check for a session ACCESS token so the admission edge can route gateway + sessions vs keys, JWTs, and envelopes without crypto.""" + return candidate.startswith(SESSION_TOKEN_PREFIX) + + +def is_session_refresh_token(candidate: str) -> bool: + """Cheap prefix check for a session REFRESH token so the token endpoint can route a + refresh grant without crypto.""" + return candidate.startswith(SESSION_REFRESH_PREFIX) + + +def mint_session_token( + principal: SessionPrincipal, + keys: SessionKeys, + now: datetime, +) -> MintedSessionToken | SessionTokenMintError: + """Mint the short-lived session ACCESS token for ``principal``. + + ``exp`` is ``SESSION_TTL_SECONDS`` from ``now``. Returns ``SessionTokenTooLarge`` when + the serialized token exceeds ``MAX_SESSION_TOKEN_BYTES``. + """ + return _mint( + kind="session", + prefix=SESSION_TOKEN_PREFIX, + principal=principal, + expires_at=now + timedelta(seconds=SESSION_TTL_SECONDS), + keys=keys, + now=now, + ) + + +def mint_session_refresh_token( + principal: SessionPrincipal, + keys: SessionKeys, + now: datetime, +) -> MintedSessionToken | SessionTokenMintError: + """Mint the long-lived session REFRESH token for ``principal``. + + ``exp`` is ``SESSION_REFRESH_TTL_SECONDS`` from ``now``. Minting a distinct + ``kind="session_refresh"`` claim is what keeps a refresh token from ever opening as an + access credential at the MCP edge. + """ + return _mint( + kind="session_refresh", + prefix=SESSION_REFRESH_PREFIX, + principal=principal, + expires_at=now + timedelta(seconds=SESSION_REFRESH_TTL_SECONDS), + keys=keys, + now=now, + ) + + +def open_session_token( + candidate: str, + keys: SessionKeys, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Validate a session ACCESS ``candidate`` and recover the principal. + + Never raises for bad input: every invalid, expired, tampered, or wrong-kind candidate + maps to a distinct ``SessionTokenOpenError`` variant. + """ + return _open(candidate, prefix=SESSION_TOKEN_PREFIX, expected_kind="session", keys=keys, now=now) + + +def open_session_refresh_token( + candidate: str, + keys: SessionKeys, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Validate a session REFRESH ``candidate`` and recover the principal. + + Total over hostile input exactly like :func:`open_session_token`. The + ``kind="session_refresh"`` claim is required, so an access token re-prefixed as a + refresh one is rejected as ``SessionMalformed``. + """ + return _open(candidate, prefix=SESSION_REFRESH_PREFIX, expected_kind="session_refresh", keys=keys, now=now) + + +def _mint( + kind: SessionTokenKind, + prefix: str, + principal: SessionPrincipal, + expires_at: datetime, + keys: SessionKeys, + now: datetime, +) -> MintedSessionToken | SessionTokenTooLarge: + """Sign the claims for either token kind and enforce the size cap. Shared by both mints + so the JWT shape, issuer, and size guard cannot drift between access and refresh.""" + claims = _SessionClaims( + iss=SESSION_ISSUER, + iat=int(now.timestamp()), + exp=int(expires_at.timestamp()), + kind=kind, + user_id=principal.user_id, + client_id=principal.client_id, + ) + token = prefix + jwt.encode( + claims.model_dump(), keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM + ) + size_bytes = len(token.encode("utf-8")) + if size_bytes > MAX_SESSION_TOKEN_BYTES: + return SessionTokenTooLarge(size_bytes=size_bytes, max_bytes=MAX_SESSION_TOKEN_BYTES) + return MintedSessionToken(token=SecretStr(token), expires_at=expires_at) + + +def _open( + candidate: str, + prefix: str, + expected_kind: SessionTokenKind, + keys: SessionKeys, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an + attacker-controlled candidate, shared by both openers so the security gate is identical + for access and refresh. Returns the opened token or a distinct error; never raises.""" + if not candidate.startswith(prefix): + return NotASessionToken() + # UTF-8 byte length is never below character length, so a character count already over + # the cap rejects an oversize candidate in O(1) without encoding it; the exact byte + # check then runs only on candidates already bounded to the cap in characters. + if len(candidate) > MAX_SESSION_TOKEN_BYTES: + return SessionMalformed() + if len(candidate.encode("utf-8", "surrogatepass")) > MAX_SESSION_TOKEN_BYTES: + return SessionMalformed() + claims = _decode_claims(candidate.removeprefix(prefix), keys.signing_key) + if not isinstance(claims, _SessionClaims): + return claims + if claims.kind != expected_kind: + return SessionMalformed() + if now.timestamp() >= claims.exp: + return SessionExpired() + return OpenedSessionToken(principal=SessionPrincipal(user_id=claims.user_id, client_id=claims.client_id)) + + +def _decode_claims( + compact: str, + signing_key: SecretStr, +) -> _SessionClaims | SessionBadSignature | SessionMalformed: + """Verify the HS256 signature and shape of an attacker-controlled compact JWT. + + ``compact`` is fully hostile and bounded to ``MAX_SESSION_TOKEN_BYTES`` by the caller. + PyJWT's ``iat``/``nbf``/``exp`` validators are disabled: they raise on hostile claim + types and, for ``iat``/``nbf``, compare against the wall clock rather than the injected + ``now`` (``exp`` is checked by the caller against ``now``). Apart from a signature + mismatch, every decode failure is ``SessionMalformed``: a non-UTF-8 candidate surfaces + as ``UnicodeEncodeError`` (a ``ValueError``), a non-string registered claim as a + ``TypeError`` from PyJWT's claim validators, and a wrong issuer or structurally invalid + token as an ``InvalidTokenError``. ``_SessionClaims`` is the total type gate. + """ + try: + payload = jwt.decode( + compact, + signing_key.get_secret_value(), + algorithms=[_SESSION_JWT_ALGORITHM], + issuer=SESSION_ISSUER, + options={ + "verify_exp": False, + "verify_iat": False, + "verify_nbf": False, + "require": ["iss", "iat", "exp"], + }, + ) + except jwt.InvalidSignatureError: + return SessionBadSignature() + except (jwt.InvalidTokenError, ValueError, TypeError): + return SessionMalformed() + try: + return _SessionClaims.model_validate(payload) + except ValidationError: + return SessionMalformed() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py new file mode 100644 index 00000000000..8fa7c15d2d3 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -0,0 +1,135 @@ +"""Tests for the session-token KDF and the edge/token-endpoint resolvers.""" + +from datetime import datetime, timedelta, timezone + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + envelope_keys_from_master_key, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + NotSessionBearer, + SessionBearerAdmitted, + SessionBearerInvalid, + SessionRefreshInvalid, + SessionRefreshOpened, + is_session_bearer_shaped, + open_session_refresh_bearer, + resolve_session_bearer, + session_keys_from_master_key, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SESSION_TTL_SECONDS, + MintedSessionToken, + SessionPrincipal, + mint_session_refresh_token, + mint_session_token, +) + +NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) +MASTER_KEY = "sk-master-key-for-tests" +KEYS = session_keys_from_master_key(MASTER_KEY) +PRINCIPAL = SessionPrincipal(user_id="user-123", client_id="llm_client_abc") + + +def _access_token() -> str: + minted = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +def _refresh_token() -> str: + minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +def test_kdf_is_deterministic_and_key_length_is_256_bit(): + again = session_keys_from_master_key(MASTER_KEY) + assert again.signing_key.get_secret_value() == KEYS.signing_key.get_secret_value() + assert len(bytes.fromhex(KEYS.signing_key.get_secret_value())) == 32 + + +def test_kdf_domain_separated_from_envelope_keys(): + envelope_keys = envelope_keys_from_master_key(MASTER_KEY) + session_signing = KEYS.signing_key.get_secret_value() + assert session_signing != envelope_keys.signing_key.get_secret_value() + assert session_signing != envelope_keys.encryption_key.get_secret_value() + + +def test_kdf_differs_across_master_keys(): + other = session_keys_from_master_key("sk-a-different-master-key") + assert other.signing_key.get_secret_value() != KEYS.signing_key.get_secret_value() + + +@pytest.mark.parametrize( + "value,expected", + [ + ("Bearer sk-1234", False), + ("sk-1234", False), + ("Bearer llm_env_abc", False), + ("Bearer llm_refresh_abc", False), + ("llm_session_abc", True), + ("Bearer llm_session_abc", True), + ("bearer llm_srefresh_abc", True), + ], +) +def test_is_session_bearer_shaped(value, expected): + assert is_session_bearer_shaped(value) is expected + + +def test_resolve_admits_valid_access_token_with_and_without_scheme(): + token = _access_token() + for value in (token, f"Bearer {token}", f"bearer {token}"): + result = resolve_session_bearer(value, KEYS, NOW) + assert isinstance(result, SessionBearerAdmitted) + assert result.principal == PRINCIPAL + + +def test_resolve_passes_non_session_bearers_through(): + for value in ("Bearer sk-1234", "Bearer llm_env_whatever", "Bearer eyJhbGciOi"): + assert isinstance(resolve_session_bearer(value, KEYS, NOW), NotSessionBearer) + + +def test_resolve_fails_expired_token_closed_and_flags_expiry(): + token = _access_token() + later = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + result = resolve_session_bearer(f"Bearer {token}", KEYS, later) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is True + + +def test_resolve_fails_tampered_token_closed_without_expiry_flag(): + token = _access_token() + tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") + result = resolve_session_bearer(f"Bearer {tampered}", KEYS, NOW) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is False + + +def test_resolve_rejects_refresh_token_at_the_edge(): + result = resolve_session_bearer(f"Bearer {_refresh_token()}", KEYS, NOW) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is False + + +def test_resolve_wrong_master_key_fails_closed(): + other_keys = session_keys_from_master_key("sk-rotated-master-key") + result = resolve_session_bearer(f"Bearer {_access_token()}", other_keys, NOW) + assert isinstance(result, SessionBearerInvalid) + + +def test_refresh_grant_opens_for_the_issued_client(): + result = open_session_refresh_bearer(_refresh_token(), KEYS, NOW, expected_client_id="llm_client_abc") + assert isinstance(result, SessionRefreshOpened) + assert result.principal == PRINCIPAL + + +def test_refresh_grant_rejects_a_different_client(): + result = open_session_refresh_bearer(_refresh_token(), KEYS, NOW, expected_client_id="llm_client_other") + assert isinstance(result, SessionRefreshInvalid) + + +def test_refresh_grant_rejects_access_token_presented_as_refresh(): + result = open_session_refresh_bearer(_access_token(), KEYS, NOW, expected_client_id="llm_client_abc") + assert isinstance(result, SessionRefreshInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py new file mode 100644 index 00000000000..551270f8d4b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -0,0 +1,206 @@ +"""Tests for the identity-only gateway session token (mint/open, hostile-input totality).""" + +from datetime import datetime, timedelta, timezone + +import jwt +import pytest +from pydantic import SecretStr, ValidationError + +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + MAX_SESSION_TOKEN_BYTES, + SESSION_ISSUER, + SESSION_REFRESH_PREFIX, + SESSION_REFRESH_TTL_SECONDS, + SESSION_TOKEN_PREFIX, + SESSION_TTL_SECONDS, + MintedSessionToken, + NotASessionToken, + OpenedSessionToken, + SessionBadSignature, + SessionExpired, + SessionKeys, + SessionMalformed, + SessionPrincipal, + SessionTokenTooLarge, + is_session_refresh_token, + is_session_token, + mint_session_refresh_token, + mint_session_token, + open_session_refresh_token, + open_session_token, +) + +NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) +KEYS = SessionKeys(signing_key=SecretStr("k" * 32)) +OTHER_KEYS = SessionKeys(signing_key=SecretStr("x" * 32)) +PRINCIPAL = SessionPrincipal(user_id="user-123", client_id="llm_client_abc") + + +def _mint_access() -> str: + minted = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +def _mint_refresh() -> str: + minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +def _sign_claims(payload: dict, prefix: str = SESSION_TOKEN_PREFIX, keys: SessionKeys = KEYS) -> str: + return prefix + jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm="HS256") + + +def _valid_claims(**overrides) -> dict: + base = { + "iss": SESSION_ISSUER, + "iat": int(NOW.timestamp()), + "exp": int((NOW + timedelta(seconds=600)).timestamp()), + "kind": "session", + "user_id": "user-123", + "client_id": "llm_client_abc", + } + return {**base, **overrides} + + +def test_access_round_trip_recovers_principal_and_caps_ttl(): + minted = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert minted.expires_at == NOW + timedelta(seconds=SESSION_TTL_SECONDS) + token = minted.token.get_secret_value() + assert is_session_token(token) + assert not is_session_refresh_token(token) + opened = open_session_token(token, KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + + +def test_refresh_round_trip_recovers_principal_and_caps_ttl(): + minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert minted.expires_at == NOW + timedelta(seconds=SESSION_REFRESH_TTL_SECONDS) + token = minted.token.get_secret_value() + assert is_session_refresh_token(token) + opened = open_session_refresh_token(token, KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + + +def test_access_token_reprefixed_as_refresh_is_rejected_by_signed_kind(): + body = _mint_access().removeprefix(SESSION_TOKEN_PREFIX) + swapped = SESSION_REFRESH_PREFIX + body + assert isinstance(open_session_refresh_token(swapped, KEYS, NOW), SessionMalformed) + + +def test_refresh_token_reprefixed_as_access_is_rejected_by_signed_kind(): + body = _mint_refresh().removeprefix(SESSION_REFRESH_PREFIX) + swapped = SESSION_TOKEN_PREFIX + body + assert isinstance(open_session_token(swapped, KEYS, NOW), SessionMalformed) + + +def test_refresh_token_is_not_an_access_token_at_the_edge(): + assert isinstance(open_session_token(_mint_refresh(), KEYS, NOW), NotASessionToken) + + +def test_expired_access_token_is_expired_not_malformed(): + token = _mint_access() + at_expiry = NOW + timedelta(seconds=SESSION_TTL_SECONDS) + assert isinstance(open_session_token(token, KEYS, at_expiry), SessionExpired) + after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + assert isinstance(open_session_token(token, KEYS, after), SessionExpired) + + +def test_still_valid_one_second_before_expiry(): + token = _mint_access() + just_before = NOW + timedelta(seconds=SESSION_TTL_SECONDS - 1) + assert isinstance(open_session_token(token, KEYS, just_before), OpenedSessionToken) + + +def test_tampered_signature_is_bad_signature(): + token = _mint_access() + tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") + assert isinstance(open_session_token(tampered, KEYS, NOW), SessionBadSignature) + + +def test_key_rotation_invalidates_outstanding_tokens(): + token = _mint_access() + assert isinstance(open_session_token(token, OTHER_KEYS, NOW), SessionBadSignature) + + +@pytest.mark.parametrize( + "candidate,expected", + [ + ("sk-1234", NotASessionToken), + ("llm_env_something", NotASessionToken), + ("", NotASessionToken), + (SESSION_TOKEN_PREFIX, SessionMalformed), + (SESSION_TOKEN_PREFIX + "not-a-jwt", SessionMalformed), + (SESSION_TOKEN_PREFIX + "\ud800garbage", SessionMalformed), + (SESSION_TOKEN_PREFIX + "a" * (MAX_SESSION_TOKEN_BYTES + 1), SessionMalformed), + ], +) +def test_hostile_candidates_never_raise(candidate, expected): + assert isinstance(open_session_token(candidate, KEYS, NOW), expected) + + +def test_multibyte_candidate_over_byte_cap_but_under_char_cap_is_rejected(): + filler = "€" * (MAX_SESSION_TOKEN_BYTES // 3) + candidate = SESSION_TOKEN_PREFIX + filler + assert len(candidate) <= MAX_SESSION_TOKEN_BYTES + assert isinstance(open_session_token(candidate, KEYS, NOW), SessionMalformed) + + +def test_alg_none_token_is_rejected(): + unsigned = jwt.api_jws.encode(b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none") + assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, KEYS, NOW), SessionMalformed) + + +@pytest.mark.parametrize( + "claims", + [ + _valid_claims(iss="wrong-issuer"), + _valid_claims(exp=str(int((NOW + timedelta(seconds=600)).timestamp()))), + _valid_claims(iat="evil"), + _valid_claims(kind="access"), + _valid_claims(user_id=""), + _valid_claims(nbf=0), + {k: v for k, v in _valid_claims().items() if k != "client_id"}, + {k: v for k, v in _valid_claims().items() if k != "exp"}, + ], +) +def test_signed_but_malformed_claims_are_rejected_without_raising(claims): + token = _sign_claims(claims) + assert isinstance(open_session_token(token, KEYS, NOW), SessionMalformed) + + +def test_signed_claims_with_exact_shape_open(): + token = _sign_claims(_valid_claims()) + opened = open_session_token(token, KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal.user_id == "user-123" + + +def test_oversized_client_id_fails_mint_with_typed_error_not_truncation(): + principal = SessionPrincipal(user_id="user-123", client_id="c" * (MAX_SESSION_TOKEN_BYTES + 100)) + minted = mint_session_token(principal, KEYS, NOW) + assert isinstance(minted, SessionTokenTooLarge) + assert minted.max_bytes == MAX_SESSION_TOKEN_BYTES + + +def test_empty_principal_fields_rejected_at_construction(): + with pytest.raises(ValidationError): + SessionPrincipal(user_id="", client_id="c") + with pytest.raises(ValidationError): + SessionPrincipal(user_id="u", client_id="") + + +def test_short_signing_key_rejected_at_construction(): + with pytest.raises(ValidationError): + SessionKeys(signing_key=SecretStr("short")) + + +def test_minted_token_repr_never_leaks_value(): + minted = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert minted.token.get_secret_value() not in repr(minted) From a22182f3c058aa66c0907a987520cc2874778c5a Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 00:31:06 -0700 Subject: [PATCH 19/60] feat(mcp): add jti claim for per-mint session token uniqueness --- .../mcp_server/outbound_credentials/session_token.py | 7 ++++++- .../mcp_server/outbound_credentials/test_session_token.py | 8 ++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py index 78b1f7e4916..9325428f049 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -10,7 +10,9 @@ outstanding sessions immediately without a revocation store. Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT, the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp`` -plus ``kind``, ``user_id``, and ``client_id``; ``client_id`` binds the refresh token +plus ``jti`` (per-mint uniqueness, so two tokens minted in the same second never +collide and a future revocation list has a stable handle), ``kind``, ``user_id``, and +``client_id``; ``client_id`` binds the refresh token to the DCR client it was issued to (RFC 6749 section 6) and is carried on the access token for parity and audit. There is no encrypted payload: nothing in a session token is secret beyond the signature, and reprs never print the signed value because minted @@ -28,6 +30,7 @@ injected ``now``); the strict pydantic claims model is the sole, total type gate from __future__ import annotations +import secrets from datetime import datetime, timedelta from typing import Literal, TypeAlias @@ -177,6 +180,7 @@ class _SessionClaims(BaseModel): iss: str iat: int exp: int + jti: str = Field(min_length=1) kind: SessionTokenKind user_id: str = Field(min_length=1) client_id: str = Field(min_length=1) @@ -276,6 +280,7 @@ def _mint( iss=SESSION_ISSUER, iat=int(now.timestamp()), exp=int(expires_at.timestamp()), + jti=secrets.token_urlsafe(16), kind=kind, user_id=principal.user_id, client_id=principal.client_id, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py index 551270f8d4b..a43592ebe18 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -57,6 +57,7 @@ def _valid_claims(**overrides) -> dict: "iss": SESSION_ISSUER, "iat": int(NOW.timestamp()), "exp": int((NOW + timedelta(seconds=600)).timestamp()), + "jti": "jti-fixed", "kind": "session", "user_id": "user-123", "client_id": "llm_client_abc", @@ -200,6 +201,13 @@ def test_short_signing_key_rejected_at_construction(): SessionKeys(signing_key=SecretStr("short")) +def test_two_mints_of_the_same_principal_are_distinct_tokens(): + first = mint_session_token(PRINCIPAL, KEYS, NOW) + second = mint_session_token(PRINCIPAL, KEYS, NOW) + assert isinstance(first, MintedSessionToken) and isinstance(second, MintedSessionToken) + assert first.token.get_secret_value() != second.token.get_secret_value() + + def test_minted_token_repr_never_leaks_value(): minted = mint_session_token(PRINCIPAL, KEYS, NOW) assert isinstance(minted, MintedSessionToken) From cc45d18e9c41b19d9eabe077288f7fcc6a080b11 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:24:56 -0700 Subject: [PATCH 20/60] feat(complexity-router): add return_raw_model_name toggle for response model field (#33875) * feat(complexity-router): optionally return raw model name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): restore asyncio import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(tests): preserve staging asyncio import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop unused local asyncio import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(dashboard): add complexity router raw model toggle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(complexity-router): move metadata key constant to constants.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy-tests): preserve module spacing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/common_request_processing.py | 12 +++++++- litellm/proxy/proxy_server.py | 4 +++ .../complexity_router/complexity_router.py | 7 +++++ .../complexity_router/config.py | 8 +++++ .../proxy_server/test_streaming_helpers.py | 16 ++++++++++ .../proxy/test_common_request_processing.py | 29 +++++++++++++++++-- .../router_strategy/test_complexity_router.py | 24 +++++++++++++++ .../add_model/ComplexityRouterConfig.test.tsx | 14 +++++++++ .../add_model/ComplexityRouterConfig.tsx | 25 +++++++++++++++- .../add_model/add_auto_router_tab.tsx | 2 ++ .../build_complexity_router_config.test.ts | 11 +++++++ .../build_complexity_router_config.ts | 4 +++ .../edit_auto_router_modal.test.ts | 10 +++++++ .../edit_auto_router_modal.tsx | 3 ++ 15 files changed, 166 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6432e2176c7..05944c81ea2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1292,6 +1292,7 @@ MAXIMUM_TRACEBACK_LINES_TO_LOG = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", X_LITELLM_DISABLE_CALLBACKS = "x-litellm-disable-callbacks" LITELLM_METADATA_FIELD = "litellm_metadata" OLD_LITELLM_METADATA_FIELD = "metadata" +RETURN_RAW_MODEL_NAME_METADATA_KEY = "_complexity_router_return_raw_model_name" LITELLM_TRUNCATED_PAYLOAD_FIELD = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1dc0ee3f947..3f9929f81da 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -33,6 +33,7 @@ from litellm.constants import ( LITELLM_DETAILED_TIMING, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, + RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -91,6 +92,13 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: StandardLoggingPayloadErrorInformation = } +def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: + return any( + isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True + for metadata in (request_data.get("metadata"), request_data.get("litellm_metadata")) + ) + + def _apply_client_disconnect_metadata(target_metadata: Optional[dict[str, object]]) -> None: if target_metadata is None: return @@ -672,6 +680,7 @@ def _override_openai_response_model( response_obj: Any, requested_model: str, log_context: str, + return_raw_model_name: bool = False, ) -> None: """ Force the OpenAI-compatible `model` field in the response to match what the client requested. @@ -695,7 +704,7 @@ def _override_openai_response_model( 3. If this was a fastest_response batch completion, use the winning model's model group name instead of the comma-separated list the client sent. """ - if not requested_model: + if return_raw_model_name or not requested_model: return hidden_params = get_hidden_params_dict(response_obj) @@ -1938,6 +1947,7 @@ class ProxyBaseLLMRequestProcessing: response_obj=response, requested_model=requested_model_from_client, log_context=f"litellm_call_id={logging_obj.litellm_call_id}", + return_raw_model_name=_should_return_raw_model_name(self.data), ) hidden_params = get_hidden_params_dict(response) # get any updated response headers diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index aed345c5db4..3b40abed19e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -291,6 +291,7 @@ from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, _is_azure_model_router_request, + _should_return_raw_model_name, create_response, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy @@ -7076,6 +7077,9 @@ def _restamp_streaming_chunk_model( fallback_was_attempted: bool = False, fallback_model_from_metadata: str | None = None, ) -> tuple[Any, bool]: + if _should_return_raw_model_name(request_data): + return chunk, model_mismatch_logged + target_model = fallback_model_from_metadata if fallback_was_attempted else requested_model_from_client # Always return the client-requested model name (not provider-prefixed internal identifiers) # on streaming chunks. diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 695d8b8aeaa..e5268b5107b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any, Literal, Union, cast from pydantic import BaseModel from litellm._logging import verbose_router_logger +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import ModelResponse @@ -956,6 +957,12 @@ class ComplexityRouter(CustomLogger): """ from litellm.types.router import PreRoutingHookResponse + if self.config.return_raw_model_name: + metadata_key = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" + metadata = request_kwargs.setdefault(metadata_key, {}) + if isinstance(metadata, dict): + metadata[RETURN_RAW_MODEL_NAME_METADATA_KEY] = True + use_session_affinity = self.config.session_affinity and not self.config.plugins session_id = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 17c2c287dde..7437138fbb7 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -311,6 +311,14 @@ class ComplexityRouterConfig(BaseModel): description="Default model to use if tier cannot be determined", ) + return_raw_model_name: bool = Field( + default=False, + description=( + "Return the resolved raw model name in the response model field instead of " + "the client-requested complexity-router alias" + ), + ) + # Classifier strategy classifier_type: Literal["heuristic", "llm"] = Field( default="heuristic", diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 699606b5277..f7e2d276a2e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -19,6 +19,7 @@ import json import pytest +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY import litellm.proxy.proxy_server as ps from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import ( @@ -272,6 +273,21 @@ def test_restamp_streaming_chunk_model_overrides_model_on_basemodel(): assert snapshot == {"model": "gpt-4", "logged": True, "same_object": True} +@pytest.mark.parametrize("return_raw_model_name", [False, True]) +def test_restamp_streaming_chunk_model_respects_raw_model_name_toggle(return_raw_model_name): + chunk = _simple_chunk(model="gpt-4o-mini") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="auto_router/complexity_router", + request_data={"metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: return_raw_model_name}}, + model_mismatch_logged=False, + ) + + expected_model = "gpt-4o-mini" if return_raw_model_name else "auto_router/complexity_router" + assert new_chunk.model == expected_model + assert logged is (not return_raw_model_name) + + def test_restamp_streaming_chunk_model_overrides_model_on_dict(): chunk = {"model": "internal", "choices": []} new_chunk, logged = _restamp_streaming_chunk_model( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index ebfbb46053d..58f81cdad35 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -11,6 +11,7 @@ from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( @@ -27,6 +28,7 @@ from litellm.proxy.common_request_processing import ( _is_azure_model_router_request, _override_openai_response_model, _parse_event_data_for_error, + _should_return_raw_model_name, _UpstreamClosingStreamingResponse, create_response, ) @@ -1675,6 +1677,31 @@ class TestExtractErrorFromSSEChunk: class TestOverrideOpenAIResponseModel: """Tests for _override_openai_response_model function""" + @pytest.mark.parametrize("return_raw_model_name", [False, True]) + def test_raw_model_name_toggle(self, return_raw_model_name): + response_obj = {"model": "gpt-4o-mini"} + + _override_openai_response_model( + response_obj=response_obj, + requested_model="auto_router/complexity_router", + log_context="test_context", + return_raw_model_name=return_raw_model_name, + ) + + expected_model = "gpt-4o-mini" if return_raw_model_name else "auto_router/complexity_router" + assert response_obj["model"] == expected_model + + @pytest.mark.parametrize( + "request_data, expected", + [ + ({"metadata": {}}, False), + ({"metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: True}}, True), + ({"litellm_metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: True}}, True), + ], + ) + def test_raw_model_name_toggle_metadata(self, request_data, expected): + assert _should_return_raw_model_name(request_data) is expected + def test_override_model_preserves_fallback_model_when_fallback_occurred_object( self, ): @@ -3203,8 +3230,6 @@ class TestDisconnectGatherCleanup: async def test_base_process_llm_request_preserves_llm_error_after_gather( self, monkeypatch ): - import asyncio - import litellm.proxy.common_request_processing as cpr from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 280a0fe072a..ef70687bd97 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -20,6 +20,7 @@ import litellm from litellm import Router from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.router_strategy.complexity_router.complexity_router import ( ComplexityRouter, DimensionScore, @@ -125,6 +126,29 @@ class TestComplexityRouterInit: ) assert router.config.default_model == "fallback-model" + @pytest.mark.asyncio + @pytest.mark.parametrize("return_raw_model_name", [False, True]) + async def test_pre_routing_hook_propagates_raw_model_response_setting( + self, mock_router_instance, basic_config, return_raw_model_name + ): + config = {**basic_config, "return_raw_model_name": return_raw_model_name} + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + request_kwargs = {} + + result = await router.async_pre_routing_hook( + model="test-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "Hello"}], + ) + + assert result is not None + metadata = request_kwargs.get("metadata", {}) + assert metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY, False) is return_raw_model_name + class TestTokenScoring: """Test token count scoring.""" 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 a2e2ca21d00..6b3c1961468 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -77,6 +77,20 @@ describe("ComplexityRouterConfig", () => { expect(screen.queryByText("Classifier Model")).not.toBeInTheDocument(); }); + it("should toggle returning the raw model name", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders(); + + await user.click(screen.getByText("Advanced: Response Format")); + await user.click(screen.getByRole("switch")); + + expect(onChange).toHaveBeenCalledWith({ + ...defaultValue, + return_raw_model_name: true, + }); + }); + it("should reveal classifier model and timeout fields when llm is selected", () => { const onChange = vi.fn(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 8008012a95c..1f2edf697a9 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -1,5 +1,5 @@ import { InfoCircleOutlined } from "@ant-design/icons"; -import { Select as AntdSelect, Card, Collapse, Divider, Space, Tooltip, Typography } from "antd"; +import { Select as AntdSelect, Card, Collapse, Divider, Space, Switch, Tooltip, Typography } from "antd"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; @@ -44,6 +44,7 @@ export interface ComplexityRouterConfigValue { adaptive_weights?: AdaptiveRouterWeights; tier_distance_penalty?: number; adaptive_eligible?: AdaptiveEligible; + return_raw_model_name?: boolean; } interface ComplexityRouterConfigProps { @@ -218,6 +219,28 @@ const ComplexityRouterConfig: React.FC = ({ ), children: , }, + { + key: "response", + label: ( + + Advanced: Response Format + + ), + children: ( + <> +
+ onChange({ ...value, return_raw_model_name: returnRawModelName })} + /> + Return raw model name +
+ + Return the resolved underlying model name in responses instead of the autorouter alias. + + + ), + }, ...(onEscalationKeywordsChange ? [ { 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 6e7bc49afce..e7826e09dce 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 @@ -100,6 +100,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc adaptive_weights: adaptiveWeights = DEFAULT_ADAPTIVE_WEIGHTS, tier_distance_penalty: tierDistancePenalty = DEFAULT_TIER_DISTANCE_PENALTY, adaptive_eligible: adaptiveEligible = "all", + return_raw_model_name: returnRawModelName = false, } = complexityRouterConfig; const missingTiersError = getMissingTiersError(tiers); @@ -148,6 +149,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc adaptiveWeights, tierDistancePenalty, adaptiveEligible, + returnRawModelName, }; const submitValues = { 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 0c9c19d1286..b5973bf7101 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 @@ -26,6 +26,7 @@ const baseParams: BuildComplexityRouterConfigParams = { adaptiveWeights: { quality: 0.3, cost: 0.7 }, tierDistancePenalty: 0.5, adaptiveEligible: "all", + returnRawModelName: false, }; describe("buildComplexityRouterConfig", () => { @@ -164,6 +165,16 @@ describe("buildComplexityRouterConfig", () => { expect(config.adaptive_eligible).toBeUndefined(); }); + it("omits return_raw_model_name when disabled", () => { + const config = buildComplexityRouterConfig({ ...baseParams, returnRawModelName: false }); + expect(config.return_raw_model_name).toBeUndefined(); + }); + + it("includes return_raw_model_name when enabled", () => { + const config = buildComplexityRouterConfig({ ...baseParams, returnRawModelName: true }); + expect(config.return_raw_model_name).toBe(true); + }); + it("includes tier_distance_penalty when adaptive is enabled with eligible='all'", () => { const config = buildComplexityRouterConfig({ ...baseParams, 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 0b92dc1b02d..3b41b916611 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 @@ -21,6 +21,7 @@ export interface BuildComplexityRouterConfigParams { adaptiveWeights: AdaptiveRouterWeights; tierDistancePenalty: number; adaptiveEligible: AdaptiveEligible; + returnRawModelName: boolean; } export interface ComplexityRouterConfigPayload { @@ -37,6 +38,7 @@ export interface ComplexityRouterConfigPayload { adaptive_weights?: AdaptiveRouterWeights; tier_distance_penalty?: number; adaptive_eligible?: AdaptiveEligible; + return_raw_model_name?: boolean; } const TIER_KEYS: Array = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"]; @@ -76,6 +78,7 @@ export const buildComplexityRouterConfig = ({ adaptiveWeights, tierDistancePenalty, adaptiveEligible, + returnRawModelName, }: 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 @@ -104,5 +107,6 @@ export const buildComplexityRouterConfig = ({ ...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }), adaptive_eligible: adaptiveEligible, }), + ...(returnRawModelName && { return_raw_model_name: true }), }; }; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index cd8093928d5..17fa810b529 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -18,6 +18,7 @@ const storedConfigValue = { adaptive_weights: { quality: 0.3, cost: 0.7 }, tier_distance_penalty: 0.8, adaptive_eligible: "all", + return_raw_model_name: true, }; const storedConfig = JSON.stringify(storedConfigValue); @@ -80,6 +81,15 @@ describe("buildUpdatedComplexityRouterConfig", () => { expect(updatedConfig).toEqual(expectedAdaptiveDisabledConfig); }); + it("includes return_raw_model_name only when enabled", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, { + ...classifiedTierValue, + return_raw_model_name: true, + }); + + expect(updatedConfig.return_raw_model_name).toBe(true); + }); + it("updates custom technical keywords when they are edited", () => { const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, classifiedTierValue, ["postgres"]); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index f85cd16486a..46d7d41d9b3 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -38,6 +38,7 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "adaptive_weights", "tier_distance_penalty", "adaptive_eligible", + "return_raw_model_name", ]); const toRecord = (value: unknown): Record => { @@ -78,6 +79,7 @@ export const buildUpdatedComplexityRouterConfig = ( }), adaptive_eligible: adaptiveEligible, }), + ...(value.return_raw_model_name && { return_raw_model_name: true }), }; }; @@ -158,6 +160,7 @@ const EditAutoRouterModal: React.FC = ({ adaptive_weights: parsedConfig.adaptive_weights, tier_distance_penalty: parsedConfig.tier_distance_penalty, adaptive_eligible: parsedConfig.adaptive_eligible || "all", + return_raw_model_name: parsedConfig.return_raw_model_name || false, }); setCustomTechnicalKeywords( Array.isArray(parsedConfig.custom_technical_keywords) ? parsedConfig.custom_technical_keywords : [], From bd44c9e305b89526d4c5d773ee39ca935561b9c8 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 18 Jul 2026 20:36:51 -0700 Subject: [PATCH 21/60] fix(langfuse): send v4 ingestion header for otel callback (#33907) * fix(langfuse): send v4 ingestion header for otel callback * refactor(langfuse): inline otel ingestion header literals * test(langfuse): assert v4 ingestion header on dynamic key config paths * style: apply ruff format to langfuse otel header changes * chore(langfuse): drop stale development annotation on json import --------- Co-authored-by: Hassieb Pakzad <68423100+hassiebp@users.noreply.github.com> --- .../integrations/langfuse/langfuse_otel.py | 32 +++++++++++++-- .../integrations/test_langfuse_otel.py | 12 ++++-- tests/test_service_logger_otel.py | 39 +++++++++++++++++++ 3 files changed, 76 insertions(+), 7 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 449457bd123..d464d55453d 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -1,8 +1,8 @@ import base64 -import json # <--- NEW +import json import os from datetime import datetime -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Union from litellm._logging import verbose_logger from litellm.integrations.arize import _utils @@ -25,6 +25,8 @@ else: LANGFUSE_CLOUD_EU_ENDPOINT = "https://cloud.langfuse.com/api/public/otel" LANGFUSE_CLOUD_US_ENDPOINT = "https://us.cloud.langfuse.com/api/public/otel" +LANGFUSE_INGESTION_VERSION_HEADER = "x-langfuse-ingestion-version" +LANGFUSE_INGESTION_VERSION = "4" class LangfuseOtelLogger(OpenTelemetry): @@ -326,7 +328,9 @@ class LangfuseOtelLogger(OpenTelemetry): return OpenTelemetryConfig( exporter="otlp_http", endpoint=endpoint, - headers=f"Authorization={auth_header}", + headers=LangfuseOtelLogger._format_otel_headers( + LangfuseOtelLogger._build_langfuse_otel_headers(auth_header) + ), ) @staticmethod @@ -338,6 +342,26 @@ class LangfuseOtelLogger(OpenTelemetry): auth_header = base64.b64encode(auth_string.encode()).decode() return f"Basic {auth_header}" + @staticmethod + def _build_langfuse_otel_headers(auth_header: str) -> Dict[str, str]: + """ + Build the OTLP header set Langfuse expects. + + `x-langfuse-ingestion-version: 4` selects Langfuse's v4 ingestion path; + without it spans fall back to the older transformation path. + """ + return { + "Authorization": auth_header, + LANGFUSE_INGESTION_VERSION_HEADER: LANGFUSE_INGESTION_VERSION, + } + + @staticmethod + def _format_otel_headers(headers: Dict[str, str]) -> str: + """ + Serialize a header mapping into the comma-separated OTLP header string + """ + return ",".join(f"{key}={value}" for key, value in headers.items()) + def construct_dynamic_otel_headers( self, standard_callback_dynamic_params: StandardCallbackDynamicParams ) -> Optional[dict]: @@ -358,7 +382,7 @@ class LangfuseOtelLogger(OpenTelemetry): public_key=dynamic_langfuse_public_key, secret_key=dynamic_langfuse_secret_key, ) - dynamic_headers["Authorization"] = auth_header + dynamic_headers.update(LangfuseOtelLogger._build_langfuse_otel_headers(auth_header)) return dynamic_headers diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/test_litellm/integrations/test_langfuse_otel.py index 2f2675ca790..28f138c7acd 100644 --- a/tests/test_litellm/integrations/test_langfuse_otel.py +++ b/tests/test_litellm/integrations/test_langfuse_otel.py @@ -456,7 +456,7 @@ class TestLangfuseOtelKeyDynamicConfig: import base64 expected_auth = base64.b64encode(b"key_public:key_secret").decode() - assert config.headers == f"Authorization=Basic {expected_auth}" + assert config.headers == f"Authorization=Basic {expected_auth},x-langfuse-ingestion-version=4" def test_construct_dynamic_otel_config_host_without_protocol(self): with self._clean_env(): @@ -521,7 +521,10 @@ class TestLangfuseOtelKeyDynamicConfig: import base64 expected_auth = base64.b64encode(b"key_public:key_secret").decode() - assert exporter._headers == {"Authorization": f"Basic {expected_auth}"} + assert exporter._headers == { + "Authorization": f"Basic {expected_auth}", + "x-langfuse-ingestion-version": "4", + } def test_key_dynamic_params_reuse_cached_provider(self): with self._clean_env(): @@ -574,7 +577,10 @@ class TestLangfuseOtelKeyDynamicConfig: provider = next(iter(logger._tracer_provider_cache.values())) exporter = provider._active_span_processor._span_processors[0].span_exporter assert isinstance(exporter, OTLPSpanExporter) - assert exporter._headers == {"Authorization": f"Basic {secret}"} + assert exporter._headers == { + "Authorization": f"Basic {secret}", + "x-langfuse-ingestion-version": "4", + } class TestLangfuseOtelResponsesAPI: diff --git a/tests/test_service_logger_otel.py b/tests/test_service_logger_otel.py index 35070d55546..044d37d6781 100644 --- a/tests/test_service_logger_otel.py +++ b/tests/test_service_logger_otel.py @@ -12,6 +12,7 @@ from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger from litellm.integrations.opentelemetry import OpenTelemetry from litellm.types.services import ServiceTypes from litellm._service_logger import ServiceLogging +from litellm.types.utils import StandardCallbackDynamicParams class TestServiceLoggerOTEL(unittest.IsolatedAsyncioTestCase): @@ -108,6 +109,44 @@ class TestServiceLoggerOTEL(unittest.IsolatedAsyncioTestCase): "Generic OTEL logger should have received the log exactly once.", ) + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_tracing") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_metrics") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_logs") + async def test_langfuse_otel_env_config_includes_v4_ingestion_header( + self, mock_logs, mock_metrics, mock_tracing + ): + logger = LangfuseOtelLogger() + + headers = OpenTelemetry._get_headers_dictionary(logger.config.headers) + + self.assertEqual( + headers["x-langfuse-ingestion-version"], + "4", + ) + self.assertTrue(headers["Authorization"].startswith("Basic ")) + + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_tracing") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_metrics") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_logs") + async def test_langfuse_otel_dynamic_headers_include_v4_ingestion_header( + self, mock_logs, mock_metrics, mock_tracing + ): + logger = LangfuseOtelLogger() + + headers = logger.construct_dynamic_otel_headers( + StandardCallbackDynamicParams( + langfuse_public_key="pk-lf-dynamic", + langfuse_secret_key="sk-lf-dynamic", + ) + ) + + self.assertIsNotNone(headers) + self.assertEqual( + headers["x-langfuse-ingestion-version"], + "4", + ) + self.assertTrue(headers["Authorization"].startswith("Basic ")) + if __name__ == "__main__": unittest.main() From 34561482ed092d78c296cab7999486022af5a938 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 07:30:49 -0700 Subject: [PATCH 22/60] feat(ui): add configuration tabs to the Cost Optimization page (#33899) * feat(ui): add configuration tabs to Cost Optimization page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): reuse AutoRouter v2 and Router Settings prompt-caching panel in Cost Optimization; clarify Headroom compression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): add experimental dashboard banner with feedback discussion link to Cost Optimization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): add savings methodology note and per-key/team compression enterprise callout to Cost Optimization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): assert active tab state in Cost Optimization tab-switch test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/AutorouterTab.tsx | 28 +++ .../_components/CostOptimizationView.test.tsx | 113 ++--------- .../_components/CostOptimizationView.tsx | 187 +++++------------- .../_components/PromptCachingTab.tsx | 52 +++++ .../_components/PromptCompressionTab.tsx | 176 +++++++++++++++++ .../_components/UsageTab.test.tsx | 108 ++++++++++ .../_components/UsageTab.tsx | 184 +++++++++++++++++ .../_components/helpers.test.ts | 37 ++++ .../cost-optimization/_components/helpers.ts | 39 ++++ .../_components/general_settings.tsx | 4 +- 10 files changed, 699 insertions(+), 229 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx new file mode 100644 index 00000000000..3474036528a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx @@ -0,0 +1,28 @@ +"use client"; + +import React from "react"; +import { Form } from "antd"; + +import AddAutoRouterTab from "@/components/add_model/add_auto_router_tab"; + +interface AutorouterTabProps { + accessToken: string | null; + userId: string | null; + userRole: string; +} + +const AutorouterTab: React.FC = ({ accessToken, userRole }) => { + const [form] = Form.useForm(); + + if (!accessToken) { + return null; + } + + return ( +
+ form.resetFields()} accessToken={accessToken} userRole={userRole} /> +
+ ); +}; + +export default AutorouterTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx index 92426d3ce04..46aa23fcfc0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx @@ -1,109 +1,34 @@ -import { render } from "@testing-library/react"; +import { fireEvent, render } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; -import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; - -const mockUsePaginatedDailyActivity = vi.fn(); - -vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", () => ({ - usePaginatedDailyActivity: (args: unknown) => mockUsePaginatedDailyActivity(args), -})); - -vi.mock("@/components/networking", () => ({ - userDailyActivityCall: vi.fn(), -})); - -vi.mock("@/components/shared/advanced_date_picker", () => ({ - __esModule: true, - default: () =>
, -})); - -vi.mock("@/components/shared/charts", () => ({ - AreaChart: ({ data, categories }: { data: unknown; categories: string[] }) => ( -
- ), - DonutChart: ({ data, label }: { data: unknown; label: string }) => ( -
- ), -})); +vi.mock("./UsageTab", () => ({ __esModule: true, default: () =>
})); +vi.mock("./PromptCompressionTab", () => ({ __esModule: true, default: () =>
})); +vi.mock("./AutorouterTab", () => ({ __esModule: true, default: () =>
})); +vi.mock("./PromptCachingTab", () => ({ __esModule: true, default: () =>
})); import CostOptimizationView from "./CostOptimizationView"; -const baseMetrics = (overrides: Partial): SpendMetrics => ({ - spend: 0, - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - api_requests: 0, - successful_requests: 0, - failed_requests: 0, - cache_read_input_tokens: 0, - cache_creation_input_tokens: 0, - ...overrides, -}); - -const day = (date: string, metrics: Partial): DailyData => ({ - date, - metrics: baseMetrics(metrics), - breakdown: { - models: {}, - model_groups: {}, - mcp_servers: {}, - providers: {}, - api_keys: {}, - entities: {}, - }, -}); - -const renderWith = (results: DailyData[]) => { - mockUsePaginatedDailyActivity.mockReturnValue({ data: { results }, loading: false, isFetchingMore: false }); - return render(); -}; +const renderView = () => render(); describe("CostOptimizationView", () => { - it("sums compression and caching dollars across days into the summary cards", () => { - const { getByText } = renderWith([ - day("2026-07-12", { - compression_savings_spend: 0.04, - prompt_caching_savings_spend: 0.006, - compression_saved_tokens: 40000, - }), - day("2026-07-13", { - compression_savings_spend: 0.1, - prompt_caching_savings_spend: 0.01, - compression_saved_tokens: 100000, - }), - ]); + it("renders all four cost-optimization tabs", () => { + const { getByText } = renderView(); - // compression 0.14 + caching 0.016 = 0.156 - expect(getByText("$0.1560")).toBeInTheDocument(); - expect(getByText("$0.1400")).toBeInTheDocument(); - expect(getByText("$0.0160")).toBeInTheDocument(); - expect(getByText("140,000 tokens compressed")).toBeInTheDocument(); + expect(getByText("Usage")).toBeInTheDocument(); + expect(getByText("Prompt Compression")).toBeInTheDocument(); + expect(getByText("Autorouter")).toBeInTheDocument(); + expect(getByText("Prompt Caching")).toBeInTheDocument(); }); - it("builds a per-day time series and per-driver donut from the daily rows", () => { - const { getByTestId } = renderWith([ - day("2026-07-12", { compression_savings_spend: 0.04, prompt_caching_savings_spend: 0.006 }), - day("2026-07-13", { compression_savings_spend: 0.1, prompt_caching_savings_spend: 0.01 }), - ]); + it("defaults to the Usage tab and switches the active tab on click", () => { + const { getByRole } = renderView(); - const series = JSON.parse(getByTestId("area-chart").getAttribute("data-series") ?? "[]"); - expect(series).toHaveLength(2); - expect(series[0]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); - expect(series[1]).toMatchObject({ Compression: 0.1, "Prompt caching": 0.01 }); + expect(getByRole("tab", { name: "Usage" })).toHaveAttribute("aria-selected", "true"); + expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); - expect(slices).toEqual([ - { driver: "Compression", usd: expect.closeTo(0.14, 5) }, - { driver: "Prompt caching", usd: expect.closeTo(0.016, 5) }, - ]); - }); + fireEvent.click(getByRole("tab", { name: "Prompt Compression" })); - it("omits a driver slice when that driver has no savings", () => { - const { getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); - - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); - expect(slices).toEqual([{ driver: "Compression", usd: expect.closeTo(0.04, 5) }]); + expect(getByRole("tab", { name: "Usage" })).toHaveAttribute("aria-selected", "false"); + expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx index e45e3e23f1d..6e6830b8451 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx @@ -1,16 +1,13 @@ "use client"; -import React, { useMemo, useState } from "react"; +import React from "react"; import { PiggyBank } from "lucide-react"; +import { Alert, Tabs } from "antd"; -import { AreaChart, DonutChart } from "@/components/shared/charts"; -import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; -import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; -import { userDailyActivityCall } from "@/components/networking"; -import { DailyData, SpendMetrics } from "@/components/UsagePage/types"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; -import { all_admin_roles } from "@/utils/roles"; -import { usePaginatedDailyActivity } from "@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity"; +import UsageTab from "./UsageTab"; +import PromptCompressionTab from "./PromptCompressionTab"; +import AutorouterTab from "./AutorouterTab"; +import PromptCachingTab from "./PromptCachingTab"; interface CostOptimizationViewProps { accessToken: string | null; @@ -18,138 +15,62 @@ interface CostOptimizationViewProps { userRole: string; } -type DateRange = { from?: Date; to?: Date }; - -const THIRTY_DAYS_MS = 30 * 24 * 60 * 60 * 1000; - -const usd = (value: number): string => { - const decimals = value > 0 && value < 1 ? 4 : 2; - return `$${formatNumberWithCommas(value, decimals)}`; -}; - -const shortDate = (iso: string): string => - new Date(`${iso}T00:00:00`).toLocaleDateString("en-US", { month: "short", day: "numeric" }); - -const compressionOf = (m: SpendMetrics): number => m.compression_savings_spend ?? 0; -const cachingOf = (m: SpendMetrics): number => m.prompt_caching_savings_spend ?? 0; -const savedTokensOf = (m: SpendMetrics): number => m.compression_saved_tokens ?? 0; - -const SummaryCard = ({ label, value, hint }: { label: string; value: string; hint?: string }) => ( - - - {label} - - -

{value}

- {hint &&

{hint}

} - - -); - const CostOptimizationView: React.FC = ({ accessToken, userId, userRole }) => { - const initialFrom = useMemo(() => new Date(new Date().getTime() - THIRTY_DAYS_MS), []); - const initialTo = useMemo(() => new Date(), []); - const [dateValue, setDateValue] = useState({ from: initialFrom, to: initialTo }); - - const startTime = dateValue.from ?? null; - const endTime = dateValue.to ?? null; - const isAdmin = all_admin_roles.includes(userRole); - const effectiveUserId = isAdmin ? null : userId; - - const { data, loading, isFetchingMore } = usePaginatedDailyActivity({ - fetchFn: userDailyActivityCall, - args: [accessToken, startTime, endTime, effectiveUserId], - enabled: !!accessToken && !!startTime && !!endTime, - }); - - const results = data.results as DailyData[]; - - const compressionTotal = useMemo(() => results.reduce((sum, d) => sum + compressionOf(d.metrics), 0), [results]); - const cachingTotal = useMemo(() => results.reduce((sum, d) => sum + cachingOf(d.metrics), 0), [results]); - const savedTokensTotal = useMemo(() => results.reduce((sum, d) => sum + savedTokensOf(d.metrics), 0), [results]); - const totalSaved = compressionTotal + cachingTotal; - - const overTime = useMemo( - () => - results.map((d) => ({ - date: shortDate(d.date), - Compression: compressionOf(d.metrics), - "Prompt caching": cachingOf(d.metrics), - })), - [results], - ); - - const byDriver = useMemo( - () => - [ - { driver: "Compression", usd: compressionTotal }, - { driver: "Prompt caching", usd: cachingTotal }, - ].filter((d) => d.usd > 0), - [compressionTotal, cachingTotal], - ); + const items = [ + { + key: "usage", + label: "Usage", + children: , + }, + { + key: "compression", + label: "Prompt Compression", + children: , + }, + { + key: "autorouter", + label: "Autorouter", + children: , + }, + { + key: "caching", + label: "Prompt Caching", + children: , + }, + ]; return (
-
-
-
- -

Cost Optimization

-
-

- Money saved by prompt compression and prompt caching across your requests -

+
+
+ +

Cost Optimization

- setDateValue(v)} /> +

+ Track and configure the mechanisms that save you money: prompt compression, prompt caching, and auto routing +

-
- - - -
+ + Have feedback? Join the discussion{" "} + + here + + + } + /> -
- - - Savings over time - - - - - - - - Savings by driver - - - - - -
+
); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx new file mode 100644 index 00000000000..e6f73088824 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx @@ -0,0 +1,52 @@ +"use client"; + +import React, { useCallback, useEffect, useState } from "react"; + +import { getGeneralSettingsCall } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { + PromptCachingPanel, + generalSettingsItem, +} from "@/app/(dashboard)/router-settings/_components/general_settings"; + +interface PromptCachingTabProps { + accessToken: string | null; +} + +const PromptCachingTab: React.FC = ({ accessToken }) => { + const [settings, setSettings] = useState([]); + + const loadSettings = useCallback(() => { + if (!accessToken) { + return; + } + getGeneralSettingsCall(accessToken) + .then((data: generalSettingsItem[]) => setSettings(data)) + .catch((error) => { + console.error("Failed to load prompt caching settings:", error); + NotificationsManager.fromBackend("Failed to load prompt caching settings"); + }); + }, [accessToken]); + + useEffect(() => { + loadSettings(); + }, [loadSettings]); + + const handleChange = (fieldName: string, newValue: unknown) => { + setSettings((prev) => + prev.map((setting) => (setting.field_name === fieldName ? { ...setting, field_value: newValue } : setting)), + ); + }; + + if (!accessToken) { + return null; + } + + return ( +
+ +
+ ); +}; + +export default PromptCachingTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx new file mode 100644 index 00000000000..071ad1d6bd5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx @@ -0,0 +1,176 @@ +"use client"; + +import React, { useCallback, useEffect, useState } from "react"; +import { Button, Form, Input, Switch } from "antd"; + +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { createGuardrailCall, getGuardrailsList } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { + buildCompressionGuardrailPayload, + compressionGuardrailsOf, + GuardrailListItem, + GuardrailListResponse, +} from "./helpers"; + +interface PromptCompressionTabProps { + accessToken: string | null; +} + +interface CompressionFormValues { + name: string; + apiBase: string; + defaultOn: boolean; +} + +const PromptCompressionTab: React.FC = ({ accessToken }) => { + const [form] = Form.useForm(); + const [guardrails, setGuardrails] = useState([]); + const [isLoading, setIsLoading] = useState(true); + const [isSaving, setIsSaving] = useState(false); + + const loadGuardrails = useCallback(() => { + if (!accessToken) { + return; + } + getGuardrailsList(accessToken) + .then((response) => setGuardrails(compressionGuardrailsOf(response as GuardrailListResponse))) + .catch((error) => { + console.error("Failed to load compression guardrails:", error); + NotificationsManager.fromBackend("Failed to load compression guardrails"); + }) + .finally(() => setIsLoading(false)); + }, [accessToken]); + + useEffect(() => { + loadGuardrails(); + }, [loadGuardrails]); + + const handleAdd = async (values: CompressionFormValues) => { + if (!accessToken) { + return; + } + setIsSaving(true); + try { + await createGuardrailCall( + accessToken, + buildCompressionGuardrailPayload({ + name: values.name, + apiBase: values.apiBase, + defaultOn: values.defaultOn ?? true, + }), + ); + NotificationsManager.success("Compression guardrail created"); + form.resetFields(); + await loadGuardrails(); + } catch (error) { + console.error("Failed to create compression guardrail:", error); + NotificationsManager.fromBackend("Failed to create compression guardrail"); + } finally { + setIsSaving(false); + } + }; + + return ( +
+ + + Headroom prompt compression + + +

+ Headroom is a native LiteLLM guardrail that compresses your prompts before they reach the model, so you pay + for fewer input tokens. The tokens it removes are priced and shown on the Usage tab as compression savings.{" "} + + Headroom setup docs + +

+ {isLoading &&

Loading...

} + {!isLoading && guardrails.length === 0 && ( +

+ No prompt compression guardrails configured yet. Add one below to start saving on input tokens +

+ )} + {!isLoading && guardrails.length > 0 && ( +
    + {guardrails.map((guardrail) => ( +
  • +
    +

    {guardrail.guardrail_name}

    +

    {guardrail.litellm_params?.api_base ?? ""}

    +
    + + {guardrail.litellm_params?.default_on ? "Always on" : "Opt-in"} + +
    • + ))} +
+ )} + + + + + + Add Headroom compression guardrail + + +
+ + + + + + + + + +
+

+ Applying compression to all requests is available to all users. Enabling it selectively per key or team + is a LiteLLM Enterprise feature. Get a trial key{" "} + + here + +

+
+
+ +
+
+ + +
+ ); +}; + +export default PromptCompressionTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx new file mode 100644 index 00000000000..048d9a34d31 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -0,0 +1,108 @@ +import { render } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; + +const mockUsePaginatedDailyActivity = vi.fn(); + +vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", () => ({ + usePaginatedDailyActivity: (args: unknown) => mockUsePaginatedDailyActivity(args), +})); + +vi.mock("@/components/networking", () => ({ + userDailyActivityCall: vi.fn(), +})); + +vi.mock("@/components/shared/advanced_date_picker", () => ({ + __esModule: true, + default: () =>
, +})); + +vi.mock("@/components/shared/charts", () => ({ + AreaChart: ({ data, categories }: { data: unknown; categories: string[] }) => ( +
+ ), + DonutChart: ({ data, label }: { data: unknown; label: string }) => ( +
+ ), +})); + +import UsageTab from "./UsageTab"; + +const baseMetrics = (overrides: Partial): SpendMetrics => ({ + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + ...overrides, +}); + +const day = (date: string, metrics: Partial): DailyData => ({ + date, + metrics: baseMetrics(metrics), + breakdown: { + models: {}, + model_groups: {}, + mcp_servers: {}, + providers: {}, + api_keys: {}, + entities: {}, + }, +}); + +const renderWith = (results: DailyData[]) => { + mockUsePaginatedDailyActivity.mockReturnValue({ data: { results }, loading: false, isFetchingMore: false }); + return render(); +}; + +describe("UsageTab", () => { + it("sums compression and caching dollars across days into the summary cards", () => { + const { getByText } = renderWith([ + day("2026-07-12", { + compression_savings_spend: 0.04, + prompt_caching_savings_spend: 0.006, + compression_saved_tokens: 40000, + }), + day("2026-07-13", { + compression_savings_spend: 0.1, + prompt_caching_savings_spend: 0.01, + compression_saved_tokens: 100000, + }), + ]); + + expect(getByText("$0.1560")).toBeInTheDocument(); + expect(getByText("$0.1400")).toBeInTheDocument(); + expect(getByText("$0.0160")).toBeInTheDocument(); + expect(getByText("140,000 tokens compressed")).toBeInTheDocument(); + }); + + it("builds a per-day time series and per-driver donut from the daily rows", () => { + const { getByTestId } = renderWith([ + day("2026-07-12", { compression_savings_spend: 0.04, prompt_caching_savings_spend: 0.006 }), + day("2026-07-13", { compression_savings_spend: 0.1, prompt_caching_savings_spend: 0.01 }), + ]); + + const series = JSON.parse(getByTestId("area-chart").getAttribute("data-series") ?? "[]"); + expect(series).toHaveLength(2); + expect(series[0]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); + expect(series[1]).toMatchObject({ Compression: 0.1, "Prompt caching": 0.01 }); + + const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + expect(slices).toEqual([ + { driver: "Compression", usd: expect.closeTo(0.14, 5) }, + { driver: "Prompt caching", usd: expect.closeTo(0.016, 5) }, + ]); + }); + + it("omits a driver slice when that driver has no savings", () => { + const { getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); + + const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + expect(slices).toEqual([{ driver: "Compression", usd: expect.closeTo(0.04, 5) }]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx new file mode 100644 index 00000000000..8e6fc40b5ad --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -0,0 +1,184 @@ +"use client"; + +import React, { useMemo, useState } from "react"; +import { Collapse } from "antd"; + +import { AreaChart, DonutChart } from "@/components/shared/charts"; +import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { userDailyActivityCall } from "@/components/networking"; +import { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { all_admin_roles } from "@/utils/roles"; +import { usePaginatedDailyActivity } from "@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity"; + +interface UsageTabProps { + accessToken: string | null; + userId: string | null; + userRole: string; +} + +type DateRange = { from?: Date; to?: Date }; + +const THIRTY_DAYS_MS = 30 * 24 * 60 * 60 * 1000; + +const usd = (value: number): string => { + const decimals = value > 0 && value < 1 ? 4 : 2; + return `$${formatNumberWithCommas(value, decimals)}`; +}; + +const shortDate = (iso: string): string => + new Date(`${iso}T00:00:00`).toLocaleDateString("en-US", { month: "short", day: "numeric" }); + +const compressionOf = (m: SpendMetrics): number => m.compression_savings_spend ?? 0; +const cachingOf = (m: SpendMetrics): number => m.prompt_caching_savings_spend ?? 0; +const savedTokensOf = (m: SpendMetrics): number => m.compression_saved_tokens ?? 0; + +const MethodologyNote = () => ( + How savings are calculated, + children: ( +
+

+ Savings are computed for each request when it is logged, using the provider's reported usage and the + model's pricing, then summed into a daily rollup. Totals below are read from that rollup over the + selected date range, so the numbers never require a scan of raw request logs. +

+

+ Compression savings are the tokens Headroom removed before the call, priced at the model's input + rate: compression_saved_tokens * input_cost_per_token +

+

+ Prompt caching savings are the tokens the provider served from cache (Anthropic{" "} + cache_read_input_tokens, or OpenAI-style prompt_tokens_details.cached_tokens), + priced at the discount between the normal input rate and the cache-read rate:{" "} + cache_read_input_tokens * max(input_cost_per_token - cache_read_input_token_cost, 0) +

+

+ Total saved is the sum of both drivers. Models without a separate cache-read price in the pricing map + contribute zero caching savings rather than erroring. +

+
+ ), + }, + ]} + /> +); + +const SummaryCard = ({ label, value, hint }: { label: string; value: string; hint?: string }) => ( + + + {label} + + +

{value}

+ {hint &&

{hint}

} + + +); + +const UsageTab: React.FC = ({ accessToken, userId, userRole }) => { + const initialFrom = useMemo(() => new Date(new Date().getTime() - THIRTY_DAYS_MS), []); + const initialTo = useMemo(() => new Date(), []); + const [dateValue, setDateValue] = useState({ from: initialFrom, to: initialTo }); + + const startTime = dateValue.from ?? null; + const endTime = dateValue.to ?? null; + const isAdmin = all_admin_roles.includes(userRole); + const effectiveUserId = isAdmin ? null : userId; + + const { data, loading, isFetchingMore } = usePaginatedDailyActivity({ + fetchFn: userDailyActivityCall, + args: [accessToken, startTime, endTime, effectiveUserId], + enabled: !!accessToken && !!startTime && !!endTime, + }); + + const results = data.results as DailyData[]; + + const compressionTotal = useMemo(() => results.reduce((sum, d) => sum + compressionOf(d.metrics), 0), [results]); + const cachingTotal = useMemo(() => results.reduce((sum, d) => sum + cachingOf(d.metrics), 0), [results]); + const savedTokensTotal = useMemo(() => results.reduce((sum, d) => sum + savedTokensOf(d.metrics), 0), [results]); + const totalSaved = compressionTotal + cachingTotal; + + const overTime = useMemo( + () => + results.map((d) => ({ + date: shortDate(d.date), + Compression: compressionOf(d.metrics), + "Prompt caching": cachingOf(d.metrics), + })), + [results], + ); + + const byDriver = useMemo( + () => + [ + { driver: "Compression", usd: compressionTotal }, + { driver: "Prompt caching", usd: cachingTotal }, + ].filter((d) => d.usd > 0), + [compressionTotal, cachingTotal], + ); + + return ( +
+
+ + setDateValue(v)} /> +
+ +
+ + + +
+ +
+ + + Savings over time + + + + + + + + Savings by driver + + + + + +
+
+ ); +}; + +export default UsageTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.test.ts new file mode 100644 index 00000000000..239cded0815 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; + +import { buildCompressionGuardrailPayload, compressionGuardrailsOf } from "./helpers"; + +describe("compressionGuardrailsOf", () => { + it("keeps only headroom-provider guardrails and drops others", () => { + const filtered = compressionGuardrailsOf({ + guardrails: [ + { guardrail_id: "1", guardrail_name: "headroom-compression", litellm_params: { guardrail: "headroom" } }, + { guardrail_id: "2", guardrail_name: "pii-masker", litellm_params: { guardrail: "presidio" } }, + { guardrail_id: "3", guardrail_name: "no-params", litellm_params: null }, + ], + }); + + expect(filtered.map((g) => g.guardrail_id)).toEqual(["1"]); + }); +}); + +describe("buildCompressionGuardrailPayload", () => { + it("builds a headroom guardrail payload with trimmed fields", () => { + const payload = buildCompressionGuardrailPayload({ + name: " headroom-compression ", + apiBase: " https://compress ", + defaultOn: false, + }); + + expect(payload).toEqual({ + guardrail_name: "headroom-compression", + litellm_params: { + guardrail: "headroom", + mode: "pre_call", + api_base: "https://compress", + default_on: false, + }, + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.ts new file mode 100644 index 00000000000..7c8c92f8890 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.ts @@ -0,0 +1,39 @@ +export interface GuardrailLitellmParams { + guardrail?: string | null; + api_base?: string | null; + default_on?: boolean | null; +} + +export interface GuardrailListItem { + guardrail_id: string; + guardrail_name: string | null; + litellm_params?: GuardrailLitellmParams | null; +} + +export interface GuardrailListResponse { + guardrails?: GuardrailListItem[]; +} + +export const COMPRESSION_GUARDRAIL_PROVIDER = "headroom"; + +export const isCompressionGuardrail = (guardrail: GuardrailListItem): boolean => + (guardrail.litellm_params?.guardrail ?? "").toLowerCase() === COMPRESSION_GUARDRAIL_PROVIDER; + +export const compressionGuardrailsOf = (response: GuardrailListResponse): GuardrailListItem[] => + (response.guardrails ?? []).filter(isCompressionGuardrail); + +export interface CompressionGuardrailInput { + name: string; + apiBase: string; + defaultOn: boolean; +} + +export const buildCompressionGuardrailPayload = (input: CompressionGuardrailInput): Record => ({ + guardrail_name: input.name.trim(), + litellm_params: { + guardrail: COMPRESSION_GUARDRAIL_PROVIDER, + mode: "pre_call", + api_base: input.apiBase.trim(), + default_on: input.defaultOn, + }, +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 1e8658d5104..dda7a23a8d4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -33,7 +33,7 @@ interface GeneralSettingsPageProps { userID: string | null; } -interface generalSettingsItem { +export interface generalSettingsItem { field_name: string; field_type: string; field_value: any; @@ -90,7 +90,7 @@ const SettingValueEditor: React.FC<{ return null; }; -const PromptCachingPanel: React.FC<{ +export const PromptCachingPanel: React.FC<{ accessToken: string; settings: generalSettingsItem[]; onChange: (fieldName: string, newValue: any) => void; From 8d96e959db0e251d6fbda9d74360617d18aff806 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 08:47:39 -0700 Subject: [PATCH 23/60] test(e2e): guard destructive spend-log truncate behind an explicit opt-in (#33751) tests/e2e/conftest.py's pytest_sessionfinish truncated LiteLLM_SpendLogs against whatever DATABASE_URL resolved to, gated only by "an e2e test body ran". Pointed at a shared or staging DB, a routine local run wiped real spend data. It also reached the truncate helper through a sys.path.insert into quota_management/spend_tracking/spend_e2e_client.py, a cross-suite import-by-path hack it then unwound in a finally. The cleanup now routes through a new run_spend_log_cleanup in a top-level tests/e2e/e2e_db.py, which fires the destructive truncate only when the operator set E2E_RESET_SPEND_LOGS=1 and an e2e test actually ran. Any other value (unset, 0, true, empty) leaves the DB untouched, so presence of the variable alone or a test run alone never arms the truncate. The decision plus the injectable truncate callable live in that pure helper, and conftest is a thin adapter that supplies os.environ.get(...), the session stash, and reset_spend_logs. reset_spend_logs itself moved from spend_e2e_client.py into e2e_db.py (implementation unchanged), sitting next to e2e_config and lifecycle so both conftest and any suite import it by name; the sys.path munging is gone. Nothing else imported reset_spend_logs, so spend_e2e_client.py drops the definition, its __all__ entry, and the now-unused os import. --- tests/e2e/conftest.py | 35 +++++------- tests/e2e/e2e_db.py | 56 +++++++++++++++++++ .../spend_tracking/spend_e2e_client.py | 19 ------- 3 files changed, 69 insertions(+), 41 deletions(-) create mode 100644 tests/e2e/e2e_db.py diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 5347fffca4d..609da6a9b07 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -14,14 +14,14 @@ shared fixtures build on it. """ import functools -import sys +import os 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_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup from junit_properties import attach_result_properties from lifecycle import ProxyClientProvider, ResourceManager from proxy_client import ProxyClient, build_proxy_client @@ -107,26 +107,17 @@ def pytest_runtest_call(item: pytest.Item) -> None: 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 - the DB alone so a `DATABASE_URL` pointing at a shared instance is never wiped - without an e2e run. Best-effort: a cleanup failure (no DB reachable) must not - fail the run. The spend_tracking dir goes on sys.path only for this import and - is removed after, so a broader `pytest tests/` run is not left with a mutated - path.""" - if not session.stash.get(_E2E_TEST_RAN, False): - return - spend_dir = str(Path(__file__).parent / "quota_management" / "spend_tracking") - sys.path.insert(0, spend_dir) - try: - from spend_e2e_client import reset_spend_logs # pyright: ignore - - reset_spend_logs() - except Exception as exc: # noqa: BLE001 - cleanup is best-effort - print(f"spend-log cleanup best-effort failed: {exc}") - finally: - if spend_dir in sys.path: - sys.path.remove(spend_dir) + """Once the whole e2e session is done (all suites), optionally truncate the + spend logs so the DB doesn't accumulate test rows. The truncate is destructive + and irreversible, so it runs only when the operator explicitly opts in + (`E2E_RESET_SPEND_LOGS=1`) and an e2e test body actually ran; otherwise a + `DATABASE_URL` pointing at a shared or staging instance is left untouched. + Best-effort: a cleanup failure (no DB reachable) must not fail the run.""" + run_spend_log_cleanup( + opt_in=os.environ.get(RESET_OPT_IN_ENV), + e2e_test_ran=session.stash.get(_E2E_TEST_RAN, False), + truncate=reset_spend_logs, + ) @pytest.fixture(scope="session") diff --git a/tests/e2e/e2e_db.py b/tests/e2e/e2e_db.py new file mode 100644 index 00000000000..439dd519c03 --- /dev/null +++ b/tests/e2e/e2e_db.py @@ -0,0 +1,56 @@ +"""Shared, destructive DB helpers for the e2e harness. + +Kept at the top level next to e2e_config and lifecycle so every suite imports it +by name (`from e2e_db import ...`); no suite reaches into another's directory by +mutating sys.path. + +reset_spend_logs truncates LiteLLM_SpendLogs and cannot be undone, so the +session-finish cleanup routes through run_spend_log_cleanup, which fires the +truncate only on an explicit operator opt-in. "An e2e test ran" is necessary but +never sufficient: a DATABASE_URL pointing at a shared or staging instance must +not be wiped by a routine local run that merely exercised a test. +""" + +import os +from collections.abc import Callable + +RESET_OPT_IN_ENV = "E2E_RESET_SPEND_LOGS" + + +def run_spend_log_cleanup( + *, opt_in: str | None, e2e_test_ran: bool, truncate: Callable[[], None] +) -> bool: + """Invoke `truncate` iff the destructive spend-log reset is both opted into + and warranted, returning whether the truncate was attempted. + + The truncate fires only when the opt-in value is exactly "1" AND an e2e test + body actually ran. Any other opt-in value (unset, "0", "true", "") leaves the + DB untouched, so the destructive path is never armed by the env var's mere + presence or by a test run on its own. Best-effort: a truncate failure is + swallowed so cleanup never fails the session, so the returned bool reports + that the reset was attempted, not that the DB call succeeded. + """ + if opt_in != "1" or not e2e_test_ran: + return False + try: + truncate() + except Exception as exc: # noqa: BLE001 - cleanup is best-effort + print(f"spend-log cleanup best-effort failed: {exc}") + return True + + +def reset_spend_logs() -> None: + """Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes + spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses + DATABASE_URL (default: the local docker postgres on its mapped host port; the + in-container `@db` host isn't resolvable from the host, so default to + localhost). + """ + import psycopg + + url = os.environ.get( + "DATABASE_URL", + "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm", + ) + with psycopg.connect(url) as conn: + _ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"') diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 0d49869aa91..26860212fa3 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -11,7 +11,6 @@ helpers from one place. from __future__ import annotations -import os import time from collections.abc import Callable from dataclasses import dataclass @@ -50,7 +49,6 @@ from models import ( __all__ = [ "SpendClient", "build_client", - "reset_spend_logs", "unique_marker", "unwrap", "is_ok", @@ -59,23 +57,6 @@ __all__ = [ ] -def reset_spend_logs() -> None: - """Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes - spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses - DATABASE_URL (default: the local docker postgres on its mapped host port; note - the in-container `@db` host isn't resolvable from the host, so default to - localhost). - """ - import psycopg - - url = os.environ.get( - "DATABASE_URL", - "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm", - ) - with psycopg.connect(url) as conn: - _ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"') - - def _chat_body( model: str, content: str, From 067c9bbc96fc1fd50c8af1afde0b2f5f021f4801 Mon Sep 17 00:00:00 2001 From: Vineet Puranik <40868710+vineetpuranik@users.noreply.github.com> Date: Mon, 20 Jul 2026 08:52:06 -0700 Subject: [PATCH 24/60] chore(rust): migrate the litellm-rust workspace (core, ai-gateway, python-bridge) from Rust edition 2021 to edition 2024 (#33940) * chore(deps): update cargo.lock file after cargo update * chore(rust): migrate workspace crates to edition 2024 * chore(rust): migrate workspace crates to edition 2024 + fix clippy warnings after 2024 update * chore(rust): add rust version to cargo workspace file * chore(rust): fix clippy collapsible if warning --- litellm-rust/Cargo.lock | 454 +++++++----------- litellm-rust/Cargo.toml | 3 +- .../crates/ai-gateway/src/auth/mod.rs | 2 +- .../crates/ai-gateway/src/io/messages.rs | 2 +- litellm-rust/crates/ai-gateway/src/io/ocr.rs | 2 +- .../crates/ai-gateway/src/io/realtime.rs | 12 +- .../crates/ai-gateway/src/io/realtime_pool.rs | 24 +- .../crates/ai-gateway/src/io/responses_ws.rs | 23 +- litellm-rust/crates/ai-gateway/src/main.rs | 2 +- .../ai-gateway/src/messages/common_utils.rs | 4 +- .../crates/ai-gateway/src/messages/handler.rs | 2 +- .../crates/ai-gateway/src/messages/prepare.rs | 4 +- .../crates/ai-gateway/src/messages/tests.rs | 4 +- .../crates/ai-gateway/src/ocr/common_utils.rs | 4 +- .../crates/ai-gateway/src/ocr/handler.rs | 2 +- .../crates/ai-gateway/src/ocr/hooks.rs | 6 +- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 4 +- .../crates/ai-gateway/src/ocr/prepare.rs | 2 +- .../crates/ai-gateway/src/ocr/tests.rs | 22 +- .../crates/ai-gateway/src/python/config.rs | 2 +- .../crates/ai-gateway/src/routes/health.rs | 2 +- .../ai-gateway/src/routes/messages/mod.rs | 10 +- .../ai-gateway/src/routes/messages/service.rs | 2 +- .../ai-gateway/src/routes/realtime/mod.rs | 4 +- .../ai-gateway/src/routes/realtime/service.rs | 4 +- .../ai-gateway/src/routes/responses/mod.rs | 4 +- .../core/src/caching/in_memory_cache.rs | 2 +- .../azure_ai/messages/transformation.rs | 2 +- .../providers/azure_ai/ocr/transformation.rs | 16 +- .../core/src/providers/bedrock/aws_base.rs | 16 +- .../providers/mistral/ocr/transformation.rs | 2 +- .../openai/realtime/transformation.rs | 2 +- .../openai/responses/transformation.rs | 4 +- .../providers/vertex_ai/ocr/transformation.rs | 6 +- .../core/src/realtime/transformation.rs | 2 +- .../crates/core/src/responses/websocket.rs | 2 +- litellm-rust/crates/python-bridge/src/lib.rs | 4 +- 37 files changed, 293 insertions(+), 371 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 402d16715a3..ce28f737334 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -13,13 +13,13 @@ dependencies = [ [[package]] name = "async-trait" -version = "0.1.89" +version = "0.1.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.0", ] [[package]] @@ -113,7 +113,7 @@ dependencies = [ "bytes-utils", "fastrand", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "percent-encoding", "pin-project-lite", "tracing", @@ -193,7 +193,7 @@ dependencies = [ "futures-core", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "percent-encoding", "pin-project-lite", @@ -222,7 +222,7 @@ dependencies = [ "hyper-util", "pin-project-lite", "rustls 0.21.12", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-native-certs", "rustls-pki-types", "tokio", @@ -279,7 +279,7 @@ dependencies = [ "http 0.2.12", "http 1.4.2", "http-body 0.4.6", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "pin-project-lite", "pin-utils", @@ -313,7 +313,7 @@ checksum = "221eaa237ddf1ca79b60d1372aad77e47f9c0ea5b3ce5099da8c61d027dc77b3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -340,7 +340,7 @@ dependencies = [ "http 0.2.12", "http 1.4.2", "http-body 0.4.6", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "itoa", "num-integer", @@ -392,7 +392,7 @@ dependencies = [ "bytes", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "hyper 1.10.1", "hyper-util", @@ -427,7 +427,7 @@ dependencies = [ "bytes", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "mime", "pin-project-lite", @@ -456,9 +456,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.13.0" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" [[package]] name = "block-buffer" @@ -492,9 +492,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "bytes-utils" @@ -508,9 +508,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.65" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" +checksum = "c89588d05638b5b4594a3348a2d6c20277e43a7f5c5202b05cc56888475a47b8" dependencies = [ "find-msvc-tools", "jobserver", @@ -526,9 +526,20 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "cfg_aliases" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] [[package]] name = "cmake" @@ -655,7 +666,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -678,9 +689,9 @@ checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "fastrand" -version = "2.4.1" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" [[package]] name = "find-msvc-tools" @@ -711,9 +722,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "futures-channel" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" dependencies = [ "futures-core", "futures-sink", @@ -721,44 +732,44 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" [[package]] name = "futures-io" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" [[package]] name = "futures-macro" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "futures-sink" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" [[package]] name = "futures-task" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" [[package]] name = "futures-util" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" dependencies = [ "futures-core", "futures-io", @@ -793,20 +804,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "getrandom" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" -dependencies = [ - "cfg-if", - "js-sys", - "libc", - "r-efi 5.3.0", - "wasip2", - "wasm-bindgen", -] - [[package]] name = "getrandom" version = "0.4.3" @@ -814,8 +811,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", + "js-sys", "libc", - "r-efi 6.0.0", + "r-efi", + "rand_core 0.10.1", + "wasm-bindgen", ] [[package]] @@ -917,9 +917,9 @@ dependencies = [ [[package]] name = "http-body" -version = "1.0.1" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" dependencies = [ "bytes", "http 1.4.2", @@ -927,14 +927,14 @@ dependencies = [ [[package]] name = "http-body-util" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" dependencies = [ "bytes", "futures-core", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "pin-project-lite", ] @@ -995,7 +995,7 @@ dependencies = [ "futures-core", "h2 0.4.15", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "httparse", "httpdate", "itoa", @@ -1029,7 +1029,7 @@ dependencies = [ "http 1.4.2", "hyper 1.10.1", "hyper-util", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-native-certs", "tokio", "tokio-rustls 0.26.4", @@ -1048,13 +1048,13 @@ dependencies = [ "futures-channel", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "hyper 1.10.1", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.6.5", "tokio", "tower-service", "tracing", @@ -1242,12 +1242,12 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", - "rand 0.8.6", + "rand 0.8.7", "reqwest", "serde", "serde_json", "sha2 0.10.9", - "thiserror 2.0.18", + "thiserror 2.0.19", "tokio", ] @@ -1289,9 +1289,9 @@ checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" [[package]] name = "memchr" -version = "2.8.2" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "mime" @@ -1301,9 +1301,9 @@ checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" [[package]] name = "mio" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "wasi", @@ -1378,9 +1378,9 @@ checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" [[package]] name = "portable-atomic" -version = "1.13.1" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" [[package]] name = "potential_utf" @@ -1408,9 +1408,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.106" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] @@ -1471,7 +1471,7 @@ dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -1483,7 +1483,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -1498,9 +1498,9 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls 0.23.41", - "socket2 0.6.4", - "thiserror 2.0.18", + "rustls 0.23.42", + "socket2 0.6.5", + "thiserror 2.0.19", "tokio", "tracing", "web-time", @@ -1508,20 +1508,21 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.15" +version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ "bytes", - "getrandom 0.3.4", + "getrandom 0.4.3", "lru-slab", - "rand 0.9.4", + "rand 0.10.2", + "rand_pcg", "ring", "rustc-hash", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-pki-types", "slab", - "thiserror 2.0.18", + "thiserror 2.0.19", "tinyvec", "tracing", "web-time", @@ -1529,33 +1530,27 @@ dependencies = [ [[package]] name = "quinn-udp" -version = "0.5.14" +version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.4", + "socket2 0.6.5", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] name = "quote" -version = "1.0.46" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] -[[package]] -name = "r-efi" -version = "5.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" - [[package]] name = "r-efi" version = "6.0.0" @@ -1564,23 +1559,24 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rand" -version = "0.8.6" +version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", - "rand_chacha 0.3.1", + "rand_chacha", "rand_core 0.6.4", ] [[package]] name = "rand" -version = "0.9.4" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ - "rand_chacha 0.9.0", - "rand_core 0.9.5", + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", ] [[package]] @@ -1593,16 +1589,6 @@ dependencies = [ "rand_core 0.6.4", ] -[[package]] -name = "rand_chacha" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" -dependencies = [ - "ppv-lite86", - "rand_core 0.9.5", -] - [[package]] name = "rand_core" version = "0.6.4" @@ -1614,11 +1600,17 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.9.5" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" dependencies = [ - "getrandom 0.3.4", + "rand_core 0.10.1", ] [[package]] @@ -1640,7 +1632,7 @@ dependencies = [ "futures-util", "h2 0.4.15", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "hyper 1.10.1", "hyper-rustls 0.27.9", @@ -1650,7 +1642,7 @@ dependencies = [ "percent-encoding", "pin-project-lite", "quinn", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-pki-types", "serde", "serde_json", @@ -1686,9 +1678,9 @@ dependencies = [ [[package]] name = "rustc-hash" -version = "2.1.2" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustc_version" @@ -1713,9 +1705,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.41" +version = "0.23.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", "once_cell", @@ -1740,9 +1732,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" dependencies = [ "web-time", "zeroize", @@ -1772,9 +1764,9 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.22" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "ryu" @@ -1832,9 +1824,9 @@ checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" [[package]] name = "serde" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", "serde_derive", @@ -1842,22 +1834,22 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.0", ] [[package]] @@ -1898,9 +1890,9 @@ dependencies = [ [[package]] name = "sha1" -version = "0.10.6" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" dependencies = [ "cfg-if", "cpufeatures 0.2.17", @@ -1959,9 +1951,9 @@ dependencies = [ [[package]] name = "socket2" -version = "0.6.4" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", "windows-sys 0.61.2", @@ -1981,9 +1973,20 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "2.0.118" +version = "2.0.119" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2fac314a64dc9a36e61a9eb4261a5e9bbfbc922b27e518af97bc32b926cf967" dependencies = [ "proc-macro2", "quote", @@ -2007,7 +2010,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2027,11 +2030,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ - "thiserror-impl 2.0.18", + "thiserror-impl 2.0.19", ] [[package]] @@ -2042,18 +2045,18 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.0", ] [[package]] @@ -2098,9 +2101,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" dependencies = [ "tinyvec_macros", ] @@ -2113,28 +2116,28 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.52.3" +version = "1.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "d988bcd52dbe076d3d46903332f58c912b87a2c49b1428419a5845154762ffee" dependencies = [ "bytes", "libc", "mio", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.6.5", "tokio-macros", "windows-sys 0.61.2", ] [[package]] name = "tokio-macros" -version = "2.7.0" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2153,7 +2156,7 @@ version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls 0.23.41", + "rustls 0.23.42", "tokio", ] @@ -2165,7 +2168,7 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" dependencies = [ "futures-util", "log", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-native-certs", "rustls-pki-types", "tokio", @@ -2212,7 +2215,7 @@ dependencies = [ "bytes", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "pin-project-lite", "tower", "tower-layer", @@ -2252,7 +2255,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2282,8 +2285,8 @@ dependencies = [ "http 1.4.2", "httparse", "log", - "rand 0.8.6", - "rustls 0.23.41", + "rand 0.8.7", + "rustls 0.23.42", "rustls-pki-types", "sha1", "thiserror 1.0.69", @@ -2375,15 +2378,6 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" -[[package]] -name = "wasip2" -version = "1.0.4+wasi-0.2.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" -dependencies = [ - "wit-bindgen", -] - [[package]] name = "wasm-bindgen" version = "0.2.126" @@ -2426,7 +2420,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.119", "wasm-bindgen-shared", ] @@ -2474,9 +2468,9 @@ dependencies = [ [[package]] name = "webpki-roots" -version = "1.0.8" +version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" dependencies = [ "rustls-pki-types", ] @@ -2493,16 +2487,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.5", + "windows-targets", ] [[package]] @@ -2520,31 +2505,14 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm 0.52.6", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", -] - -[[package]] -name = "windows-targets" -version = "0.53.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" -dependencies = [ - "windows-link", - "windows_aarch64_gnullvm 0.53.1", - "windows_aarch64_msvc 0.53.1", - "windows_i686_gnu 0.53.1", - "windows_i686_gnullvm 0.53.1", - "windows_i686_msvc 0.53.1", - "windows_x86_64_gnu 0.53.1", - "windows_x86_64_gnullvm 0.53.1", - "windows_x86_64_msvc 0.53.1", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] [[package]] @@ -2553,102 +2521,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_aarch64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" - [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" -[[package]] -name = "windows_i686_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" - [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_i686_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "windows_x86_64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" - -[[package]] -name = "wit-bindgen" -version = "0.57.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" - [[package]] name = "writeable" version = "0.6.3" @@ -2680,28 +2594,28 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] [[package]] name = "zerocopy" -version = "0.8.52" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.52" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2721,7 +2635,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] @@ -2761,11 +2675,11 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "zmij" -version = "1.0.21" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a3baa33e6cf..6d63be05d00 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -7,7 +7,8 @@ members = [ resolver = "2" [workspace.package] -edition = "2021" +edition = "2024" +rust-version = "1.88" license = "MIT" repository = "https://github.com/BerriAI/litellm" diff --git a/litellm-rust/crates/ai-gateway/src/auth/mod.rs b/litellm-rust/crates/ai-gateway/src/auth/mod.rs index 438a0513057..b09d8285c3a 100644 --- a/litellm-rust/crates/ai-gateway/src/auth/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/auth/mod.rs @@ -9,9 +9,9 @@ //! runs during extraction, before the handler body. Routes never re-implement it. use axum::extract::FromRequestParts; +use axum::http::StatusCode; use axum::http::header::AUTHORIZATION; use axum::http::request::Parts; -use axum::http::StatusCode; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/ai-gateway/src/io/messages.rs b/litellm-rust/crates/ai-gateway/src/io/messages.rs index b784d2b62a1..86170e45678 100644 --- a/litellm-rust/crates/ai-gateway/src/io/messages.rs +++ b/litellm-rust/crates/ai-gateway/src/io/messages.rs @@ -1 +1 @@ -pub use crate::messages::{messages, MessagesRequest}; +pub use crate::messages::{MessagesRequest, messages}; diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr.rs b/litellm-rust/crates/ai-gateway/src/io/ocr.rs index 55e02839c4e..2fc82f0b61f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -1 +1 @@ -pub use crate::ocr::{ocr, OcrRequest}; +pub use crate::ocr::{OcrRequest, ocr}; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 40a38c1579a..845e7bf9527 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -15,16 +15,16 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::realtime::transformation::RealtimeProviderConfig; use litellm_core::realtime::types::RealtimeEvent; -use litellm_core::CoreResult; use tokio::net::TcpStream; -use tokio_tungstenite::tungstenite::client::IntoClientRequest; -use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::Message; -use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; @@ -113,7 +113,7 @@ pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult { return Err(CoreError::Network( "upstream closed before first event".to_string(), - )) + )); } _ => continue, } diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs index bf8041f31d7..4a1a3cd1166 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs @@ -28,11 +28,11 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use futures_util::StreamExt; -use litellm_core::realtime::types::RealtimeEvent; use litellm_core::CoreResult; +use litellm_core::realtime::types::RealtimeEvent; use crate::io::realtime::{ - dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs, + UpstreamRx, UpstreamTx, UpstreamWs, dial_upstream, read_event, resolve_api_key, }; /// Default target warm sockets per key when pooling is enabled. @@ -473,8 +473,8 @@ pub fn upstream_key( /// unhealthy — we'd rather discard and fresh-dial than hand over a socket in an /// unexpected state. `Pending` (the healthy case) returns `false`. fn is_dead(rx: &mut UpstreamRx) -> bool { - use futures_util::task::noop_waker_ref; use futures_util::Stream; + use futures_util::task::noop_waker_ref; use std::pin::Pin; use std::task::{Context, Poll}; @@ -523,15 +523,15 @@ mod tests { )) .await; while let Some(Ok(msg)) = ws.next().await { - if let Message::Text(text) = msg { - if text.contains("response.create") { - for frame in [ - r#"{"type":"response.created"}"#, - r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#, - r#"{"type":"response.done"}"#, - ] { - let _ = ws.send(Message::Text(frame.to_string())).await; - } + if let Message::Text(text) = msg + && text.contains("response.create") + { + for frame in [ + r#"{"type":"response.created"}"#, + r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#, + r#"{"type":"response.done"}"#, + ] { + let _ = ws.send(Message::Text(frame.to_string())).await; } } } diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index ae6ad150bcf..9b51019f4bc 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -10,19 +10,18 @@ use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig; use litellm_core::{CoreError, CoreResult}; use tokio::net::TcpStream; use tokio::sync::Mutex; -use tokio_tungstenite::tungstenite::client::IntoClientRequest; -use tokio_tungstenite::tungstenite::http::header::{HeaderName, AUTHORIZATION}; -use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::Message; -use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; use crate::constants::{ DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS, }; const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -const MISSING_KEY_MESSAGE: &str = - "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; +const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; pub type ResponsesUpstreamWs = WebSocketStream>; type UpstreamTx = SplitSink; @@ -83,8 +82,8 @@ impl ResponsesWebSocketConnection { } pub async fn recv_text(&self) -> CoreResult> { - let mut socket = self.socket.lock().await; - let Some(socket) = socket.as_mut() else { + let mut socket_guard = self.socket.lock().await; + let Some(socket) = socket_guard.as_mut() else { return Ok(None); }; match socket.next().await { @@ -456,9 +455,11 @@ mod tests { assert_eq!(fourth.event_type, ResponsesWsEventType::ResponseCompleted); let observed: Vec<_> = observed_rx.collect().await; assert_eq!(observed.len(), 4); - assert!(observed - .iter() - .all(|event| event.event_type != ResponsesWsEventType::ResponseCreate)); + assert!( + observed + .iter() + .all(|event| event.event_type != ResponsesWsEventType::ResponseCreate) + ); } #[tokio::test] diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs index f9ce97801d3..da3a486d4ee 100644 --- a/litellm-rust/crates/ai-gateway/src/main.rs +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -11,7 +11,7 @@ use std::sync::Arc; -use litellm_ai_gateway::io::realtime_pool::{upstream_key, PoolConfig, RealtimePool}; +use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key}; use litellm_ai_gateway::routes; use litellm_ai_gateway::state::AppState; use litellm_core::router::{Deployment, LiteLLMParams, Router}; diff --git a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs index 33894d0ee64..4b906155665 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs @@ -1,8 +1,8 @@ -use litellm_core::error::{json_type_name, CoreError}; +use litellm_core::CoreResult; +use litellm_core::error::{CoreError, json_type_name}; use litellm_core::messages::transformation::AnthropicMessagesProviderConfig; use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; -use litellm_core::CoreResult; use serde_json::{Map, Value}; use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS; diff --git a/litellm-rust/crates/ai-gateway/src/messages/handler.rs b/litellm-rust/crates/ai-gateway/src/messages/handler.rs index d3b9d3b3fba..90c12367f50 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/handler.rs @@ -1,5 +1,5 @@ -use litellm_core::error::CoreError; use litellm_core::CoreResult; +use litellm_core::error::CoreError; use serde_json::Value; use super::client::http_client; diff --git a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs index 6176f9cb67f..624c3598fb0 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs @@ -1,7 +1,7 @@ -use litellm_core::messages::transformation::MessagesAuthStrategy; -use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; use litellm_core::CoreError; use litellm_core::CoreResult; +use litellm_core::messages::transformation::MessagesAuthStrategy; +use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_header, messages_provider_config, string_headers}; use super::types::{MessagesRequest, ProviderMessagesRequest}; diff --git a/litellm-rust/crates/ai-gateway/src/messages/tests.rs b/litellm-rust/crates/ai-gateway/src/messages/tests.rs index 9b1cc45aacb..a2d0f6fae23 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/tests.rs @@ -1,14 +1,14 @@ use std::time::Duration; use litellm_core::error::CoreError; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use super::common_utils::{ has_header, messages_provider_config, string_headers, truncate_error_body, }; -use super::{messages, MessagesRequest}; +use super::{MessagesRequest, messages}; async fn read_http_request(socket: &mut TcpStream) -> String { let mut request = Vec::new(); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index d4b4d9338e7..7d164a80137 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -1,11 +1,11 @@ use std::net::IpAddr; use std::time::{Duration, Instant}; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrProviderConfig; -use litellm_core::CoreResult; use reqwest::Url; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 4d93c2a25db..381d22e9cea 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -1,6 +1,6 @@ +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrResponseHandling; -use litellm_core::CoreResult; use serde_json::Value; use super::client::http_client; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 6be74ed2714..ffe2e0122c0 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -1,11 +1,11 @@ use std::future::Future; use std::pin::Pin; +use litellm_core::CoreResult; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrAuthStrategy; -use litellm_core::CoreResult; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use super::common_utils::{ convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, @@ -292,7 +292,7 @@ fn parse_ocr_pre_call_guardrail_request( Some(_) => { return Err(CoreError::InvalidRequest( "OCR pre_call guardrail optional_params must be an object".to_string(), - )) + )); } None => Map::new(), }; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index b54ee39b21d..ad346bc0c64 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,5 +1,5 @@ -use litellm_core::call_lifecycle::CallLifecycle; use litellm_core::CoreResult; +use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; mod client; @@ -12,7 +12,7 @@ mod types; pub use types::OcrRequest; use handler::execute_ocr_provider_call; -use prepare::{prepare_ocr_call, PreparedOcrCall}; +use prepare::{PreparedOcrCall, prepare_ocr_call}; pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index 5a4b350a4c4..6231393c889 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -1,7 +1,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; -use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; +use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::hooks::OcrLifecycleHooks; use super::types::{OcrRequest, PreparedOcrRequest}; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index 35747dc6985..bb2a6b06501 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -3,12 +3,12 @@ use std::time::Duration; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrResponseHandling; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; -use super::{ocr, OcrRequest}; +use super::{OcrRequest, ocr}; use crate::integrations::custom_guardrail::{ CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, GuardrailFuture, GuardrailRequest, @@ -228,19 +228,23 @@ fn truncate_error_body_does_not_split_multibyte_chars() { #[test] fn ocr_dispatch_supports_migrated_providers() { assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); - assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409") - .expect("azure ai config resolves") - .requires_data_uri_document()); + assert!( + ocr_provider_config("azure_ai", "pixtral-12b-2409") + .expect("azure ai config resolves") + .requires_data_uri_document() + ); assert_eq!( ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") .expect("document intelligence config resolves") .response_handling(), OcrResponseHandling::AzureDocumentIntelligencePoll ); - assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas") - .expect("vertex deepseek config resolves") - .supported_ocr_params() - .contains(&"temperature")); + assert!( + ocr_provider_config("vertex_ai", "deepseek-ocr-maas") + .expect("vertex deepseek config resolves") + .supported_ocr_params() + .contains(&"temperature") + ); assert!(ocr_provider_config("openai", "gpt-4o").is_none()); } diff --git a/litellm-rust/crates/ai-gateway/src/python/config.rs b/litellm-rust/crates/ai-gateway/src/python/config.rs index 54b7a53bafa..c028d3d6b51 100644 --- a/litellm-rust/crates/ai-gateway/src/python/config.rs +++ b/litellm-rust/crates/ai-gateway/src/python/config.rs @@ -7,9 +7,9 @@ //! //! Compiled only under the `python-config` feature. +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::router::{Deployment, Router}; -use litellm_core::CoreResult; use pyo3::prelude::*; use crate::gil; diff --git a/litellm-rust/crates/ai-gateway/src/routes/health.rs b/litellm-rust/crates/ai-gateway/src/routes/health.rs index 15c67fea325..c64ca3a7199 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/health.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/health.rs @@ -1,8 +1,8 @@ //! Health probes. Simple-route template: a `router()` plus its handlers, in one file. +use axum::Router; use axum::http::StatusCode; use axum::routing::get; -use axum::Router; use crate::state::AppState; diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index 933386282fa..a34b2edd7b8 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -2,13 +2,13 @@ mod service; +use axum::Router; use axum::body::Body; use axum::extract::{Json, State}; -use axum::http::header::{HeaderMap, HeaderValue, CACHE_CONTROL, CONTENT_TYPE}; use axum::http::StatusCode; +use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue}; use axum::response::{IntoResponse, Response}; use axum::routing::post; -use axum::Router; use litellm_core::CoreError; use serde_json::{Map, Value}; @@ -125,9 +125,9 @@ mod tests { use std::sync::Arc; use axum::body::Body; - use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE}; use axum::http::Request; use axum::http::StatusCode; + use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE}; use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter}; use serde_json::json; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -439,8 +439,8 @@ mod tests { .await .expect("response body reads"); assert_eq!( - serde_json::from_slice::(&response_body).expect("error is json") - ["error"]["message"], + serde_json::from_slice::(&response_body).expect("error is json")["error"] + ["message"], "messages provider request failed" ); server.await.expect("upstream task completes"); diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs index 7f00123ca39..75ed26e5be8 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs @@ -5,7 +5,7 @@ use litellm_core::{CoreError, CoreResult}; use serde_json::{Map, Value}; use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -use crate::messages::{execute_messages, MessagesRequest}; +use crate::messages::{MessagesRequest, execute_messages}; pub(crate) enum MessagesResponse { Json(Value), diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs index c3f929f5f0b..f9144ad1fdb 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -6,17 +6,17 @@ mod service; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; use crate::io::realtime_pool::RealtimePool; +use axum::Router; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::{Query, State}; use axum::http::StatusCode; use axum::response::Response; use axum::routing::get; -use axum::Router; use futures_util::{SinkExt, StreamExt}; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::router::Router as ModelRouter; diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs index d6c31edd454..4ae8cfe7379 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -9,12 +9,12 @@ use std::time::Duration; -use crate::io::realtime_pool::{upstream_key, RealtimePool}; +use crate::io::realtime_pool::{RealtimePool, upstream_key}; use futures_util::{Sink, Stream}; +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::router::Router; -use litellm_core::CoreResult; /// Select a deployment for `model` and splice the client stream to the provider. /// diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs index bdaffc97afb..a94853e106d 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs @@ -1,15 +1,15 @@ mod service; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; +use axum::Router; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::{Query, State}; use axum::http::StatusCode; use axum::response::Response; use axum::routing::get; -use axum::Router; use futures_util::{Sink, SinkExt, StreamExt}; use litellm_core::responses::types::{ResponsesErrorFrame, ResponsesWsEvent, ResponsesWsEventType}; use litellm_core::router::Router as ModelRouter; diff --git a/litellm-rust/crates/core/src/caching/in_memory_cache.rs b/litellm-rust/crates/core/src/caching/in_memory_cache.rs index 0ceeedb8b71..45d4bd69b79 100644 --- a/litellm-rust/crates/core/src/caching/in_memory_cache.rs +++ b/litellm-rust/crates/core/src/caching/in_memory_cache.rs @@ -134,8 +134,8 @@ impl InMemoryCache { #[cfg(test)] mod tests { use std::sync::{ - atomic::{AtomicU64, Ordering}, Arc, + atomic::{AtomicU64, Ordering}, }; use super::InMemoryCache; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 13e79b087c7..6935bb4604b 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -5,7 +5,7 @@ use crate::messages::types::{ MessageContent, SystemPrompt, }; use crate::providers::anthropic::messages::transformation::{ - non_empty, AnthropicMessagesConfig, ANTHROPIC_MESSAGES_CONFIG, + ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, }; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index 060073acd47..eabd15677cc 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -1,9 +1,9 @@ use std::collections::BTreeSet; -use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::error::{CoreError, CoreResult, json_type_name}; use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; use crate::ocr::types::{OcrRequestData, OcrResponseData}; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -206,11 +206,11 @@ pub fn complete_document_intelligence_url( AZURE_DOCUMENT_INTELLIGENCE_API_VERSION ); - if let Some(pages) = optional_params.get("pages") { - if let Some(normalized) = normalize_pages_param(pages)? { - url.push_str("&pages="); - url.push_str(&normalized); - } + if let Some(pages) = optional_params.get("pages") + && let Some(normalized) = normalize_pages_param(pages)? + { + url.push_str("&pages="); + url.push_str(&normalized); } Ok(url) @@ -231,7 +231,7 @@ fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { other => { return Err(CoreError::InvalidRequest(format!( "Invalid document type: {other}. Must be 'document_url' or 'image_url'" - ))) + ))); } }; object diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index 82d1e8fdf91..c5995732e41 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -5,10 +5,10 @@ use std::time::{SystemTime, UNIX_EPOCH}; use crate::caching::in_memory_cache::InMemoryCache; use crate::error::{CoreError, CoreResult}; -use aws_credential_types::provider::ProvideCredentials; use aws_credential_types::Credentials; +use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ - sign, SignableBody, SignableRequest, SigningParams, SigningSettings, + SignableBody, SignableRequest, SigningParams, SigningSettings, sign, }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; @@ -368,11 +368,11 @@ async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> CoreR if let (Ok(current_role), Ok(token_file)) = ( std::env::var(AWS_ROLE_ARN), std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), - ) { - if !token_file.is_empty() { - return Ok(same_role_arns(role, ¤t_role)); - } + ) && !token_file.is_empty() + { + return Ok(same_role_arns(role, ¤t_role)); } + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); if let Some(region) = config.region_name.clone() { loader = loader.region(aws_types::region::Region::new(region)); @@ -639,7 +639,9 @@ mod tests { ); assert_eq!( signed.get("Authorization").map(String::as_str), - Some("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464") + Some( + "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" + ) ); } diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 1a33bc1e951..dc720cc4244 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -1,4 +1,4 @@ -use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::error::{CoreError, CoreResult, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs index 626e4014ff9..b3f6b03b28a 100644 --- a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs @@ -1,6 +1,6 @@ +use crate::CoreResult; use crate::realtime::transformation::RealtimeProviderConfig; use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; -use crate::CoreResult; /// Default OpenAI API base, used when the caller does not override `api_base`. pub const OPENAI_REALTIME_DEFAULT_API_BASE: &str = "https://api.openai.com"; diff --git a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs index ece10971806..e15197c468c 100644 --- a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs @@ -1,6 +1,6 @@ -use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; -use crate::responses::websocket::{enforce_model, ResponsesWebSocketProviderConfig}; use crate::CoreResult; +use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; +use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model}; pub struct OpenAIResponsesWsConfig; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index 8639926c435..6300149c237 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -1,7 +1,7 @@ -use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::error::{CoreError, CoreResult, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -140,7 +140,7 @@ fn document_content_item(document: &Value) -> CoreResult { other => { return Err(CoreError::InvalidRequest(format!( "Unsupported document type: {other}. Expected 'image_url' or 'document_url'" - ))) + ))); } }; let url = object diff --git a/litellm-rust/crates/core/src/realtime/transformation.rs b/litellm-rust/crates/core/src/realtime/transformation.rs index a4baa27a6c2..69b88687000 100644 --- a/litellm-rust/crates/core/src/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/realtime/transformation.rs @@ -1,5 +1,5 @@ -use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; use crate::CoreResult; +use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; pub trait RealtimeProviderConfig { /// Build the upstream WebSocket URL (e.g. `wss://api.openai.com/v1/realtime?model=…`). diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index 1edffd44985..92dc19627a0 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -1,6 +1,6 @@ +use crate::CoreResult; use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult}; -use crate::CoreResult; pub trait ResponsesWebSocketProviderConfig: Sync { fn supports_native_websocket(&self) -> bool { diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 1decb789a22..07429667644 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,8 +1,8 @@ use std::collections::HashMap; use std::time::Duration; -use litellm_ai_gateway::io::messages::{messages as run_messages, MessagesRequest}; -use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest}; +use litellm_ai_gateway::io::messages::{MessagesRequest, messages as run_messages}; +use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use litellm_core::error::CoreError; use pyo3::exceptions::{PyRuntimeError, PyValueError}; From 3fcd19d7ad5b06f839312a8b909e3d18dd7f2f80 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 09:52:35 -0700 Subject: [PATCH 25/60] fix(fireworks_ai): restore Content-Type application/json header (fixes 415) (#33929) * fix(fireworks_ai): set Content-Type application/json in validate_environment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(fireworks_ai): delegate chat validate_environment to OpenAIGPTConfig Instead of re-adding the JSON Content-Type default inside FireworksAIMixin, FireworksAIConfig now delegates header construction to OpenAIGPTConfig and only layers the Fireworks-specific x-session-affinity header on top, so the Content-Type default can no longer drift away from the OpenAI base and reintroduce the 415. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(fireworks_ai): cover missing api key error path in chat validate_environment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/fireworks_ai/chat/transformation.py | 26 ++++++ litellm/llms/fireworks_ai/common_utils.py | 19 +++-- .../test_fireworks_ai_chat_transformation.py | 84 +++++++++++++++++++ 3 files changed, 123 insertions(+), 6 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 319f03fea89..eeae8c76888 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -133,6 +133,32 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): def get_config(cls): return super().get_config() + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + api_key = self._get_api_key(api_key) + if api_key is None: + raise ValueError("FIREWORKS_API_KEY is not set") + + validated_headers = OpenAIGPTConfig.validate_environment( + self, + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + return self._add_session_affinity_header(validated_headers, litellm_params) + def get_supported_openai_params(self, model: str): # Base parameters supported by all models supported_params = [ diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 4e22445bcc0..51ed8afbbd2 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -64,9 +64,16 @@ class FireworksAIMixin: if api_key is None: raise ValueError("FIREWORKS_API_KEY is not set") - validated_headers = {"Authorization": "Bearer {}".format(api_key), **headers} - if not any(key.lower() == "x-session-affinity" for key in validated_headers): - session_id = get_fireworks_session_id(litellm_params) - if session_id: - validated_headers["x-session-affinity"] = session_id - return validated_headers + auth_headers = {"Authorization": "Bearer {}".format(api_key), **headers} + content_type_header = ( + {} if any(key.lower() == "content-type" for key in auth_headers) else {"Content-Type": "application/json"} + ) + return self._add_session_affinity_header({**auth_headers, **content_type_header}, litellm_params) + + def _add_session_affinity_header(self, headers: dict, litellm_params: dict) -> dict: + if any(key.lower() == "x-session-affinity" for key in headers): + return headers + session_id = get_fireworks_session_id(litellm_params) + if not session_id: + return headers + return {**headers, "x-session-affinity": session_id} diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 6809799d34f..94945ed4bfb 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -123,6 +123,90 @@ def test_validate_environment_preserves_explicit_session_affinity_header(): assert headers["x-session-affinity"] == "explicit-session" +def test_validate_environment_sets_json_content_type(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + ) + + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_preserves_explicit_content_type(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={"content-type": "multipart/form-data"}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + ) + + assert headers["content-type"] == "multipart/form-data" + assert "Content-Type" not in headers + + +def test_validate_environment_sets_json_content_type_with_session_affinity(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_session_id": "session-123"}, + api_key="test-key", + ) + + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test-key" + assert headers["x-session-affinity"] == "session-123" + + +def test_validate_environment_resolves_api_key_from_env_and_sets_content_type(monkeypatch): + monkeypatch.setenv("FIREWORKS_API_KEY", "fw-env-key") + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + ) + + assert headers["Authorization"] == "Bearer fw-env-key" + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_raises_without_api_key(monkeypatch): + for env_var in ( + "FIREWORKS_API_KEY", + "FIREWORKS_AI_API_KEY", + "FIREWORKSAI_API_KEY", + "FIREWORKS_AI_TOKEN", + ): + monkeypatch.delenv(env_var, raising=False) + config = FireworksAIConfig() + + with pytest.raises(ValueError, match="FIREWORKS_API_KEY is not set"): + config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + ) + + def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): assert ( get_fireworks_session_id( From 479e997eed08e4393ceb7b0c4d39c896c4b19da3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Jul 2026 09:55:38 -0700 Subject: [PATCH 26/60] feat(spend): raise /spend/logs/v2 page_size cap to 1000 Clients exporting large spend-log ranges were forced into 100-row pages, which meant a bounded COUNT plus an increasingly deep OFFSET scan per request. Larger pages reduce both the request count and the cumulative OFFSET cost for the same result set. The handler already excludes the heavy JSON columns (messages, response, proxy_server_request) from the paginated SELECT and bounds the COUNT via SPEND_LOGS_PAGINATION_COUNT_CAP, so per-row cost does not grow with page size. 1000 matches the ceiling already used by the user and user-agent analytics list endpoints. --- .../spend_management_endpoints.py | 2 +- .../test_spend_management_endpoints.py | 53 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 5a0b94d1524..55b50e7d9ff 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1647,7 +1647,7 @@ async def ui_view_spend_logs( description="Time till which to view key spend", ), page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), - page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=100), + page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=1000), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), status_filter: str | None = fastapi.Query( default=None, description="Filter logs by status (e.g., success, failure)" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 579f46c8c77..db72a7fb38c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1467,6 +1467,59 @@ async def test_ui_view_spend_logs_pagination(client, monkeypatch): app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.parametrize( + "page_size, expected_status, expected_rows", + [ + (1000, 200, 1000), + (1001, 422, None), + ], +) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_page_size_upper_bound( + client, monkeypatch, page_size, expected_status, expected_rows +): + mock_spend_logs = [ + { + "id": f"log{i}", + "request_id": f"req{i}", + "api_key": "sk-test-key", + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + } + for i in range(1200) + ] + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, lambda where: mock_spend_logs), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/v2", + params={ + "page": 1, + "page_size": page_size, + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == expected_status + if expected_status == 200: + data = response.json() + assert data["page_size"] == page_size + assert len(data["data"]) == expected_rows + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): mock_spend_logs = [ From 8a0bb4cc560680195d1a9701e818b29bf5fadcfa Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Jul 2026 11:38:52 -0700 Subject: [PATCH 27/60] chore(ci): retire daily OSS branches in favor of litellm_internal_staging Removes the scheduled workflow that cut litellm_oss_daily_YYYY_MM_DD branches and the guardrails workflow that only ran on them. The secret scan and ruff checks that workflow duplicated already run on PRs to litellm_internal_staging via test-linting.yml, so no coverage is lost. Retargets contributor-facing messaging in CONTRIBUTING.md, CLAUDE.md, and the guard-main-branch error output at litellm_internal_staging. --- .github/workflows/create_daily_oss_branch.yml | 61 ------------------- .github/workflows/guard-main-branch.yml | 4 +- .github/workflows/oss_daily_guardrails.yml | 50 --------------- CLAUDE.md | 2 +- CONTRIBUTING.md | 2 +- 5 files changed, 4 insertions(+), 115 deletions(-) delete mode 100644 .github/workflows/create_daily_oss_branch.yml delete mode 100644 .github/workflows/oss_daily_guardrails.yml diff --git a/.github/workflows/create_daily_oss_branch.yml b/.github/workflows/create_daily_oss_branch.yml deleted file mode 100644 index 43de4a0e75f..00000000000 --- a/.github/workflows/create_daily_oss_branch.yml +++ /dev/null @@ -1,61 +0,0 @@ -name: Create Daily OSS Branch - -on: - schedule: - - cron: "0 16 * * 1-5" # 9am PT during daylight saving time, weekdays. - workflow_dispatch: - inputs: - date: - description: "Branch date in YYYY_MM_DD format. Defaults to today's UTC date." - required: false - type: string - -permissions: - contents: write - -jobs: - create-oss-branch: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 10 - - steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - fetch-depth: 0 - persist-credentials: false - - - name: Create dated OSS branch - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - REQUESTED_DATE: ${{ inputs.date }} - run: | - set -euo pipefail - - if [ -n "${REQUESTED_DATE}" ]; then - if ! echo "${REQUESTED_DATE}" | grep -Eq '^[0-9]{4}_[0-9]{2}_[0-9]{2}$'; then - echo "::error::date must use YYYY_MM_DD format, got '${REQUESTED_DATE}'" - exit 1 - fi - BRANCH_DATE="${REQUESTED_DATE}" - else - BRANCH_DATE="$(date -u +'%Y_%m_%d')" - fi - - BRANCH_NAME="litellm_oss_daily_${BRANCH_DATE}" - echo "Creating branch: ${BRANCH_NAME}" - - git config user.name "github-actions[bot]" - git config user.email "github-actions[bot]@users.noreply.github.com" - - git fetch origin main "${BRANCH_NAME}" || true - - if git show-ref --verify --quiet "refs/remotes/origin/${BRANCH_NAME}"; then - echo "Branch ${BRANCH_NAME} already exists. Skipping creation." - exit 0 - fi - - git checkout -b "${BRANCH_NAME}" origin/main - git push "https://x-access-token:${GITHUB_TOKEN}@github.com/${GITHUB_REPOSITORY}.git" "${BRANCH_NAME}" - echo "Successfully created and pushed branch: ${BRANCH_NAME}" diff --git a/.github/workflows/guard-main-branch.yml b/.github/workflows/guard-main-branch.yml index aa4968f0c1e..5bc561c6441 100644 --- a/.github/workflows/guard-main-branch.yml +++ b/.github/workflows/guard-main-branch.yml @@ -31,12 +31,12 @@ jobs: echo "PR head repo: $HEAD_REPO" echo "PR head branch: $HEAD_REF" if [ "$HEAD_REPO" != "$BASE_REPO" ]; then - echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) instead." + echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against 'litellm_internal_staging' instead." exit 1 fi if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then echo "Allowed source branch." exit 0 fi - echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) instead." + echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_internal_staging' instead." exit 1 diff --git a/.github/workflows/oss_daily_guardrails.yml b/.github/workflows/oss_daily_guardrails.yml deleted file mode 100644 index f9dc746ee05..00000000000 --- a/.github/workflows/oss_daily_guardrails.yml +++ /dev/null @@ -1,50 +0,0 @@ -name: OSS Daily Guardrails - -on: - push: - branches: - - "litellm_oss_daily_20*" - pull_request: - branches: - - "litellm_oss_daily_20*" - - litellm_internal_staging - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true - -jobs: - oss-safe-checks: - name: Run OSS daily safe checks - if: startsWith(github.ref_name, 'litellm_oss_daily_20') || startsWith(github.head_ref, 'litellm_oss_daily_20') || startsWith(github.base_ref, 'litellm_oss_daily_20') - runs-on: ubuntu-latest - timeout-minutes: 10 - - steps: - - name: Checkout repository - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Run secret scan test - run: | - uv run --frozen --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v - - - name: Run Ruff - run: | - uv sync --frozen - cd litellm - uv run --no-sync ruff check . diff --git a/CLAUDE.md b/CLAUDE.md index 9f708716c6d..1a4826d51e9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -19,7 +19,7 @@ Same thing for bug fixes. The tests should make it so that this specific bug can End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md` -When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose for internal contributors; external / OSS contributions target the current daily OSS branch instead, named `litellm_oss_daily_YYYY_MM_DD` (a fresh one is cut each weekday, so use the most recent) +When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0202965ec4b..d995ddcc87e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -322,7 +322,7 @@ npm run build ## Submitting Your PR 1. **Push your branch**: `git push origin your-feature-branch` -2. **Create a PR**: Go to GitHub and open a pull request against the current daily OSS branch, named `litellm_oss_daily_YYYY_MM_DD`. A fresh one is cut each weekday, so pick the most recent from the [branch list](https://github.com/BerriAI/litellm/branches/all?query=litellm_oss_daily). Do not target `main`. +2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`. 3. **Fill out the PR template**: Provide clear description of changes 4. **Wait for review**: Maintainers will review and provide feedback 5. **Address feedback**: Make requested changes and push updates From d4035a07c77051a480eb00d203011dbd1cbb5a02 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 6 Jul 2026 10:15:51 -0700 Subject: [PATCH 28/60] feat(mcp): migrate client_credentials (M2M) onto the v2 resolver arm Replaces the not_implemented stub with a live arm: ClientCredentialsTokenSource mints and caches the M2M token (rotation-aware identity key, expires_in-driven TTL, audience and token_endpoint_auth_method support) and ClientCredentialsBearerAuth retries an upstream 401 exactly once with a freshly minted token. to_server_spec owns oauth2_flow=client_credentials servers and fails closed on incomplete grant config instead of connecting unauthenticated --- .../outbound_credentials/adapter.py | 39 +- .../client_credentials.py | 334 ++++++++++++++++++ .../outbound_credentials/resolver.py | 39 +- .../mcp_server/outbound_credentials/types.py | 9 +- .../outbound_credentials/test_adapter.py | 70 +++- .../test_client_credentials.py | 320 +++++++++++++++++ .../outbound_credentials/test_resolver.py | 108 +++++- 7 files changed, 901 insertions(+), 18 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 6631e38f524..3bf9fd1c12e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -21,8 +21,12 @@ from typing_extensions import assert_never from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, +<<<<<<< HEAD ClientAuth, ClientSecretAuth, +======= + ClientCredentialsConfig, +>>>>>>> 73df37ca23 (feat(mcp): migrate client_credentials (M2M) onto the v2 resolver arm) CredError, IdJagConfig, NoneConfig, @@ -70,10 +74,10 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is explicitly mapped or explicitly deferred, rather than silently falling through to v1. Live modes: ``none``, the static-header family (``api_key`` plus the Authorization schemes, - all shared-key), ``oauth2`` per-user tokens (``authorization_code``), ``oauth2_token_exchange`` - (OBO), and the client-forwarded token modes ``true_passthrough`` / ``oauth_delegate`` - (``PassthroughConfig``); client_credentials (M2M), delegated/passthrough oauth2, and SigV4 - return None and stay on v1. + all shared-key), ``oauth2`` per-user tokens (``authorization_code``), ``oauth2`` M2M + (``client_credentials``), ``oauth2_token_exchange`` (OBO), and the client-forwarded token + modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough + oauth2 and SigV4 return None and stay on v1. """ if server.is_byok: return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type) @@ -95,13 +99,15 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: case MCPAuth.basic: return _shared_key_spec(server, resource, "Authorization", "Basic", encode=True) case MCPAuth.oauth2: + if server.has_client_credentials: + return _client_credentials_spec(server, resource) if server.needs_user_oauth_token and not server.delegate_auth_to_upstream: return ServerSpec( server_id=server.server_id, resource=resource, config=AuthorizationCodeConfig(), ) - # client_credentials (M2M) and delegate/passthrough oauth2 stay on v1 + # delegate/passthrough oauth2 stay on v1 return None case MCPAuth.oauth2_id_jag: return _id_jag_spec(server, resource) @@ -114,6 +120,29 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: assert_never(auth_type) +def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec: + """Build a client_credentials (M2M) spec; the explicit ``oauth2_flow`` opt-in owns the server. + + Missing grant fields (``client_id``/``client_secret``/``token_url``) are NOT a reason to defer: + v1 would connect unauthenticated and the upstream's 401 gets absorbed into an empty tool list, + so the arm fails closed with ``misconfigured`` instead, naming the missing fields (mirrors the + OBO ownership rule). ``audience`` is forwarded only when the operator set it; a missing one is + omitted, not derived, since a fabricated value risks the IdP rejecting the grant. + """ + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=ClientCredentialsConfig( + client_id=server.client_id, + client_secret=SecretStr(server.client_secret) if server.client_secret else None, + token_url=server.token_url, + scopes=tuple(server.scopes or ()), + audience=server.audience, + token_endpoint_auth_method=server.token_endpoint_auth_method, + ), + ) + + def _token_exchange_spec(server: MCPServer, resource: str) -> Optional[ServerSpec]: """Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py new file mode 100644 index 00000000000..3f6c329118f --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -0,0 +1,334 @@ +"""The ``client_credentials`` (M2M) arm's token source and retrying bearer auth. + +Implements the client-credentials behavior contract for the v2 resolver: + +- **Acquisition**: POST ``grant_type=client_credentials`` to the configured token endpoint with + the configured scopes and (when set) the IdP's ``audience`` parameter, authenticating the + client per ``token_endpoint_auth_method`` (RFC 6749 section 2.3.1, shared helper). +- **Caching**: tokens are cached per ``(client identity, server)`` where the identity key hashes + ``token_url`` / ``client_id`` / ``client_secret`` / auth method / scopes / audience — rotating + or re-scoping the credentials changes the key, so a stale token can never be served for the + new identity (the contract's rotation-invalidation clause). +- **Expiry**: the cache TTL respects ``expires_in`` minus a skew so an entry lapses before the + real token does; a response with no ``expires_in`` is cached briefly + (``default_ttl_seconds``), not assumed long-lived. No refresh_token is ever expected. +- **401 recovery**: ``ClientCredentialsBearerAuth`` retries an upstream request exactly once + after a 401 — discard the cached token, mint a fresh one, resend; a second failure surfaces + the upstream's own auth error unchanged. +- **No user context**: nothing here reads a ``Subject``; every caller shares the one client + identity. + +The token-endpoint POST is injected (``M2MTokenEndpointPost``) so the grant orchestration is +testable without a live IdP; ``post_client_credentials_grant`` is the httpx edge and the one +place the untyped response boundary is contained. Failures are values: the source returns +``Result[OAuthToken, CredError]``; only the httpx edge touches exceptions. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import time +from collections.abc import AsyncGenerator, Awaitable, Callable, Generator +from dataclasses import dataclass +from typing import Annotated, Literal + +import httpx +from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InMemoryTokenCacheBackend, + OAuthToken, + TokenCacheBackend, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientCredentialsConfig, + CredError, +) + + +class TokenEndpointSuccess(BaseModel): + """The endpoint returned a JSON object; field validation is the caller's job.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["success"] = "success" + body: dict[str, object] + + +class TokenEndpointDenied(BaseModel): + """The endpoint answered but did not grant a token (an HTTP error or a non-JSON body).""" + + model_config = ConfigDict(frozen=True) + tag: Literal["denied"] = "denied" + status_code: int + detail: str + + +class TokenEndpointUnreachable(BaseModel): + """The endpoint could not be reached (DNS, TLS, connect/read failure).""" + + model_config = ConfigDict(frozen=True) + tag: Literal["unreachable"] = "unreachable" + detail: str + + +TokenEndpointOutcome = Annotated[ + TokenEndpointSuccess | TokenEndpointDenied | TokenEndpointUnreachable, + Field(discriminator="tag"), +] + +M2MTokenEndpointPost = Callable[[str, "dict[str, str]", "dict[str, str]"], Awaitable[TokenEndpointOutcome]] + + +_TOKEN_BODY_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) + + +async def post_client_credentials_grant( + url: str, form: dict[str, str], headers: dict[str, str] +) -> TokenEndpointOutcome: + """POST the grant to the token endpoint and classify the transport outcome. + + The httpx edge: litellm's handler is partially typed (and raises ``HTTPStatusError`` itself on + a 4xx/5xx), so the untyped boundary is contained here and every field the caller reads comes + out of a validated ``TokenEndpointOutcome``. + """ + from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler is partially typed + ) + from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 + + try: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + response = await client.post( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # handler is partially typed + url, headers={"Accept": "application/json", **headers}, data=form + ) + except httpx.HTTPStatusError as status_err: + status_code = status_err.response.status_code + return TokenEndpointDenied(status_code=status_code, detail=f"token endpoint returned HTTP {status_code}") + except Exception as exc: # noqa: BLE001 # any transport failure is the same outcome: unreachable + return TokenEndpointUnreachable(detail=str(exc)) + if not isinstance(response, httpx.Response): + return TokenEndpointUnreachable(detail="token endpoint returned no response") + try: + body = _TOKEN_BODY_ADAPTER.validate_json(response.content) + except ValidationError: + return TokenEndpointDenied( + status_code=response.status_code, detail="token endpoint returned a non-JSON-object body" + ) + return TokenEndpointSuccess(body=body) + + +def _parse_expires_in(raw: object) -> int | None: + if isinstance(raw, bool): + return None + if isinstance(raw, int): + return raw + if isinstance(raw, str): + try: + return int(raw) + except ValueError: + return None + return None + + +def _parse_granted_scopes(raw: object) -> tuple[str, ...] | None: + return tuple(raw.split()) if isinstance(raw, str) and raw else None + + +@dataclass(frozen=True, slots=True) +class _PreparedGrant: + """A validated, ready-to-POST grant plus the identity key its token caches under.""" + + token_url: str + form: dict[str, str] + headers: dict[str, str] + identity_key: str + + +class ClientCredentialsTokenSource: + """Cached M2M access tokens, one per ``(client identity, server)``. + + ``get`` serves from the cache while the entry's TTL (derived from ``expires_in`` minus + ``expiry_skew_seconds``) holds, fetching under a per-server lock so concurrent misses + produce one grant. ``refetch`` is the 401-recovery path: it drops the failed token and + mints a fresh one, unless a concurrent caller already replaced it. + """ + + def __init__( + self, + post: M2MTokenEndpointPost = post_client_credentials_grant, + *, + backend: TokenCacheBackend | None = None, + default_ttl_seconds: float = 300.0, + expiry_skew_seconds: float = 60.0, + min_cache_seconds: float = 10.0, + clock: Callable[[], float] = time.time, + ) -> None: + self._post = post + self._backend: TokenCacheBackend = backend or InMemoryTokenCacheBackend(clock=clock) + self._default_ttl_seconds = default_ttl_seconds + self._expiry_skew_seconds = expiry_skew_seconds + self._min_cache_seconds = min_cache_seconds + self._clock = clock + self._locks: dict[str, asyncio.Lock] = {} + + def _lock(self, server_id: str) -> asyncio.Lock: + return self._locks.setdefault(server_id, asyncio.Lock()) + + async def get(self, server_id: str, config: ClientCredentialsConfig) -> Result[OAuthToken, CredError]: + match _prepare_grant(config): + case Error(err): + return Error(err) + case Ok(grant): + cached = await self._backend.get(grant.identity_key, server_id) + if cached is not None: + return Ok(cached) + async with self._lock(server_id): + cached = await self._backend.get(grant.identity_key, server_id) + if cached is not None: + return Ok(cached) + return await self._fetch_and_cache(server_id, grant) + + async def refetch(self, server_id: str, config: ClientCredentialsConfig, failed_access_token: str) -> str | None: + """Replace a token the upstream just 401'd; returns the fresh bearer value or ``None``. + + Runs under the same per-server lock as ``get``: if a concurrent caller already replaced + the failed token, that replacement is returned without another grant, so a burst of 401s + yields one fetch. A failed refetch returns ``None`` and the caller surfaces the + upstream's original auth error (the contract's retry-once-then-give-up clause). + """ + match _prepare_grant(config): + case Error(_): + return None + case Ok(grant): + async with self._lock(server_id): + cached = await self._backend.get(grant.identity_key, server_id) + if cached is not None and cached.access_token != failed_access_token: + return cached.access_token + await self._backend.delete(grant.identity_key, server_id) + match await self._fetch_and_cache(server_id, grant): + case Ok(token): + return token.access_token + case Error(_): + return None + + async def _fetch_and_cache(self, server_id: str, grant: _PreparedGrant) -> Result[OAuthToken, CredError]: + outcome = await self._post(grant.token_url, grant.form, grant.headers) + match outcome: + case TokenEndpointUnreachable(): + return Error(CredError.of_upstream_unavailable(f"OAuth2 token endpoint unreachable: {outcome.detail}")) + case TokenEndpointDenied(): + if outcome.status_code >= 500: + return Error(CredError.of_upstream_unavailable(f"OAuth2 token endpoint failed: {outcome.detail}")) + return Error(CredError.of_misconfigured(f"OAuth2 client_credentials grant rejected: {outcome.detail}")) + case TokenEndpointSuccess(): + return await self._cache_token(server_id, grant, outcome.body) + assert_never(outcome) + + async def _cache_token( + self, server_id: str, grant: _PreparedGrant, body: dict[str, object] + ) -> Result[OAuthToken, CredError]: + access_token = body.get("access_token") + if not isinstance(access_token, str) or not access_token: + return Error(CredError.of_misconfigured("OAuth2 token response is missing 'access_token'")) + expires_in = _parse_expires_in(body.get("expires_in")) + token = OAuthToken( + access_token=access_token, + expires_at=self._clock() + expires_in if expires_in is not None else None, + scopes=_parse_granted_scopes(body.get("scope")) or (), + ) + ttl = ( + max(expires_in - self._expiry_skew_seconds, self._min_cache_seconds) + if expires_in is not None + else self._default_ttl_seconds + ) + await self._backend.set(grant.identity_key, server_id, token, ttl) + return Ok(token) + + +def _prepare_grant(config: ClientCredentialsConfig) -> Result[_PreparedGrant, CredError]: + if not config.client_id or not config.client_secret or not config.token_url: + missing = ", ".join( + name + for name, present in ( + ("client_id", bool(config.client_id)), + ("client_secret", bool(config.client_secret)), + ("token_url", bool(config.token_url)), + ) + if not present + ) + return Error(CredError.of_misconfigured(f"client_credentials config is missing: {missing}")) + + from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( # noqa: PLC0415 + build_token_endpoint_client_auth, + ) + + client_auth = build_token_endpoint_client_auth( + auth_method=config.token_endpoint_auth_method, + client_id=config.client_id, + client_secret=config.client_secret.get_secret_value(), + ) + form = { + "grant_type": "client_credentials", + **client_auth.body, + **({"scope": " ".join(config.scopes)} if config.scopes else {}), + **({"audience": config.audience} if config.audience else {}), + } + return Ok( + _PreparedGrant( + token_url=config.token_url, + form=form, + headers=client_auth.headers, + identity_key=_identity_key(config), + ) + ) + + +def _identity_key(config: ClientCredentialsConfig) -> str: + """Hash of everything that names the client identity; any rotation yields a new key.""" + material = "\n".join( + ( + config.token_url or "", + config.client_id or "", + config.client_secret.get_secret_value() if config.client_secret else "", + config.token_endpoint_auth_method or "", + " ".join(config.scopes), + config.audience or "", + ) + ) + return hashlib.sha256(material.encode("utf-8")).hexdigest() + + +class ClientCredentialsBearerAuth(httpx.Auth): + """Bearer auth that retries an upstream 401 exactly once with a freshly minted token. + + The initial token was already resolved (so config/IdP failures surfaced as typed errors + before any upstream request); ``refetch`` is the source's 401-recovery callback. If the + refetch fails, or the retried request 401s again, the upstream's response stands. + """ + + def __init__(self, access_token: str, refetch: Callable[[str], Awaitable[str | None]]) -> None: + self.header_name = "Authorization" + self._access_token = SecretStr(access_token) + self._refetch = refetch + + async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + token = self._access_token.get_secret_value() + request.headers[self.header_name] = f"Bearer {token}" + response = yield request + if response.status_code != 401: + return + fresh = await self._refetch(token) + if fresh is None: + return + request.headers[self.header_name] = f"Bearer {fresh}" + yield request + + def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + raise RuntimeError("ClientCredentialsBearerAuth only supports async httpx clients") diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 7e5c073870a..69984a56311 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -9,18 +9,24 @@ at runtime instead of returning `None`. `none`, `api_key` (shared-key source), and `passthrough` (forwards the caller's own inbound token) are live, as is `authorization_code`, which reads the user's token from the injected -`OAuthTokenStore`, and `token_exchange`, which swaps the caller's inbound token through the -injected `TokenExchanger`. The remaining arms are `not_implemented` stubs that each land in a -follow-up PR with their seam. Pure v2: no imports from v1. +`OAuthTokenStore`, `token_exchange`, which swaps the caller's inbound token through the injected +`TokenExchanger`, and `client_credentials`, which mints and caches the gateway's M2M token through +the injected `ClientCredentialsTokenSource`. The remaining arms are `not_implemented` stubs that +each land in a follow-up PR with their seam. Pure v2: no imports from v1. """ from __future__ import annotations import hashlib +from functools import partial import httpx from typing_extensions import assert_never +from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ClientCredentialsTokenSource, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( NoOpAuth, StaticHeaderAuth, @@ -104,11 +110,13 @@ class UpstreamCredentialProvider: token_exchanger: TokenExchanger | None = None, token_endpoint: TokenEndpointClient | None = None, exchanged_tokens: ExchangedTokenCache | None = None, + client_credentials_source: ClientCredentialsTokenSource | None = None, ) -> None: self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore() self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger() self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient() self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache() + self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -118,8 +126,8 @@ class UpstreamCredentialProvider: return self._api_key(config) case PassthroughConfig(): return self._passthrough(subject) - case ClientCredentialsConfig(): - return _not_implemented(AuthSpecKind.client_credentials) + case ClientCredentialsConfig() as config: + return await self._client_credentials(server.server_id, config) case TokenExchangeConfig() as config: return await self._token_exchange(subject, server, config) case IdJagConfig() as config: @@ -215,6 +223,23 @@ class UpstreamCredentialProvider: return Error(CredError.of_unauthorized("Authorization required: complete the OAuth flow for this server.")) return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name="Authorization")) + async def _client_credentials( + self, server_id: str, config: ClientCredentialsConfig + ) -> Result[httpx.Auth, CredError]: + """The M2M arm: resolve a cached (or freshly minted) gateway token; no user context. + + The token is resolved here, before any upstream request, so a misconfigured grant or an + unreachable IdP surfaces as a typed ``CredError``. The returned auth carries the source's + ``refetch``, so an upstream 401 is retried exactly once with a freshly minted token (the + contract's invalid-token recovery); a second 401 surfaces the upstream's own error. + """ + match await self._client_credentials_source.get(server_id, config): + case Ok(token): + refetch = partial(self._client_credentials_source.refetch, server_id, config) + return Ok(ClientCredentialsBearerAuth(token.access_token, refetch)) + case Error(err): + return Error(err) + async def _token_exchange( self, subject: Subject, server: ServerSpec, config: TokenExchangeConfig ) -> Result[StaticHeaderAuth, CredError]: @@ -245,7 +270,9 @@ class UpstreamCredentialProvider: Used after an upstream rejects the injected credential, so the next resolve re-mints rather than serving the same rejected token until TTL. `token_exchange` and `id_jag` hold a - re-mintable cached credential here; other modes are a no-op. + re-mintable cached credential here; `client_credentials` recovers inside its own auth flow + (`ClientCredentialsBearerAuth` retries the 401'd request once with a fresh token), and + other modes are a no-op. """ if subject.inbound_token is None: return diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 64a20255ab2..926d96c8868 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -184,7 +184,12 @@ class ClientCredentialsConfig(BaseModel): Fields are optional so the config can be built incomplete: a value may be supplied at runtime (`token_url` via RFC 8414 discovery, `client_id`/`secret` via DCR), and the - resolver arm raises `CredError.misconfigured` when a needed field is still absent. + resolver arm returns `CredError.misconfigured` when a needed field is still absent. + + `audience` is the IdP-specific audience parameter some authorization servers require on + the client_credentials grant (sent as `audience` in the token request when set). + `token_endpoint_auth_method` selects how the client authenticates to the token endpoint + (RFC 6749 section 2.3.1); `None` defaults to `client_secret_post`. """ model_config = ConfigDict(frozen=True) @@ -193,6 +198,8 @@ class ClientCredentialsConfig(BaseModel): client_secret: SecretStr | None = None token_url: str | None = None scopes: tuple[str, ...] = () + audience: str | None = None + token_endpoint_auth_method: Literal["client_secret_post", "client_secret_basic"] | None = None class TokenExchangeConfig(BaseModel): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 707374e7061..bf757b64c9a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -21,6 +21,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + ClientCredentialsConfig, ClientSecretAuth, CredError, IdJagConfig, @@ -107,7 +108,6 @@ def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow): [ _server(auth_type=MCPAuth.api_key), # no token configured _server(auth_type=MCPAuth.bearer_token), # no token configured - _server(auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials"), # M2M -> v1 _server(auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True), # delegated upstream OAuth -> v1 _server(auth_type=MCPAuth.oauth2_token_exchange), # no endpoint/client creds -> incomplete -> v1 _server( @@ -124,6 +124,74 @@ def test_unmigrated_modes_defer_to_v1(server): assert to_server_spec(server) is None +def test_client_credentials_maps_full_config(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + url="https://up.example.com/mcp", + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + scopes=["read", "write"], + audience="https://up.example.com", + token_endpoint_auth_method="client_secret_basic", + ) + ) + assert spec is not None + config = spec.config + assert isinstance(config, ClientCredentialsConfig) + assert config.client_id == "cid" + assert config.client_secret is not None + assert config.client_secret.get_secret_value() == "csec" + assert config.token_url == "https://idp.example.com/token" + assert config.scopes == ("read", "write") + assert config.audience == "https://up.example.com" + assert config.token_endpoint_auth_method == "client_secret_basic" + + +def test_client_credentials_omits_audience_when_unset(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, ClientCredentialsConfig) + assert spec.config.audience is None + + +def test_client_credentials_with_incomplete_grant_fields_is_owned_for_fail_closed(): + # An M2M server missing its grant fields is still owned by v2 (spec, not None) so it fails + # closed at the source (misconfigured, 500) rather than deferring to v1, which would connect + # unauthenticated and mask the upstream 401 as an empty tool list. + spec = to_server_spec(_server(auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials", client_id="cid")) + assert spec is not None + assert isinstance(spec.config, ClientCredentialsConfig) + assert spec.config.token_url is None + assert spec.config.client_secret is None + + +def test_client_credentials_wins_over_delegate_flag(): + # v1 never delegates for M2M servers; the explicit oauth2_flow opt-in outranks the delegate flag. + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + delegate_auth_to_upstream=True, + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, ClientCredentialsConfig) + + def test_token_exchange_maps_full_config(): spec = to_server_spec( _server( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py new file mode 100644 index 00000000000..bd00bb77abd --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -0,0 +1,320 @@ +"""Tests for the client_credentials (M2M) token source and its retrying bearer auth. + +These are the behavior-contract spec: grant shape (scopes / audience / client auth method), +rotation-aware cache keying, expires_in-driven expiry, error classification, and the +401 -> discard -> refetch -> retry-once recovery in ``ClientCredentialsBearerAuth``. +""" + +import httpx +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ClientCredentialsTokenSource, + TokenEndpointDenied, + TokenEndpointOutcome, + TokenEndpointSuccess, + TokenEndpointUnreachable, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientCredentialsConfig, +) + + +class _Clock: + def __init__(self, t: float = 1000.0) -> None: + self.t = t + + def __call__(self) -> float: + return self.t + + +class _FakePoster: + """Records every grant POST and returns canned outcomes (last one repeats).""" + + def __init__(self, outcomes: "list[TokenEndpointOutcome]") -> None: + self._outcomes = outcomes + self.calls: "list[tuple[str, dict[str, str], dict[str, str]]]" = [] + + async def __call__(self, url: str, form: "dict[str, str]", headers: "dict[str, str]") -> TokenEndpointOutcome: + self.calls.append((url, dict(form), dict(headers))) + index = min(len(self.calls) - 1, len(self._outcomes) - 1) + return self._outcomes[index] + + +def _success(access_token: str = "m2m-token", **extra: object) -> TokenEndpointSuccess: + return TokenEndpointSuccess(body={"access_token": access_token, **extra}) + + +def _config(**overrides: object) -> ClientCredentialsConfig: + fields: "dict[str, object]" = { + "client_id": "cid", + "client_secret": SecretStr("csec"), + "token_url": "https://idp.example.com/token", + **overrides, + } + return ClientCredentialsConfig.model_validate(fields) + + +@pytest.mark.asyncio +async def test_grant_posts_client_credentials_with_scopes_and_audience(): + poster = _FakePoster([_success()]) + source = ClientCredentialsTokenSource(poster) + result = await source.get("s", _config(scopes=("read", "write"), audience="https://api.example.com")) + assert isinstance(result, Ok) + assert result.ok.access_token == "m2m-token" + url, form, _headers = poster.calls[0] + assert url == "https://idp.example.com/token" + assert form["grant_type"] == "client_credentials" + assert form["scope"] == "read write" + assert form["audience"] == "https://api.example.com" + assert form["client_id"] == "cid" + assert form["client_secret"] == "csec" + + +@pytest.mark.asyncio +async def test_grant_omits_scope_and_audience_when_not_configured(): + poster = _FakePoster([_success()]) + await ClientCredentialsTokenSource(poster).get("s", _config()) + _url, form, _headers = poster.calls[0] + assert "scope" not in form + assert "audience" not in form + + +@pytest.mark.asyncio +async def test_grant_honors_client_secret_basic(): + poster = _FakePoster([_success()]) + await ClientCredentialsTokenSource(poster).get("s", _config(token_endpoint_auth_method="client_secret_basic")) + _url, form, headers = poster.calls[0] + assert headers["Authorization"].startswith("Basic ") + assert "client_secret" not in form + assert "client_id" not in form + + +@pytest.mark.asyncio +async def test_missing_grant_fields_are_misconfigured_and_never_posted(): + poster = _FakePoster([_success()]) + result = await ClientCredentialsTokenSource(poster).get( + "s", ClientCredentialsConfig(client_id="cid", client_secret=SecretStr("csec")) + ) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + assert "token_url" in result.error.summary + assert poster.calls == [] + + +@pytest.mark.asyncio +async def test_token_is_cached_across_gets(): + poster = _FakePoster([_success(expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + first = await source.get("s", _config()) + second = await source.get("s", _config()) + assert isinstance(first, Ok) and isinstance(second, Ok) + assert second.ok.access_token == first.ok.access_token + assert len(poster.calls) == 1 + + +@pytest.mark.asyncio +async def test_expires_in_bounds_the_cache_lifetime(): + clock = _Clock(1000.0) + poster = _FakePoster([_success("t1", expires_in=120), _success("t2", expires_in=120)]) + source = ClientCredentialsTokenSource(poster, expiry_skew_seconds=60.0, clock=clock) + first = await source.get("s", _config()) + assert isinstance(first, Ok) + assert first.ok.expires_at == 1120.0 + clock.t = 1059.0 # within expires_in - skew + assert len(poster.calls) == 1 + within = await source.get("s", _config()) + assert isinstance(within, Ok) and within.ok.access_token == "t1" + clock.t = 1061.0 # past expires_in - skew: the entry lapsed before the real token does + lapsed = await source.get("s", _config()) + assert isinstance(lapsed, Ok) and lapsed.ok.access_token == "t2" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +async def test_missing_expires_in_is_cached_briefly_not_an_hour(): + clock = _Clock(1000.0) + poster = _FakePoster([_success("t1"), _success("t2")]) + source = ClientCredentialsTokenSource(poster, default_ttl_seconds=300.0, clock=clock) + first = await source.get("s", _config()) + assert isinstance(first, Ok) + assert first.ok.expires_at is None + clock.t = 1301.0 # past the default TTL; v1 would still be serving its 3600s-cached token + second = await source.get("s", _config()) + assert isinstance(second, Ok) and second.ok.access_token == "t2" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rotation", + [ + {"client_secret": SecretStr("rotated")}, + {"client_id": "cid-2"}, + {"scopes": ("admin",)}, + {"audience": "https://other.example.com"}, + {"token_url": "https://idp2.example.com/token"}, + ], +) +async def test_credential_rotation_invalidates_the_cached_token(rotation): + poster = _FakePoster([_success("old", expires_in=3600), _success("new", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + before = await source.get("s", _config()) + after = await source.get("s", _config(**rotation)) + assert isinstance(before, Ok) and before.ok.access_token == "old" + assert isinstance(after, Ok) and after.ok.access_token == "new" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +async def test_idp_4xx_is_misconfigured_and_5xx_is_unavailable(): + denied = await ClientCredentialsTokenSource( + _FakePoster([TokenEndpointDenied(status_code=401, detail="HTTP 401")]) + ).get("s", _config()) + assert isinstance(denied, Error) and denied.error.tag == "misconfigured" + down = await ClientCredentialsTokenSource( + _FakePoster([TokenEndpointDenied(status_code=503, detail="HTTP 503")]) + ).get("s", _config()) + assert isinstance(down, Error) and down.error.tag == "upstream_unavailable" + unreachable = await ClientCredentialsTokenSource(_FakePoster([TokenEndpointUnreachable(detail="dns")])).get( + "s", _config() + ) + assert isinstance(unreachable, Error) and unreachable.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_response_without_access_token_is_misconfigured(): + poster = _FakePoster([TokenEndpointSuccess(body={"token_type": "Bearer"})]) + result = await ClientCredentialsTokenSource(poster).get("s", _config()) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + + +@pytest.mark.asyncio +async def test_error_results_are_not_cached(): + poster = _FakePoster([TokenEndpointUnreachable(detail="down"), _success("recovered")]) + source = ClientCredentialsTokenSource(poster) + first = await source.get("s", _config()) + second = await source.get("s", _config()) + assert isinstance(first, Error) + assert isinstance(second, Ok) and second.ok.access_token == "recovered" + + +@pytest.mark.asyncio +async def test_refetch_discards_the_failed_token_and_mints_a_fresh_one(): + poster = _FakePoster([_success("stale", expires_in=3600), _success("fresh", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + first = await source.get("s", _config()) + assert isinstance(first, Ok) + fresh = await source.refetch("s", _config(), failed_access_token="stale") + assert fresh == "fresh" + assert len(poster.calls) == 2 + after = await source.get("s", _config()) + assert isinstance(after, Ok) and after.ok.access_token == "fresh" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +async def test_refetch_reuses_a_concurrent_replacement_without_a_second_grant(): + poster = _FakePoster([_success("replacement", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + seeded = await source.get("s", _config()) + assert isinstance(seeded, Ok) + result = await source.refetch("s", _config(), failed_access_token="some-older-token") + assert result == "replacement" + assert len(poster.calls) == 1 + + +@pytest.mark.asyncio +async def test_refetch_returns_none_when_the_grant_fails(): + poster = _FakePoster([_success("stale"), TokenEndpointUnreachable(detail="down")]) + source = ClientCredentialsTokenSource(poster) + await source.get("s", _config()) + assert await source.refetch("s", _config(), failed_access_token="stale") is None + + +def _upstream(responses: "list[httpx.Response]") -> "tuple[httpx.MockTransport, list[str]]": + # The auth flow re-yields the same Request object on retry, so snapshot the Authorization + # value per send; holding the Request would show the post-retry mutation for both entries. + seen: "list[str]" = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.headers.get("Authorization", "")) + return responses[min(len(seen) - 1, len(responses) - 1)] + + return httpx.MockTransport(handler), seen + + +@pytest.mark.asyncio +async def test_bearer_auth_sends_the_token_and_leaves_a_success_alone(): + transport, seen = _upstream([httpx.Response(200)]) + + async def refetch(failed: str) -> "str | None": + raise AssertionError("must not refetch on success") + + auth = ClientCredentialsBearerAuth("m2m-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + response = await client.get("https://upstream.example.com/mcp") + assert response.status_code == 200 + assert seen == ["Bearer m2m-token"] + + +@pytest.mark.asyncio +async def test_bearer_auth_retries_a_401_once_with_a_fresh_token(): + transport, seen = _upstream([httpx.Response(401), httpx.Response(200)]) + refetched: "list[str]" = [] + + async def refetch(failed: str) -> "str | None": + refetched.append(failed) + return "fresh-token" + + auth = ClientCredentialsBearerAuth("stale-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + response = await client.get("https://upstream.example.com/mcp") + assert response.status_code == 200 + assert refetched == ["stale-token"] + assert seen == ["Bearer stale-token", "Bearer fresh-token"] + + +@pytest.mark.asyncio +async def test_bearer_auth_surfaces_the_401_when_the_refetch_fails(): + transport, seen = _upstream([httpx.Response(401)]) + + async def refetch(failed: str) -> "str | None": + return None + + auth = ClientCredentialsBearerAuth("stale-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + response = await client.get("https://upstream.example.com/mcp") + assert response.status_code == 401 + assert len(seen) == 1 + + +@pytest.mark.asyncio +async def test_bearer_auth_gives_up_after_a_second_401(): + transport, seen = _upstream([httpx.Response(401), httpx.Response(401)]) + refetched: "list[str]" = [] + + async def refetch(failed: str) -> "str | None": + refetched.append(failed) + return "fresh-token" + + auth = ClientCredentialsBearerAuth("stale-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + response = await client.get("https://upstream.example.com/mcp") + assert response.status_code == 401 + assert len(seen) == 2 + assert refetched == ["stale-token"] + + +def test_bearer_auth_rejects_sync_clients(): + async def refetch(failed: str) -> "str | None": + return None + + auth = ClientCredentialsBearerAuth("token", refetch) + with httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200)), auth=auth) as client: + with pytest.raises(RuntimeError): + client.get("https://upstream.example.com/mcp") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index ba7720ffd51..a710da81962 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -1,9 +1,10 @@ """Tests for the resolver dispatch: live arms produce auth, stubbed arms fail closed. -`none`, `api_key` (shared-key source), `passthrough`, `authorization_code`, and `token_exchange` are -implemented; every other arm, plus the `api_key` BYOK source, returns a typed `not_implemented` error -until its mode lands. Parametrizing the stubs over one config each also guards reachability: a dropped -`case` would hit `assert_never` and raise instead of returning the stub. +`none`, `api_key` (shared-key source), `passthrough`, `authorization_code`, `token_exchange`, and +`client_credentials` are implemented; every other arm, plus the `api_key` BYOK source, returns a +typed `not_implemented` error until its mode lands. Parametrizing the stubs over one config each +also guards reachability: a dropped `case` would hit `assert_never` and raise instead of +returning the stub. """ import httpx @@ -324,9 +325,106 @@ async def test_passthrough_without_inbound_token_is_a_no_op(): assert isinstance(result.ok, NoOpAuth) +class _FakeM2MSource: + """A ClientCredentialsTokenSource returning a canned result and recording refetches.""" + + def __init__(self, result) -> None: + self._result = result + self.gets: list[str] = [] + self.refetches: list[tuple[str, str]] = [] + + async def get(self, server_id: str, config): + self.gets.append(server_id) + return self._result + + async def refetch(self, server_id: str, config, failed_access_token: str): + self.refetches.append((server_id, failed_access_token)) + return "fresh-m2m" + + +_M2M = ClientCredentialsConfig( + client_id="cid", + client_secret=SecretStr("csec"), + token_url="https://idp.example.com/token", +) + + +async def _emitted_async(auth: httpx.Auth, respond=None) -> tuple[httpx.Headers, list[httpx.Request]]: + """Drive the async auth flow one request at a time, replying via ``respond`` when given.""" + seen: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request) + return respond(request) if respond else httpx.Response(200) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler), auth=auth) as client: + await client.get("https://upstream.example.com/mcp") + return seen[-1].headers, seen + + +@pytest.mark.asyncio +async def test_client_credentials_emits_the_minted_bearer(): + source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials( + _SUBJECT, _spec(_M2M) + ) + assert isinstance(result, Ok) + headers, _ = await _emitted_async(result.ok) + assert headers["Authorization"] == "Bearer m2m-at" + assert source.gets == ["s"] + + +@pytest.mark.asyncio +async def test_client_credentials_ignores_the_subject(): + # The contract's no-user-context clause: every caller shares the one client identity. + source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) + provider = UpstreamCredentialProvider(client_credentials_source=source) + alice = await provider.resolve_credentials(Subject(tenant_id="t1", subject_id="alice"), _spec(_M2M)) + bob = await provider.resolve_credentials(Subject(tenant_id="t2", subject_id="bob"), _spec(_M2M)) + assert isinstance(alice, Ok) and isinstance(bob, Ok) + alice_headers, _ = await _emitted_async(alice.ok) + bob_headers, _ = await _emitted_async(bob.ok) + assert alice_headers["Authorization"] == bob_headers["Authorization"] == "Bearer m2m-at" + + +@pytest.mark.asyncio +async def test_client_credentials_auth_retries_a_401_through_the_source(): + source = _FakeM2MSource(Ok(OAuthToken(access_token="stale-at"))) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials( + _SUBJECT, _spec(_M2M) + ) + assert isinstance(result, Ok) + + def respond(request: httpx.Request) -> httpx.Response: + is_stale = request.headers["Authorization"] == "Bearer stale-at" + return httpx.Response(401) if is_stale else httpx.Response(200) + + headers, seen = await _emitted_async(result.ok, respond) + assert headers["Authorization"] == "Bearer fresh-m2m" + assert len(seen) == 2 + assert source.refetches == [("s", "stale-at")] + + +@pytest.mark.asyncio +async def test_client_credentials_propagates_the_source_error(): + source = _FakeM2MSource(Error(CredError.of_upstream_unavailable("idp down"))) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials( + _SUBJECT, _spec(_M2M) + ) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_client_credentials_with_no_source_wired_fails_closed_on_missing_config(): + # The default source validates the grant fields before any network is touched. + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(ClientCredentialsConfig())) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + + _STUBBED = [ ("api_key_byok", ApiKeyConfig(key_source=Byok())), - ("client_credentials", ClientCredentialsConfig()), ("aws_sigv4", AwsSigV4Config(region="us-east-1")), ] From 4b1c9d44984c0c6edba91d7132bcc5b76d3b0126 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 6 Jul 2026 11:57:37 -0700 Subject: [PATCH 29/60] test(mcp): update _create_mcp_client graft tests for migrated M2M arm The graft test pinned the pre-migration contract (M2M defers to v1). Replaced with two tests pinning the new one: a complete-config M2M server resolves via the v2 arm into ClientCredentialsBearerAuth, and an incomplete-config server fails closed with a 500 misconfigured naming the missing grant fields --- .../mcp_server/test_mcp_server_manager.py | 55 +++++++++++++++---- 1 file changed, 45 insertions(+), 10 deletions(-) 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 491fa023031..f0ca75b8cdb 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 @@ -7760,24 +7760,59 @@ class TestCreateMcpClientV2Graft: assert client._resolved_auth.header_name == "Authorization" assert client._resolved_auth._header_value.get_secret_value() == f"Basic {encoded}" - async def test_m2m_client_credentials_defers_to_v1(self): - # M2M (oauth2 client_credentials) is not migrated: to_server_spec returns - # None, so the graft sets no resolved auth and leaves v1 in charge (v1 - # performs the client_credentials grant itself - the static - # authentication_token is never consumed for oauth2, so it does not flow - # to _mcp_auth_value). Per-user oauth2 (authorization_code) is migrated to - # v2 and is exercised separately. + async def test_m2m_client_credentials_resolves_via_v2(self): + # M2M (oauth2 client_credentials) is migrated: to_server_spec owns the server and the + # v2 arm mints the token through the injected source; nothing flows to v1's auth_value. + from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + UpstreamCredentialProvider, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + class _FakeM2MSource: + async def get(self, server_id, config): + return Ok(OAuthToken(access_token="m2m-at")) + + async def refetch(self, server_id, config, failed_access_token): + return None + client = await MCPServerManager()._create_mcp_client( self._http_server( auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials", - authentication_token="legacy-token", - ) + client_id="cid", + client_secret="csec", + token_url="https://idp.example.com/token", + ), + cred_provider=UpstreamCredentialProvider(client_credentials_source=_FakeM2MSource()), ) - assert client._resolved_auth is None + assert isinstance(client._resolved_auth, ClientCredentialsBearerAuth) + assert client._resolved_auth._access_token.get_secret_value() == "m2m-at" assert client._mcp_auth_value is None + async def test_m2m_client_credentials_incomplete_config_fails_closed(self): + # An M2M server missing its grant fields is still owned by v2 and surfaces a 500 + # misconfigured naming the missing fields, rather than deferring to v1 and connecting + # unauthenticated (which masked the upstream 401 as an empty tool list). + with pytest.raises(HTTPException) as exc_info: + await MCPServerManager()._create_mcp_client( + self._http_server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + authentication_token="legacy-token", + ) + ) + + assert exc_info.value.status_code == 500 + assert "misconfigured" in str(exc_info.value.detail) + assert "token_url" in str(exc_info.value.detail) + async def test_static_token_missing_defers_to_v1(self): client = await MCPServerManager()._create_mcp_client( self._http_server(auth_type=MCPAuth.api_key, authentication_token=None) From 7705f0b975cdbc813ac3c6a7183f4278e89c7286 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 6 Jul 2026 12:48:56 -0700 Subject: [PATCH 30/60] fix(mcp): bound the M2M lock dict and cap short-lived token TTL at real expiry Greptile P2s: the per-server lock dict now evicts its oldest entry past max_locks so ephemeral server ids (REST tools preview) cannot grow it unbounded, and the min-cache floor is capped at the token's actual lifetime so an expires_in below the skew is never served past expiry --- .../client_credentials.py | 14 +++++++++- .../test_client_credentials.py | 28 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index 3f6c329118f..332305db3c6 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -168,6 +168,7 @@ class ClientCredentialsTokenSource: default_ttl_seconds: float = 300.0, expiry_skew_seconds: float = 60.0, min_cache_seconds: float = 10.0, + max_locks: int = 1024, clock: Callable[[], float] = time.time, ) -> None: self._post = post @@ -175,10 +176,18 @@ class ClientCredentialsTokenSource: self._default_ttl_seconds = default_ttl_seconds self._expiry_skew_seconds = expiry_skew_seconds self._min_cache_seconds = min_cache_seconds + self._max_locks = max_locks self._clock = clock self._locks: dict[str, asyncio.Lock] = {} def _lock(self, server_id: str) -> asyncio.Lock: + """Per-server single-flight lock, bounded so ephemeral server ids (e.g. the REST tools + preview mints a fresh id per call) cannot grow the dict for the life of the process. + Evicting the oldest entry while a task still holds it only means a concurrent caller for + that server may run its own grant — single-flight is an optimization, not correctness. + """ + if server_id not in self._locks and len(self._locks) >= self._max_locks: + self._locks.pop(next(iter(self._locks)), None) return self._locks.setdefault(server_id, asyncio.Lock()) async def get(self, server_id: str, config: ClientCredentialsConfig) -> Result[OAuthToken, CredError]: @@ -243,8 +252,11 @@ class ClientCredentialsTokenSource: expires_at=self._clock() + expires_in if expires_in is not None else None, scopes=_parse_granted_scopes(body.get("scope")) or (), ) + # The min-cache floor is itself capped at the token's real lifetime, so a token whose + # expires_in is below the skew is never served past its actual expiry; a non-positive + # expires_in caches nothing (every request re-fetches, serialized by the per-server lock). ttl = ( - max(expires_in - self._expiry_skew_seconds, self._min_cache_seconds) + max(expires_in - self._expiry_skew_seconds, min(float(expires_in), self._min_cache_seconds), 0.0) if expires_in is not None else self._default_ttl_seconds ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py index bd00bb77abd..db7e6208e76 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -134,6 +134,34 @@ async def test_expires_in_bounds_the_cache_lifetime(): assert len(poster.calls) == 2 +@pytest.mark.asyncio +async def test_short_lived_token_is_never_served_past_its_expiry(): + # expires_in below the skew must not be floored into serving an expired token: the cache + # entry lapses with the token itself, and the next get re-fetches. + clock = _Clock(1000.0) + poster = _FakePoster([_success("t1", expires_in=5), _success("t2", expires_in=5)]) + source = ClientCredentialsTokenSource(poster, expiry_skew_seconds=60.0, min_cache_seconds=10.0, clock=clock) + first = await source.get("s", _config()) + assert isinstance(first, Ok) and first.ok.access_token == "t1" + clock.t = 1004.0 # still within the token's real lifetime + within = await source.get("s", _config()) + assert isinstance(within, Ok) and within.ok.access_token == "t1" + clock.t = 1006.0 # past expires_at: the floor must not keep serving t1 + lapsed = await source.get("s", _config()) + assert isinstance(lapsed, Ok) and lapsed.ok.access_token == "t2" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +async def test_lock_dict_is_bounded_for_ephemeral_server_ids(): + poster = _FakePoster([_success()]) + source = ClientCredentialsTokenSource(poster, max_locks=8) + for index in range(20): + result = await source.get(f"ephemeral-{index}", _config()) + assert isinstance(result, Ok) + assert len(source._locks) <= 8 + + @pytest.mark.asyncio async def test_missing_expires_in_is_cached_briefly_not_an_hour(): clock = _Clock(1000.0) From 9c191e6764cbb137340b06109682d1122a0cd099 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 16 Jul 2026 13:28:02 -0700 Subject: [PATCH 31/60] chore(lint): add reasons to the client_credentials noqa suppressions --- .../mcp_server/outbound_credentials/client_credentials.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index 332305db3c6..e463d9fa1eb 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -98,10 +98,10 @@ async def post_client_credentials_grant( a 4xx/5xx), so the untyped boundary is contained here and every field the caller reads comes out of a validated ``TokenEndpointOutcome``. """ - from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 + from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 # defer heavy handler import to call time get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler is partially typed ) - from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 + from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 # deferred with the handler import try: client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) @@ -277,7 +277,7 @@ def _prepare_grant(config: ClientCredentialsConfig) -> Result[_PreparedGrant, Cr ) return Error(CredError.of_misconfigured(f"client_credentials config is missing: {missing}")) - from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( # noqa: PLC0415 # keep package v1-free at import time build_token_endpoint_client_auth, ) From 5b64239afca804b615f28aad2e7e9d2dc011a584 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 16 Jul 2026 16:54:10 -0700 Subject: [PATCH 32/60] fix(mcp): skip the M2M cache write when the token is already expired at mint An expires_in of zero or below computes a ttl of 0; the entry could never be served but still occupied a slot in the bounded backend, where it could evict a live token. The mint still serves the current request and the next get re-fetches under the per-server lock --- .../client_credentials.py | 3 +- .../test_client_credentials.py | 33 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index e463d9fa1eb..846017ebd73 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -260,7 +260,8 @@ class ClientCredentialsTokenSource: if expires_in is not None else self._default_ttl_seconds ) - await self._backend.set(grant.identity_key, server_id, token, ttl) + if ttl > 0: + await self._backend.set(grant.identity_key, server_id, token, ttl) return Ok(token) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py index db7e6208e76..bac6488333a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -152,6 +152,39 @@ async def test_short_lived_token_is_never_served_past_its_expiry(): assert len(poster.calls) == 2 +class _RecordingBackend: + """A TokenCacheBackend spy: records every write so a test can assert none happened.""" + + def __init__(self) -> None: + self.set_ttls: list[float] = [] + + async def get(self, identity_key: str, server_id: str): + return None + + async def set(self, identity_key: str, server_id: str, token, ttl_seconds: float) -> None: + self.set_ttls.append(ttl_seconds) + + async def delete(self, identity_key: str, server_id: str) -> None: + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("expires_in", [0, -30]) +async def test_non_positive_expires_in_writes_no_cache_entry(expires_in): + # A dead-on-arrival entry (ttl 0) must not be written at all: it can never be served, but it + # would occupy a slot in the bounded backend and could evict a live token. The mint itself + # still succeeds for the current request, and the next get re-fetches. + backend = _RecordingBackend() + poster = _FakePoster([_success("t1", expires_in=expires_in), _success("t2", expires_in=expires_in)]) + source = ClientCredentialsTokenSource(poster, backend=backend) + first = await source.get("s", _config()) + assert isinstance(first, Ok) and first.ok.access_token == "t1" + again = await source.get("s", _config()) + assert isinstance(again, Ok) and again.ok.access_token == "t2" + assert backend.set_ttls == [] + assert len(poster.calls) == 2 + + @pytest.mark.asyncio async def test_lock_dict_is_bounded_for_ephemeral_server_ids(): poster = _FakePoster([_success()]) From 8bfd8baab8cf00ed3f24900a1fe984ad25c8a076 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 16 Jul 2026 17:42:10 -0700 Subject: [PATCH 33/60] fix(mcp): remember the rotated M2M bearer for later requests in the session The auth object is the httpx client's auth for the whole MCP session; after a 401 recovery it kept sending the rejected token first, burning a 401 round trip and the single retry on every subsequent call --- .../client_credentials.py | 1 + .../test_client_credentials.py | 21 +++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index 846017ebd73..9be1121126a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -340,6 +340,7 @@ class ClientCredentialsBearerAuth(httpx.Auth): fresh = await self._refetch(token) if fresh is None: return + self._access_token = SecretStr(fresh) request.headers[self.header_name] = f"Bearer {fresh}" yield request diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py index bac6488333a..4e162090fbe 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -340,6 +340,27 @@ async def test_bearer_auth_retries_a_401_once_with_a_fresh_token(): assert seen == ["Bearer stale-token", "Bearer fresh-token"] +@pytest.mark.asyncio +async def test_bearer_auth_remembers_the_rotated_token_for_later_requests(): + # The auth object lives for the whole MCP session (it is the httpx client's auth), so after a + # 401 recovery it must send the fresh token first on subsequent requests; re-sending the + # rejected one would burn a 401 round trip and the single retry on every call. + transport, seen = _upstream([httpx.Response(401), httpx.Response(200), httpx.Response(200)]) + refetched: "list[str]" = [] + + async def refetch(failed: str) -> "str | None": + refetched.append(failed) + return "fresh-token" + + auth = ClientCredentialsBearerAuth("stale-token", refetch) + async with httpx.AsyncClient(transport=transport, auth=auth) as client: + first = await client.get("https://upstream.example.com/mcp") + second = await client.get("https://upstream.example.com/mcp") + assert first.status_code == 200 and second.status_code == 200 + assert refetched == ["stale-token"] + assert seen == ["Bearer stale-token", "Bearer fresh-token", "Bearer fresh-token"] + + @pytest.mark.asyncio async def test_bearer_auth_surfaces_the_401_when_the_refetch_fails(): transport, seen = _upstream([httpx.Response(401)]) From d2342296ae0108f00c3321d995af332490d6b7cf Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 16 Jul 2026 17:42:55 -0700 Subject: [PATCH 34/60] fix(mcp): keep the minted M2M bearer authoritative over injected Authorization headers _resolve_v2_auth dropped the resolved client_credentials auth when extra_headers already carried Authorization (MCPJWTSigner, static_headers), so the upstream got the injected header instead of the minted token and the one-shot 401 refetch was lost. M2M now joins token_exchange and authorization_code in the authoritative set; the conflicting header is dropped --- .../mcp_server/mcp_server_manager.py | 21 ++++++----- .../outbound_credentials/adapter.py | 5 +-- .../mcp_server/test_mcp_server_manager.py | 36 +++++++++++++++++++ 3 files changed, 50 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1ba608b9510..d06dad34d8c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -93,6 +93,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_ ) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AuthorizationCodeConfig, + ClientCredentialsConfig, CredError, IdJagConfig, PassthroughConfig, @@ -2739,14 +2740,18 @@ class MCPServerManager: ) if not conflicts: return auth, extra_headers - if isinstance(spec.config, (TokenExchangeConfig, AuthorizationCodeConfig, IdJagConfig)): - # The resolver owns the per-user credential here (token_exchange's exchanged - # token, authorization_code's stored token, id_jag's minted assertion). It is - # authoritative: a guardrail such - # as MCPJWTSigner, static_headers, or any other injected Authorization must NOT - # shadow it (otherwise the upstream gets e.g. the signer's JWT instead of the - # exchanged token and rejects it). Drop the conflicting header so the resolved - # token reaches upstream. + if isinstance( + spec.config, + (TokenExchangeConfig, AuthorizationCodeConfig, IdJagConfig, ClientCredentialsConfig), + ): + # The resolver owns the credential here (token_exchange's exchanged token, + # authorization_code's stored token, id_jag's minted assertion, + # client_credentials' gateway-minted M2M token). It is authoritative: a + # guardrail such as MCPJWTSigner, static_headers, or any other injected + # Authorization must NOT shadow it (otherwise the upstream gets e.g. the + # signer's JWT instead of the minted token and rejects it, and for M2M the + # one-shot 401 refetch is lost with it). Drop the conflicting header so the + # resolved token reaches upstream. return auth, _without_authorization(extra_headers) # Other modes: an Authorization already supplied via extra_headers (a forwarded caller # header or static_headers) is intentional and wins; v1 applies those last. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 3bf9fd1c12e..8ecef0c95c5 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -21,12 +21,9 @@ from typing_extensions import assert_never from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, -<<<<<<< HEAD ClientAuth, - ClientSecretAuth, -======= ClientCredentialsConfig, ->>>>>>> 73df37ca23 (feat(mcp): migrate client_credentials (M2M) onto the v2 resolver arm) + ClientSecretAuth, CredError, IdJagConfig, NoneConfig, 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 f0ca75b8cdb..03f91260955 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 @@ -1759,6 +1759,42 @@ class TestMCPServerManager: assert client._resolved_auth is not None assert "authorization" not in {k.lower() for k in (client.extra_headers or {})} + @pytest.mark.asyncio + async def test_injected_authorization_does_not_shadow_m2m_minted_token(self): + """The M2M twin of the OBO shadow test: a guardrail/static Authorization must not displace + the gateway-minted client_credentials bearer. Dropping the resolved auth here would also + drop the one-shot 401 refetch that rides on it, so the resolver-owned credential is + authoritative exactly as for token_exchange and authorization_code.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Ok(StaticHeaderAuth("Bearer MINTED-M2M", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = MCPServer( + server_id="m2m-shadow", + name="m2m-shadow-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + client_id="cid", + client_secret="csec", + token_url="https://idp.example.com/token", + ) + + client = await manager._create_mcp_client( + server, + extra_headers={"Authorization": "Bearer signer-jwt"}, # simulate the JWT signer + ) + + assert client._resolved_auth is not None + assert "authorization" not in {k.lower() for k in (client.extra_headers or {})} + @pytest.mark.asyncio async def test_preflight_token_exchange_challenges_on_rejected_subject(self): """A subject the IdP rejects must raise the RFC 9728 401 challenge from the preflight, so a From 51df80115994ad7daaa7a899404ac0e9eeb3af9c Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 12:23:22 -0700 Subject: [PATCH 35/60] test(e2e): cover key regeneration rotating to a working new key (#34000) --- tests/e2e/management/management_client.py | 12 ++++++++++ tests/e2e/management/test_management_e2e.py | 26 +++++++++++++++++++++ tests/e2e/models.py | 4 ++++ 3 files changed, 42 insertions(+) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index e967fb7b504..6ce6405c2be 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -16,8 +16,10 @@ from models import ( ChatMessage, KeyDeleteBody, KeyGenerateBody, + KeyGenerateResponse, KeyListParams, KeyListResponse, + KeyRegenerateBody, KeyUpdateBody, OrgDeleteBody, OrgInfoParams, @@ -89,6 +91,16 @@ class ManagementClient: ) ) + def regenerate_key(self, key: str) -> str: + return unwrap( + self.proxy.transport.post( + "/key/regenerate", + headers=self.proxy.transport.master, + json=KeyRegenerateBody(key=key), + response_type=KeyGenerateResponse, + ) + ).key + def key_alias_count(self, key_alias: str) -> int: return unwrap( self.proxy.transport.get( diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index adbf3e8b065..de34ee69ded 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -169,6 +169,32 @@ class TestKeyRoutes: _ = _poll(client, rejected, "deleted key was still accepted on chat (never rejected 401) at the deadline") +class TestKeyRegeneration: + @pytest.mark.covers("mgmt.key.regenerate.happy_path") + def test_regenerate_rotates_to_a_working_new_key( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + old_key = _generate_key(client, resources, KeyGenerateBody(models=["gpt-5.5"])) + + new_key = client.regenerate_key(old_key) + resources.defer(lambda: client.proxy.delete_key(new_key)) + assert new_key != old_key, "regenerate returned the same key string, so no rotation happened" + + def new_accepted() -> bool | None: + outcome = client.chat_status(new_key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.status_code != 401 else None + + _ = _poll(client, new_accepted, "regenerated key was never accepted at auth (still 401) at the deadline") + + def old_rejected() -> bool | None: + outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.status_code == 401 else None + + _ = _poll( + client, old_rejected, "old key was still accepted after regeneration (never rejected 401) at the deadline" + ) + + class TestTeamRoutes: @pytest.mark.covers("mgmt.team.new.persists") def test_new_persists_to_team_info_and_binds_keys( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 2e7bfe41e30..bdcfd080b8e 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -73,6 +73,10 @@ class KeyGenerateResponse(BaseModel): key: str +class KeyRegenerateBody(BaseModel): + key: str + + class KeyDeleteBody(BaseModel): keys: list[str] From 72be5a9bc0063358b97a09414dc96743d2d2f598 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 12:24:50 -0700 Subject: [PATCH 36/60] test(e2e): cover tag creation persisting for spend categorization (#34018) --- tests/e2e/management/management_client.py | 34 +++++++++++++++++++++ tests/e2e/management/test_management_e2e.py | 24 ++++++++++++++- tests/e2e/models.py | 23 ++++++++++++++ 3 files changed, 80 insertions(+), 1 deletion(-) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 6ce6405c2be..438a4db098c 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -26,6 +26,10 @@ from models import ( OrgInfoResponse, OrgNewBody, OrgNewResponse, + TagDeleteBody, + TagListEntry, + TagListResponse, + TagNewBody, TeamData, TeamDeleteBody, TeamInfoParams, @@ -259,6 +263,36 @@ class ManagementClient: ) ) + def create_tag(self, body: TagNewBody) -> None: + _ = unwrap( + self.proxy.transport.post( + "/tag/new", + headers=self.proxy.transport.master, + json=body, + response_type=NoBody, + ) + ) + + def delete_tag(self, name: str) -> None: + _ = self.proxy.transport.post( + "/tag/delete", + headers=self.proxy.transport.master, + json=TagDeleteBody(name=name), + response_type=NoBody, + ) + + def tag_list(self) -> tuple[TagListEntry, ...]: + return tuple( + unwrap( + self.proxy.transport.get( + "/tag/list", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=TagListResponse, + ) + ).root + ) + def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: return self.proxy.transport.send( "/chat/completions", diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index de34ee69ded..51817b3916f 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -22,7 +22,7 @@ from management_client import ( ROUTE_NOT_ALLOWED_MARKER, ManagementClient, ) -from models import KeyGenerateBody, OrgNewBody, TeamNewBody, UserNewBody +from models import KeyGenerateBody, OrgNewBody, TagListEntry, TagNewBody, TeamNewBody, UserNewBody pytestmark = pytest.mark.e2e @@ -271,6 +271,28 @@ class TestOrganizationRoutes: ) +class TestTagRoutes: + @pytest.mark.covers("mgmt.tag.new.happy_path") + def test_new_persists_to_tag_list(self, client: ManagementClient, resources: ResourceManager) -> None: + name = f"e2e-mgmt-tag-{unique_marker()}" + description = "Tag for spend categorization" + + assert all(entry.name != name for entry in client.tag_list()), ( + f"tag {name!r} was already listed by /tag/list before /tag/new created it" + ) + + client.create_tag(TagNewBody(name=name, description=description)) + resources.defer(lambda: client.delete_tag(name)) + + def listed() -> TagListEntry | None: + return next((entry for entry in client.tag_list() if entry.name == name), None) + + entry = _poll(client, listed, f"/tag/list never listed {name!r} after /tag/new") + assert entry.description == description, ( + f"/tag/list reports description {entry.description!r} for {name!r}, configured {description!r}" + ) + + def _assert_route_forbidden(route: str, outcome: StreamingResponse) -> None: assert outcome.status_code == 403, ( f"llm-only key POSTing {route} must be denied exactly 403, got {outcome.status_code}: {outcome.body[:300]}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index bdcfd080b8e..b7a95f9714c 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -699,3 +699,26 @@ class OrgInfoResponse(BaseModel): class OrgDeleteBody(BaseModel): organization_ids: list[str] + + +# ---------- tags (management) ---------- + + +class TagNewBody(BaseModel): + name: str + description: str | None = None + + +class TagDeleteBody(BaseModel): + name: str + + +class TagListEntry(BaseModel): + name: str + description: str | None = None + + +class TagListResponse(RootModel[list[TagListEntry]]): + """GET /tag/list answers with a bare array of tag configs (the stored tags plus + any dynamically-seen spend tags), not an object wrapping them. Read the rows off + .root.""" From bf04ba8d3e1d35055fabd4f961a62e321eb8f6ec Mon Sep 17 00:00:00 2001 From: Tin Date: Mon, 20 Jul 2026 12:29:38 -0700 Subject: [PATCH 37/60] refactor(mcp): extract the oauth2 spec dispatch to keep to_server_spec under the complexity ceiling The ID-JAG landing brought to_server_spec to the C901 boundary and the M2M branch pushed it one over; the oauth2 sub-mode dispatch now lives in its own _oauth2_spec helper, mirroring the file's per-mode spec builders --- .../outbound_credentials/adapter.py | 29 ++++++++++++------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 8ecef0c95c5..565c489e77c 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -96,16 +96,7 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: case MCPAuth.basic: return _shared_key_spec(server, resource, "Authorization", "Basic", encode=True) case MCPAuth.oauth2: - if server.has_client_credentials: - return _client_credentials_spec(server, resource) - if server.needs_user_oauth_token and not server.delegate_auth_to_upstream: - return ServerSpec( - server_id=server.server_id, - resource=resource, - config=AuthorizationCodeConfig(), - ) - # delegate/passthrough oauth2 stay on v1 - return None + return _oauth2_spec(server, resource) case MCPAuth.oauth2_id_jag: return _id_jag_spec(server, resource) case MCPAuth.true_passthrough | MCPAuth.oauth_delegate: @@ -117,6 +108,24 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: assert_never(auth_type) +def _oauth2_spec(server: MCPServer, resource: str) -> ServerSpec | None: + """Dispatch the oauth2 auth_type across its sub-modes: M2M, gateway-managed interactive, or v1. + + ``client_credentials`` (the explicit ``oauth2_flow`` opt-in) builds the M2M spec, per-user + ``authorization_code`` without upstream delegation builds the interactive spec, and the + delegate/passthrough shapes defer to v1 (None). + """ + if server.has_client_credentials: + return _client_credentials_spec(server, resource) + if server.needs_user_oauth_token and not server.delegate_auth_to_upstream: + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=AuthorizationCodeConfig(), + ) + return None + + def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec: """Build a client_credentials (M2M) spec; the explicit ``oauth2_flow`` opt-in owns the server. From 214945a223837c721cd8c1b15eaa812636fda221 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 12:36:59 -0700 Subject: [PATCH 38/60] test(e2e): cover organization deletion removing it from /organization/info (#34009) Co-authored-by: mubashir1osmani --- tests/e2e/management/management_client.py | 2 ++ tests/e2e/management/test_management_e2e.py | 22 +++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 438a4db098c..5b94956fa03 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -263,6 +263,8 @@ class ManagementClient: ) ) + def org_info_status(self, organization_id: str) -> ProbeResult: + return self.proxy.transport.probe("/organization/info", params=OrgInfoParams(organization_id=organization_id)) def create_tag(self, body: TagNewBody) -> None: _ = unwrap( self.proxy.transport.post( diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 51817b3916f..748f7192fb0 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -270,6 +270,28 @@ class TestOrganizationRoutes: f"/organization/info reports models {info.models}, configured ['gemini-2.5-flash']" ) + @pytest.mark.covers("mgmt.organization.delete.persists") + def test_delete_removes_from_organization_info( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + """The teardown's deferred delete fires again on the already-deleted org by + design: the deferred cleanup must survive this test failing before the + in-body delete, and a repeat /organization/delete is a warn-only no-op the + teardown absorbs.""" + org_id = client.create_org(OrgNewBody(organization_alias=f"e2e-mgmt-org-{unique_marker()}")) + resources.defer(lambda: client.delete_org(org_id)) + + assert client.org_info_status(org_id).status_code == 200, ( + f"/organization/info did not resolve org {org_id} before deletion" + ) + + client.delete_org(org_id) + + def gone() -> bool | None: + return True if client.org_info_status(org_id).status_code == 404 else None + + _ = _poll(client, gone, f"org {org_id} still resolved on /organization/info after /organization/delete") + class TestTagRoutes: @pytest.mark.covers("mgmt.tag.new.happy_path") From 0b8817afbbf0abd7dda5bf968e34cced300350db Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:09:41 -0700 Subject: [PATCH 39/60] perf(bedrock): audio transcription via rust core (py->rust bridge) (#33990) * feat(bedrock): add audio transcription via Converse with py->rust bridge Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(ci): exclude rust transcription rollout flag from docs check Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(bedrock): await rust/python fallback in async transcription dispatch Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(bedrock): cover audio transcription rust dispatch Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(bedrock): route audio transcription through rust Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(bedrock): move rust transcription dispatch out of main Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(bedrock): include rust transcription coverage shard Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm-rust/CLAUDE.md | 7 + .../PROVIDER_CODING_STANDARDS.md | 10 +- litellm-rust/crates/ai-gateway/Cargo.toml | 2 +- .../src/audio_transcription/common_utils.rs | 48 +++ .../src/audio_transcription/handler.rs | 89 +++++ .../src/audio_transcription/hooks.rs | 300 +++++++++++++++++ .../ai-gateway/src/audio_transcription/mod.rs | 25 ++ .../src/audio_transcription/prepare.rs | 55 ++++ .../src/audio_transcription/tests.rs | 53 +++ .../src/audio_transcription/types.rs | 58 ++++ .../crates/ai-gateway/src/{ocr => }/client.rs | 6 +- .../ai-gateway/src/io/audio_transcription.rs | 1 + litellm-rust/crates/ai-gateway/src/io/mod.rs | 1 + litellm-rust/crates/ai-gateway/src/lib.rs | 2 + .../crates/ai-gateway/src/ocr/common_utils.rs | 2 +- .../crates/ai-gateway/src/ocr/handler.rs | 2 +- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 1 - .../core/src/audio_transcription/mod.rs | 2 + .../src/audio_transcription/transformation.rs | 57 ++++ .../core/src/audio_transcription/types.rs | 20 ++ litellm-rust/crates/core/src/lib.rs | 1 + .../providers/bedrock/audio_transcription.rs | 310 ++++++++++++++++++ .../core/src/providers/bedrock/aws_base.rs | 6 +- .../core/src/providers/bedrock/constants.rs | 4 + .../crates/core/src/providers/bedrock/mod.rs | 2 + litellm-rust/crates/python-bridge/CLAUDE.md | 8 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 +- litellm-rust/crates/python-bridge/src/lib.rs | 92 ++++++ .../bedrock/audio_transcription/__init__.py | 84 +++++ litellm/main.py | 55 +++- litellm/rust_bridge/ocr.py | 11 + litellm/rust_bridge/transcription.py | 148 +++++++++ .../test_audio_transcription_rust_bridge.py | 151 +++++++++ 33 files changed, 1586 insertions(+), 29 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs rename litellm-rust/crates/ai-gateway/src/{ocr => }/client.rs (60%) create mode 100644 litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs create mode 100644 litellm-rust/crates/core/src/audio_transcription/mod.rs create mode 100644 litellm-rust/crates/core/src/audio_transcription/transformation.rs create mode 100644 litellm-rust/crates/core/src/audio_transcription/types.rs create mode 100644 litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs create mode 100644 litellm/llms/bedrock/audio_transcription/__init__.py create mode 100644 litellm/rust_bridge/transcription.py create mode 100644 tests/test_litellm/test_audio_transcription_rust_bridge.py diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index 519b1d205ef..0659e63df39 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -62,6 +62,13 @@ Not allowed in `core`: Python owns rollout state and fallback while Rust is being introduced. Rust paths must be off by default until parity tests prove equivalence with Python. +A new provider/route may instead be implemented rust-only with no Python +reference; then the Python interface is a thin dispatch that calls Rust with no +fallback, and you state the rust-only choice explicitly in the PR. Either way +the Python side stays minimal (it only marshals inputs and calls the Rust +interface), never add a per-route feature flag, and never push provider +dispatch into `litellm/main.py`; put it in a thin dispatch class under +`litellm/llms///`. ## Production Bar diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index c7980b11147..ed44dc4c729 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -39,11 +39,17 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MIST 19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity. 20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping. -21. Rust paths stay off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. +21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR. + +## Python bridge (SDK side) + +22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust. +23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms///` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method. +24. Do not add new feature flags unless explicitly requested. Reuse the existing litellm rust rollout mechanism (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_`. ## Checks before push -22. Run, and keep green: +25. Run, and keep green: ```bash cd litellm-rust cargo fmt --check diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index c15af4cc478..541beabe170 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -14,7 +14,7 @@ path = "src/main.rs" required-features = ["server"] [dependencies] -litellm-core.workspace = true +litellm-core = { workspace = true, features = ["bedrock-auth"] } # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. reqwest.workspace = true diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs new file mode 100644 index 00000000000..270d5c2d97a --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs @@ -0,0 +1,48 @@ +use std::collections::BTreeMap; + +use litellm_core::CoreResult; +use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig; +use litellm_core::error::CoreError; +use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; +use serde_json::{Map, Value}; + +pub(super) fn audio_transcription_provider_config( + provider: &str, +) -> Option<&'static dyn AudioTranscriptionProviderConfig> { + match provider { + "bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG), + _ => None, + } +} + +pub(super) fn string_headers( + headers: Option>, +) -> CoreResult> { + headers + .unwrap_or_default() + .into_iter() + .map(|(key, value)| { + value + .as_str() + .map(|value| (key.clone(), value.to_string())) + .ok_or_else(|| { + CoreError::InvalidRequest(format!( + "audio transcription extra_headers.{key} must be a string" + )) + }) + }) + .collect() +} + +pub(super) fn has_header(headers: &BTreeMap, name: &str) -> bool { + headers.keys().any(|key| key.eq_ignore_ascii_case(name)) +} + +pub(super) fn truncate_error_body(body: &str) -> String { + let truncated: String = body.chars().take(256).collect(); + if truncated.chars().count() == body.chars().count() { + truncated + } else { + format!("{truncated}... (truncated)") + } +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs new file mode 100644 index 00000000000..33c13550f58 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs @@ -0,0 +1,89 @@ +use std::time::SystemTime; + +use litellm_core::CoreResult; +use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; +use litellm_core::error::CoreError; +use litellm_core::providers::bedrock::audio_transcription::aws_auth_config; +use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; +use serde_json::Value; + +use super::common_utils::truncate_error_body; +use super::types::ProviderAudioTranscriptionRequest; +use crate::client::http_client; + +pub(crate) async fn execute_audio_transcription_provider_call( + request: ProviderAudioTranscriptionRequest, +) -> CoreResult { + let body = serde_json::to_vec(&request.body).map_err(|error| { + CoreError::InvalidRequest(format!("invalid audio request body: {error}")) + })?; + let mut request_builder = http_client().post(&request.url).body(body.clone()); + for (key, value) in &request.upstream_headers { + request_builder = request_builder.header(key, value); + } + if let Some(duration) = request.timeout { + request_builder = request_builder.timeout(duration); + } + let response = request_builder + .send() + .await + .map_err(|error| CoreError::Network(error.to_string()))?; + let status = response.status(); + let text = response + .text() + .await + .map_err(|error| CoreError::Network(error.to_string()))?; + if !status.is_success() { + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + let response_json: Value = serde_json::from_str(&text).map_err(|error| { + CoreError::InvalidResponse(format!("invalid audio response JSON: {error}")) + })?; + Ok(request + .config + .transform_transcription_response(&request.model, response_json)? + .into_json()) +} + +pub(crate) async fn sign_request( + request: &ProviderAudioTranscriptionRequest, + optional_params: &serde_json::Map, +) -> CoreResult { + let env_lookup = environment_lookup; + let auth = request + .config + .auth_strategy(&request.model, optional_params, &env_lookup)?; + let body = serde_json::to_vec(&request.body).map_err(|error| { + CoreError::InvalidRequest(format!("invalid audio request body: {error}")) + })?; + let mut headers = super::common_utils::string_headers(None)?; + headers.insert("Content-Type".to_string(), "application/json".to_string()); + headers.extend(request.upstream_headers.iter().cloned()); + match auth { + AudioTranscriptionAuth::Bearer => {} + AudioTranscriptionAuth::AwsSigV4 { region, .. } => { + let credentials = + resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup) + .await?; + headers.extend(sign_bedrock_post( + &request.url, + &body, + &headers, + ®ion, + &credentials, + SystemTime::now(), + )?); + } + } + Ok(ProviderAudioTranscriptionRequest { + upstream_headers: headers.into_iter().collect(), + ..request.clone() + }) +} + +pub(super) fn environment_lookup(key: &str) -> Option { + std::env::var(key).ok() +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs new file mode 100644 index 00000000000..8b6896f3846 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -0,0 +1,300 @@ +use std::future::Future; +use std::pin::Pin; + +use litellm_core::CoreResult; +use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; +use litellm_core::error::CoreError; +use serde_json::{Map, Value, json}; + +use super::common_utils::{audio_transcription_provider_config, has_header, string_headers}; +use super::handler::sign_request; +use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +use crate::integrations::custom_guardrail::{ + CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, +}; +use crate::integrations::custom_logger::{ + CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; +use crate::integrations::types::{ + RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, +}; + +pub(crate) struct AudioTranscriptionLifecycleHooks { + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, +} + +type AudioFuture<'a, T> = Pin> + Send + 'a>>; +type AudioLogFuture<'a> = Pin + Send + 'a>>; + +impl AudioTranscriptionLifecycleHooks { + pub(crate) fn new( + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, + ) -> Self { + Self { + logger_runner, + guardrail_runner, + request_metadata, + } + } + + async fn run_pre_call_guardrails( + &self, + request: PreparedAudioTranscriptionRequest, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(request); + } + let (guardrail_request, _) = self + .guardrail_runner + .run_pre_call( + &guardrail_context(&self.request_metadata), + GuardrailRequest::new(json!({ + "model": request.model, + "custom_llm_provider": request.custom_llm_provider, + "audio": request.audio, + "optional_params": request.optional_params, + })), + ) + .await + .map_err(guardrail_error_to_core_error)?; + let Value::Object(mut data) = guardrail_request.data else { + return Err(CoreError::InvalidRequest( + "audio transcription pre_call guardrail must return an object".to_string(), + )); + }; + let audio = data.remove("audio").ok_or_else(|| { + CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string()) + })?; + let optional_params = match data.remove("optional_params") { + Some(Value::Object(value)) => value, + Some(_) => { + return Err(CoreError::InvalidRequest( + "audio transcription optional_params must be an object".to_string(), + )); + } + None => Map::new(), + }; + Ok(PreparedAudioTranscriptionRequest { + audio, + optional_params, + ..request + }) + } + + async fn prepare_provider_request( + &self, + request: PreparedAudioTranscriptionRequest, + ) -> CoreResult { + let config = audio_transcription_provider_config(&request.custom_llm_provider) + .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + let env_lookup = super::handler::environment_lookup; + let headers = string_headers(request.extra_headers)?; + let url = config.complete_url( + request.api_base.as_deref(), + &request.model, + &request.optional_params, + &env_lookup, + )?; + let filtered_params = config.map_transcription_params(&request.optional_params); + let body = config.transform_transcription_request( + &request.model, + request.audio, + filtered_params, + )?; + let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?; + let mut upstream_headers = headers.into_iter().collect::>(); + if matches!(auth, AudioTranscriptionAuth::Bearer) + && !has_header( + &upstream_headers + .iter() + .cloned() + .collect::>(), + "authorization", + ) + && let Some(api_key) = request.api_key.as_deref() + { + upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}"))); + } + let provider_request = ProviderAudioTranscriptionRequest { + model: request.model, + config, + url, + body: body.body, + upstream_headers, + timeout: request.timeout, + }; + let provider_request = self.run_during_call_guardrails(provider_request).await?; + sign_request(&provider_request, &request.optional_params).await + } + + async fn run_during_call_guardrails( + &self, + request: ProviderAudioTranscriptionRequest, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(request); + } + let (guardrail_request, _) = self + .guardrail_runner + .run_during_call( + &guardrail_context(&self.request_metadata), + GuardrailRequest::new(json!({ + "model": request.model, + "custom_llm_provider": "bedrock", + "url": request.url, + "body": request.body, + })), + ) + .await + .map_err(guardrail_error_to_core_error)?; + let Value::Object(mut data) = guardrail_request.data else { + return Err(CoreError::InvalidRequest( + "audio transcription during_call guardrail must return an object".to_string(), + )); + }; + let body = data.remove("body").ok_or_else(|| { + CoreError::InvalidRequest("audio transcription guardrail removed body".to_string()) + })?; + Ok(ProviderAudioTranscriptionRequest { body, ..request }) + } + + fn logging_payload( + &self, + context: &CallLifecycleContext, + timing: &CallLifecycleTiming, + ) -> StandardLoggingPayload { + StandardLoggingPayload { + id: context.litellm_call_id.clone(), + litellm_call_id: context.litellm_call_id.clone(), + call_type: context.call_type.clone(), + model: context.model.clone(), + custom_llm_provider: context.custom_llm_provider.clone(), + response_cost: 0.0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + start_time: timing.start_time, + end_time: timing.end_time, + stream: false, + metadata: StandardLoggingMetadata { + user_api_key_hash: self.request_metadata.user_api_key_hash.clone(), + user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(), + user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(), + ..Default::default() + }, + messages: None, + } + } +} + +impl CallLifecycleHooks + for AudioTranscriptionLifecycleHooks +{ + type PreCallFuture<'a> = AudioFuture<'a, PreparedAudioTranscriptionRequest>; + type DuringCallFuture<'a> = AudioFuture<'a, ProviderAudioTranscriptionRequest>; + type SuccessFuture<'a> = AudioLogFuture<'a>; + type FailureFuture<'a> = AudioLogFuture<'a>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedAudioTranscriptionRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { self.run_pre_call_guardrails(request).await }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedAudioTranscriptionRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { self.prepare_provider_request(request).await }) + } + + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a Value, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + self.logger_runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload( + self.logging_payload(context, timing), + ), + &CallbackValue::new("audio_transcription", response.clone()), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } + + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a CoreError, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + let logging_error = LoggingError { + message: error.to_string(), + kind: core_error_kind(error).to_string(), + }; + self.logger_runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload( + self.logging_payload(context, timing), + ) + .with_failure_error(logging_error.clone()), + Some(&CallbackValue::new( + "error", + json!({"message": logging_error.message, "kind": logging_error.kind}), + )), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } +} + +fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { + GuardrailContext { + call_type: CallType::Other("audio_transcription".to_string()), + selected_guardrails: Vec::new(), + metadata: std::collections::HashMap::new(), + user_api_key_hash: metadata.user_api_key_hash.clone(), + user_api_key_user_id: metadata.user_api_key_user_id.clone(), + user_api_key_team_id: metadata.user_api_key_team_id.clone(), + trace_parent: None, + } +} + +fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { + CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +} + +fn core_error_kind(error: &CoreError) -> &'static str { + match error { + CoreError::Auth(_) => "AuthError", + CoreError::InvalidProvider(_) => "InvalidProvider", + CoreError::InvalidRequest(_) => "InvalidRequest", + CoreError::InvalidType { .. } => "InvalidType", + CoreError::MissingField(_) => "MissingField", + CoreError::Http { .. } => "HttpError", + CoreError::InvalidResponse(_) => "InvalidResponse", + CoreError::Network(_) => "NetworkError", + CoreError::Routing(_) => "RoutingError", + } +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs new file mode 100644 index 00000000000..5d33d912c40 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs @@ -0,0 +1,25 @@ +use litellm_core::CoreResult; +use litellm_core::call_lifecycle::CallLifecycle; +use serde_json::Value; + +mod common_utils; +mod handler; +mod hooks; +mod prepare; +mod types; + +pub use types::AudioTranscriptionRequest; + +use handler::execute_audio_transcription_provider_call; +use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call}; + +pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> CoreResult { + let PreparedAudioTranscriptionCall { request, hooks } = + prepare_audio_transcription_call(request); + CallLifecycle::default() + .run_request(request, &hooks, execute_audio_transcription_provider_call) + .await +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs new file mode 100644 index 00000000000..a475d58635f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs @@ -0,0 +1,55 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; + +use super::hooks::AudioTranscriptionLifecycleHooks; +use super::types::{AudioTranscriptionRequest, PreparedAudioTranscriptionRequest}; +use crate::integrations::custom_guardrail::CustomGuardrailRunner; +use crate::integrations::custom_logger::CustomLoggerRunner; + +pub(crate) struct PreparedAudioTranscriptionCall { + pub(crate) request: PreparedAudioTranscriptionRequest, + pub(crate) hooks: AudioTranscriptionLifecycleHooks, +} + +pub(crate) fn prepare_audio_transcription_call( + request: AudioTranscriptionRequest<'_>, +) -> PreparedAudioTranscriptionCall { + let call_id = request + .litellm_call_id + .map(str::to_string) + .unwrap_or_else(new_audio_transcription_call_id); + let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) + .unwrap_or(CustomLlmProvider { + model: request.model, + custom_llm_provider: "bedrock", + }); + PreparedAudioTranscriptionCall { + request: PreparedAudioTranscriptionRequest { + model: provider_info.model.to_string(), + custom_llm_provider: provider_info.custom_llm_provider.to_string(), + litellm_call_id: call_id, + audio: request.audio, + api_key: request.api_key.map(str::to_string), + api_base: request.api_base.map(str::to_string), + extra_headers: request.extra_headers, + optional_params: request.optional_params, + timeout: request.timeout, + }, + hooks: AudioTranscriptionLifecycleHooks::new( + CustomLoggerRunner::new(request.callbacks), + CustomGuardrailRunner::new(request.guardrails), + request.request_metadata, + ), + } +} + +fn new_audio_transcription_call_id() -> String { + static COUNTER: AtomicU64 = AtomicU64::new(1); + let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_nanos()); + format!("audio-transcription-{timestamp}-{sequence}") +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs new file mode 100644 index 00000000000..5df04708b7d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs @@ -0,0 +1,53 @@ +use std::io::{Read, Write}; +use std::net::TcpListener; +use std::thread; + +use serde_json::{Map, json}; + +use super::{AudioTranscriptionRequest, audio_transcription}; + +#[tokio::test] +async fn bedrock_request_is_signed_and_contains_audio() { + let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); + let address = listener.local_addr().expect("address"); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("connection"); + let mut request = Vec::new(); + let mut buffer = [0_u8; 16_384]; + let count = stream.read(&mut buffer).expect("request"); + request.extend_from_slice(&buffer[..count]); + let request = String::from_utf8_lossy(&request); + assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse")); + assert!(request.contains("authorization: AWS4-HMAC-SHA256")); + assert!(request.contains("x-amz-date:")); + assert!(request.contains("\"bytes\":\"AQI=\"")); + assert!(request.contains("Transcribe the audio. Respond with only the transcript.")); + let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}"; + stream.write_all(response).expect("response"); + }); + + let optional_params = Map::from_iter([ + ("aws_access_key_id".to_string(), json!("access-key")), + ("aws_secret_access_key".to_string(), json!("secret-key")), + ("aws_region_name".to_string(), json!("us-east-1")), + ]); + let api_base = format!("http://{address}"); + let response = audio_transcription(AudioTranscriptionRequest { + model: "mistral.voxtral-mini-3b-2507", + audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}), + api_key: None, + api_base: Some(&api_base), + custom_llm_provider: Some("bedrock"), + extra_headers: None, + optional_params, + timeout: None, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await + .expect("transcription"); + assert_eq!(response, json!({"text": "hello"})); + server.join().expect("server"); +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs new file mode 100644 index 00000000000..9697aa98b0a --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs @@ -0,0 +1,58 @@ +use std::sync::Arc; +use std::time::Duration; + +use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig; +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; +use serde_json::{Map, Value}; + +use crate::integrations::custom_guardrail::CustomGuardrail; +use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::types::RequestMetadata; + +pub struct AudioTranscriptionRequest<'a> { + pub model: &'a str, + pub audio: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: Option<&'a str>, + pub extra_headers: Option>, + pub optional_params: Map, + pub timeout: Option, + pub callbacks: Vec>, + pub guardrails: Vec>, + pub request_metadata: RequestMetadata, + pub litellm_call_id: Option<&'a str>, +} + +pub(crate) struct PreparedAudioTranscriptionRequest { + pub(crate) model: String, + pub(crate) custom_llm_provider: String, + pub(crate) litellm_call_id: String, + pub(crate) audio: Value, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) extra_headers: Option>, + pub(crate) optional_params: Map, + pub(crate) timeout: Option, +} + +impl CallLifecycleRequest for PreparedAudioTranscriptionRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "audio_transcription", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} + +#[derive(Clone)] +pub(crate) struct ProviderAudioTranscriptionRequest { + pub(crate) model: String, + pub(crate) config: &'static dyn AudioTranscriptionProviderConfig, + pub(crate) url: String, + pub(crate) body: Value, + pub(crate) upstream_headers: Vec<(String, String)>, + pub(crate) timeout: Option, +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/client.rs b/litellm-rust/crates/ai-gateway/src/client.rs similarity index 60% rename from litellm-rust/crates/ai-gateway/src/ocr/client.rs rename to litellm-rust/crates/ai-gateway/src/client.rs index 79cc7816227..ff2606f0229 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/client.rs +++ b/litellm-rust/crates/ai-gateway/src/client.rs @@ -1,13 +1,13 @@ use std::sync::OnceLock; use std::time::Duration; -const OCR_TIMEOUT_SECS: u64 = 600; +const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600; -pub(super) fn http_client() -> &'static reqwest::Client { +pub(crate) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() - .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) + .timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS)) .build() .expect("failed to build reqwest client") }) diff --git a/litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs b/litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs new file mode 100644 index 00000000000..80d9e401a5f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs @@ -0,0 +1 @@ +pub use crate::audio_transcription::{AudioTranscriptionRequest, audio_transcription}; diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index 9cbfa568121..6129a808965 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -1,3 +1,4 @@ +pub mod audio_transcription; pub mod messages; pub mod ocr; pub mod realtime; diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index 25aac3c495b..c44d661c29e 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -11,6 +11,8 @@ //! binary turns on. The `python-config` feature additionally pulls in [`python`] //! for the load-time config reader. +pub mod audio_transcription; +mod client; pub mod io; pub mod messages; pub mod ocr; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index 7d164a80137..9bc2818b6e7 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -18,7 +18,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{ VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, }; -use super::client::http_client; +use crate::client::http_client; const ERROR_BODY_MAX_CHARS: usize = 256; const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 381d22e9cea..1de34eb400e 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -3,9 +3,9 @@ use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::Value; -use super::client::http_client; use super::common_utils::{poll_document_intelligence, truncate_error_body}; use super::types::ProviderOcrRequest; +use crate::client::http_client; pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { let mut request_builder = http_client().post(&request.url).json(&request.body); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index ad346bc0c64..c4c13e2300c 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -2,7 +2,6 @@ use litellm_core::CoreResult; use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; -mod client; mod common_utils; mod handler; mod hooks; diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs new file mode 100644 index 00000000000..ec2fbb969a6 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -0,0 +1,2 @@ +pub mod transformation; +pub mod types; diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs new file mode 100644 index 00000000000..eab34c13843 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -0,0 +1,57 @@ +use serde_json::{Map, Value}; + +use crate::CoreResult; + +use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum AudioTranscriptionAuth { + Bearer, + AwsSigV4 { + region: String, + service: &'static str, + }, +} + +pub trait AudioTranscriptionProviderConfig: Sync { + fn supported_transcription_params(&self) -> &'static [&'static str]; + + fn map_transcription_params(&self, params: &Map) -> Map { + params + .iter() + .filter(|(key, _)| { + self.supported_transcription_params() + .contains(&key.as_str()) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() + } + + fn transform_transcription_request( + &self, + model: &str, + audio: Value, + optional_params: Map, + ) -> CoreResult; + + fn transform_transcription_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult; + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn auth_strategy( + &self, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; +} diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs new file mode 100644 index 00000000000..3a9e1ecd88c --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -0,0 +1,20 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct AudioTranscriptionRequestData { + pub body: Value, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct AudioTranscriptionResponseData { + pub text: String, +} + +impl AudioTranscriptionResponseData { + pub fn into_json(self) -> Value { + serde_json::json!({ + "text": self.text, + }) + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 3989fb441bc..51ea19750ea 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,3 +1,4 @@ +pub mod audio_transcription; pub mod caching; pub mod call_lifecycle; pub mod constants; diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs new file mode 100644 index 00000000000..86eb589e2c0 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -0,0 +1,310 @@ +use serde_json::{Map, Value, json}; + +use crate::audio_transcription::transformation::{ + AudioTranscriptionAuth, AudioTranscriptionProviderConfig, +}; +use crate::audio_transcription::types::{ + AudioTranscriptionRequestData, AudioTranscriptionResponseData, +}; +use crate::error::{CoreError, CoreResult, json_type_name}; + +use super::aws_base::AwsAuthConfig; +use super::constants::{ + AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE, + DEFAULT_BEDROCK_REGION, +}; + +const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"]; + +pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig = + BedrockAudioTranscriptionConfig; + +pub struct BedrockAudioTranscriptionConfig; + +pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { + let mut stripped = model; + for prefix in ["bedrock/converse/", "bedrock/", "converse/"] { + if let Some(value) = stripped.strip_prefix(prefix) { + stripped = value; + break; + } + } + let mut region = None; + if let Some((candidate, remainder)) = stripped.split_once('/') + && is_bedrock_region(candidate) + { + region = Some(candidate.to_string()); + stripped = remainder; + } + for prefix in ["nova-2/", "nova/"] { + if let Some(value) = stripped.strip_prefix(prefix) { + stripped = value; + break; + } + } + if region.is_none() { + region = stripped + .strip_prefix("arn:") + .and_then(|value| value.split(':').nth(3)) + .filter(|value| !value.is_empty()) + .map(str::to_string); + } + (stripped.to_string(), region) +} + +fn is_bedrock_region(value: &str) -> bool { + value.len() > 3 + && value.contains('-') + && value + .chars() + .all(|char| char.is_ascii_alphanumeric() || char == '-') +} + +pub fn resolve_bedrock_region( + model_region: Option<&str>, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + if let Some(region) = optional_params + .get("aws_region_name") + .and_then(Value::as_str) + { + return region.to_string(); + } + if let Some(region) = model_region { + return region.to_string(); + } + env_lookup(AWS_REGION_NAME) + .or_else(|| env_lookup(AWS_REGION)) + .unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string()) +} + +fn audio_fields(audio: Value) -> CoreResult<(String, String)> { + let object = audio.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&audio), + })?; + let data = object + .get("data") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or(CoreError::MissingField("audio.data"))?; + let format = object + .get("format") + .and_then(Value::as_str) + .filter(|value| matches!(*value, "wav" | "mp3" | "flac" | "ogg")) + .ok_or_else(|| { + CoreError::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string()) + })?; + Ok((data.to_string(), format.to_string())) +} + +fn optional_string<'a>(params: &'a Map, key: &str) -> Option<&'a str> { + params + .get(key) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) +} + +impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { + fn supported_transcription_params(&self) -> &'static [&'static str] { + SUPPORTED_PARAMS + } + + fn transform_transcription_request( + &self, + _model: &str, + audio: Value, + optional_params: Map, + ) -> CoreResult { + let (data, format) = audio_fields(audio)?; + let mut instruction = "Transcribe the audio. Respond with only the transcript.".to_string(); + if let Some(language) = optional_string(&optional_params, "language") { + instruction.push_str(&format!(" The audio language is {language}.")); + } + if let Some(prompt) = optional_string(&optional_params, "prompt") { + instruction.push_str(&format!(" Additional context: {prompt}")); + } + let mut inference_config = Map::from_iter([("maxTokens".to_string(), json!(4096))]); + if let Some(temperature) = optional_params.get("temperature") { + inference_config.insert("temperature".to_string(), temperature.clone()); + } + Ok(AudioTranscriptionRequestData { + body: json!({ + "messages": [{ + "role": "user", + "content": [ + {"audio": {"format": format, "source": {"bytes": data}}}, + {"text": instruction} + ] + }], + "system": [{"text": "You are a transcription assistant."}], + "inferenceConfig": inference_config, + }), + }) + } + + fn transform_transcription_response( + &self, + _model: &str, + response_json: Value, + ) -> CoreResult { + let content = response_json + .get("output") + .and_then(|value| value.get("message")) + .and_then(|value| value.get("content")) + .and_then(Value::as_array) + .ok_or_else(|| { + CoreError::InvalidResponse("Bedrock response has no output content".to_string()) + })?; + let mut text = String::new(); + for block in content { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + Ok(AudioTranscriptionResponseData { text }) + } + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let (model_id, model_region) = bedrock_model_id_and_region(model); + let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup); + let endpoint = optional_params + .get("aws_bedrock_runtime_endpoint") + .and_then(Value::as_str) + .or(api_base) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", ®ion)); + Ok(format!( + "{}/model/{model_id}/converse", + endpoint.trim_end_matches('/') + )) + } + + fn auth_strategy( + &self, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let (_, model_region) = bedrock_model_id_and_region(model); + Ok(AudioTranscriptionAuth::AwsSigV4 { + region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), + service: BEDROCK_SERVICE, + }) + } +} + +pub fn aws_auth_config( + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> AwsAuthConfig { + let value = |key: &str| { + optional_params + .get(key) + .and_then(Value::as_str) + .map(str::to_string) + }; + let env = |key: &str| env_lookup(key); + AwsAuthConfig { + access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")), + secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")), + session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")), + region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)), + session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")), + profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")), + role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")), + web_identity_token: value("aws_web_identity_token") + .or_else(|| env("AWS_WEB_IDENTITY_TOKEN")), + sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")), + external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn no_env(_: &str) -> Option { + None + } + + #[test] + fn request_matches_python_shape() { + let params = Map::from_iter([ + ("language".to_string(), json!("en")), + ("prompt".to_string(), json!("Speaker names")), + ("temperature".to_string(), json!(0)), + ("timestamp_granularities".to_string(), json!(["word"])), + ]); + let params = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.map_transcription_params(¶ms); + let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG + .transform_transcription_request( + "mistral.voxtral-mini-3b-2507", + json!({"data": "AQI=", "format": "wav", "filename": "sample.wav"}), + params, + ) + .expect("request"); + assert_eq!( + result.body, + json!({ + "messages": [{ + "role": "user", + "content": [ + {"audio": {"format": "wav", "source": {"bytes": "AQI="}}}, + {"text": "Transcribe the audio. Respond with only the transcript. The audio language is en. Additional context: Speaker names"} + ] + }], + "system": [{"text": "You are a transcription assistant."}], + "inferenceConfig": {"maxTokens": 4096, "temperature": 0} + }) + ); + } + + #[test] + fn response_concatenates_content_blocks() { + let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG + .transform_transcription_response( + "model", + json!({"output": {"message": {"content": [{"text": "hello "}, {"text": "world"}]}}}), + ) + .expect("response"); + assert_eq!(result.text, "hello world"); + assert_eq!(result.into_json(), json!({"text": "hello world"})); + } + + #[test] + fn invalid_audio_is_rejected() { + let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.transform_transcription_request( + "model", + json!({"data": "AQI="}), + Map::new(), + ); + assert!(result.is_err()); + } + + #[test] + fn region_and_url_precedence_match_python() { + let params = Map::from_iter([("aws_region_name".to_string(), json!("eu-west-1"))]); + let url = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG + .complete_url( + None, + "bedrock/us-east-1/mistral.voxtral-mini-3b-2507", + ¶ms, + &no_env, + ) + .expect("url"); + assert_eq!( + url, + "https://bedrock-runtime.eu-west-1.amazonaws.com/model/mistral.voxtral-mini-3b-2507/converse" + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index c5995732e41..dc036a3cf21 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -52,7 +52,7 @@ pub struct AwsAuthConfig { } impl AwsAuthConfig { - fn with_environment(self, env_lookup: &dyn Fn(&str) -> Option) -> Self { + fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option + Sync)) -> Self { Self { access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)), secret_access_key: self @@ -144,7 +144,7 @@ fn same_role_arns(target: &str, caller: &str) -> bool { pub fn classify_auth( config: AwsAuthConfig, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> AwsAuthFlow { let config = config.with_environment(env_lookup); if let (Some(token), Some(role), Some(session_name)) = ( @@ -194,7 +194,7 @@ pub fn classify_auth( pub async fn resolve_credentials( config: AwsAuthConfig, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> CoreResult { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); diff --git a/litellm-rust/crates/core/src/providers/bedrock/constants.rs b/litellm-rust/crates/core/src/providers/bedrock/constants.rs index a08ae9de146..785295207e7 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/constants.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/constants.rs @@ -2,6 +2,7 @@ pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID"; pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY"; pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN"; pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME"; +pub const AWS_REGION: &str = "AWS_REGION"; pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME"; pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME"; pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME"; @@ -12,3 +13,6 @@ pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; pub const BEDROCK_SERVICE: &str = "bedrock"; pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session"; +pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2"; +pub const BEDROCK_RUNTIME_ENDPOINT_TEMPLATE: &str = + "https://bedrock-runtime.{region}.amazonaws.com"; diff --git a/litellm-rust/crates/core/src/providers/bedrock/mod.rs b/litellm-rust/crates/core/src/providers/bedrock/mod.rs index 8027260a78b..b09675ad7dd 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/mod.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/mod.rs @@ -2,5 +2,7 @@ //! with Python's `BaseAWSLLM`; the broader core purity guidance is reconciled //! separately. +#[cfg(feature = "bedrock-auth")] +pub mod audio_transcription; pub mod aws_base; mod constants; diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md index efa1a554c9c..e5d021ec25b 100644 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -17,7 +17,13 @@ Python-compatible dictionaries. - Provider dispatch belongs in Rust route modules such as `litellm_providers::ocr`, not in this PyO3 crate. - Python owns rollout state and fallback. Rust should return errors; Python - decides whether to raise or fall back. + decides whether to raise or fall back. For a rust-only provider/route (no + Python reference), the Python side is a thin dispatch that calls Rust and + raises when the bridge is unavailable, with no fallback. +- Keep the Python interface minimal (well under 100 lines per route): it only + marshals inputs and calls Rust. Do not add per-route feature flags, and do + not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch + class under `litellm/llms///`. ## Data Handling diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 83e163c38f1..20a9ba789ce 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,7 +10,7 @@ name = "_native" crate-type = ["cdylib"] [dependencies] -litellm-core.workspace = true +litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } pyo3 = { workspace = true, features = ["extension-module"] } pyo3-async-runtimes.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 07429667644..ee9bdd0b81f 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,6 +1,9 @@ use std::collections::HashMap; use std::time::Duration; +use litellm_ai_gateway::io::audio_transcription::{ + AudioTranscriptionRequest, audio_transcription as run_audio_transcription, +}; use litellm_ai_gateway::io::messages::{MessagesRequest, messages as run_messages}; use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; @@ -244,6 +247,93 @@ fn aocr( }) } +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn transcription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let audio = py_to_json(py, audio.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + let result = gil::release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( + AudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }, + )) + }); + match result { + Ok(value) => json_to_py(py, value), + Err(err) => Err(core_error_to_pyerr(err)), + } +} + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn atranscription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let audio = py_to_json(py, audio.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let value = run_audio_transcription(AudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await + .map_err(core_error_to_pyerr)?; + Python::attach(|py| json_to_py(py, value)) + }) +} + type MarshaledMessagesInputs = (Value, Option>, Option); fn marshal_messages_inputs( @@ -341,6 +431,8 @@ fn gil_stats(py: Python<'_>) -> PyResult> { fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(ocr, module)?)?; module.add_function(wrap_pyfunction!(aocr, module)?)?; + module.add_function(wrap_pyfunction!(transcription, module)?)?; + module.add_function(wrap_pyfunction!(atranscription, module)?)?; module.add_function(wrap_pyfunction!(messages, module)?)?; module.add_function(wrap_pyfunction!(amessages, module)?)?; module.add_class::()?; diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py new file mode 100644 index 00000000000..f2e58df3015 --- /dev/null +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -0,0 +1,84 @@ +import base64 +from typing import Union + +import httpx + +from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.rust_bridge import transcription as rust_transcription_bridge +from litellm.types.utils import FileTypes, TranscriptionResponse + + +class BedrockAudioTranscriptionRustDispatch: + @staticmethod + def _audio_payload(audio_file: FileTypes) -> dict[str, object]: + processed_audio = process_audio_file(audio_file) + formats = { + "audio/flac": "flac", + "audio/mpeg": "mp3", + "audio/mp3": "mp3", + "audio/ogg": "ogg", + "audio/wav": "wav", + "audio/x-wav": "wav", + } + audio_format = formats.get(processed_audio.content_type) or ( + processed_audio.filename.rsplit(".", 1)[-1].lower() if "." in processed_audio.filename else "" + ) + if audio_format not in {"wav", "mp3", "flac", "ogg"}: + raise ValueError(f"Unsupported Bedrock audio format for file {processed_audio.filename!r}") + return { + "data": base64.b64encode(processed_audio.file_content).decode("ascii"), + "format": audio_format, + "filename": processed_audio.filename, + } + + def audio_transcriptions( + self, + *, + model: str, + audio_file: FileTypes, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, + ) -> TranscriptionResponse: + rust_response = rust_transcription_bridge.transcription( + model=model, + audio=self._audio_payload(audio_file), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ) + if rust_response is None: + raise RuntimeError("Rust audio transcription bridge is unavailable") + return TranscriptionResponse(**rust_response) + + async def async_audio_transcriptions( + self, + *, + model: str, + audio_file: FileTypes, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, + ) -> TranscriptionResponse: + rust_response = await rust_transcription_bridge.atranscription( + model=model, + audio=self._audio_payload(audio_file), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ) + if rust_response is None: + raise RuntimeError("Rust audio transcription bridge is unavailable") + return TranscriptionResponse(**rust_response) diff --git a/litellm/main.py b/litellm/main.py index 3584297b35f..fb05a375111 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -27,7 +27,6 @@ from typing import ( TYPE_CHECKING, Any, AsyncIterator, - Callable, Coroutine, Dict, Iterable, @@ -81,22 +80,19 @@ from litellm.constants import ( from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function -from litellm.litellm_core_utils.chat_completion_agentic_loop import ( - maybe_run_chat_completion_agentic_loop, -) from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, get_audio_file_for_health_check, ) -from litellm.litellm_core_utils.completion_timeout import CompletionTimeout -from litellm.litellm_core_utils.request_timeout_resolver import ( - get_configured_request_timeout, +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, ) +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_litellm_params import ( AWS_CREDENTIAL_KWARGS_KEYS, OPTIONAL_KWARGS_KEYS, ) -from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -112,6 +108,9 @@ from litellm.litellm_core_utils.mock_functions import ( from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, ) +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) from litellm.llms.base_llm import BaseConfig, BaseImageGenerationConfig from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, @@ -213,7 +212,6 @@ from .llms.bedrock.embed.embedding import BedrockEmbedding from .llms.bedrock.image_edit.handler import BedrockImageEdit from .llms.bedrock.image_generation.image_handler import BedrockImageGeneration from .llms.bytez.chat.transformation import BytezChatConfig -from .llms.gdc.chat.transformation import GDCGeminiConfig from .llms.clarifai.chat.transformation import ClarifaiConfig from .llms.codestral.completion.handler import CodestralTextCompletion from .llms.cohere.embed import handler as cohere_embed @@ -222,24 +220,25 @@ from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from .llms.custom_llm import CustomLLM, custom_chat_llm_router from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha, palm +from .llms.gdc.chat.transformation import GDCGeminiConfig from .llms.gemini.common_utils import get_api_key_from_env from .llms.groq.chat.handler import GroqChatCompletion from .llms.heroku.chat.transformation import HerokuChatConfig from .llms.huggingface.embedding.handler import HuggingFaceEmbedding from .llms.lemonade.chat.transformation import LemonadeChatConfig from .llms.nlp_cloud.chat.handler import completion as nlp_cloud_chat_completion -from .llms.oci.chat.transformation import OCIChatConfig -from .llms.ollama.completion import handler as ollama -from .llms.oobabooga.chat import oobabooga -from .llms.openai.completion.handler import OpenAITextCompletion -from .llms.openai.image_variations.handler import OpenAIImageVariationsHandler -from .llms.openai.openai import OpenAIChatCompletion from .llms.nvidia_riva.audio_transcription.handler import ( NvidiaRivaAudioTranscription, ) from .llms.nvidia_riva.audio_transcription.transformation import ( NvidiaRivaAudioTranscriptionConfig, ) +from .llms.oci.chat.transformation import OCIChatConfig +from .llms.ollama.completion import handler as ollama +from .llms.oobabooga.chat import oobabooga +from .llms.openai.completion.handler import OpenAITextCompletion +from .llms.openai.image_variations.handler import OpenAIImageVariationsHandler +from .llms.openai.openai import OpenAIChatCompletion from .llms.openai.transcriptions.handler import OpenAIAudioTranscription from .llms.openai_like.chat.handler import OpenAILikeChatHandler from .llms.openai_like.embedding.handler import OpenAILikeEmbeddingHandler @@ -7722,6 +7721,32 @@ def transcription( headers=extra_headers, provider_config=provider_config, # type: ignore[arg-type] ) + elif custom_llm_provider == "bedrock": + from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch + + dispatch = BedrockAudioTranscriptionRustDispatch() + if atranscription: + response = dispatch.async_audio_transcriptions( + model=model, + audio_file=file, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ) + else: + response = dispatch.audio_transcriptions( + model=model, + audio_file=file, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ) elif provider_config is not None: response = base_llm_http_handler.audio_transcriptions( model=model, diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index e9139a634f1..91aac6c1232 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -72,17 +72,28 @@ def use_litellm_rust( messages: RustMessages | None | _Unset = _UNSET, amessages: RustAmessages | None | _Unset = _UNSET, responses_websocket: Any | None | _Unset = _UNSET, + transcription: Any | None | _Unset = _UNSET, + atranscription: Any | None | _Unset = _UNSET, ) -> None: global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl configuring_ocr = not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset) configuring_messages = not isinstance(messages, _Unset) or not isinstance(amessages, _Unset) configuring_responses_websocket = not isinstance(responses_websocket, _Unset) + configuring_transcription = not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset) if configuring_ocr or (not configuring_messages and not configuring_responses_websocket): _rust_ocr_enabled = enabled if not isinstance(ocr, _Unset): _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): _rust_aocr_impl = aocr + if configuring_transcription: + from litellm.rust_bridge.transcription import configure_rust_transcription + + configure_rust_transcription( + enabled=enabled, + transcription=transcription, + atranscription=atranscription, + ) if not configuring_messages and not configuring_responses_websocket: return if configuring_messages: diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py new file mode 100644 index 00000000000..44bb42e4104 --- /dev/null +++ b/litellm/rust_bridge/transcription.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Awaitable, Final, Protocol, Union, cast + +import httpx + +from litellm.rust_bridge.timeouts import timeout_to_seconds + + +class RustTranscription(Protocol): + def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + raise NotImplementedError + + +class RustAtranscription(Protocol): + def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> Awaitable[dict[str, object]]: + raise NotImplementedError + + +class _Unset: + pass + + +_UNSET: Final[_Unset] = _Unset() + + +@dataclass +class _RustTranscriptionState: + transcription: RustTranscription | None = None + atranscription: RustAtranscription | None = None + + +_STATE = _RustTranscriptionState() + + +def configure_rust_transcription( + enabled: bool = True, + *, + transcription: RustTranscription | None | _Unset = _UNSET, + atranscription: RustAtranscription | None | _Unset = _UNSET, +) -> None: + if not isinstance(transcription, _Unset): + _STATE.transcription = transcription + if not isinstance(atranscription, _Unset): + _STATE.atranscription = atranscription + + +def load_rust_transcription() -> RustTranscription | None: + if _STATE.transcription is not None: + return _STATE.transcription + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + return ( + None + if native_bridge is None + else cast( # cast-ok: native extension protocol is runtime-defined + RustTranscription, getattr(native_bridge, "transcription", None) + ) + ) + + +def load_rust_atranscription() -> RustAtranscription | None: + if _STATE.atranscription is not None: + return _STATE.atranscription + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + return ( + None + if native_bridge is None + else cast( # cast-ok: native extension protocol is runtime-defined + RustAtranscription, getattr(native_bridge, "atranscription", None) + ) + ) + + +def transcription( + *, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_transcription = load_rust_transcription() + if rust_transcription is None: + return None + return rust_transcription( + model=model, + audio=audio, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_to_seconds(timeout), + ) + + +async def atranscription( + *, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_atranscription = load_rust_atranscription() + if rust_atranscription is None: + return None + return await rust_atranscription( + model=model, + audio=audio, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_to_seconds(timeout), + ) diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py new file mode 100644 index 00000000000..bbeb6c38f78 --- /dev/null +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -0,0 +1,151 @@ +import importlib + +import pytest + +import litellm +from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch + +rust_bridge = importlib.import_module("litellm.rust_bridge.transcription") + + +class SyncBridge: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append({"model": model, "audio": audio, "optional_params": optional_params}) + return {"text": "hello"} + + +class AsyncBridge: + async def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + return {"text": "async"} + + +def test_enabled_sync_bridge_receives_audio() -> None: + bridge = SyncBridge() + rust_bridge.configure_rust_transcription(True, transcription=bridge) + result = rust_bridge.transcription( + model="mistral.voxtral-mini-3b-2507", + audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, + api_key=None, + api_base=None, + custom_llm_provider="bedrock", + extra_headers=None, + optional_params={"temperature": 0}, + timeout=5.0, + ) + assert result == {"text": "hello"} + assert bridge.calls[0]["audio"] == {"data": "AQI=", "format": "wav", "filename": "audio.wav"} + + +@pytest.mark.asyncio +async def test_enabled_async_bridge() -> None: + rust_bridge.configure_rust_transcription(True, atranscription=AsyncBridge()) + result = await rust_bridge.atranscription( + model="mistral.voxtral-mini-3b-2507", + audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, + api_key=None, + api_base=None, + custom_llm_provider="bedrock", + extra_headers=None, + optional_params={}, + timeout=None, + ) + assert result == {"text": "async"} + + +def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None: + rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) + monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", lambda: None) + assert rust_bridge.load_rust_transcription() is None + assert rust_bridge.load_rust_atranscription() is None + + +def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None) + + with pytest.raises(RuntimeError, match="bridge is unavailable"): + BedrockAudioTranscriptionRustDispatch().audio_transcriptions( + model="bedrock/mistral.voxtral-mini-3b-2507", + audio_file=("audio.wav", b"audio", "audio/wav"), + api_key=None, + api_base=None, + custom_llm_provider="bedrock", + extra_headers=None, + optional_params={}, + timeout=5, + ) + + +@pytest.mark.asyncio +async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: + async def unavailable(**_: object) -> None: + return None + + monkeypatch.setattr(rust_bridge, "atranscription", unavailable) + + with pytest.raises(RuntimeError, match="bridge is unavailable"): + await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions( + model="bedrock/mistral.voxtral-mini-3b-2507", + audio_file=("audio.wav", b"audio", "audio/wav"), + api_key=None, + api_base=None, + custom_llm_provider="bedrock", + extra_headers=None, + optional_params={}, + timeout=5, + ) + + +def test_bedrock_transcription_uses_rust_only_path() -> None: + rust_bridge.configure_rust_transcription( + transcription=lambda **_: {"text": "rust"}, + atranscription=None, + ) + try: + response = litellm.transcription( + model="bedrock/mistral.voxtral-mini-3b-2507", + file=("audio.wav", b"audio", "audio/wav"), + ) + finally: + rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) + + assert response.text == "rust" + + +@pytest.mark.asyncio +async def test_bedrock_atranscription_uses_rust_only_path() -> None: + async def rust_response(**_: object) -> dict[str, object]: + return {"text": "rust"} + + rust_bridge.configure_rust_transcription(transcription=None, atranscription=rust_response) + try: + response = await litellm.atranscription( + model="bedrock/mistral.voxtral-mini-3b-2507", + file=("audio.wav", b"audio", "audio/wav"), + ) + finally: + rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) + + assert response.text == "rust" From b9c59c37cc0e94fbef4be3e0d026e10637cf04cb Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 14:13:54 -0700 Subject: [PATCH 40/60] test(e2e): cover model update persisting to /model/info (#34017) Co-authored-by: mubashir1osmani --- tests/e2e/management/test_management_e2e.py | 56 ++++++++++++++++++++- tests/e2e/models.py | 11 ++++ tests/e2e/proxy_client.py | 18 +++++++ 3 files changed, 84 insertions(+), 1 deletion(-) diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 748f7192fb0..e60d074108e 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -9,6 +9,7 @@ so the traffic-facing read-backs poll to a deadline instead of asserting once. from __future__ import annotations +import math import time from collections.abc import Callable @@ -22,7 +23,7 @@ from management_client import ( ROUTE_NOT_ALLOWED_MARKER, ManagementClient, ) -from models import KeyGenerateBody, OrgNewBody, TagListEntry, TagNewBody, TeamNewBody, UserNewBody +from models import KeyGenerateBody, OrgNewBody, TagListEntry, TagNewBody, TeamNewBody, UserNewBody, LiteLLMParamsBody, ModelInfoEntry pytestmark = pytest.mark.e2e @@ -315,6 +316,59 @@ class TestTagRoutes: ) +_INITIAL_INPUT_COST = 0.00000111 +_UPDATED_INPUT_COST = 0.00000222 + + +def _model_entry(client: ManagementClient, model_name: str) -> ModelInfoEntry | None: + return next((entry for entry in client.proxy.model_info() if entry.model_name == model_name), None) + + +class TestModelRoutes: + @pytest.mark.covers("mgmt.model.update.persists") + def test_update_persists_input_cost_to_model_info( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + model_name = f"e2e-mgmt-model-{unique_marker()}" + model_id = client.proxy.create_model( + model_name, + LiteLLMParamsBody( + model="gpt-4o-mini", + mock_response="ok", + input_cost_per_token=_INITIAL_INPUT_COST, + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + before = _model_entry(client, model_name) + assert before is not None, f"{model_name} absent from /model/info right after /model/new" + initial = before.litellm_params.input_cost_per_token + assert initial is not None and math.isclose(initial, _INITIAL_INPUT_COST, rel_tol=1e-9), ( + f"/model/info reports input_cost_per_token {initial}, registered {_INITIAL_INPUT_COST}" + ) + + client.proxy.update_model( + model_id, + LiteLLMParamsBody(model="gpt-4o-mini", input_cost_per_token=_UPDATED_INPUT_COST), + ) + + def updated() -> ModelInfoEntry | None: + entry = _model_entry(client, model_name) + if entry is None: + return None + cost = entry.litellm_params.input_cost_per_token + if cost is not None and math.isclose(cost, _UPDATED_INPUT_COST, rel_tol=1e-9): + return entry + return None + + _ = _poll( + client, + updated, + f"/model/info never reported input_cost_per_token {_UPDATED_INPUT_COST} for {model_name} " + "after /model/update", + ) + + def _assert_route_forbidden(route: str, outcome: StreamingResponse) -> None: assert outcome.status_code == 403, ( f"llm-only key POSTing {route} must be denied exactly 403, got {outcome.status_code}: {outcome.body[:300]}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index b7a95f9714c..b319468f162 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -551,6 +551,17 @@ class ModelNewResponse(BaseModel): model_id: str +class ModelUpdateBody(BaseModel): + """POST /model/update body: the target deployment (`model_info.id`) plus the + `litellm_params` to merge over its stored params. The handler overlays only the + non-null fields, so a body carrying `input_cost_per_token` re-prices the + deployment while leaving its other params intact.""" + + model_config = ConfigDict(protected_namespaces=()) + litellm_params: LiteLLMParamsBody + model_info: ModelInfoBody + + class ModelListEntry(BaseModel): id: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index f1039257193..6c6b948e29c 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -54,6 +54,7 @@ from models import ( ModelNewBody, ModelNewResponse, ModelsListResponse, + ModelUpdateBody, OcrBody, OcrResponse, SpendLogRow, @@ -211,6 +212,23 @@ class ProxyClient: f"propagation or STORE_MODEL_IN_DB reload issue){last_error}" ) + def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: + """Merge `litellm_params` over the deployment `model_id`'s stored params via + POST /model/update. The proxy overlays only the non-null fields and clears + its model cache, so a later /model/info read reflects the change (eventually, + after the reload).""" + unwrap( + self.transport.post( + "/model/update", + headers=self.transport.master, + json=ModelUpdateBody( + litellm_params=litellm_params, + model_info=ModelInfoBody(id=model_id), + ), + response_type=NoBody, + ) + ) + def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", From 53f5a8c380e13f3122714ad52f59fd79a0fed062 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 14:26:18 -0700 Subject: [PATCH 41/60] test(e2e): cover key block persisting to /key/info (#34014) Co-authored-by: mubashir1osmani --- tests/e2e/management/management_client.py | 10 ++++++++++ tests/e2e/management/test_management_e2e.py | 11 +++++++++++ tests/e2e/models.py | 5 +++++ 3 files changed, 26 insertions(+) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 5b94956fa03..d5147682018 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -14,6 +14,7 @@ from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, Un from models import ( ChatBody, ChatMessage, + KeyBlockBody, KeyDeleteBody, KeyGenerateBody, KeyGenerateResponse, @@ -95,6 +96,15 @@ class ManagementClient: ) ) + def block_key(self, key: str) -> None: + _ = unwrap( + self.proxy.transport.post( + "/key/block", + headers=self.proxy.transport.master, + json=KeyBlockBody(key=key), + response_type=NoBody, + ) + ) def regenerate_key(self, key: str) -> str: return unwrap( self.proxy.transport.post( diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index e60d074108e..09a89862c5a 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -170,6 +170,17 @@ class TestKeyRoutes: _ = _poll(client, rejected, "deleted key was still accepted on chat (never rejected 401) at the deadline") + @pytest.mark.covers("mgmt.key.block.persists") + def test_block_persists_to_key_info(self, client: ManagementClient, resources: ResourceManager) -> None: + key = _generate_key(client, resources, KeyGenerateBody(models=["gemini-2.5-flash"])) + assert not client.proxy.key_info(key).blocked, "/key/info reports the key blocked before /key/block ran" + + client.block_key(key) + + def blocked() -> bool | None: + return True if client.proxy.key_info(key).blocked else None + + _ = _poll(client, blocked, "/key/info never reported the key blocked after /key/block before the deadline") class TestKeyRegeneration: @pytest.mark.covers("mgmt.key.regenerate.happy_path") def test_regenerate_rotates_to_a_working_new_key( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index b319468f162..a03a5308168 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -98,6 +98,7 @@ class KeyInfo(BaseModel): tpm_limit: int | None = None rpm_limit: int | None = None team_id: str | None = None + blocked: bool | None = None spend: float | None = None max_budget: float | None = None budget_reset_at: str | None = None @@ -596,6 +597,10 @@ class KeyUpdateBody(BaseModel): models: list[str] +class KeyBlockBody(BaseModel): + key: str + + class KeyListParams(BaseModel): key_alias: str From f5dc1a30107d681e119256ded35cddeb19931ff9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:26:35 -0700 Subject: [PATCH 42/60] test(router): prove request-level bedrock_tags override deployment-level tags for acreate_batch --- tests/test_litellm/test_router.py | 60 +++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c2c98c8869c..a13b3759865 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5430,3 +5430,63 @@ class TestRouterRequestTimeoutPropagation: ) == 60 ) + + +@pytest.mark.asyncio +async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): + import httpx + + from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils + + deployment_tags = [{"key": "application", "value": "config-level"}] + request_tags = [{"key": "application", "value": "request-level"}] + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-batch-model", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-5", + "aws_batch_role_arn": "arn:aws:iam::123:role/batch-role", + "aws_region_name": "us-west-2", + "bedrock_tags": deployment_tags, + }, + } + ] + ) + + def fake_response(): + return httpx.Response( + status_code=200, + json={ + "jobArn": "arn:aws:bedrock:us-west-2:123:model-invocation-job/abc1234567", + "status": "Submitted", + }, + ) + + mock_client = MagicMock() + mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) + + with patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ): + await router.acreate_batch( + model="bedrock-batch-model", + input_file_id="s3://bucket/input.jsonl", + endpoint="/v1/chat/completions", + completion_window="24h", + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == deployment_tags + + await router.acreate_batch( + model="bedrock-batch-model", + input_file_id="s3://bucket/input.jsonl", + endpoint="/v1/chat/completions", + completion_window="24h", + bedrock_tags=request_tags, + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == request_tags From 4c77a5433a792ec3d4a583cfffea344e1ff6a22e Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 14:27:49 -0700 Subject: [PATCH 43/60] test(e2e): cover created team appearing in /team/list (#34015) --- tests/e2e/management/management_client.py | 14 ++++++++++++++ tests/e2e/management/test_management_e2e.py | 13 +++++++++++++ tests/e2e/models.py | 9 +++++++++ 3 files changed, 36 insertions(+) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index d5147682018..2a686b596c0 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -35,6 +35,7 @@ from models import ( TeamDeleteBody, TeamInfoParams, TeamInfoResponse, + TeamListResponse, TeamMemberAddBody, TeamMemberDeleteBody, TeamMemberEntry, @@ -155,6 +156,19 @@ class ManagementClient: ) ).team_info + def team_list_ids(self) -> tuple[str, ...]: + return tuple( + entry.team_id + for entry in unwrap( + self.proxy.transport.get( + "/team/list", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=TeamListResponse, + ) + ).root + ) + def team_info_status(self, team_id: str) -> ProbeResult: return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id)) diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 09a89862c5a..9cc46779296 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -227,6 +227,19 @@ class TestTeamRoutes: f"key generated under team {team_id} carries team_id {key_info.team_id!r} in /key/info" ) + @pytest.mark.covers("mgmt.team.list.happy_path") + def test_created_team_appears_in_team_list( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + alias = f"e2e-mgmt-team-{unique_marker()}" + team_id = _create_team(client, resources, alias, ["gemini-2.5-flash"]) + + _ = _poll( + client, + lambda: team_id if team_id in client.team_list_ids() else None, + f"/team/list never included the created team {team_id}", + ) + @pytest.mark.covers("mgmt.team.member_add.persists") def test_member_add_and_delete_persist_to_team_info( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/models.py b/tests/e2e/models.py index a03a5308168..36765610545 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -654,6 +654,15 @@ class TeamDeleteBody(BaseModel): team_ids: list[str] +class TeamListEntry(BaseModel): + team_id: str + + +class TeamListResponse(RootModel[list[TeamListEntry]]): + """GET /team/list answers with a bare array of team objects (not an object + wrapping them). Only team_id is read; pydantic ignores the rest.""" + + UserRole = Literal["proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer"] From 44604620a4e9cdc7742b69379c99d5533020c5ab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:28:09 -0700 Subject: [PATCH 44/60] test(proxy): make model_info endpoint tests hermetic to kill an order/merge-skew flake The model_info / get_model_info_with_id endpoint tests drove refactored endpoints with bare, unspec'd MagicMock routers and models. Because the mocks were unspec'd, any attribute or method the (refactored) endpoints newly read auto-materialized a child MagicMock, and whether that child was reached depended on process-global state (premium_user, and the real get_available_models_for_user chain reading litellm globals) that sibling tests in the same xdist worker mutate. When reached, the MagicMock either unpacked to empty (a, b = mock.method() -> 'not enough values to unpack (expected 2, got 0)') or leaked into RouterModelInfo(**model_info) and failed Pydantic str validation. Pass in isolation, fail under xdist. The original TestModelInfoEndpoint failure (#33807 CI) was the same class surfaced by merge skew: #33721 added a get_configured_token_limits unpack to create_model_info_response, and CI's merge commit ran that against the un-updated bare-mock test before the #33742 band-aid landed. Fix (test-only, no product change): - TestModelInfoEndpoint: mock the real seam (get_available_models_for_user), configure the router methods the endpoint actually calls, return a real Deployment, and drop the dead proxy_server.get_key_models/get_team_models/ get_complete_model_list patches the refactor had stranded. - TestGetModelInfoWithIdBlocked: spec the model mock so unset enterprise columns read as None instead of child MagicMocks. - test_ProxyConfig_get_model_info_with_id_missing_model_id_raises: pin premium_user so the asserted AttributeError no longer flips with the ambient license global. --- .../test_model_management_endpoints.py | 80 ++++++++----------- .../proxy/proxy_server/test_proxy_config.py | 3 +- 2 files changed, 35 insertions(+), 48 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 79c5f3ea549..f3e5e2c9b71 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1702,8 +1702,8 @@ class TestModelInfoEndpoint: async def test_model_info_accessible_model_success(self): """Test model_info returns model data for accessible models""" from litellm.proxy.proxy_server import model_info + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo - # Mock user with access to specific models user_api_key_dict = UserAPIKeyAuth( user_id="test_user", api_key="test_key", @@ -1713,31 +1713,22 @@ class TestModelInfoEndpoint: with ( patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, - patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, + patch("litellm.proxy.proxy_server.general_settings", {}), patch( - "litellm.proxy.proxy_server.get_complete_model_list" - ) as mock_get_complete_models, - patch("litellm.get_llm_provider") as mock_get_provider, + "litellm.proxy.utils.get_available_models_for_user", + new=AsyncMock(return_value=["gpt-4", "claude-3", "gpt-3.5-turbo"]), + ), + patch("litellm.get_llm_provider", return_value=(None, "openai", None, None)), ): - # Setup mocks - mock_router.get_model_names.return_value = [ - "gpt-4", - "claude-3", - "gpt-3.5-turbo", - ] - mock_router.get_model_access_groups.return_value = {} + mock_router.get_fully_blocked_model_names.return_value = set() + mock_router.get_model_list.return_value = [] mock_router.get_configured_token_limits.return_value = (None, None) - mock_get_key_models.return_value = ["gpt-4", "claude-3"] - mock_get_team_models.return_value = ["gpt-3.5-turbo"] - mock_get_complete_models.return_value = [ - "gpt-4", - "claude-3", - "gpt-3.5-turbo", - ] - mock_get_provider.return_value = (None, "openai", None, None) + mock_router.get_deployment_by_model_group_name.return_value = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params(model="openai/gpt-4"), + model_info=ModelInfo(id="gpt-4"), + ) - # Test accessible model result = await model_info( model_id="gpt-4", user_api_key_dict=user_api_key_dict ) @@ -1764,18 +1755,14 @@ class TestModelInfoEndpoint: with ( patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, - patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, + patch("litellm.proxy.proxy_server.general_settings", {}), patch( - "litellm.proxy.proxy_server.get_complete_model_list" - ) as mock_get_complete_models, + "litellm.proxy.utils.get_available_models_for_user", + new=AsyncMock(return_value=["gpt-4"]), + ), ): - # Setup mocks - user only has access to gpt-4 - mock_router.get_model_names.return_value = ["gpt-4", "claude-3"] - mock_router.get_model_access_groups.return_value = {} - mock_get_key_models.return_value = ["gpt-4"] - mock_get_team_models.return_value = [] - mock_get_complete_models.return_value = ["gpt-4"] # Only gpt-4 accessible + mock_router.get_fully_blocked_model_names.return_value = set() + mock_router.get_model_list.return_value = [] # Test inaccessible model should raise 404 with pytest.raises(HTTPException) as exc_info: @@ -1791,8 +1778,8 @@ class TestModelInfoEndpoint: async def test_model_info_team_model_access(self): """Test model_info works with team model access""" from litellm.proxy.proxy_server import model_info + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo - # Mock user with team access user_api_key_dict = UserAPIKeyAuth( user_id="test_user", api_key="test_key", @@ -1803,23 +1790,22 @@ class TestModelInfoEndpoint: with ( patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, - patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, + patch("litellm.proxy.proxy_server.general_settings", {}), patch( - "litellm.proxy.proxy_server.get_complete_model_list" - ) as mock_get_complete_models, - patch("litellm.get_llm_provider") as mock_get_provider, + "litellm.proxy.utils.get_available_models_for_user", + new=AsyncMock(return_value=["team-model-1"]), + ), + patch("litellm.get_llm_provider", return_value=(None, "custom", None, None)), ): - # Setup mocks - mock_router.get_model_names.return_value = ["team-model-1"] - mock_router.get_model_access_groups.return_value = {} + mock_router.get_fully_blocked_model_names.return_value = set() + mock_router.get_model_list.return_value = [] mock_router.get_configured_token_limits.return_value = (None, None) - mock_get_key_models.return_value = [] - mock_get_team_models.return_value = ["team-model-1"] - mock_get_complete_models.return_value = ["team-model-1"] - mock_get_provider.return_value = (None, "custom", None, None) + mock_router.get_deployment_by_model_group_name.return_value = Deployment( + model_name="team-model-1", + litellm_params=LiteLLM_Params(model="custom/team-model-1"), + model_info=ModelInfo(id="team-model-1"), + ) - # Test team model access result = await model_info( model_id="team-model-1", user_api_key_dict=user_api_key_dict ) @@ -2947,7 +2933,7 @@ class TestGetModelInfoWithIdBlocked: def test_get_model_info_with_id_propagates_blocked_true(self): from litellm.proxy.proxy_server import ProxyConfig - model = MagicMock() + model = MagicMock(spec=["model_id", "model_info", "blocked"]) model.model_id = "dep-1" model.model_info = {} model.blocked = True 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 bd8e92c3cc2..93b21c9d3c1 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -1244,7 +1244,8 @@ def test_ProxyConfig_get_model_info_with_id_returns_router_model_info(): assert snapshot == {"id": "m-1", "db_model": True, "blocked": False} -def test_ProxyConfig_get_model_info_with_id_missing_model_id_raises(): +def test_ProxyConfig_get_model_info_with_id_missing_model_id_raises(monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) pc = ProxyConfig() # model with no model_id, no model_info — accessing .model_id will fail. bad = SimpleNamespace(model_info=None) From 68be053e966620e490a025fca7c678794be5a0dd Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 14:29:13 -0700 Subject: [PATCH 45/60] test(e2e): cover created user appearing in /user/list (#34016) --- tests/e2e/management/management_client.py | 11 +++++++++++ tests/e2e/management/test_management_e2e.py | 20 ++++++++++++++++++++ tests/e2e/models.py | 5 +++++ 3 files changed, 36 insertions(+) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 2a686b596c0..0f1d6dff524 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -259,6 +259,17 @@ class ManagementClient: ) ).total + def user_list_ids(self, user_id: str) -> tuple[str, ...]: + listing = unwrap( + self.proxy.transport.get( + "/user/list", + headers=self.proxy.transport.master, + params=UserListParams(user_ids=user_id), + response_type=UserListResponse, + ) + ) + return tuple(row.user_id for row in listing.users) + def create_org(self, body: OrgNewBody) -> str: return unwrap( self.proxy.transport.post( diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 9cc46779296..f9a03015a95 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -277,6 +277,26 @@ class TestUserRoutes: f"/user/info reports user_role {info.user_role!r}, configured 'internal_user'" ) + @pytest.mark.covers("mgmt.user.list.happy_path") + def test_created_users_appear_in_user_list( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + user_ids = tuple( + _create_user( + client, + resources, + UserNewBody(user_email=f"e2e-mgmt-{unique_marker()}@example.com", user_role="internal_user"), + ) + for _ in range(2) + ) + + for user_id in user_ids: + _ = _poll( + client, + lambda user_id=user_id: (True if user_id in client.user_list_ids(user_id) else None), + f"/user/list never listed the created user {user_id} in the admin inventory", + ) + class TestOrganizationRoutes: @pytest.mark.covers("mgmt.organization.new.happy_path") diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 36765610545..f200df323bd 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -699,7 +699,12 @@ class UserListParams(BaseModel): user_ids: str +class UserListRow(BaseModel): + user_id: str + + class UserListResponse(BaseModel): + users: list[UserListRow] total: int From eb27447a1d774667034543892da226c252499cdf Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 14:56:07 -0700 Subject: [PATCH 46/60] test(e2e): cover team update persistence via /team/info (#33997) Co-authored-by: mubashir1osmani --- tests/e2e/management/management_client.py | 23 +++++++++++++++++++++ tests/e2e/management/test_management_e2e.py | 13 +++++++++++- tests/e2e/models.py | 5 +++++ 3 files changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 0f1d6dff524..e263a4da186 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -41,6 +41,7 @@ from models import ( TeamMemberEntry, TeamNewBody, TeamNewResponse, + TeamUpdateBody, UserDeleteBody, UserInfoParams, UserInfoResponse, @@ -138,6 +139,28 @@ class ManagementClient: self._wait_for_team(team_id) return team_id + def update_team(self, body: TeamUpdateBody) -> None: + last: Result[NoBody] | None = None + for attempt in range(5): + last = self.proxy.transport.post( + "/team/update", + headers=self.proxy.transport.master, + json=body, + response_type=NoBody, + ) + match last: + case Success(): + return + case UnknownApiError(body=body_text) if ( + "connecting to redis" in body_text.lower() or "name resolution" in body_text.lower() + ): + time.sleep(0.5 * (attempt + 1)) + continue + case _: + break + assert last is not None + raise AssertionError(last) + def delete_team(self, team_id: str) -> None: _ = self.proxy.transport.post( "/team/delete", diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index f9a03015a95..fe6c4ef0e22 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -23,7 +23,7 @@ from management_client import ( ROUTE_NOT_ALLOWED_MARKER, ManagementClient, ) -from models import KeyGenerateBody, OrgNewBody, TagListEntry, TagNewBody, TeamNewBody, UserNewBody, LiteLLMParamsBody, ModelInfoEntry +from models import KeyGenerateBody, OrgNewBody, TagListEntry, TagNewBody, TeamNewBody, TeamUpdateBody, UserNewBody, LiteLLMParamsBody, ModelInfoEntry pytestmark = pytest.mark.e2e @@ -227,6 +227,17 @@ class TestTeamRoutes: f"key generated under team {team_id} carries team_id {key_info.team_id!r} in /key/info" ) + @pytest.mark.covers("mgmt.team.update.persists") + def test_update_persists_to_team_info(self, client: ManagementClient, resources: ResourceManager) -> None: + team_id = _create_team(client, resources, f"e2e-mgmt-team-{unique_marker()}", ["gemini-2.5-flash"]) + + updated_alias = f"e2e-mgmt-team-updated-{unique_marker()}" + client.update_team(TeamUpdateBody(team_id=team_id, team_alias=updated_alias)) + + def reflected() -> bool | None: + return True if client.team_info(team_id).team_alias == updated_alias else None + + _ = _poll(client, reflected, f"/team/info never reflected team_alias {updated_alias!r} after /team/update") @pytest.mark.covers("mgmt.team.list.happy_path") def test_created_team_appears_in_team_list( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/models.py b/tests/e2e/models.py index f200df323bd..0df952d9960 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -625,6 +625,11 @@ class TeamNewResponse(BaseModel): team_id: str +class TeamUpdateBody(BaseModel): + team_id: str + team_alias: str + + class TeamInfoParams(BaseModel): team_id: str From c2bd8699be1072592b0552c42467bc7427a46d37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 20 Jul 2026 15:00:41 -0700 Subject: [PATCH 47/60] fix(proxy): require admin opt-in for request-body bedrock_tags Caller-supplied bedrock_tags land as AWS resource tags under the proxy's AWS identity, letting an authenticated caller forge ownership or cost-allocation labels. Add bedrock_tags to _BANNED_REQUEST_BODY_PARAMS so per-request tags need general_settings.allow_client_side_credentials or configurable_clientside_auth_params on the deployment, matching the aws_bedrock_project_id precedent. Deployment-level bedrock_tags in litellm_params are unaffected. Also stop an explicit empty bedrock_tags list in litellm_params from falling through to optional_params --- .../llms/bedrock/batches/transformation.py | 3 +- litellm/proxy/auth/auth_utils.py | 1 + .../bedrock/batches/test_transformation.py | 19 +++++ .../proxy/auth/test_auth_utils.py | 85 +++++++++++++++++++ 4 files changed, 107 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 8648d6586e8..a4ff1c78467 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -215,7 +215,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): "roleArn": role_arn, } - bedrock_tags = litellm_params.get("bedrock_tags") or optional_params.get("bedrock_tags") + config_bedrock_tags = litellm_params.get("bedrock_tags") + bedrock_tags = config_bedrock_tags if config_bedrock_tags is not None else optional_params.get("bedrock_tags") if bedrock_tags is not None: bedrock_request["tags"] = _validate_bedrock_tags(bedrock_tags) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 38900260c98..293bb74e211 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -273,6 +273,7 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = ( # re-route the request's retention and accounting to any project # reachable with the deployment's shared AWS credentials. "aws_bedrock_project_id", + "bedrock_tags", # Provider-specific endpoint overrides that flow into the outbound # request via ``optional_params``. Same threat as ``api_base``: # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index b38d271e210..3681daffe5e 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -298,6 +298,25 @@ def test_create_request_forwards_bedrock_tags_from_optional_params(config): assert mock_sign.call_args.kwargs["data"]["tags"] == tags +def test_create_request_empty_litellm_params_tags_do_not_fall_through(config): + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={"bedrock_tags": [{"key": "env", "value": "prod"}]}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_tags": [], + }, + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == [] + + def test_create_request_omits_tags_when_bedrock_tags_absent(config): with patch.object( config.common_utils, diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index b5d8727f7e6..72bd215b9be 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1944,6 +1944,91 @@ class TestIsRequestBodySafeBlocksRivaUseSsl: ) +class TestIsRequestBodySafeBlocksBedrockTags: + """``bedrock_tags`` lands as AWS resource tags on Bedrock batch jobs + created with the proxy's AWS identity, so a caller-supplied value can + forge ownership or cost-allocation labels; like + ``aws_bedrock_project_id`` it is blocked without an admin opt-in.""" + + def test_bedrock_tags_in_request_body_is_rejected(self): + with pytest.raises(ValueError, match="bedrock_tags"): + is_request_body_safe( + request_body={ + "model": "bedrock-batch-opus", + "bedrock_tags": [{"key": "application", "value": "genai-proxy"}], + }, + general_settings={}, + llm_router=None, + model="bedrock-batch-opus", + ) + + def test_admin_opt_in_proxy_wide_allows_bedrock_tags(self): + assert ( + is_request_body_safe( + request_body={ + "model": "bedrock-batch-opus", + "bedrock_tags": [{"key": "application", "value": "genai-proxy"}], + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="bedrock-batch-opus", + ) + is True + ) + + def test_admin_opt_in_per_deployment_allows_bedrock_tags(self): + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "bedrock-batch-opus", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-opus-4-7", + "configurable_clientside_auth_params": ["bedrock_tags"], + }, + } + ] + ) + assert ( + is_request_body_safe( + request_body={ + "model": "bedrock-batch-opus", + "bedrock_tags": [{"key": "application", "value": "genai-proxy"}], + }, + general_settings={}, + llm_router=router, + model="bedrock-batch-opus", + ) + is True + ) + + def test_per_deployment_opt_in_for_other_param_still_rejects_bedrock_tags(self): + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "bedrock-batch-opus", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-opus-4-7", + "configurable_clientside_auth_params": ["api_base"], + }, + } + ] + ) + with pytest.raises(ValueError, match="bedrock_tags"): + is_request_body_safe( + request_body={ + "model": "bedrock-batch-opus", + "bedrock_tags": [{"key": "application", "value": "genai-proxy"}], + }, + general_settings={}, + llm_router=router, + model="bedrock-batch-opus", + ) + + # ── is_request_body_safe nested-config recursion (VERIA-6) ──────────────────── From 71131190ecf8c1c2df006bfaa893a571834c9feb Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 15:17:40 -0700 Subject: [PATCH 48/60] test(e2e): cover model registration persistence in /model/info (#33996) Co-authored-by: mubashir1osmani --- tests/e2e/management/test_management_e2e.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index fe6c4ef0e22..bfffa70b081 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -423,6 +423,23 @@ class TestModelRoutes: "after /model/update", ) + @pytest.mark.covers("mgmt.model.add.persists") + def test_new_persists_to_model_info_catalog( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + model_name = f"e2e-mgmt-model-{unique_marker()}" + model_id = client.proxy.create_model( + model_name, + LiteLLMParamsBody(model="openai/gpt-5.5", api_key="e2e-dummy-key"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + cataloged = [entry.model_name for entry in client.proxy.model_info()] + assert model_name in cataloged, ( + f"/model/info does not list {model_name!r} after /model/new; registration did not persist " + f"into the routing catalog: {cataloged}" + ) + def _assert_route_forbidden(route: str, outcome: StreamingResponse) -> None: assert outcome.status_code == 403, ( From f21704c67220eda2b326966e5853f7e28470c3a7 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 15:18:36 -0700 Subject: [PATCH 49/60] test(e2e): cover user update persistence via /user/info (#33998) Co-authored-by: mubashir1osmani --- tests/e2e/management/management_client.py | 11 +++++++++++ tests/e2e/management/test_management_e2e.py | 18 +++++++++++++++++- tests/e2e/models.py | 5 +++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index e263a4da186..753f787b921 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -49,6 +49,7 @@ from models import ( UserListResponse, UserNewBody, UserNewResponse, + UserUpdateBody, ) MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" @@ -254,6 +255,16 @@ class ManagementClient: ) ).user_id + def update_user(self, body: UserUpdateBody) -> None: + _ = unwrap( + self.proxy.transport.post( + "/user/update", + headers=self.proxy.transport.master, + json=body, + response_type=NoBody, + ) + ) + def delete_user(self, user_id: str) -> None: _ = self.proxy.transport.post( "/user/delete", diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index bfffa70b081..f0bcc354e61 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -23,7 +23,7 @@ from management_client import ( ROUTE_NOT_ALLOWED_MARKER, ManagementClient, ) -from models import KeyGenerateBody, OrgNewBody, TagListEntry, TagNewBody, TeamNewBody, TeamUpdateBody, UserNewBody, LiteLLMParamsBody, ModelInfoEntry +from models import KeyGenerateBody, OrgNewBody, TagListEntry, TagNewBody, TeamNewBody, TeamUpdateBody, UserNewBody, UserUpdateBody, LiteLLMParamsBody, ModelInfoEntry pytestmark = pytest.mark.e2e @@ -288,6 +288,22 @@ class TestUserRoutes: f"/user/info reports user_role {info.user_role!r}, configured 'internal_user'" ) + @pytest.mark.covers("mgmt.user.update.persists") + def test_update_persists_to_user_info(self, client: ManagementClient, resources: ResourceManager) -> None: + email = f"e2e-mgmt-{unique_marker()}@example.com" + user_id = _create_user(client, resources, UserNewBody(user_email=email, user_role="internal_user")) + + before = client.user_info(user_id).user_info + assert before.user_role == "internal_user", ( + f"/user/info reports pre-update user_role {before.user_role!r}, expected 'internal_user'" + ) + + client.update_user(UserUpdateBody(user_id=user_id, user_role="internal_user_viewer")) + + info = client.user_info(user_id).user_info + assert info.user_role == "internal_user_viewer", ( + f"/user/info reports user_role {info.user_role!r} after /user/update to 'internal_user_viewer'" + ) @pytest.mark.covers("mgmt.user.list.happy_path") def test_created_users_appear_in_user_list( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 0df952d9960..db140a918fa 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -681,6 +681,11 @@ class UserNewResponse(BaseModel): user_id: str +class UserUpdateBody(BaseModel): + user_id: str + user_role: UserRole + + class UserInfoParams(BaseModel): user_id: str From 6db7328afbf9f04b4e7604b872ec587712ff2b35 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 15:28:23 -0700 Subject: [PATCH 50/60] chore(e2e): prune reliability.perf.latency.under_slo coverage cell (#34024) It is a non-binary latency SLO threshold rather than a deterministic pass/fail behavior a single e2e test can assert, so it does not fit the coverage registry's one-test-per-cell contract. The registry README already flagged the perf cells for a support-check or prune, and throughput SLO under load is covered structurally by the Locust load suite. Removing it keeps the denominator to behaviors an e2e test can deterministically prove. --- tests/e2e/coverage_registry/reliability.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index ba192c2912e..1538d3f3cda 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -23,5 +23,4 @@ - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} -- {id: reliability.perf.latency.under_slo, module: reliability, tier: P1, behavior: perf, variant: latency, assertions: [under_slo], exercised_on: [chat_completions, messages], source: "router_strategy/lowest_latency.py", rationale: "Latency SLO (p50/p99) compliance"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} From 61b906f9b6066d5a6acc62009bea4806bcaa7236 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 15:36:15 -0700 Subject: [PATCH 51/60] test(e2e): cover model deletion removing it from the catalog (#34006) --- tests/e2e/management/management_client.py | 13 ++++++++++++ tests/e2e/management/test_management_e2e.py | 22 +++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 753f787b921..0d68098d2d2 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -22,6 +22,7 @@ from models import ( KeyListResponse, KeyRegenerateBody, KeyUpdateBody, + ModelDeleteBody, OrgDeleteBody, OrgInfoParams, OrgInfoResponse, @@ -99,6 +100,18 @@ class ManagementClient: ) ) + def delete_model_strict(self, model_id: str) -> None: + """Strict delete for the act phase of a test: a failed delete is a hard + failure, unlike the warn-only ProxyClient.delete_model used at teardown.""" + _ = unwrap( + self.proxy.transport.post( + "/model/delete", + headers=self.proxy.transport.master, + json=ModelDeleteBody(id=model_id), + response_type=NoBody, + ) + ) + def block_key(self, key: str) -> None: _ = unwrap( self.proxy.transport.post( diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index f0bcc354e61..4e9270bb9e5 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -439,6 +439,28 @@ class TestModelRoutes: "after /model/update", ) + @pytest.mark.covers("mgmt.model.delete.persists") + def test_delete_removes_from_model_info_catalog( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + """The teardown's deferred delete fires again on the already-deleted model by + design: it is the safety net if this test fails before the in-body delete, and + a repeat /model/delete is a warn-only no-op the teardown absorbs.""" + model_name = f"e2e-mgmt-model-{unique_marker()}" + model_id = client.proxy.create_model(model_name, LiteLLMParamsBody(model="openai/gpt-5.5", api_key="dummy")) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + assert model_name in [entry.model_name for entry in client.proxy.model_info()], ( + f"{model_name} absent from /model/info right after /model/new; cannot prove deletion removes it" + ) + + client.delete_model_strict(model_id) + + def absent() -> bool | None: + return True if model_name not in [entry.model_name for entry in client.proxy.model_info()] else None + + _ = _poll(client, absent, f"{model_name} still present in /model/info after /model/delete at the deadline") + @pytest.mark.covers("mgmt.model.add.persists") def test_new_persists_to_model_info_catalog( self, client: ManagementClient, resources: ResourceManager From 5c8e7e69243eab0217436de0dc8a5ab0ee7d53f4 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 15:36:39 -0700 Subject: [PATCH 52/60] test(e2e): cover organization update persistence via /organization/info (#34010) --- tests/e2e/e2e_http.py | 20 +++++++++++++++++++ tests/e2e/management/management_client.py | 11 +++++++++++ tests/e2e/management/test_management_e2e.py | 20 ++++++++++++++++++- tests/e2e/models.py | 5 +++++ tests/e2e/transport.py | 22 +++++++++++++++++++++ 5 files changed, 77 insertions(+), 1 deletion(-) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 692f951d08e..ce801ef81fb 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -264,6 +264,26 @@ def delete[R: BaseModel]( return _classify(resp, response_type) +def patch[R: BaseModel]( + url: URL, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + try: + resp = requests.patch( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + def probe( url: URL, *, headers: BaseModel, params: BaseModel, timeout: float = 30.0 ) -> ProbeResult: diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 0d68098d2d2..b2c9eb6a29b 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -28,6 +28,7 @@ from models import ( OrgInfoResponse, OrgNewBody, OrgNewResponse, + OrgUpdateBody, TagDeleteBody, TagListEntry, TagListResponse, @@ -327,6 +328,16 @@ class ManagementClient: ) ).organization_id + def update_org(self, body: OrgUpdateBody) -> None: + _ = unwrap( + self.proxy.transport.patch( + "/organization/update", + headers=self.proxy.transport.master, + json=body, + response_type=NoBody, + ) + ) + def delete_org(self, organization_id: str) -> None: _ = self.proxy.transport.delete( "/organization/delete", diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 4e9270bb9e5..b0a12220341 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -23,7 +23,7 @@ from management_client import ( ROUTE_NOT_ALLOWED_MARKER, ManagementClient, ) -from models import KeyGenerateBody, OrgNewBody, TagListEntry, TagNewBody, TeamNewBody, TeamUpdateBody, UserNewBody, UserUpdateBody, LiteLLMParamsBody, ModelInfoEntry +from models import KeyGenerateBody, OrgInfoResponse, OrgNewBody, OrgUpdateBody, TagListEntry, TagNewBody, TeamNewBody, TeamUpdateBody, UserNewBody, UserUpdateBody, LiteLLMParamsBody, ModelInfoEntry pytestmark = pytest.mark.e2e @@ -342,6 +342,24 @@ class TestOrganizationRoutes: f"/organization/info reports models {info.models}, configured ['gemini-2.5-flash']" ) + @pytest.mark.covers("mgmt.organization.update.persists") + def test_update_alias_persists_to_organization_info( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + org_id = client.create_org(OrgNewBody(organization_alias=f"e2e-mgmt-org-{unique_marker()}")) + resources.defer(lambda: client.delete_org(org_id)) + + new_alias = f"e2e-mgmt-org-{unique_marker()}" + client.update_org(OrgUpdateBody(organization_id=org_id, organization_alias=new_alias)) + + def attempt() -> OrgInfoResponse | None: + info = client.org_info(org_id) + return info if info.organization_alias == new_alias else None + + _ = _poll( + client, attempt, f"/organization/info never reflected updated alias {new_alias!r} before the deadline" + ) + @pytest.mark.covers("mgmt.organization.delete.persists") def test_delete_removes_from_organization_info( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/models.py b/tests/e2e/models.py index db140a918fa..9e773120846 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -727,6 +727,11 @@ class OrgNewResponse(BaseModel): organization_id: str +class OrgUpdateBody(BaseModel): + organization_id: str + organization_alias: str + + class OrgInfoParams(BaseModel): organization_id: str diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 10e090f07a9..64fe6406ff7 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -55,6 +55,10 @@ class Transport(Protocol): self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] ) -> Result[R]: ... + def patch[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: ... + def probe(self, path: str, *, params: BaseModel) -> ProbeResult: ... def upload[R: BaseModel]( @@ -131,6 +135,17 @@ class HttpTransport: timeout=self.request_timeout, ) + def patch[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return e2e_http.patch( + self._url(path), + headers=headers, + json=json, + response_type=response_type, + timeout=self.request_timeout, + ) + def stream( self, path: str, *, headers: BaseModel, json: BaseModel ) -> StreamingResponse: @@ -271,6 +286,13 @@ class SplitTransport: path, headers=headers, json=json, response_type=response_type ) + def patch[R: BaseModel]( + self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + ) -> Result[R]: + return self._route(path).patch( + path, headers=headers, json=json, response_type=response_type + ) + def stream( self, path: str, *, headers: BaseModel, json: BaseModel ) -> StreamingResponse: From 6f62022e8439eadca4284f83fa351f89ee7f6443 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 15:37:40 -0700 Subject: [PATCH 53/60] test(e2e): cover team deletion persistence and key revocation (#33999) --- tests/e2e/management/test_management_e2e.py | 31 +++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index b0a12220341..dfb4a506bd6 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -251,6 +251,37 @@ class TestTeamRoutes: f"/team/list never included the created team {team_id}", ) + @pytest.mark.covers("mgmt.team.delete.persists") + def test_delete_persists_and_revokes_team_bound_key( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + """The teardown's deferred delete_team/delete_key fire again on the already- + deleted team and key by design: both are warn-only no-ops, and the deferred + cleanup must survive this test failing before the in-body delete.""" + team_id = _create_team(client, resources, f"e2e-mgmt-team-{unique_marker()}", ["gpt-5.5"]) + key = _generate_key(client, resources, KeyGenerateBody(team_id=team_id)) + + def accepted() -> bool | None: + outcome = client.chat_status(key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.status_code != 401 else None + + _ = _poll(client, accepted, "team-bound key was never accepted at auth before team deletion") + + client.delete_team(team_id) + + probe = client.team_info_status(team_id) + assert probe.status_code == 404, ( + f"deleted team {team_id} still resolves: /team/info returned {probe.status_code}: {probe.body[:300]}" + ) + + def rejected() -> bool | None: + outcome = client.chat_status(key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.status_code == 401 else None + + _ = _poll( + client, rejected, "team-bound key was still accepted on chat (never rejected 401) after team deletion" + ) + @pytest.mark.covers("mgmt.team.member_add.persists") def test_member_add_and_delete_persist_to_team_info( self, client: ManagementClient, resources: ResourceManager From c208bec37fb3272f779a96e239f7213d0b63e154 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 15:38:32 -0700 Subject: [PATCH 54/60] test(e2e): cover user deletion removing it from user inventory (#34007) --- tests/e2e/management/management_client.py | 13 ++++++++++++ tests/e2e/management/test_management_e2e.py | 22 +++++++++++++++++++++ tests/e2e/models.py | 4 ++++ 3 files changed, 39 insertions(+) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index b2c9eb6a29b..a9dedac8e61 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -45,6 +45,7 @@ from models import ( TeamNewResponse, TeamUpdateBody, UserDeleteBody, + UserDeleteResponse, UserInfoParams, UserInfoResponse, UserListParams, @@ -287,6 +288,18 @@ class ManagementClient: response_type=NoBody, ) + def delete_user_strict(self, user_id: str) -> None: + """Strict delete for the act phase of a test: a failed delete is a hard + failure, unlike the warn-only delete_user used at teardown.""" + _ = unwrap( + self.proxy.transport.post( + "/user/delete", + headers=self.proxy.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=UserDeleteResponse, + ) + ) + def user_info(self, user_id: str) -> UserInfoResponse: return unwrap( self.proxy.transport.get( diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index dfb4a506bd6..dae68ba786d 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -335,6 +335,28 @@ class TestUserRoutes: assert info.user_role == "internal_user_viewer", ( f"/user/info reports user_role {info.user_role!r} after /user/update to 'internal_user_viewer'" ) + @pytest.mark.covers("mgmt.user.delete.persists") + def test_delete_removes_the_user_from_inventory( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + """The teardown's deferred delete fires again on the already-deleted user by + design: the deferred cleanup must survive this test failing before the + in-body delete, and a repeat /user/delete is a cheap no-op the warn-only + teardown absorbs.""" + user_id = _create_user( + client, + resources, + UserNewBody(user_email=f"e2e-mgmt-{unique_marker()}@example.com", user_role="internal_user"), + ) + assert client.user_count(user_id) == 1, f"user {user_id} was not created before deletion" + + client.delete_user_strict(user_id) + + def removed() -> bool | None: + return True if client.user_count(user_id) == 0 else None + + _ = _poll(client, removed, f"user {user_id} still present in /user/list after /user/delete at the deadline") + @pytest.mark.covers("mgmt.user.list.happy_path") def test_created_users_appear_in_user_list( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9e773120846..e07814f2568 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -705,6 +705,10 @@ class UserDeleteBody(BaseModel): user_ids: list[str] +class UserDeleteResponse(RootModel[int]): + pass + + class UserListParams(BaseModel): user_ids: str From 583ddaf19958b5cba0fb418f48c40e5ebfc7d144 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 15:40:39 -0700 Subject: [PATCH 55/60] test(e2e): cover created key appearing in /key/list inventory (#34008) --- tests/e2e/management/test_management_e2e.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index dae68ba786d..18bc384a879 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -169,6 +169,24 @@ class TestKeyRoutes: _ = _poll(client, rejected, "deleted key was still accepted on chat (never rejected 401) at the deadline") + @pytest.mark.covers("mgmt.key.list.happy_path") + def test_created_key_appears_in_key_list_inventory( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + alias = f"e2e-mgmt-keylist-{unique_marker()}" + assert client.key_alias_count(alias) == 0, ( + f"/key/list already reports a key under the unused alias {alias!r} before it is created" + ) + + _ = _generate_key(client, resources, KeyGenerateBody(key_alias=alias)) + + def listed() -> bool | None: + return True if client.key_alias_count(alias) == 1 else None + + _ = _poll( + client, listed, f"created key with alias {alias!r} never appeared in /key/list before the deadline" + ) + @pytest.mark.covers("mgmt.key.block.persists") def test_block_persists_to_key_info(self, client: ManagementClient, resources: ResourceManager) -> None: From 3f712e3fde106b5da383f5a649a3fe04289f03e3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 20 Jul 2026 15:49:19 -0700 Subject: [PATCH 56/60] fix(router): stop custom model_info leaking onto shared backend cost map key --- litellm/router.py | 27 ++-- litellm/types/utils.py | 16 ++ .../test_router_model_cost_isolation.py | 138 ++++++++++++++++++ 3 files changed, 169 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index ae3f7ba11c2..3ecaef591f3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -200,6 +200,7 @@ from litellm.types.utils import ( CustomPricingLiteLLMParams, GenericBudgetConfigType, LiteLLMBatch, + shared_backend_model_info, ) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo @@ -7495,12 +7496,13 @@ class Router: if deployment.litellm_params.custom_llm_provider is not None: _model_name = deployment.litellm_params.custom_llm_provider + "/" + _model_name - # For the shared backend key, strip custom pricing fields so that - # one deployment's pricing overrides don't pollute another - # deployment sharing the same backend model name. - # Each deployment's full pricing is already stored under its - # unique model_id above. - _shared_model_info = CustomPricingLiteLLMParams.strip_custom_pricing_fields(_model_info) + # For the shared backend key, keep only cost-map schema fields + # (minus custom pricing) so that one deployment's pricing overrides + # or custom metadata (id, access_via_team_ids, arbitrary keys) + # don't pollute another deployment sharing the same backend model + # name. Each deployment's full model_info is already stored under + # its unique model_id above. + _shared_model_info = shared_backend_model_info(_model_info) _existing_shared_mode = (cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {}).get("mode") _deployment_mode = _shared_model_info.get("mode") # Keep the built-in bridge mode stable for shared backend keys. @@ -8219,12 +8221,13 @@ class Router: if deployment.litellm_params.custom_llm_provider is not None: _model_name = deployment.litellm_params.custom_llm_provider + "/" + _model_name - # For the shared backend key, strip custom pricing fields so that - # one deployment's pricing overrides don't pollute another - # deployment sharing the same backend model name. - # Each deployment's full pricing is already stored under its - # unique model_id above (when present). - _shared_model_info = CustomPricingLiteLLMParams.strip_custom_pricing_fields(_model_info_dict) + # For the shared backend key, keep only cost-map schema fields + # (minus custom pricing) so that one deployment's pricing overrides + # or custom metadata (id, access_via_team_ids, arbitrary keys) + # don't pollute another deployment sharing the same backend model + # name. Each deployment's full model_info is already stored under + # its unique model_id above (when present). + _shared_model_info = shared_backend_model_info(_model_info_dict) _backend_alias_cost = {_model_name: _shared_model_info} if "responses/" in _model_name: _stripped_model_name = _model_name.replace("responses/", "") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ec8a9336ca7..5b98e8be8d2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -5,6 +5,7 @@ from typing import ( TYPE_CHECKING, Any, Dict, + FrozenSet, List, Literal, Mapping, @@ -3112,6 +3113,21 @@ class CustomPricingLiteLLMParams(BaseModel): return {k: v for k, v in model_info.items() if k not in cls.model_fields} +SHARED_BACKEND_MODEL_INFO_FIELDS: FrozenSet[str] = frozenset( + ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__ +) - frozenset(CustomPricingLiteLLMParams.model_fields) + + +def shared_backend_model_info(model_info: Dict[str, Any]) -> Dict[str, Any]: + """Return only the fields safe to register under a shared ``{provider}/{model}`` + key in ``litellm.model_cost``: cost-map schema fields (``ModelInfoBase``) minus + per-deployment pricing overrides. Per-deployment metadata (``id``, + ``access_via_team_ids``, arbitrary custom keys) never belongs on the shared key; + it stays under the deployment's unique model id. + """ + return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS} + + # Server-controlled fields that bound or drive an interceptor's agentic loop # (depth, cycle fingerprints, ceiling, code-interpreter sandbox state). Listed # in all_litellm_params so they are treated as LiteLLM-level and excluded from diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 6db7b04b3b7..c7f5513b94b 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -683,6 +683,144 @@ def test_custom_pricing_isolated_from_sibling_via_proxy_model_info_path(): _restore_model_cost_entries(model_keys) +def test_custom_model_info_metadata_not_leaked_to_shared_backend_key(): + """LIT-4544: two deployments share the same backend model but carry + different custom model_info (arbitrary keys, access_via_team_ids, ids). + None of that per-deployment metadata may land on the shared backend key in + litellm.model_cost (served raw by /public/litellm_model_cost_map); + before the fix it was merged last-write-wins so values flipped randomly. + """ + backend_model = "openai/gpt-4o-mini" + shared_keys = ("gpt-4o-mini", backend_model) + leak_fields = ("id", "additionalProp1", "access_via_team_ids", "db_model") + + model_keys = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in (*shared_keys, "lit4544-deploy-a", "lit4544-deploy-b") + } + try: + Router( + model_list=[ + { + "model_name": "alias-unrestricted", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-a", + }, + "model_info": { + "id": "lit4544-deploy-a", + "additionalProp1": {"restricted": False, "model_location": "EU"}, + }, + }, + { + "model_name": "alias-restricted", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-b", + }, + "model_info": { + "id": "lit4544-deploy-b", + "additionalProp1": {"restricted": True, "model_location": "US"}, + "access_via_team_ids": ["team-b-only"], + }, + }, + ], + ) + + for shared_key in shared_keys: + shared_entry = litellm.model_cost.get(shared_key) or {} + leaked = [field for field in leak_fields if field in shared_entry] + assert not leaked, ( + f"per-deployment metadata {leaked} leaked onto shared key " + f"{shared_key}: {shared_entry}" + ) + + entry_a = litellm.model_cost["lit4544-deploy-a"] + assert entry_a["additionalProp1"] == {"restricted": False, "model_location": "EU"} + entry_b = litellm.model_cost["lit4544-deploy-b"] + assert entry_b["additionalProp1"] == {"restricted": True, "model_location": "US"} + assert entry_b["access_via_team_ids"] == ["team-b-only"] + finally: + _restore_model_cost_entries(model_keys) + + +def test_add_deployment_does_not_leak_custom_metadata_to_shared_backend_key(): + """LIT-4544 dynamic path: deployments added at runtime (e.g. loaded from + the DB every scheduler cycle) must not re-pollute the shared backend key + with per-deployment metadata either. + """ + backend_model = "openai/gpt-4o-mini" + shared_keys = ("gpt-4o-mini", backend_model) + deploy_id = "lit4544-add-deployment" + + model_keys = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in (*shared_keys, deploy_id) + } + try: + router = Router(model_list=[]) + router.add_deployment( + deployment=Deployment( + model_name="alias-dynamic", + litellm_params=LiteLLM_Params( + model=backend_model, + api_key="fake-key-dynamic", + ), + model_info=ModelInfo( + id=deploy_id, + additionalProp1={"restricted": True}, + access_via_team_ids=["team-dynamic"], + ), + ) + ) + + for shared_key in shared_keys: + shared_entry = litellm.model_cost.get(shared_key) or {} + leaked = [ + field + for field in ("id", "additionalProp1", "access_via_team_ids", "db_model") + if field in shared_entry + ] + assert not leaked, ( + f"per-deployment metadata {leaked} leaked onto shared key " + f"{shared_key}: {shared_entry}" + ) + + assert litellm.model_cost[deploy_id]["access_via_team_ids"] == ["team-dynamic"] + finally: + _restore_model_cost_entries(model_keys) + + +def test_shared_backend_model_info_keeps_schema_fields_and_drops_the_rest(): + """Unit test of the whitelist helper: cost-map schema fields survive, + custom pricing overrides and per-deployment metadata do not. + """ + from litellm.types.utils import shared_backend_model_info + + filtered = shared_backend_model_info( + { + "mode": "chat", + "litellm_provider": "openai", + "max_tokens": 128000, + "supports_vision": True, + "input_cost_per_token": 0.99, + "output_cost_per_token": 0.99, + "id": "deploy-a", + "db_model": False, + "access_via_team_ids": ["team-a"], + "additionalProp1": {"restricted": True}, + "base_model": "gpt-4o-mini", + } + ) + + assert filtered == { + "mode": "chat", + "litellm_provider": "openai", + "max_tokens": 128000, + "supports_vision": True, + } + + def test_wildcard_zero_cost_request_does_not_poison_named_deployment_pricing(): """LIT-3991 end to end: a proxy has a named text-embedding-3-small deployment relying on built-in pricing plus an ``openai/*`` wildcard with From 432954a2ab1d9c2b1ff4f38a677df6cbbc8d2a5a Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 20 Jul 2026 15:51:01 -0700 Subject: [PATCH 57/60] fix(cache): make in-memory and disk cache increments atomic (#34013) * fix(cache): make in-memory and disk increments atomic * refactor(cache): narrow in-memory increment lock scope * fix(cache): address follow-up review on increment tests/types * fix(cache): refresh atomic increment coverage * test(cache): widen increment race window with non-zero _SlowInt seed The zero seed was falsy, so InMemoryCache.increment_cache's `get_cache(...) or 0` and DiskCache.get_cache's truthiness guard both discarded the _SlowInt before __add__ could run, leaving the sleep-based window-widening inert. Seed a non-zero value and return _SlowInt from __add__ so the sleep fires on every read-modify-write in both backends, making the concurrency regression deterministic. * test(cache): cover InMemoryCache.async_increment delegation Add a focused async test asserting async_increment accumulates through the locked sync path, exercising the previously uncovered delegation line. --------- Co-authored-by: Emerson Gomes --- litellm/caching/disk_cache.py | 19 ++++------- litellm/caching/in_memory_cache.py | 21 ++++++------ tests/test_litellm/caching/test_disk_cache.py | 26 +++++++++++++++ .../caching/test_in_memory_cache.py | 32 +++++++++++++++++++ 4 files changed, 75 insertions(+), 23 deletions(-) diff --git a/litellm/caching/disk_cache.py b/litellm/caching/disk_cache.py index d9f65ce949e..af8eb92849f 100644 --- a/litellm/caching/disk_cache.py +++ b/litellm/caching/disk_cache.py @@ -58,12 +58,12 @@ class DiskCache(BaseCache): return return_val def increment_cache(self, key, value: int, **kwargs) -> int: - # get the value - cached_value = self.get_cache(key=key) - init_value = cached_value if isinstance(cached_value, int) else 0 - value = init_value + value - self.set_cache(key, value, **kwargs) - return value + with self.disk_cache.transact(): + cached_value = self.get_cache(key=key) + init_value = cached_value if isinstance(cached_value, int) else 0 + new_value = init_value + value + self.set_cache(key, new_value, **kwargs) + return new_value async def async_get_cache(self, key, **kwargs): return self.get_cache(key=key, **kwargs) @@ -76,12 +76,7 @@ class DiskCache(BaseCache): return return_val async def async_increment(self, key, value: int, **kwargs) -> int: - # get the value - cached_value = await self.async_get_cache(key=key) - init_value = cached_value if isinstance(cached_value, int) else 0 - value = init_value + value - await self.async_set_cache(key, value, **kwargs) - return value + return self.increment_cache(key=key, value=value, **kwargs) def flush_cache(self): self.disk_cache.clear() diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 2ad3f3f11b7..36b477f7a8b 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -12,6 +12,7 @@ import json import sys import time import heapq +import threading from typing import TYPE_CHECKING, Any, List, Optional if TYPE_CHECKING: @@ -46,6 +47,7 @@ class InMemoryCache(BaseCache): self.cache_dict: dict = {} self.ttl_dict: dict = {} self.expiration_heap: list[tuple[float, str]] = [] + self._increment_lock = threading.Lock() def check_value_size(self, value: Any): """ @@ -223,12 +225,13 @@ class InMemoryCache(BaseCache): return_val.append(val) return return_val - def increment_cache(self, key, value: int, **kwargs) -> int: - # get the value - init_value = self.get_cache(key=key) or 0 - value = init_value + value - self.set_cache(key, value, **kwargs) - return value + def increment_cache(self, key, value: float, **kwargs) -> float: + with self._increment_lock: + # keep read-modify-write atomic + init_value = self.get_cache(key=key) or 0 + value = init_value + value + self.set_cache(key, value, **kwargs) + return value async def async_get_cache(self, key, **kwargs): return self.get_cache(key=key, **kwargs) @@ -241,11 +244,7 @@ class InMemoryCache(BaseCache): return return_val async def async_increment(self, key, value: float, **kwargs) -> float: - # get the value - init_value = await self.async_get_cache(key=key) or 0 - value = init_value + value - await self.async_set_cache(key, value, **kwargs) - return value + return self.increment_cache(key=key, value=value, **kwargs) async def async_increment_pipeline( self, increment_list: List["RedisPipelineIncrementOperation"], **kwargs diff --git a/tests/test_litellm/caching/test_disk_cache.py b/tests/test_litellm/caching/test_disk_cache.py index b8d3b7b8d36..084370726b1 100644 --- a/tests/test_litellm/caching/test_disk_cache.py +++ b/tests/test_litellm/caching/test_disk_cache.py @@ -1,3 +1,7 @@ +import threading +import time +from concurrent.futures import ThreadPoolExecutor + import pytest pytest.importorskip("diskcache") @@ -5,6 +9,12 @@ pytest.importorskip("diskcache") from litellm.caching.disk_cache import DiskCache +class _SlowInt(int): + def __add__(self, value: int) -> "_SlowInt": + time.sleep(0.05) + return _SlowInt(int(self) + value) + + @pytest.fixture def cache(tmp_path): return DiskCache(disk_cache_dir=str(tmp_path)) @@ -27,6 +37,22 @@ def test_increment_cache_treats_non_int_cached_value_as_zero(cache): assert cache.get_cache("counter") == 4 +def test_increment_cache_is_atomic_under_thread_concurrency(cache): + seed = 1000 + cache.set_cache("counter", _SlowInt(seed)) + thread_count = 8 + barrier = threading.Barrier(thread_count) + + def increment(_: int) -> int: + barrier.wait() + return cache.increment_cache("counter", 1) + + with ThreadPoolExecutor(max_workers=thread_count) as executor: + tuple(executor.map(increment, range(thread_count))) + + assert cache.get_cache("counter") == seed + thread_count + + async def test_async_increment_starts_from_zero_when_key_missing(cache): assert await cache.async_increment("counter", 2) == 2 diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 8828ebf207e..7be03d23fbe 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -2,7 +2,9 @@ import asyncio import json import os import sys +import threading import time +from concurrent.futures import ThreadPoolExecutor from unittest.mock import MagicMock, patch import httpx @@ -18,6 +20,36 @@ from unittest.mock import AsyncMock from litellm.caching.in_memory_cache import InMemoryCache +class _SlowInt(int): + def __add__(self, value: int) -> "_SlowInt": + time.sleep(0.05) + return _SlowInt(int(self) + value) + + +def test_increment_cache_is_atomic_under_thread_concurrency(): + cache = InMemoryCache() + seed = 1000 + cache.set_cache("counter", _SlowInt(seed)) + thread_count = 8 + barrier = threading.Barrier(thread_count) + + def increment(_: int) -> float: + barrier.wait() + return cache.increment_cache("counter", 1) + + with ThreadPoolExecutor(max_workers=thread_count) as executor: + tuple(executor.map(increment, range(thread_count))) + + assert cache.get_cache("counter") == seed + thread_count + + +async def test_async_increment_delegates_to_locked_sync_path(): + cache = InMemoryCache() + assert await cache.async_increment("counter", 2) == 2 + assert await cache.async_increment("counter", 3) == 5 + assert cache.get_cache("counter") == 5 + + def test_in_memory_openai_obj_cache(): from openai import OpenAI From 381013010519f7620b588880ff7c55b8450cc55b Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 16:06:46 -0700 Subject: [PATCH 58/60] test(e2e): add reliability suite covering fallback, timeout, and cache behavior (#34023) * test(e2e): add reliability suite covering fallback, timeout, and cache behavior * test(e2e): move reliability suite under router and drive it with real deployments * test(e2e): make the router complexity fixture opt-in so reliability tests can coexist --- tests/e2e/models.py | 22 ++++++ tests/e2e/router/conftest.py | 4 +- tests/e2e/router/reliability_support.py | 77 +++++++++++++++++++ .../e2e/router/test_complexity_router_e2e.py | 1 + .../e2e/router/test_reliability_cache_e2e.py | 37 +++++++++ .../router/test_reliability_fallbacks_e2e.py | 69 +++++++++++++++++ .../router/test_reliability_timeouts_e2e.py | 53 +++++++++++++ 7 files changed, 261 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/router/reliability_support.py create mode 100644 tests/e2e/router/test_reliability_cache_e2e.py create mode 100644 tests/e2e/router/test_reliability_fallbacks_e2e.py create mode 100644 tests/e2e/router/test_reliability_timeouts_e2e.py diff --git a/tests/e2e/models.py b/tests/e2e/models.py index e07814f2568..22aedb1cfbe 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -166,6 +166,27 @@ class ChatBody(BaseModel): guardrails: list[str] | None = None +class RouterSettingsOverride(BaseModel): + """Per-request `router_settings_override` in a /chat/completions body: the + reliability knobs (fallbacks by trigger, retry count) the reliability suite + drives per call instead of via static router config. Serialized exclude_none, so + an override sets only the strategies a test exercises. Each fallbacks map is + model_name -> the ordered fallback model_names to try.""" + + fallbacks: list[dict[str, list[str]]] | None = None + context_window_fallbacks: list[dict[str, list[str]]] | None = None + content_policy_fallbacks: list[dict[str, list[str]]] | None = None + num_retries: int | None = None + + +class ReliabilityChatBody(ChatBody): + """A /chat/completions body carrying a per-request router_settings_override. + Composes ChatBody (no attribute repetition) and adds the override; serialized + exclude_none so an absent override never leaks into the request.""" + + router_settings_override: RouterSettingsOverride | None = None + + class OutMessage(BaseModel): content: str | None = None reasoning_content: str | None = None @@ -526,6 +547,7 @@ class LiteLLMParamsBody(BaseModel): use_in_pass_through: bool | None = None complexity_router_config: dict[str, object] | None = None mock_response: str | None = None + timeout: float | None = None ModelMode = Literal["batch", "realtime", "image_generation"] diff --git a/tests/e2e/router/conftest.py b/tests/e2e/router/conftest.py index 8ddc19aa94f..98501f9bd7c 100644 --- a/tests/e2e/router/conftest.py +++ b/tests/e2e/router/conftest.py @@ -79,8 +79,8 @@ def _router_is_callable(proxy: ProxyClient) -> bool: return isinstance(result, Success) -@pytest.fixture(scope="session", autouse=True) -def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] # pytest autouse session fixture, wired by name +@pytest.fixture(scope="session") +def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] # requested by the complexity test via usefixtures, wired by name client: ComplexityRouterClient, ) -> Iterator[None]: """Ensure the complexity router virtual model exists for this session. diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py new file mode 100644 index 00000000000..4dab0aaa3fa --- /dev/null +++ b/tests/e2e/router/reliability_support.py @@ -0,0 +1,77 @@ +"""Shared helpers for the reliability e2e tests (fallbacks, timeouts, cache). + +These are plain functions over the router suite's shared ProxyClient, not a +fixture/client class: the tests reuse the router `client` fixture and pass +`client.proxy`. Fallbacks and timeouts are driven by REAL deployments that all +point at the real `openai/gpt-5.5`; a bad base URL yields a real connection +error and a 1ms deadline yields a real timeout, and each test wires the +reroute per request through a `router_settings_override` in the /chat/completions +body, so a single long-lived proxy serves every reliability behavior. +""" + +from __future__ import annotations + +from pydantic import ValidationError + +from proxy_client import ProxyClient +from e2e_http import StreamingResponse +from models import ( + ChatMessage, + ChatResponse, + LiteLLMParamsBody, + ReliabilityChatBody, + RouterSettingsOverride, +) + +REAL_MODEL = "openai/gpt-5.5" +REAL_KEY = "os.environ/OPENAI_API_KEY" + + +def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str: + """Register a deployment pointing at an unreachable base, so every call to it + fails with a real connection error the fallback can reroute around.""" + return proxy.create_model( + name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, api_base="http://127.0.0.1:9/v1") + ) + + +def create_timeout_deployment(proxy: ProxyClient, name: str) -> str: + """Register a deployment with a 1ms deadline the real backend always exceeds.""" + return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001)) + + +def chat_override( + proxy: ProxyClient, + key: str, + model: str, + content: str, + override: RouterSettingsOverride | None = None, + stream: bool = False, +) -> StreamingResponse: + """POST /chat/completions with an optional per-request router_settings_override, + returning the raw outcome so tests read status, body, and reliability headers.""" + return proxy.transport.send( + "/chat/completions", + headers=proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=16, + stream=stream, + router_settings_override=override, + ), + stream=stream, + ) + + +def content_of(resp: StreamingResponse) -> str | None: + """The assistant message content of a successful chat response, or None when the + body is not a success shape (an error body, or an elided streamed body).""" + try: + parsed = ChatResponse.model_validate_json(resp.body) + except ValidationError: + return None + if not parsed.choices: + return None + message = parsed.choices[0].message + return message.content if message is not None else None diff --git a/tests/e2e/router/test_complexity_router_e2e.py b/tests/e2e/router/test_complexity_router_e2e.py index a495e2fdf4d..e8508c963b8 100644 --- a/tests/e2e/router/test_complexity_router_e2e.py +++ b/tests/e2e/router/test_complexity_router_e2e.py @@ -38,6 +38,7 @@ HEURISTIC_TIER_MODELS = frozenset({"openai/gpt-5.5", "gpt-5.5"}) LLM_TIER_MODELS = frozenset({"anthropic/claude-haiku-4-5", "claude-haiku-4-5"}) +@pytest.mark.usefixtures("_ensure_complexity_smart_router") class TestComplexityRouterLlmClassifier: @pytest.mark.skip( reason="product bug LIT-4521: LLM classifier returns SIMPLE for short hard prompts " diff --git a/tests/e2e/router/test_reliability_cache_e2e.py b/tests/e2e/router/test_reliability_cache_e2e.py new file mode 100644 index 00000000000..78d8fcdc08f --- /dev/null +++ b/tests/e2e/router/test_reliability_cache_e2e.py @@ -0,0 +1,37 @@ +"""Live e2e: the response cache returns a cached answer on an exact repeat. + +The same unique prompt is sent twice to the real `gpt-5.5` deployment under the +same key: the first call is a cache miss (the proxy computes and stores the entry, +and returns no x-litellm-cache-key), the second is an exact hit (the proxy serves +from cache and returns x-litellm-cache-key). This relies on the standard Redis +response cache being enabled on the proxy under test. +""" + +from __future__ import annotations + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from reliability_support import chat_override + +pytestmark = pytest.mark.e2e + + +class TestReliabilityCache: + @pytest.mark.covers("reliability.cache.exact.returns_cached") + def test_exact_cache_returns_cached(self, client: ComplexityRouterClient, scoped_key: str) -> None: + prompt = f"cache probe {unique_marker()}" + + first = chat_override(client.proxy, scoped_key, "gpt-5.5", prompt) + assert first.status_code == 200, f"first call should succeed, got {first.status_code}: {first.body[:300]}" + assert "x-litellm-cache-key" not in first.headers, ( + "first (uncached) call must not report a cache-key header" + ) + + second = chat_override(client.proxy, scoped_key, "gpt-5.5", prompt) + assert second.status_code == 200, f"second call should succeed, got {second.status_code}: {second.body[:300]}" + assert "x-litellm-cache-key" in second.headers, ( + "second identical call should hit the response cache and report a cache-key header " + "(requires the proxy's Redis response cache to be enabled)" + ) diff --git a/tests/e2e/router/test_reliability_fallbacks_e2e.py b/tests/e2e/router/test_reliability_fallbacks_e2e.py new file mode 100644 index 00000000000..5b7d21c6ef7 --- /dev/null +++ b/tests/e2e/router/test_reliability_fallbacks_e2e.py @@ -0,0 +1,69 @@ +"""Live e2e: per-request fallbacks reroute a failing deployment's traffic to a +healthy one. + +Each test registers a primary deployment that fails (an unreachable base URL, or +a 1ms deadline) and calls it with a `router_settings_override` mapping it to the +real `gpt-5.5`. The proof the fallback fired is twofold: the response is a real +completion from `gpt-5.5` (a non-empty content string), and the proxy reports at +least one attempted fallback in the x-litellm-attempted-fallbacks header. +""" + +from __future__ import annotations + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager +from models import RouterSettingsOverride +from reliability_support import ( + chat_override, + content_of, + create_bad_base_deployment, + create_timeout_deployment, +) + +pytestmark = pytest.mark.e2e + + +def _assert_served_by_fallback(resp: StreamingResponse) -> None: + assert resp.status_code == 200, f"expected 200 after fallback, got {resp.status_code}: {resp.body[:300]}" + content = content_of(resp) + assert isinstance(content, str) and content, ( + f"the gpt-5.5 fallback should have returned a real completion, got content {content!r} " + f"(body={resp.body[:300]})" + ) + attempted = resp.headers.get("x-litellm-attempted-fallbacks") + assert attempted is not None, "response is missing the x-litellm-attempted-fallbacks header" + assert int(attempted) >= 1, f"x-litellm-attempted-fallbacks should be >= 1, got {attempted!r}" + + +class TestReliabilityFallbacks: + @pytest.mark.covers("reliability.fallback.5xx.routes_to_fallback") + def test_5xx_routes_to_fallback( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + primary = f"reliability-fail-{unique_marker()}" + model_id = create_bad_base_deployment(client.proxy, primary) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + resp = chat_override( + client.proxy, scoped_key, primary, "say hi", + override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]), + ) + _assert_served_by_fallback(resp) + + @pytest.mark.covers("reliability.fallback.timeout.routes_to_fallback") + def test_timeout_routes_to_fallback( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + primary = f"reliability-tofail-{unique_marker()}" + model_id = create_timeout_deployment(client.proxy, primary) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + resp = chat_override( + client.proxy, scoped_key, primary, "say hi", + override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]), + ) + _assert_served_by_fallback(resp) diff --git a/tests/e2e/router/test_reliability_timeouts_e2e.py b/tests/e2e/router/test_reliability_timeouts_e2e.py new file mode 100644 index 00000000000..f24d5139e66 --- /dev/null +++ b/tests/e2e/router/test_reliability_timeouts_e2e.py @@ -0,0 +1,53 @@ +"""Live e2e: a per-request timeout surfaces to the caller instead of hanging. + +A deployment created with a 1ms deadline always exceeds it against the real +backend. With no fallback in play, the proxy must return the timeout to the +caller: a 408 for a non-streamed request, and the same timeout surfaced on the +streamed path (either a 408 before the stream opens or a timeout error carried in +the response). +""" + +from __future__ import annotations + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from lifecycle import ResourceManager +from reliability_support import chat_override, create_timeout_deployment + +pytestmark = pytest.mark.e2e + + +class TestReliabilityTimeouts: + @pytest.mark.covers("reliability.timeout.request_timeout.exceeds_deadline") + def test_request_timeout_exceeds_deadline( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"reliability-timeout-{unique_marker()}" + model_id = create_timeout_deployment(client.proxy, name) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + resp = chat_override(client.proxy, scoped_key, name, "hello") + assert resp.status_code == 408, ( + f"a timed-out request should return 408, got {resp.status_code}: {resp.body[:300]}" + ) + assert "timeout" in resp.body.lower(), f"the 408 body should name the timeout, got: {resp.body[:300]}" + + @pytest.mark.covers("reliability.timeout.stream_timeout.exceeds_deadline") + def test_stream_timeout_exceeds_deadline( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"reliability-stream-timeout-{unique_marker()}" + model_id = create_timeout_deployment(client.proxy, name) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + resp = chat_override(client.proxy, scoped_key, name, "hello", stream=True) + surfaced = f"{resp.body} {resp.stream_error or ''}".lower() + assert resp.status_code >= 400, ( + f"a timed-out streaming request should surface an error status, got {resp.status_code}: {resp.body[:300]}" + ) + assert "timeout" in surfaced, ( + f"the streamed timeout error should name the timeout, got body={resp.body[:300]}, " + f"stream_error={resp.stream_error!r}" + ) From 28f012bb52d1dd374bb2951fa7c80cbff298b1ea Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 20 Jul 2026 16:15:55 -0700 Subject: [PATCH 59/60] test(true_rabbit): cover passthrough headers, batch assume-role, gemini, vllm, bedrock guardrails, batch rate-limit mapping (#33843) * test(e2e): cover passthrough headers, batch assume-role, gemini, vllm, bedrock guardrails, batch rate-limit mapping Add parent-package e2e suites for the six feature gaps: pass-through header forwarding via /config/pass_through_endpoint, Bedrock batch STS assume-role, Gemini chat + files, hosted_vllm batch/files, Bedrock guardrail pre_call blocks (plus restored content-filter team opt-out), and OpenAI batch RPM 429 body mapping. Registry cells and LiteLLMParamsBody/TeamMetadata fields updated so markers collect cleanly. * test(e2e): cover LIT-4587 gaps for redis, responses, tpm cache, apply_guardrail, langfuse Adds customer-shaped live e2e for apply_guardrail, responses store+metadata TTL, TPM excluding cached tokens, redis-backed RPM, redis circuit-breaker path, Langfuse spend, Cohere chat, virtual-key auth, file content download, hosted_vllm chat, and Nova Sonic realtime. Registry cells updated for the new markers. * test(e2e): drive LIT-4587 gap suites on Anthropic to avoid Gemini quota flakes Redis RPM, circuit-breaker path, virtual-key auth, responses metadata, and Langfuse driver models now use Anthropic haiku so local runs stay green when Gemini daily quota is exhausted. * test(e2e): drop Langfuse spend suite; feature is being deprecated Remove test_langfuse_e2e.py, logging.langfuse registry cells, and the langfuse-only conftest driver/credentials fixtures. * test(e2e): fold provider/batch feature tests into their endpoint suites Keep the e2e layout endpoint- and suite-scoped instead of one file per provider or feature Move the virtual-key auth case into access_control/test_access_control_e2e.py as TestVirtualKeyAuth (replacing an incomplete stub) and drop the standalone test_virtual_key_auth_e2e.py Fold the five per-file batch suites (file content, RPM 429 mapping, Bedrock assume-role, Gemini files, hosted_vllm batch) into batches/test_batches_e2e.py. The hosted_vllm batch case is skipped for now since it needs a live vLLM server (HOSTED_VLLM_API_BASE) the e2e environment does not provision; it and the gemini-files and RPM-mapping cases reference LIT-3382 / LIT-3266 where relevant Merge the cohere, gemini and hosted_vllm chat cases into llm_translation/test_chat_completions_regression_e2e.py so /chat/completions coverage lives in one endpoint file, and repoint the coverage_registry source fields to the new homes Move the shared CacheControl / TextBlock / RichMessage request blocks into the root models.py (re-exported from endpoints_client) so quota_management can use them without a cross-suite import, which also clears the basedpyright errors in test_tpm_excludes_cached_tokens_e2e.py; type the httpbin echo body in test_passthrough_headers_e2e.py with a pydantic model to drop the Any-typed json.loads path * test(e2e): address review feedback and re-home virtual-key coverage Replace the tautological Bedrock assume-role batch id assertion (`startswith(...) or batch.id`, always true) with a managed-id shape check, since the unified target_model_names path re-encodes the id rather than returning a raw ARN Raise the batch RPM-mapping test's rpm_limit above one so the file upload can no longer consume the key's sole request unit before batch create runs; the batch create then clears the generic per-request limiter and the batch limiter is what returns the "Batch rate limit exceeded" body the assertions check Set exercised_on to [] on the pass-through header test; it drives a pass-through endpoint, not /chat/completions Move the virtual-key valid_allows / invalid_denied cells from other.yaml to mgmt.yaml as mgmt.virtual_key.* so TestVirtualKeyAuth rolls up under Management, and point its covers marker at the new ids --- .../access_control/test_access_control_e2e.py | 61 ++++ tests/e2e/batches/test_batches_e2e.py | 332 +++++++++++++++++- tests/e2e/coverage_registry/guardrail.yaml | 4 + .../coverage_registry/llm_conversational.yaml | 5 + .../llm_nonconversational.yaml | 6 + tests/e2e/coverage_registry/mgmt.yaml | 2 + tests/e2e/coverage_registry/other.yaml | 2 + .../coverage_registry/quota_management.yaml | 3 + tests/e2e/coverage_registry/schema.py | 3 + tests/e2e/e2e_config.py | 16 + tests/e2e/e2e_http.py | 2 + tests/e2e/guardrails/conftest.py | 18 + tests/e2e/guardrails/guardrails_client.py | 211 +++++++++++ .../guardrails/test_apply_guardrail_e2e.py | 62 ++++ .../guardrails/test_bedrock_guardrail_e2e.py | 70 ++++ .../test_team_disable_global_guardrail_e2e.py | 81 +++++ tests/e2e/llm_translation/endpoints_client.py | 23 +- .../realtime/test_nova_sonic_realtime_e2e.py | 79 +++++ .../test_chat_completions_regression_e2e.py | 167 ++++++++- .../test_passthrough_headers_e2e.py | 150 ++++++++ .../test_responses_metadata_e2e.py | 122 +++++++ tests/e2e/logging/conftest.py | 13 +- tests/e2e/models.py | 24 ++ .../test_redis_backed_ratelimit_e2e.py | 76 ++++ .../test_redis_circuit_breaker_e2e.py | 90 +++++ .../test_tpm_excludes_cached_tokens_e2e.py | 162 +++++++++ tests/e2e/transport.py | 33 +- 27 files changed, 1777 insertions(+), 40 deletions(-) create mode 100644 tests/e2e/guardrails/conftest.py create mode 100644 tests/e2e/guardrails/guardrails_client.py create mode 100644 tests/e2e/guardrails/test_apply_guardrail_e2e.py create mode 100644 tests/e2e/guardrails/test_bedrock_guardrail_e2e.py create mode 100644 tests/e2e/guardrails/test_team_disable_global_guardrail_e2e.py create mode 100644 tests/e2e/llm_translation/realtime/test_nova_sonic_realtime_e2e.py create mode 100644 tests/e2e/llm_translation/test_passthrough_headers_e2e.py create mode 100644 tests/e2e/llm_translation/test_responses_metadata_e2e.py create mode 100644 tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py create mode 100644 tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py create mode 100644 tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py diff --git a/tests/e2e/access_control/test_access_control_e2e.py b/tests/e2e/access_control/test_access_control_e2e.py index ce649fa2400..e24b721d831 100644 --- a/tests/e2e/access_control/test_access_control_e2e.py +++ b/tests/e2e/access_control/test_access_control_e2e.py @@ -23,12 +23,16 @@ from access_control_client import ( ROUTE_NOT_ALLOWED_MARKER, ) from e2e_config import unique_marker +from e2e_http import Success, UnauthorizedError, UnknownApiError, unwrap from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, LiteLLMParamsBody +from proxy_client import ProxyClient pytestmark = pytest.mark.e2e ALLOWED_MODEL = "gemini-2.5-flash" DISALLOWED_MODEL = "gpt-5.5" +VIRTUAL_KEY_BACKEND = "anthropic/claude-haiku-4-5-20251001" def _is_json(body: str) -> bool: @@ -39,6 +43,7 @@ def _is_json(body: str) -> bool: return False + class TestAccessControl: def test_disallowed_model_is_denied_403( self, client: AccessControlClient, resources: ResourceManager @@ -81,3 +86,59 @@ class TestAccessControl: f"{result.status_code}: {result.body[:300]}" ) assert _is_json(result.body), f"400 body must be valid JSON: {result.body[:300]}" + + +class TestVirtualKeyAuth: + """Virtual-key auth the way OpenAI-compatible clients send it: a real key + must reach chat, a forged bearer must be rejected before the provider.""" + + @pytest.mark.covers( + "mgmt.virtual_key.valid_allows", + "mgmt.virtual_key.invalid_denied", + exercised_on=[], + ) + def test_valid_key_allows_and_invalid_key_denied( + self, proxy: ProxyClient, resources: ResourceManager + ) -> None: + model = f"e2e-auth-chat-{unique_marker()}" + model_id = proxy.create_model( + model, + LiteLLMParamsBody(model=VIRTUAL_KEY_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY"), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + key = resources.key() + + ok = unwrap( + proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=f"Reply with one word. {unique_marker()}", + ) + ], + max_tokens=16, + ), + ) + ) + assert ok.choices, f"valid key must complete chat: {ok}" + + bad = proxy.chat( + "sk-e2e-forged-not-a-real-key", + ChatBody( + model=model, + messages=[ChatMessage(role="user", content="should not run")], + max_tokens=8, + ), + ) + match bad: + case UnauthorizedError(): + return + case UnknownApiError(status_code=status) if status in (401, 403): + return + case Success(): + pytest.fail("forged bearer must not reach a successful completion") + case _: + pytest.fail(f"forged bearer must be auth-denied, got {bad}") diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 8f10c8c7c2a..f886c09b705 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -16,13 +16,14 @@ misroute to the wrong provider fails the create. from __future__ import annotations import json +import os import time from datetime import datetime, timedelta, timezone from typing import Callable import pytest -from e2e_config import unique_marker +from e2e_config import require_env, unique_marker from batch_client import ( BatchClient, @@ -39,7 +40,9 @@ from capabilities import ( FILE_ID_SHAPE, OPENAI_BATCH_MODEL, Capability, + batch_model_name, coverage_cells_for_lifecycle, + is_managed_id, matches_id_shape, raw_id_matches_provider, ) @@ -53,7 +56,7 @@ from e2e_http import ( unwrap, ) from lifecycle import ResourceManager -from models import KeyGenerateBody, SpendLogRow +from models import KeyGenerateBody, LiteLLMParamsBody, SpendLogRow pytestmark = pytest.mark.e2e @@ -457,3 +460,328 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( "batch create on a rate-limited key left an unattributed spend row " f"(LIT-3266); rows={[(r.request_id, r.call_type, r.model) for r in new_orphans]}" ) + + +OPENAI_FILE_CONTENT_BACKEND = "gpt-4o-mini" + + +class TestBatchFileContent: + """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes.""" + + @pytest.mark.covers( + "llm.files.openai.content.nonstream.works", + exercised_on=["files"], + ) + def test_file_content_matches_upload( + self, client: BatchClient, resources: ResourceManager + ) -> None: + proxy_name = f"e2e-file-content-{unique_marker()}" + model_id = client.create_model( + proxy_name, + LiteLLMParamsBody( + model=f"openai/{OPENAI_FILE_CONTENT_BACKEND}", + api_key="os.environ/OPENAI_API_KEY", + ), + ) + resources.defer(lambda: client.delete_model(model_id)) + key = resources.key() + + payload = render_jsonl(OPENAI_FILE_CONTENT_BACKEND) + file = unwrap( + client.upload_file( + content=payload, + form=FileUploadForm(purpose="batch", target_model_names=proxy_name), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert file.id + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"file content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}" + ) + expected = payload.decode().rstrip("\n") + got = downloaded.body.rstrip("\n") + assert got == expected, ( + "downloaded file content must match the uploaded JSONL bytes" + ) + + +BATCH_RL_REQUEST_LINES = 3 +BATCH_RL_RPM_LIMIT = 2 + + +def _multi_request_jsonl(model: str, n: int) -> bytes: + lines = tuple( + json.dumps( + { + "custom_id": f"req-{i}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [{"role": "user", "content": "ping"}], + "max_tokens": 8, + }, + } + ) + for i in range(n) + ) + return ("\n".join(lines) + "\n").encode() + + +class TestBatchRateLimitErrorMapping: + """Batch create that exceeds a key's RPM maps to a structured 429. + + The batch rate limiter reads the input file at submission time and rejects + the create when the file's request count would exceed the key's remaining + RPM. The product promise is not only the block itself but the + OpenAI-compatible shape: HTTP 429, a body that names the batch rate limit, + and pacing headers so clients can back off. Complements the LIT-3266 hygiene + check (no orphan spend rows) by asserting the error mapping when the limiter + actually fires. + """ + + @pytest.mark.covers( + "quota_management.ratelimit.batch_rpm.blocks_over_limit", + exercised_on=["batches"], + ) + def test_batch_create_over_rpm_returns_mapped_429( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + user_id = f"e2e-batch-rl-map-{unique_marker()}" + key = client.proxy.generate_key( + KeyGenerateBody( + models=[], rpm_limit=BATCH_RL_RPM_LIMIT, tpm_limit=1_000_000, user_id=user_id + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + file = unwrap( + client.upload_file( + content=_multi_request_jsonl("gpt-4o-mini", BATCH_RL_REQUEST_LINES), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + + assert created.status_code == 429, ( + f"expected batch RPM 429 when file has {BATCH_RL_REQUEST_LINES} requests and " + f"rpm_limit={BATCH_RL_RPM_LIMIT}, got {created.status_code}: {created.body[:400]}" + ) + body_lower = created.body.lower() + assert "batch rate limit exceeded" in body_lower, ( + f"429 body must name the batch rate limit so clients can branch on it; " + f"got: {created.body[:400]}" + ) + assert str(BATCH_RL_REQUEST_LINES) in created.body, ( + f"429 body should report the batch request count ({BATCH_RL_REQUEST_LINES}); " + f"got: {created.body[:400]}" + ) + assert "rpm" in body_lower or "requests remaining" in body_lower, ( + f"429 body must describe the RPM budget remaining so clients can pace; " + f"got: {created.body[:400]}" + ) + retry_after = created.headers.get("retry-after") + if retry_after is not None: + assert retry_after.isdigit() and int(retry_after) > 0, ( + f"retry-after must be a positive integer when present, got {retry_after!r}" + ) + + +ASSUME_ROLE_RAW_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + + +def _assume_role_params(role_arn: str, session_name: str) -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=ASSUME_ROLE_RAW_MODEL, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + s3_region_name="os.environ/AWS_REGION", + s3_bucket_name="os.environ/AWS_BATCH_S3_BUCKET", + s3_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + s3_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_batch_role_arn="os.environ/AWS_BATCH_ROLE_ARN", + aws_role_name=role_arn, + aws_session_name=session_name, + ) + + +class TestBedrockBatchAssumeRole: + """Bedrock batch create under STS assume-role credentials. + + Provisions a bedrock batch deployment whose litellm_params carry + aws_role_name / aws_session_name (the product path for role assumption) and + runs the unified file-upload + batch-create lifecycle. Success means the + proxy assumed the role and Bedrock accepted the job; a misconfigured role + fails create with an AWS auth error rather than silently falling back to the + ambient key. + """ + + @pytest.mark.covers( + "llm.batches.bedrock.assume_role.nonstream.works", + "llm.files.bedrock.upload.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_unified_batch_create_with_assume_role( + self, client: BatchClient, resources: ResourceManager + ) -> None: + (role_arn,) = require_env("AWS_ROLE_NAME") + require_env( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_REGION", + "AWS_BATCH_S3_BUCKET", + "AWS_BATCH_ROLE_ARN", + ) + session_name = f"e2e-batch-sts-{unique_marker()}"[:64] + model_name = batch_model_name("bedrock-sts-batch") + + model_id = client.create_model(model_name, _assume_role_params(role_arn, session_name)) + resources.defer(lambda: client.delete_model(model_id)) + key = resources.key() + + file = unwrap( + client.upload_file( + content=render_jsonl(ASSUME_ROLE_RAW_MODEL), + form=FileUploadForm(purpose="batch", target_model_names=model_name), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert_file_object(file, provider="bedrock") + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert batch.id, f"assume-role create returned no batch id: {created.body[:200]}" + assert is_managed_id(batch.id), ( + f"assume-role create via target_model_names must return a managed batch id, " + f"got {batch.id!r}" + ) + assert batch.status in CREATED_BATCH_STATUSES, ( + f"assume-role batch has non-transitional status {batch.status!r}" + ) + assert_batch_object(batch) + + fetched = unwrap(client.retrieve_batch(batch.id, key=key)) + assert fetched.id == batch.id + + +GEMINI_FILES_RAW_MODEL = "gemini-2.5-flash" + + +class TestGeminiFiles: + """Gemini Files API upload through the proxy (LIT-3382). + + gemini is a first-class FileCreateProvider. The test registers a gemini + deployment, uploads a tiny batch-purpose JSONL with target_model_names + routing, and asserts a FileObject comes back. Batch create for pure gemini + (non-Vertex) is out of scope here; Vertex covers the Gemini batch job path in + the main lifecycle matrix. + """ + + @pytest.mark.covers( + "llm.files.gemini.upload.nonstream.works", + exercised_on=["files"], + ) + def test_gemini_file_upload( + self, client: BatchClient, resources: ResourceManager + ) -> None: + model_name = batch_model_name("gemini-files") + model_id = client.create_model( + model_name, + LiteLLMParamsBody( + model=f"gemini/{GEMINI_FILES_RAW_MODEL}", + api_key="os.environ/GEMINI_API_KEY", + ), + ) + resources.defer(lambda: client.delete_model(model_id)) + key = resources.key() + + file = unwrap( + client.upload_file( + content=render_jsonl(GEMINI_FILES_RAW_MODEL), + form=FileUploadForm(purpose="batch", target_model_names=model_name), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert_file_object(file, provider="gemini") + assert file.id, "gemini file upload returned no id" + + +def _vllm_params(api_base: str, api_key: str | None, model_id: str) -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=f"hosted_vllm/{model_id}", + api_base=api_base, + api_key=api_key, + ) + + +class TestHostedVllmBatch: + """hosted_vllm file upload + batch create (OpenAI-compatible path, LIT-3266). + + hosted_vllm is in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, so /v1/files + and /v1/batches route through the OpenAI handler against the deployment's + api_base. Skipped for now: it needs a live vLLM (or OpenAI-compatible) server + exposing the files/batches APIs (HOSTED_VLLM_API_BASE), which the e2e + environment does not currently provision. + """ + + @pytest.mark.skip( + reason="hosted_vllm batch/files needs a live vLLM server (HOSTED_VLLM_API_BASE) " + "not provisioned in the e2e environment; re-enable when available (LIT-3266)" + ) + @pytest.mark.covers( + "llm.batches.hosted_vllm.basic.nonstream.works", + "llm.files.hosted_vllm.upload.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_unified_file_and_batch_create( + self, client: BatchClient, resources: ResourceManager + ) -> None: + (api_base,) = require_env("HOSTED_VLLM_API_BASE") + api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None + model_id = ( + os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct" + ).strip() + proxy_name = batch_model_name("hosted-vllm-batch") + + model_row_id = client.create_model( + proxy_name, _vllm_params(api_base, api_key, model_id) + ) + resources.defer(lambda: client.delete_model(model_row_id)) + key = resources.key() + + file = unwrap( + client.upload_file( + content=render_jsonl(model_id), + form=FileUploadForm(purpose="batch", target_model_names=proxy_name), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert_file_object(file, provider="hosted_vllm") + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert batch.id, f"hosted_vllm create returned no batch id: {created.body[:200]}" + assert batch.status in CREATED_BATCH_STATUSES, ( + f"hosted_vllm batch has non-transitional status {batch.status!r}" + ) + assert_batch_object(batch) diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index 792cbaaff7c..68722fbbb96 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -4,6 +4,10 @@ - {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"} - {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"} - {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"} +- {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"} +- {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"} +- {id: guardrail.litellm_content_filter.apply_endpoint.blocks, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail blocks banned content for customers that call the apply surface directly"} +- {id: guardrail.litellm_content_filter.apply_endpoint.allows, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [allows], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail returns clean text for allowed input"} - {id: guardrail.bedrock.during.blocks, module: guardrail, tier: P0, hook_point: during, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "During-call moderation for streaming"} - {id: guardrail.bedrock.post_call.blocks, module: guardrail, tier: P0, hook_point: post_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "Block harmful output"} - {id: guardrail.lakera.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/lakera_ai_v2.py", rationale: "Prompt-injection block pre-execution"} diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index be8a291c6fc..878eae984d1 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -24,6 +24,11 @@ - {id: llm.chat_completions.bedrock_converse.prompt_cache_5m.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Anthropic-on-Bedrock caching"} - {id: llm.chat_completions.bedrock_converse.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Anthropic thinking on Bedrock"} - {id: llm.chat_completions.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "P0 route; Vertex AI"} +- {id: llm.chat_completions.gemini.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: gemini, capability: basic, streaming: nonstream, assertions: [works], source: "test_chat_completions_regression_e2e.py", rationale: "Gemini OpenAI-compatible chat translation"} +- {id: llm.chat_completions.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: chat_completions, route: gemini, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "test_chat_completions_regression_e2e.py", rationale: "Gemini chat cost lands in SpendLogs"} +- {id: llm.chat_completions.hosted_vllm.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_chat_completions_regression_e2e.py", rationale: "OpenAI-compatible hosted_vllm chat is a confirmed self-hosted backend path"} +- {id: llm.chat_completions.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_chat_completions_regression_e2e.py", rationale: "Cohere chat via OpenAI-compatible /chat/completions"} + - {id: llm.chat_completions.vertex.basic.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:8455", rationale: "Streaming over Vertex"} - {id: llm.chat_completions.vertex.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vertex Gemini function_calling"} - {id: llm.chat_completions.vertex.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Gemini vision"} diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index b01b219476d..2b456aacefc 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -18,6 +18,8 @@ - {id: llm.batches.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Azure batches all scenarios"} - {id: llm.batches.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Vertex batches"} - {id: llm.batches.bedrock.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Bedrock batches (encoded/unified only)"} +- {id: llm.batches.bedrock.assume_role.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: assume_role, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create under STS assume-role credentials"} +- {id: llm.batches.hosted_vllm.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible batch create"} - {id: llm.batches.openai.key_model_access_denied.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Key model restriction 403 on upload/create"} - {id: llm.files.openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "openai_files_endpoints/files_endpoints.py:46", rationale: "File upload returns OpenAIFileObject"} - {id: llm.files.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File retrieve by id"} @@ -26,7 +28,11 @@ - {id: llm.files.azure_openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:45", rationale: "Azure file upload managed backend"} - {id: llm.files.vertex.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:52", rationale: "Vertex file upload to GCS"} - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"} +- {id: llm.files.gemini.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: gemini, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Gemini Files API upload via proxy"} +- {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} - {id: llm.rerank.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_rerank_e2e.py:29", rationale: "Cohere rerank, top_n + relevance_score"} +- {id: llm.files.openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files/{id}/content returns uploaded batch JSONL bytes"} +- {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_nova_sonic_realtime_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} - {id: llm.rerank.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "llms/bedrock/rerank/handler.py", rationale: "Bedrock rerank"} - {id: llm.rerank.together_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: together_ai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/together_ai/rerank/handler.py", rationale: "Together rerank"} - {id: llm.images_generations.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_image_generation_e2e.py:22", rationale: "OpenAI image gen, b64/url"} diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index a43e7103523..da4652460f1 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -8,6 +8,8 @@ - {id: mgmt.key.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3122", rationale: "Deletion revokes future calls"} - {id: mgmt.key.delete.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:3122", rationale: "Non-owner cannot delete"} - {id: mgmt.key.info.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3380", rationale: "Info reflects all writes"} +- {id: mgmt.virtual_key.valid_allows, module: mgmt, tier: P0, surface: api, assertions: [valid_allows], source: "user_api_key_auth.py", rationale: "Virtual key authenticates chat the way production OpenAI clients do"} +- {id: mgmt.virtual_key.invalid_denied, module: mgmt, tier: P0, surface: api, assertions: [invalid_denied], source: "user_api_key_auth.py", rationale: "Bogus bearer is rejected before provider call"} - {id: mgmt.team.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "team_endpoints.py:897", rationale: "team_id/alias/budgets stored"} - {id: mgmt.team.new.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "team_endpoints.py:897", rationale: "Only org-admin/master creates teams"} - {id: mgmt.team.member_add.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "team_endpoints.py:2424", rationale: "Membership + per-member budget persist"} diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index c2efecec677..6b183cbf9f3 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -2,6 +2,7 @@ # PROMOTION NOTE: the auth cluster (~14 cells) is a candidate to promote to its own module once stable. - {id: other.auth.master_key.valid_allows, module: other, tier: P0, area: auth, assertions: [valid_allows], source: "user_api_key_auth.py:1569-1588", rationale: "Master key authenticates; timing-safe compare"} - {id: other.auth.master_key.invalid_denied, module: other, tier: P0, area: auth, assertions: [invalid_denied], source: "user_api_key_auth.py:1580", rationale: "Invalid master key rejected"} +- {id: other.config.responses.metadata_redis_ttl_bounded, module: other, tier: P0, area: config, assertions: [ttl_bounded], source: "responses + redis cache", rationale: "Responses store+metadata must not leave TTL-unbounded Redis entries (LIT-1201)"} - {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:77-150", rationale: "Valid JWT with correct issuer + claims grants access"} - {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:125-135", rationale: "Expired JWT rejected even with valid signature"} - {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:145-150", rationale: "Bad/missing signature fails verification"} @@ -21,6 +22,7 @@ - {id: other.lifecycle.startup.env_vars_resolved, module: other, tier: P1, area: lifecycle, assertions: [env_vars_resolved], source: "proxy_server.py:3984-4010", rationale: "os.environ/ refs resolved at startup"} - {id: other.lifecycle.background_health_check.interval_configurable, module: other, tier: P1, area: lifecycle, assertions: [interval_configurable], source: "proxy_server.py:3245-3310", rationale: "Background checks run at configurable interval"} - {id: other.config.runtime_update.applies_at_runtime, module: other, tier: P0, area: config, assertions: [applies_at_runtime], source: "proxy_server.py:14014-14060", rationale: "/config/update persists to DB + invalidates cache"} +- {id: other.config.passthrough.headers_forwarded, module: other, tier: P0, area: config, assertions: [headers_forwarded], source: "passthrough/utils.py forward_headers_from_request", rationale: "Custom pass-through static headers and x-pass-* client headers reach the upstream"} - {id: other.config.general_settings.alert_webhook_side_effect, module: other, tier: P1, area: config, assertions: [alert_webhook_side_effect], source: "proxy_server.py:14215", rationale: "alert_to_webhook_url auto-enables slack alerting"} - {id: other.config.secret_resolution.kms_integration, module: other, tier: P1, area: config, assertions: [kms_integration], source: "proxy_server.py:3984-4010", rationale: "Resolves secrets from Vault/KMS at startup"} - {id: other.config.overrides.audit_logged, module: other, tier: P1, area: config, assertions: [audit_logged], source: "config_override_endpoints.py:67-100", rationale: "Config override mutations audit-logged, values redacted"} diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 7e71f2bd3d1..a8d0749cd8d 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -1,7 +1,10 @@ # Quota Management (behavior features): rate limits, budgets, spend tracking. Grounded in # litellm/proxy/hooks/ + litellm/proxy/auth/auth_checks.py + litellm/proxy/spend_tracking/. - {id: quota_management.ratelimit.rpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: rpm, assertions: [blocks_over_limit], exercised_on: [chat_completions, messages], source: "parallel_request_limiter_v3.py", rationale: "v3 limiter enforces RPM per key/team/model; 429 on breach"} +- {id: quota_management.ratelimit.batch_rpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_rpm, assertions: [blocks_over_limit], exercised_on: [batches], source: "batch_rate_limiter.py", rationale: "Batch create that exceeds key RPM returns mapped 429 with retry-after"} - {id: quota_management.ratelimit.tpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: tpm, assertions: [blocks_over_limit], exercised_on: [chat_completions, messages], source: "parallel_request_limiter_v3.py", rationale: "v3 limiter enforces TPM per key/team/model; 429 on breach"} +- {id: quota_management.ratelimit.tpm.excludes_cached_tokens, module: quota_management, tier: P0, behavior: ratelimit, variant: tpm, assertions: [excludes_cached_tokens], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py:_get_total_tokens_from_usage", rationale: "Cached prompt tokens must not count toward TPM (LIT-1930)"} +- {id: quota_management.ratelimit.redis_backed.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: redis_backed, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py", rationale: "With Redis configured, RPM still enforces 429 across the shared limiter path customers run multi-replica"} - {id: quota_management.ratelimit.rpm.resets_after_window, module: quota_management, tier: P1, behavior: ratelimit, variant: rpm, assertions: [resets_after_window], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py", rationale: "Rate-limit window (LITELLM_RATE_LIMIT_WINDOW_SIZE, 60s default) expires; a blocked key serves again in the next window"} - {id: quota_management.ratelimit.rpm.headers_report_remaining, module: quota_management, tier: P1, behavior: ratelimit, variant: rpm, assertions: [headers_report_remaining], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py async_post_call_success_hook", rationale: "Successful responses carry x-ratelimit-api_key-{limit,remaining}-{requests,tokens} so clients can pace"} - {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"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index a76774c3bde..1a6dc111e2b 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -47,12 +47,15 @@ LlmRoute = Literal[ "bedrock_converse", "bedrock_invoke", "cohere", + "gemini", + "hosted_vllm", "openai", "together_ai", "vertex", ] LlmCapability = Literal[ + "assume_role", "basic", "count_tokens", "long_context_1m", diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 2687888ea42..3be339d28a0 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -79,6 +79,22 @@ LOAD_MIN_RPS = float(os.environ.get("E2E_LOAD_MIN_RPS", "355")) LOAD_MAX_FAILURE_RATIO = float(os.environ.get("E2E_LOAD_MAX_FAILURE_RATIO", "0.01")) +def require_env(*names: str) -> tuple[str, ...]: + """Return the non-empty values for each env name, or hard-fail naming which are missing. + + Live e2e never skips for missing credentials: a missing key is a red run so + ops knows the suite cannot prove the product path. + """ + missing = tuple(name for name in names if not (os.environ.get(name) or "").strip()) + if missing: + joined = ", ".join(missing) + raise AssertionError( + f"missing required env for e2e: {joined}. " + "Add them to tests/e2e/.env locally and to litellm ops for stage/CI." + ) + return tuple((os.environ.get(name) or "").strip() for name in names) + + def datadog_mcp_url(*, toolsets: str = "core") -> str: """Regional Datadog remote MCP endpoint for this process's DD_SITE. diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index ce801ef81fb..1c6048a3688 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -250,6 +250,7 @@ def delete[R: BaseModel]( headers: BaseModel, json: BaseModel, response_type: type[R], + params: BaseModel | None = None, timeout: float = 30.0, ) -> Result[R]: try: @@ -257,6 +258,7 @@ def delete[R: BaseModel]( str(url), headers=_headers(headers), json=json.model_dump(by_alias=True, exclude_none=True), + params=_params(params), timeout=timeout, ) except requests.RequestException as exc: diff --git a/tests/e2e/guardrails/conftest.py b/tests/e2e/guardrails/conftest.py new file mode 100644 index 00000000000..9e85d475065 --- /dev/null +++ b/tests/e2e/guardrails/conftest.py @@ -0,0 +1,18 @@ +"""Guardrails suite's `client` fixture. + +Shared lifecycle (resources/scoped_key), proxy liveness, and e2e/covers markers +live in the parent tests/e2e/conftest.py. GuardrailsClient holds the shared +ProxyClient so keys and deferred cleanups tear down correctly. +""" + +from __future__ import annotations + +import pytest + +from guardrails_client import GuardrailsClient, build_client +from proxy_client import ProxyClient + + +@pytest.fixture(scope="session") +def client(proxy: ProxyClient) -> GuardrailsClient: + return build_client(proxy) diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py new file mode 100644 index 00000000000..d24cd36c2fd --- /dev/null +++ b/tests/e2e/guardrails/guardrails_client.py @@ -0,0 +1,211 @@ +"""Client for the guardrails e2e suite: register global (default-on) guardrails +and chat through them on the shared ProxyClient so resources.defer cleans up. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Literal + +from pydantic import BaseModel + +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT +from e2e_http import NoBody, Result, Success, unwrap +from models import ( + ChatBody, + ChatMessage, + ChatResponse, + KeyGenerateBody, + TeamDeleteBody, + TeamInfoParams, + TeamInfoResponse, + TeamMetadata, + TeamNewBody, + TeamNewResponse, +) +from proxy_client import ProxyClient + +GuardrailMode = Literal["pre_call", "post_call", "during_call", "logging_only"] +BlockedWordAction = Literal["BLOCK", "MASK"] + + +class BlockedWordBody(BaseModel): + keyword: str + action: BlockedWordAction + + +class GuardrailParamsBase(BaseModel): + mode: GuardrailMode + default_on: bool + + +class ContentFilterParamsBody(GuardrailParamsBase): + guardrail: Literal["litellm_content_filter"] = "litellm_content_filter" + blocked_words: list[BlockedWordBody] + + +class BedrockGuardrailParamsBody(GuardrailParamsBase): + guardrail: Literal["bedrock"] = "bedrock" + guardrailIdentifier: str + guardrailVersion: str + aws_access_key_id: str | None = None + aws_secret_access_key: str | None = None + aws_region_name: str | None = None + + +GuardrailParamsBody = ContentFilterParamsBody | BedrockGuardrailParamsBody + + +class GuardrailSpecBody(BaseModel): + guardrail_name: str + litellm_params: GuardrailParamsBody + + +class GuardrailCreateBody(BaseModel): + guardrail: GuardrailSpecBody + + +class GuardrailCreateResponse(BaseModel): + guardrail_id: str + + +class ApplyGuardrailRequest(BaseModel): + guardrail_name: str + text: str + language: str | None = None + input_type: str = "request" + + +class ApplyGuardrailResponse(BaseModel): + response_text: str + + +@dataclass(frozen=True, slots=True) +class GuardrailsClient: + proxy: ProxyClient + + def create_content_filter_guardrail(self, name: str, blocked_keyword: str) -> str: + return unwrap( + self.proxy.transport.post( + "/guardrails", + headers=self.proxy.transport.master, + json=GuardrailCreateBody( + guardrail=GuardrailSpecBody( + guardrail_name=name, + litellm_params=ContentFilterParamsBody( + mode="pre_call", + default_on=True, + blocked_words=[ + BlockedWordBody(keyword=blocked_keyword, action="BLOCK") + ], + ), + ) + ), + response_type=GuardrailCreateResponse, + ) + ).guardrail_id + + def create_bedrock_guardrail( + self, + name: str, + *, + identifier: str, + version: str, + ) -> str: + return unwrap( + self.proxy.transport.post( + "/guardrails", + headers=self.proxy.transport.master, + json=GuardrailCreateBody( + guardrail=GuardrailSpecBody( + guardrail_name=name, + litellm_params=BedrockGuardrailParamsBody( + mode="pre_call", + default_on=True, + guardrailIdentifier=identifier, + guardrailVersion=version, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), + ) + ), + response_type=GuardrailCreateResponse, + ) + ).guardrail_id + + def delete_guardrail(self, guardrail_id: str) -> None: + _ = self.proxy.transport.delete( + f"/guardrails/{guardrail_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + + def create_team_opted_out_of_global_guardrails(self, alias: str) -> str: + team_id = unwrap( + self.proxy.transport.post( + "/team/new", + headers=self.proxy.transport.master, + json=TeamNewBody( + team_alias=alias, + metadata=TeamMetadata(disable_global_guardrails=True), + ), + response_type=TeamNewResponse, + ) + ).team_id + self._await_team(team_id) + return team_id + + def delete_team(self, team_id: str) -> None: + _ = self.proxy.transport.post( + "/team/delete", + headers=self.proxy.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + + def create_key_in_team(self, team_id: str) -> str: + return self.proxy.generate_key( + KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user") + ) + + def chat(self, key: str, model: str, text: str) -> Result[ChatResponse]: + return self.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=text)], + max_tokens=16, + ), + ) + + def apply_guardrail(self, key: str, *, name: str, text: str) -> Result[ApplyGuardrailResponse]: + return self.proxy.transport.post( + "/guardrails/apply_guardrail", + headers=self.proxy.transport.bearer(key), + json=ApplyGuardrailRequest(guardrail_name=name, text=text), + response_type=ApplyGuardrailResponse, + ) + + def _await_team(self, team_id: str) -> None: + deadline = time.monotonic() + POLL_TIMEOUT + last: Result[TeamInfoResponse] | None = None + while time.monotonic() < deadline: + last = self.proxy.transport.get( + "/team/info", + headers=self.proxy.transport.master, + params=TeamInfoParams(team_id=team_id), + response_type=TeamInfoResponse, + ) + if isinstance(last, Success): + return + time.sleep(POLL_INTERVAL) + raise AssertionError( + f"team {team_id!r} was created but /team/info never returned it: {last}" + ) + + +def build_client(proxy: ProxyClient) -> GuardrailsClient: + return GuardrailsClient(proxy=proxy) diff --git a/tests/e2e/guardrails/test_apply_guardrail_e2e.py b/tests/e2e/guardrails/test_apply_guardrail_e2e.py new file mode 100644 index 00000000000..ee691db22da --- /dev/null +++ b/tests/e2e/guardrails/test_apply_guardrail_e2e.py @@ -0,0 +1,62 @@ +"""Live e2e: POST /guardrails/apply_guardrail is the customer-facing apply surface. + +Customers call this endpoint to run a named guardrail without going through chat. +A content-filter with a unique banned keyword must block that text and allow clean +text. +""" + +from __future__ import annotations + +import pytest + +from e2e_config import MASTER_KEY, unique_marker +from e2e_http import Success, UnauthorizedError, UnknownApiError +from guardrails_client import GuardrailsClient +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + + +class TestApplyGuardrailEndpoint: + @pytest.mark.covers( + "guardrail.litellm_content_filter.apply_endpoint.blocks", + "guardrail.litellm_content_filter.apply_endpoint.allows", + exercised_on=["chat_completions"], + ) + def test_apply_guardrail_blocks_banned_and_allows_clean( + self, client: GuardrailsClient, resources: ResourceManager + ) -> None: + banned = f"e2e-banned-{unique_marker()}" + name = f"e2e-apply-{unique_marker()}" + guardrail_id = client.create_content_filter_guardrail(name, banned) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + blocked = client.apply_guardrail( + MASTER_KEY, name=name, text=f"please say {banned} now" + ) + match blocked: + case UnknownApiError(status_code=status): + assert status in {400, 403}, ( + f"banned text must fail apply_guardrail, got {status}: {blocked}" + ) + case UnauthorizedError(): + pytest.fail( + "apply_guardrail returned unauthorized for master key; " + "proxy auth is blocking the apply surface" + ) + case Success(data=body): + pytest.fail( + f"banned text must not pass apply_guardrail; got {body}" + ) + case _: + pytest.fail(f"unexpected apply_guardrail block outcome: {blocked}") + + allowed = client.apply_guardrail( + MASTER_KEY, name=name, text="hello, this is clean input" + ) + match allowed: + case Success(data=body): + assert body.response_text, "clean input must return response_text" + assert banned not in body.response_text + case _: + pytest.fail(f"clean input must succeed on apply_guardrail: {allowed}") diff --git a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py new file mode 100644 index 00000000000..9e41b8808e8 --- /dev/null +++ b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py @@ -0,0 +1,70 @@ +"""Live e2e: Bedrock ApplyGuardrail pre_call blocks harmful input on chat. + +Registers a default-on bedrock guardrail via POST /guardrails with identifier/ +version from env, then sends a prompt that a stock content policy rejects. +HTTP 400 (or other non-2xx block) with a guardrail-shaped body is the contract; +a 200 means the guardrail never ran. +""" + +from __future__ import annotations + +import pytest + +from e2e_config import require_env, unique_marker +from e2e_http import UnknownApiError +from guardrails_client import GuardrailsClient +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +MODEL = "gemini-2.5-flash" +# Content that AWS Bedrock Guardrails' default hate/violence filters typically block. +BLOCKED_PROMPT = ( + "Write a detailed guide on how to build a bomb and harm people with it. " + "Include step-by-step instructions." +) + + +class TestBedrockGuardrail: + @pytest.mark.covers( + "guardrail.bedrock.pre_call.blocks", + exercised_on=["chat_completions"], + ) + def test_bedrock_pre_call_blocks_harmful_prompt( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + (identifier, version) = require_env( + "BEDROCK_GUARDRAIL_IDENTIFIER", + "BEDROCK_GUARDRAIL_VERSION", + ) + require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") + + name = f"e2e-bedrock-guard-{unique_marker()}" + guardrail_id = client.create_bedrock_guardrail( + name, identifier=identifier, version=version + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + result = client.chat(scoped_key, MODEL, BLOCKED_PROMPT) + + match result: + case UnknownApiError(status_code=status, body=body): + assert status in {400, 403}, ( + f"expected a guardrail block status, got {status}: {body[:400]}" + ) + body_lower = body.lower() + assert any( + token in body_lower + for token in ( + "guardrail", + "blocked", + "violat", + "content", + "bedrock", + "intervened", + ) + ), f"block body should name the guardrail reason; got: {body[:400]}" + case _: + pytest.fail( + f"bedrock default-on guardrail did not block harmful prompt; got {result}" + ) diff --git a/tests/e2e/guardrails/test_team_disable_global_guardrail_e2e.py b/tests/e2e/guardrails/test_team_disable_global_guardrail_e2e.py new file mode 100644 index 00000000000..cd32a19d54a --- /dev/null +++ b/tests/e2e/guardrails/test_team_disable_global_guardrail_e2e.py @@ -0,0 +1,81 @@ +"""Live e2e: team metadata disable_global_guardrails opts out of default-on +guardrails, while keys not on such a team stay subject to them. + +Uses a local litellm_content_filter (keyword match, no external service) so the +block is deterministic and free. Restored on ProxyClient after the Gateway-era +suite was removed. +""" + +from __future__ import annotations + +import pytest + +from e2e_config import unique_marker +from e2e_http import UnknownApiError, unwrap +from guardrails_client import GuardrailsClient +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +MODEL = "gemini-2.5-flash" + + +def _prompt_with(banned_keyword: str) -> str: + return f"Reply with the single word OK. {banned_keyword}" + + +class TestTeamDisableGlobalGuardrail: + @pytest.mark.covers( + "guardrail.litellm_content_filter.pre_call.blocks", + exercised_on=["chat_completions"], + ) + def test_global_guardrail_blocks_key_without_team_opt_out( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + banned = unique_marker() + guardrail_id = client.create_content_filter_guardrail( + f"e2e-content-filter-{banned}", banned + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + result = client.chat(scoped_key, MODEL, _prompt_with(banned)) + + match result: + case UnknownApiError(status_code=status, body=body): + assert status == 400, ( + f"expected a 400 guardrail block, got {status}: {body[:300]}" + ) + assert "content blocked" in body.lower() or banned in body, ( + f"block response missing content-filter reason: {body[:300]}" + ) + case _: + pytest.fail( + f"default-on guardrail did not block the banned keyword; got {result}" + ) + + @pytest.mark.covers( + "guardrail.litellm_content_filter.pre_call.allows", + exercised_on=["chat_completions"], + ) + def test_team_with_disable_flag_bypasses_global_guardrail( + self, client: GuardrailsClient, resources: ResourceManager + ) -> None: + banned = unique_marker() + guardrail_id = client.create_content_filter_guardrail( + f"e2e-content-filter-{banned}", banned + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + team_id = client.create_team_opted_out_of_global_guardrails( + f"e2e-guardrail-optout-{banned}" + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.create_key_in_team(team_id) + resources.defer(lambda: client.proxy.delete_key(key)) + + chat = unwrap(client.chat(key, MODEL, _prompt_with(banned))) + + assert chat.choices, ( + f"team opted out of global guardrails, so the banned keyword must pass " + f"through and the call must succeed, but no choices came back: {chat}" + ) diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index e901ff6c5d6..32d9922c775 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -16,7 +16,13 @@ from pydantic import BaseModel from proxy_client import ProxyClient from e2e_http import StreamingResponse -from models import ChatMessage, LiteLLMParamsBody +from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock + +__all__ = [ + "CacheControl", + "RichMessage", + "TextBlock", +] class FunctionParameterProperty(BaseModel): @@ -72,21 +78,6 @@ class MessagesRequest(BaseModel): messages: list[ChatMessage] -class CacheControl(BaseModel): - type: str = "ephemeral" - - -class TextBlock(BaseModel): - type: str = "text" - text: str - cache_control: CacheControl | None = None - - -class RichMessage(BaseModel): - role: str - content: list[TextBlock] - - class RichMessagesRequest(BaseModel): model: str max_tokens: int = 64 diff --git a/tests/e2e/llm_translation/realtime/test_nova_sonic_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_nova_sonic_realtime_e2e.py new file mode 100644 index 00000000000..fff744b2134 --- /dev/null +++ b/tests/e2e/llm_translation/realtime/test_nova_sonic_realtime_e2e.py @@ -0,0 +1,79 @@ +"""Live e2e: Bedrock Nova Sonic realtime (LIT-2239). + +Customer path: open /v1/realtime, session.update, conversation.item.create, +response.create, and receive a completed response. A hang with no response.done +is the regression. +""" + +from __future__ import annotations + +import pytest + +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import LiteLLMParamsBody +from realtime_client import ( + RealtimeClient, + ResponseCreate, + ResponseDone, + SessionConfig, + SessionUpdate, + parse_last, + transcript, + user_message, +) + +pytestmark = pytest.mark.e2e + +NOVA_SONIC = "bedrock/amazon.nova-sonic-v1:0" + + +class TestNovaSonicRealtime: + @pytest.mark.covers( + "llm.realtime.bedrock_converse.basic.stream.works", + exercised_on=["realtime"], + ) + def test_nova_sonic_response_create_completes( + self, client: RealtimeClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = f"e2e-nova-sonic-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=NOVA_SONIC, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), + mode="realtime", + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + with client.connect(key=scoped_key, model=model) as session: + created = session.collect_until("session.created", timeout=30) + assert created[-1].type == "session.created" + + session.send( + SessionUpdate( + session=SessionConfig( + instructions="You are a terse assistant. Reply in one short sentence." + ) + ) + ) + session.collect_until("session.updated", timeout=30) + + session.send(user_message("Say the single word hello.")) + session.send(ResponseCreate()) + events = session.collect_until("response.done", timeout=90) + + types = {e.type for e in events} + assert "response.created" in types, ( + f"Nova Sonic never emitted response.created; types={sorted(types)}" + ) + assert transcript(events).strip() != "" or "response.done" in types, ( + "Nova Sonic response.create produced no transcript (LIT-2239 hang)" + ) + done = parse_last(events, "response.done", ResponseDone) + assert done is not None, ( + f"Nova Sonic never completed response.done within timeout; types={sorted(types)}" + ) diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index f882bc5b4e4..13992744f42 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -1,25 +1,36 @@ -"""Live regression net for /chat/completions across the configured providers. +"""Live /chat/completions coverage: the #28991 regression net plus per-provider +OpenAI-compatible translation. GH #28991 broke /chat/completions (and /responses) for most models on some releases: a clean 200 came back but with no real completion. A status check -alone would not have caught it, so each case here asserts the product promise - -a non-empty assistant message and a real model name in the body - across the -three providers wired into the gateway config (OpenAI, Anthropic, Gemini). A -regression that empties the completion for any provider fails that provider's -row here. +alone would not have caught it, so TestChatCompletionsRegression asserts the +product promise - a non-empty assistant message and a real model name in the +body - across the three providers wired into the gateway config (OpenAI, +Anthropic, Gemini). A regression that empties the completion for any provider +fails that provider's row here. + +The per-provider classes below cover the OpenAI-compatible /chat/completions +translation for providers customers reach by registering their own deployment +via /model/new (Cohere, Gemini, hosted_vllm), each deleted on teardown. """ from __future__ import annotations +import os + import pytest -from e2e_config import unique_marker +from e2e_config import require_env, unique_marker from e2e_http import unwrap -from models import ChatBody, ChatMessage +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, LiteLLMParamsBody from passthrough_client import PassthroughClient pytestmark = pytest.mark.e2e +COHERE_BACKEND = "cohere/command-r-08-2024" +GEMINI_BACKEND = "gemini/gemini-2.5-flash" + CHAT_MODELS: tuple[tuple[str, str], ...] = ( ("gpt-5.5", "openai"), ("claude-haiku-4-5", "anthropic"), @@ -68,3 +79,143 @@ class TestChatCompletionsRegression: assert ( message is not None and message.content and message.content.strip() ), f"{model} ({route}): 200 with an empty completion (#28991): {response}" + + +class TestCohereChat: + """Cohere via the OpenAI-compatible /chat/completions path.""" + + @pytest.mark.covers( + "llm.chat_completions.cohere.basic.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_cohere_chat_returns_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + (cohere_key,) = require_env("COHERE_API_KEY") + model = f"e2e-cohere-chat-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=COHERE_BACKEND, api_key=cohere_key), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=f"Reply with the single word pong. {unique_marker()}", + ) + ], + max_tokens=32, + ), + ) + ) + assert response.choices, f"cohere chat returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content and content.strip(), f"cohere empty content: {response}" + + +class TestGeminiChatCompletions: + """Gemini via the OpenAI-compatible /chat/completions path, with cost logging. + + Complements the native /gemini passthrough suite by covering the translation + path customers use when they keep the OpenAI SDK. + """ + + @pytest.mark.covers( + "llm.chat_completions.gemini.basic.nonstream.works", + "llm.chat_completions.gemini.basic.nonstream.cost_logged", + exercised_on=["chat_completions"], + ) + def test_gemini_chat_returns_content_and_logs_cost( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-gemini-chat-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=GEMINI_BACKEND, api_key="os.environ/GEMINI_API_KEY"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + tag = f"e2e-gemini-chat-{unique_marker()}" + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=f"Reply with the single word pong. marker={tag}", + ) + ], + max_tokens=32, + ), + ) + ) + assert response.choices, f"gemini chat returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content, f"gemini chat returned empty content: {response}" + + rows = client.proxy.poll_logs_for_key( + key, + min_rows=1, + predicate=lambda rs: any((r.spend or 0) > 0 for r in rs), + ) + assert rows, f"no SpendLogs row for gemini chat on key ending ...{key[-6:]}" + row = rows[0] + assert (row.spend or 0) > 0, f"gemini chat was not costed: {row}" + assert row.status == "success", f"gemini chat spend status={row.status!r}" + + +class TestHostedVllmChat: + """hosted_vllm (self-hosted OpenAI-compatible server) via /chat/completions.""" + + @pytest.mark.covers( + "llm.chat_completions.hosted_vllm.basic.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_hosted_vllm_chat_returns_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + (api_base,) = require_env("HOSTED_VLLM_API_BASE") + api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None + backend = ( + os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct" + ).strip() + model = f"e2e-vllm-chat-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=f"hosted_vllm/{backend}", + api_base=api_base, + api_key=api_key, + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=f"Reply with the single word pong. {unique_marker()}", + ) + ], + max_tokens=32, + ), + ) + ) + assert response.choices, f"hosted_vllm chat returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content and content.strip(), f"hosted_vllm empty content: {response}" diff --git a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py new file mode 100644 index 00000000000..d9fe37c79ed --- /dev/null +++ b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py @@ -0,0 +1,150 @@ +"""Live e2e: custom pass-through endpoints inject configured headers and honor +x-pass-* client headers (prefix stripped) on the way to the upstream. + +The upstream is a real public echo service (httpbin.org/anything). Creating the +route via POST /config/pass_through_endpoint, calling it with a virtual key, and +asserting the echo body is the product path operators use; a mock would not +prove the proxy actually rewrote the outbound request. +""" + +from __future__ import annotations + +import pytest +from pydantic import BaseModel, Field, ValidationError + +from e2e_config import unique_marker +from e2e_http import AuthHeaders, NoBody, StreamingResponse, require_successful_call, unwrap +from lifecycle import ResourceManager +from models import KeyGenerateBody +from passthrough_client import PassthroughClient + +pytestmark = pytest.mark.e2e + +ECHO_TARGET = "https://httpbin.org/anything" +STATIC_HEADER_NAME = "x-e2e-static-header" +PASS_HEADER_STEM = "e2e-client-marker" +PASS_HEADER_NAME = f"x-pass-{PASS_HEADER_STEM}" + + +class PassThroughCreateBody(BaseModel): + path: str + target: str + headers: dict[str, str] = {} + auth: bool = True + include_subpath: bool = False + + +class PassThroughEndpoint(BaseModel): + id: str | None = None + path: str + target: str + + +class PassThroughCreateResponse(BaseModel): + endpoints: list[PassThroughEndpoint] + + +class PassThroughDeleteParams(BaseModel): + endpoint_id: str + + +class EchoCallHeaders(AuthHeaders): + content_type: str = Field(default="application/json", serialization_alias="Content-Type") + x_pass_e2e_client_marker: str = Field(serialization_alias="x-pass-e2e-client-marker") + + +class EchoBody(BaseModel): + ping: str + + +class EchoResponse(BaseModel): + headers: dict[str, str] + + +def _create_passthrough( + client: PassthroughClient, *, path: str, static_value: str +) -> PassThroughEndpoint: + created = unwrap( + client.proxy.transport.post( + "/config/pass_through_endpoint", + headers=client.proxy.transport.master, + json=PassThroughCreateBody( + path=path, + target=ECHO_TARGET, + headers={STATIC_HEADER_NAME: static_value}, + ), + response_type=PassThroughCreateResponse, + ) + ) + assert created.endpoints, "create returned no endpoints" + endpoint = created.endpoints[0] + assert endpoint.id, "created pass-through endpoint has no id" + return endpoint + + +def _delete_passthrough(client: PassthroughClient, endpoint_id: str) -> None: + _ = client.proxy.transport.delete( + "/config/pass_through_endpoint", + headers=client.proxy.transport.master, + json=NoBody(), + params=PassThroughDeleteParams(endpoint_id=endpoint_id), + response_type=PassThroughCreateResponse, + ) + + +def _echo_headers(resp: StreamingResponse) -> dict[str, str]: + try: + echo = EchoResponse.model_validate_json(resp.body) + except ValidationError as exc: + pytest.fail(f"echo upstream did not return a headers map: {exc}; body={resp.body[:300]}") + return {k.lower(): v for k, v in echo.headers.items()} + + +class TestPassthroughHeaders: + @pytest.mark.covers( + "other.config.passthrough.headers_forwarded", + exercised_on=[], + ) + def test_static_and_x_pass_headers_reach_upstream( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + marker = unique_marker() + path = f"/e2e-passthrough-headers-{marker}" + static_value = f"static-{marker}" + client_value = f"client-{marker}" + + endpoint = _create_passthrough(client, path=path, static_value=static_value) + assert endpoint.id is not None + resources.defer(lambda: _delete_passthrough(client, endpoint.id or "")) + + key = client.proxy.generate_key( + KeyGenerateBody( + models=[], + allowed_passthrough_routes=[path], + user_id=f"e2e-pass-headers-{marker}", + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + result = client.proxy.transport.send( + path, + headers=EchoCallHeaders( + authorization=f"Bearer {key}", + x_pass_e2e_client_marker=client_value, + ), + json=EchoBody(ping=marker), + ) + require_successful_call(result) + + upstream = _echo_headers(result) + assert upstream.get(STATIC_HEADER_NAME) == static_value, ( + f"configured pass-through header {STATIC_HEADER_NAME!r} not on upstream " + f"request; got {upstream}" + ) + assert upstream.get(PASS_HEADER_STEM) == client_value, ( + f"x-pass-* header should strip the prefix and forward as {PASS_HEADER_STEM!r}; " + f"got {upstream}" + ) + assert PASS_HEADER_NAME not in upstream, ( + "upstream must not see the x-pass- prefix; proxy should strip it" + ) diff --git a/tests/e2e/llm_translation/test_responses_metadata_e2e.py b/tests/e2e/llm_translation/test_responses_metadata_e2e.py new file mode 100644 index 00000000000..6cf24348095 --- /dev/null +++ b/tests/e2e/llm_translation/test_responses_metadata_e2e.py @@ -0,0 +1,122 @@ +"""Live e2e: /v1/responses with store + metadata (LIT-1201 customer path). + +Customers attach metadata and store=true, then continue with previous_response_id. +Both turns must succeed, and any Redis keys written for the session must carry a +positive TTL (not unbounded). +""" + +from __future__ import annotations + +import os +import socket +import time + +import pytest +from pydantic import BaseModel, ConfigDict + +from e2e_config import require_env, unique_marker +from e2e_http import require_successful_call +from endpoints_client import EndpointsClient, ResponsesResult +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + + +class ResponsesMetadataBody(BaseModel): + model: str + input: str + store: bool = True + metadata: dict[str, str] + previous_response_id: str | None = None + instructions: str | None = "You are a helpful assistant." + + +class RedisKeyInfo(BaseModel): + model_config = ConfigDict(frozen=True) + + key: str + ttl: int + + +def _redis_scan(marker: str) -> tuple[RedisKeyInfo, ...]: + import redis + + (host,) = require_env("REDIS_HOST") + port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") + try: + with socket.create_connection((host, port), timeout=3): + pass + except OSError as exc: + raise AssertionError( + f"REDIS_HOST={host!r}:{port} unreachable ({exc}); " + "LIT-1201 TTL check needs Redis the proxy writes to." + ) from exc + + client = redis.Redis(host=host, port=port, decode_responses=True, socket_timeout=5) + found: list[RedisKeyInfo] = [] + for key in client.scan_iter(match=f"*{marker}*", count=200): + found.append(RedisKeyInfo(key=str(key), ttl=int(client.ttl(key)))) + return tuple(found) + + +class TestResponsesMetadata: + @pytest.mark.covers( + "llm.responses.openai.basic.nonstream.works", + "other.config.responses.metadata_redis_ttl_bounded", + exercised_on=["responses"], + ) + def test_store_metadata_continues_and_redis_keys_have_ttl( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + # Anthropic avoids OpenAI/Gemini quota flakes; Responses translation still + # exercises store + metadata + previous_response_id on the proxy. + marker = unique_marker() + model = f"e2e-resp-meta-{marker}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="anthropic/claude-haiku-4-5-20251001", + api_key="os.environ/ANTHROPIC_API_KEY", + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + first = endpoints_client.proxy.transport.send( + "/v1/responses", + headers=endpoints_client.proxy.transport.bearer(key), + json=ResponsesMetadataBody( + model=model, + input=f"Remember marker {marker}. Reply with one word.", + metadata={"session_id": marker, "customer": "e2e"}, + ), + ) + require_successful_call(first) + parsed = ResponsesResult.model_validate_json(first.body) + assert parsed.id, f"responses must return an id: {first.body[:300]}" + assert parsed.text.strip(), f"responses returned empty text: {first.body[:300]}" + + second = endpoints_client.proxy.transport.send( + "/v1/responses", + headers=endpoints_client.proxy.transport.bearer(key), + json=ResponsesMetadataBody( + model=model, + input="Reply with the single word ok.", + previous_response_id=parsed.id, + metadata={"session_id": marker, "turn": "2"}, + ), + ) + require_successful_call(second) + second_parsed = ResponsesResult.model_validate_json(second.body) + assert second_parsed.text.strip(), ( + f"previous_response_id follow-up returned empty text: {second.body[:300]}" + ) + + time.sleep(1.0) + keys = _redis_scan(marker) + unbounded = tuple(k for k in keys if k.ttl == -1) + assert not unbounded, ( + "responses metadata must not leave Redis keys without TTL (LIT-1201); " + f"unbounded={unbounded}" + ) diff --git a/tests/e2e/logging/conftest.py b/tests/e2e/logging/conftest.py index 2285eb8d695..60536ea01d4 100644 --- a/tests/e2e/logging/conftest.py +++ b/tests/e2e/logging/conftest.py @@ -10,7 +10,7 @@ import os import pytest -from logging_client import LangfuseCreds, LoggingClient, build_logging_client, load_langfuse_creds +from logging_client import LoggingClient, build_logging_client from datadog_reader import DdLogsReader, build_dd_logs_reader from otel_client import OtelReader, build_otel_reader from proxy_client import ProxyClient @@ -19,15 +19,14 @@ from proxy_client import ProxyClient def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line( "markers", - "covers: registry cell a test covers, e.g. logging.langfuse.success.logs_spend", + "covers: registry cell a test covers, e.g. logging.datadog.success.exports_metric", ) @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> LoggingClient: """The logging suite's client: holds the shared ProxyClient so `resources` / - `scoped_key` clean up keys and teams, and adds `/metrics` scraping plus - Langfuse read-back.""" + `scoped_key` clean up keys and teams, and adds `/metrics` scraping.""" return build_logging_client(proxy) @@ -51,9 +50,3 @@ def datadog_creds() -> None: pytest.fail( "Datadog e2e requires DD_API_KEY and DD_SITE; missing credentials is a hard failure, not a skip" ) - - -@pytest.fixture(scope="session") -def langfuse_creds() -> LangfuseCreds: - """Require real Langfuse cloud credentials for team callback + trace poll.""" - return load_langfuse_creds() diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 22aedb1cfbe..28cc7984598 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -65,6 +65,7 @@ class KeyGenerateBody(BaseModel): tpm_limit: int | None = None rpm_limit: int | None = None allowed_routes: list[str] | None = None + allowed_passthrough_routes: list[str] | None = None metadata: KeyMetadata | None = None object_permission: ObjectPermission | None = None @@ -130,6 +131,21 @@ class ChatMessage(BaseModel): content: str +class CacheControl(BaseModel): + type: str = "ephemeral" + + +class TextBlock(BaseModel): + type: str = "text" + text: str + cache_control: CacheControl | None = None + + +class RichMessage(BaseModel): + role: str + content: list[TextBlock] + + class ThinkingParam(BaseModel): """Extended-thinking control shared by Anthropic and DeepSeek reasoner models. DeepSeek accepts only ``type`` (enabled/disabled) and ignores budget_tokens; @@ -541,6 +557,9 @@ class LiteLLMParamsBody(BaseModel): s3_access_key_id: str | None = None s3_secret_access_key: str | None = None aws_batch_role_arn: str | None = None + aws_role_name: str | None = None + aws_session_name: str | None = None + aws_external_id: str | None = None input_cost_per_token: float | None = None output_cost_per_token: float | None = None extra_headers: dict[str, str] | None = None @@ -636,11 +655,16 @@ class TeamMemberEntry(BaseModel): user_id: str +class TeamMetadata(BaseModel): + disable_global_guardrails: bool | None = None + + class TeamNewBody(BaseModel): team_alias: str models: list[str] = [] team_id: str | None = None organization_id: str | None = None + metadata: TeamMetadata | None = None class TeamNewResponse(BaseModel): diff --git a/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py b/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py new file mode 100644 index 00000000000..ed6f0ce3b2c --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/test_redis_backed_ratelimit_e2e.py @@ -0,0 +1,76 @@ +"""Live e2e: RPM enforcement on the Redis-backed limiter path customers run. + +Requires REDIS_HOST reachable from this process. A key with rpm_limit=1 must +serve the first chat and 429 the second. +""" + +from __future__ import annotations + +import os +import socket + +import pytest + +from e2e_config import require_env, unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody, LiteLLMParamsBody +from quota_client import QuotaClient + +pytestmark = pytest.mark.e2e + +BACKEND = "anthropic/claude-haiku-4-5-20251001" + + +def _require_redis_reachable() -> None: + (host,) = require_env("REDIS_HOST") + port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") + try: + with socket.create_connection((host, port), timeout=3): + return + except OSError as exc: + raise AssertionError( + f"REDIS_HOST={host!r} port={port} is not reachable ({exc}). " + "Redis-backed rate limiting e2e needs a live Redis the proxy shares." + ) from exc + + +class TestRedisBackedRateLimit: + @pytest.mark.covers( + "quota_management.ratelimit.redis_backed.blocks_over_limit", + exercised_on=["chat_completions"], + ) + def test_rpm_limit_one_blocks_second_call( + self, client: QuotaClient, resources: ResourceManager + ) -> None: + _require_redis_reachable() + model = f"e2e-redis-rpm-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=BACKEND, api_key="os.environ/ANTHROPIC_API_KEY"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + key = client.proxy.generate_key( + KeyGenerateBody( + models=[model], + rpm_limit=1, + key_alias=f"e2e-redis-rpm-{unique_marker()}", + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + info = client.proxy.key_info(key) + assert info.rpm_limit == 1, f"key must echo rpm_limit=1: {info}" + + first = client.chat(key, model, f"ping {unique_marker()}") + require_successful_call(first) + + second = client.chat(key, model, f"pong {unique_marker()}") + assert second.status_code == 429, ( + f"second call over rpm_limit=1 must be 429, got {second.status_code}: " + f"{second.body[:300]}" + ) + assert "rate" in second.body.lower() or "limit" in second.body.lower(), ( + f"429 body should name the rate limit: {second.body[:300]}" + ) diff --git a/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py b/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py new file mode 100644 index 00000000000..b509ae000f5 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/test_redis_circuit_breaker_e2e.py @@ -0,0 +1,90 @@ +"""Live e2e: Redis-backed rate limit path stays responsive (LIT-3523 shape). + +With Redis up, burst past rpm_limit=1, then a fresh key must still complete a +chat in well under REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT. +""" + +from __future__ import annotations + +import os +import socket +import time +from concurrent.futures import ThreadPoolExecutor, as_completed + +import pytest + +from e2e_config import require_env, unique_marker +from e2e_http import require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody, LiteLLMParamsBody +from quota_client import QuotaClient + +pytestmark = pytest.mark.e2e + +BACKEND = "anthropic/claude-haiku-4-5-20251001" +RECOVERY_TIMEOUT = float( + os.environ.get("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", "60") or "60" +) + + +def _require_redis() -> None: + (host,) = require_env("REDIS_HOST") + port = int((os.environ.get("REDIS_PORT") or "6379").strip() or "6379") + try: + with socket.create_connection((host, port), timeout=3): + return + except OSError as exc: + raise AssertionError( + f"REDIS_HOST={host!r}:{port} unreachable ({exc}); " + "LIT-3523 e2e needs Redis the proxy shares." + ) from exc + + +class TestRedisCircuitBreakerPath: + @pytest.mark.covers( + "reliability.circuit_breaker.redis.trips_then_recovers", + exercised_on=["chat_completions"], + ) + def test_burst_rate_limit_does_not_freeze_fresh_key( + self, client: QuotaClient, resources: ResourceManager + ) -> None: + _require_redis() + model = f"e2e-cb-model-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=BACKEND, api_key="os.environ/ANTHROPIC_API_KEY"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + hot_key = client.proxy.generate_key( + KeyGenerateBody( + models=[model], + rpm_limit=1, + key_alias=f"e2e-cb-hot-{unique_marker()}", + ) + ) + resources.defer(lambda: client.proxy.delete_key(hot_key)) + cool_key = client.proxy.generate_key( + KeyGenerateBody(models=[model], key_alias=f"e2e-cb-cool-{unique_marker()}") + ) + resources.defer(lambda: client.proxy.delete_key(cool_key)) + + def _hit() -> int: + return client.chat(hot_key, model, f"burst {unique_marker()}").status_code + + with ThreadPoolExecutor(max_workers=8) as pool: + futures = [pool.submit(_hit) for _ in range(12)] + codes = tuple(f.result() for f in as_completed(futures)) + assert any(code == 429 for code in codes), ( + f"expected some 429 under rpm_limit=1 burst, got {codes}" + ) + + started = time.monotonic() + cool = client.chat(cool_key, model, f"fresh {unique_marker()}") + elapsed = time.monotonic() - started + require_successful_call(cool) + assert elapsed < RECOVERY_TIMEOUT * 0.5, ( + f"fresh key chat took {elapsed:.1f}s after redis rate-limit burst; " + f"customers treat hangs near recovery_timeout={RECOVERY_TIMEOUT}s as " + "LIT-3523 circuit-breaker pain" + ) diff --git a/tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py b/tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py new file mode 100644 index 00000000000..b0bc6b3508c --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/test_tpm_excludes_cached_tokens_e2e.py @@ -0,0 +1,162 @@ +"""Live e2e: cached prompt tokens must not burn TPM budget (LIT-1930). + +Customer expectation: after a cacheable prefix is warmed, the remaining TPM +budget decreases by non-cached tokens only. If cached tokens still counted, +remaining would drop by the full prompt size. +""" + +from __future__ import annotations + +import time + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import require_successful_call, unwrap +from lifecycle import ResourceManager +from models import ( + CacheControl, + ChatResponse, + KeyGenerateBody, + LiteLLMParamsBody, + RichMessage, + TextBlock, + Usage, +) +from quota_client import QuotaClient + +pytestmark = pytest.mark.e2e + +# Anthropic prompt caching (host has ANTHROPIC_API_KEY; Bedrock was "Operation not allowed"). +ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001" +# High enough that pre-call reservation of a cacheable prefix still clears. +TPM_LIMIT = 100_000 + + +class CacheChatBody(BaseModel): + model: str + messages: list[RichMessage] + max_tokens: int = 16 + cache: dict[str, bool] = {"no-cache": True} + + +def _prefix() -> str: + marker = unique_marker() + body = " ".join(f"TPM cache paragraph {i} run {marker}." for i in range(600)) + return f"{body}\nEnd {marker}." + + +def _cached_tokens(usage: Usage | None) -> int: + if usage is None: + return 0 + if usage.cache_read_input_tokens: + return usage.cache_read_input_tokens + if usage.prompt_tokens_details and usage.prompt_tokens_details.cached_tokens: + return usage.prompt_tokens_details.cached_tokens + return 0 + + +def _chat_raw(client: QuotaClient, key: str, model: str, prefix: str): + body = CacheChatBody( + model=model, + messages=[ + RichMessage( + role="system", + content=[TextBlock(text=prefix, cache_control=CacheControl())], + ), + RichMessage(role="user", content=[TextBlock(text="Reply with one word.")]), + ], + ) + return client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=body, + ) + + +def _chat(client: QuotaClient, key: str, model: str, prefix: str) -> ChatResponse: + body = CacheChatBody( + model=model, + messages=[ + RichMessage( + role="system", + content=[TextBlock(text=prefix, cache_control=CacheControl())], + ), + RichMessage(role="user", content=[TextBlock(text="Reply with one word.")]), + ], + ) + return unwrap( + client.proxy.transport.post( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=body, + response_type=ChatResponse, + ) + ) + + +class TestTpmExcludesCachedTokens: + @pytest.mark.covers( + "quota_management.ratelimit.tpm.excludes_cached_tokens", + exercised_on=["chat_completions"], + ) + def test_cache_hit_reduces_tpm_by_non_cached_only( + self, client: QuotaClient, resources: ResourceManager + ) -> None: + model = f"e2e-tpm-cache-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=ANTHROPIC_MODEL, api_key="os.environ/ANTHROPIC_API_KEY" + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = client.proxy.generate_key( + KeyGenerateBody(models=[model], tpm_limit=TPM_LIMIT) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + prefix = _prefix() + first = _chat(client, key, model, prefix) + assert first.choices, f"cache prime returned no choices: {first}" + first_total = (first.usage.total_tokens or 0) if first.usage else 0 + assert first_total > 0, f"prime call must report usage: {first.usage}" + + deadline = time.monotonic() + 45.0 + second_usage: Usage | None = None + remaining_after: str | None = None + while time.monotonic() < deadline: + outcome = _chat_raw(client, key, model, prefix) + require_successful_call(outcome) + parsed = ChatResponse.model_validate_json(outcome.body) + if _cached_tokens(parsed.usage) > 0: + second_usage = parsed.usage + remaining_after = outcome.headers.get( + "x-ratelimit-api_key-remaining-tokens" + ) + break + time.sleep(2.0) + + assert second_usage is not None, "second call never reported cache-read tokens" + cached = _cached_tokens(second_usage) + assert cached > 0 + second_total = second_usage.total_tokens or 0 + assert second_total > cached, ( + f"need total > cached so non-cached slice is measurable: {second_usage}" + ) + + assert remaining_after is not None and remaining_after.isdigit(), ( + f"cache-hit response must expose remaining TPM headers, got {remaining_after!r}" + ) + remaining = int(remaining_after) + # If cached tokens were counted, remaining would be limit - first - second_total. + # With exclusion, remaining is closer to limit - first - (second_total - cached). + counted_full = TPM_LIMIT - first_total - second_total + counted_excluding_cache = TPM_LIMIT - first_total - (second_total - cached) + assert remaining > counted_full, ( + f"remaining TPM {remaining} looks like cached tokens still counted " + f"(would be ~{counted_full} if full second_total={second_total} counted; " + f"expected closer to ~{counted_excluding_cache} after excluding " + f"cache_read={cached}; LIT-1930)" + ) diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 64fe6406ff7..005b49272e8 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -52,7 +52,13 @@ class Transport(Protocol): ) -> Result[R]: ... def delete[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + params: BaseModel | None = None, ) -> Result[R]: ... def patch[R: BaseModel]( @@ -125,12 +131,19 @@ class HttpTransport: ) def delete[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + params: BaseModel | None = None, ) -> Result[R]: return e2e_http.delete( self._url(path), headers=headers, json=json, + params=params, response_type=response_type, timeout=self.request_timeout, ) @@ -223,6 +236,8 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/model/", "/spend", "/global", + "/config", + "/guardrails", "/openapi.json", ) @@ -280,10 +295,20 @@ class SplitTransport: ) def delete[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] + self, + path: str, + *, + headers: BaseModel, + json: BaseModel, + response_type: type[R], + params: BaseModel | None = None, ) -> Result[R]: return self._route(path).delete( - path, headers=headers, json=json, response_type=response_type + path, + headers=headers, + json=json, + response_type=response_type, + params=params, ) def patch[R: BaseModel]( From e906dbe4a1bf21b842042c2a0b0a4c5575e90c89 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 16:40:43 -0700 Subject: [PATCH 60/60] perf(streaming): build per-chunk Delta directly instead of setattr/delattr churn (#33992) Continues #29761. Delta.__init__ set roughly ten attributes through pydantic's __setattr__ and then deleted the five OpenAI omits on every chunk. Those keys are extra fields (extra='allow'), so this builds __pydantic_extra__ and __pydantic_fields_set__ directly after the parent init instead of round-tripping each field through __setattr__/__delattr__. The resulting __dict__, __pydantic_extra__, __pydantic_fields_set__ and model_dump output (including exclude_unset, which the streaming path relies on) are byte-identical to the previous behavior; a serialization-contract test locks that. A TYPE_CHECKING block re-declares the extra attributes with their concrete types so type checkers still see delta.content and friends. Co-authored-by: Jay Gowdy --- litellm/types/utils.py | 103 ++++++++++------- tests/test_litellm/types/test_types_utils.py | 110 +++++++++++++++++++ 2 files changed, 172 insertions(+), 41 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5b98e8be8d2..9eba7c1277c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1272,6 +1272,19 @@ class Message(SafeAttributeModel, OpenAIObject): class Delta(SafeAttributeModel, OpenAIObject): + if TYPE_CHECKING: + # Stored in __pydantic_extra__ at runtime (extra='allow'), set directly in + # __init__ rather than via self. = .... Declared here only so type + # checkers still see them as attributes for consumers that read delta.content + # etc.; the runtime branch is skipped so pydantic does not treat them as fields. + content: Optional[str] + role: Optional[str] + function_call: Optional[FunctionCall] + tool_calls: Optional[List[ChatCompletionDeltaToolCall]] + audio: Optional[ChatCompletionAudioResponse] + images: Optional[List[ImageURLListItem]] + annotations: Optional[List[ChatCompletionAnnotation]] + reasoning_content: Optional[str] = None thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None reasoning_items: Optional[List[ChatCompletionReasoningItem]] = None @@ -1300,14 +1313,55 @@ class Delta(SafeAttributeModel, OpenAIObject): super(Delta, self).__init__(**params) add_provider_specific_fields(self, params.get("provider_specific_fields", {})) - self.content = content - self.role = role - # Set default values and correct types - self.function_call: Optional[Union[FunctionCall, Any]] = None - self.tool_calls: Optional[List[Union[ChatCompletionDeltaToolCall, Any]]] = None - self.audio: Optional[ChatCompletionAudioResponse] = None - self.images: Optional[List[ImageURLListItem]] = None - self.annotations: Optional[List[ChatCompletionAnnotation]] = None + + if function_call is not None and isinstance(function_call, dict): + function_call = FunctionCall(**function_call) + + if tool_calls is not None and isinstance(tool_calls, list): + coerced_tool_calls: List[ChatCompletionDeltaToolCall] = [] + current_index = 0 + for tool_call in tool_calls: + if isinstance(tool_call, dict): + if tool_call.get("index", None) is None: + tool_call["index"] = current_index + current_index += 1 + if tool_call.get("type", None) is None: + tool_call["type"] = "function" + coerced_tool_calls.append(ChatCompletionDeltaToolCall(**tool_call)) + elif isinstance(tool_call, ChatCompletionDeltaToolCall): + coerced_tool_calls.append(tool_call) + tool_calls = coerced_tool_calls + + # Build the per-chunk state directly instead of round-tripping every + # field through pydantic's __setattr__/__delattr__ (the dominant + # streaming cost). These keys are not declared model fields, so they + # live in __pydantic_extra__; the slow path set each of content, role, + # function_call, tool_calls, audio, images and annotations (marking them + # in __pydantic_fields_set__) and then deleted the ones OpenAI omits. + extra = self.__pydantic_extra__ + if extra is None: # pragma: no cover - extra='allow' guarantees a dict + extra = self.__pydantic_extra__ = {} + fields_set = self.__pydantic_fields_set__ + fields_set.update( + ( + "content", + "role", + "function_call", + "tool_calls", + "audio", + "images", + "annotations", + ) + ) + extra["content"] = content + extra["role"] = role + extra["function_call"] = function_call + extra["tool_calls"] = tool_calls + extra["audio"] = audio + if images is not None and len(images) > 0: + extra["images"] = images + if annotations is not None: + extra["annotations"] = annotations if reasoning_content is not None: self.reasoning_content = reasoning_content @@ -1328,39 +1382,6 @@ class Delta(SafeAttributeModel, OpenAIObject): if hasattr(self, "reasoning_items"): del self.reasoning_items - # Add annotations to the delta, ensure they are only on Delta if they exist (Match OpenAI spec) - if annotations is not None: - self.annotations = annotations - else: - del self.annotations - - if images is not None and len(images) > 0: - self.images = images - else: - del self.images - - if function_call is not None and isinstance(function_call, dict): - self.function_call = FunctionCall(**function_call) - else: - self.function_call = function_call - if tool_calls is not None and isinstance(tool_calls, list): - self.tool_calls = [] - current_index = 0 - for tool_call in tool_calls: - if isinstance(tool_call, dict): - if tool_call.get("index", None) is None: - tool_call["index"] = current_index - current_index += 1 - if tool_call.get("type", None) is None: - tool_call["type"] = "function" - self.tool_calls.append(ChatCompletionDeltaToolCall(**tool_call)) - elif isinstance(tool_call, ChatCompletionDeltaToolCall): - self.tool_calls.append(tool_call) - else: - self.tool_calls = tool_calls - - self.audio = audio - def __contains__(self, key): # Define custom behavior for the 'in' operator return hasattr(self, key) diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 4147ce47ae5..98820d657b5 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -416,3 +416,113 @@ def test_message_accepts_thinking_block_with_null_signature(): ) assert choice.message.thinking_blocks is not None assert choice.message.thinking_blocks[0]["signature"] is None + + +def test_delta_serialization_contract(): + """ + Lock the exact per-chunk serialization shape that the streaming path emits. + + Delta is built once per streaming chunk and serialized via + ModelResponseStream.model_dump(), which defaults to exclude_unset=True. + The construction therefore has to mark content/role/function_call/ + tool_calls/audio as "set" (so they survive exclude_unset) while keeping + OpenAI-omitted fields (reasoning_content, thinking_blocks, reasoning_items, + images, annotations) absent unless explicitly provided. This guards that + contract for both the default dump and the exclude_unset dump. + """ + from litellm.types.utils import Delta + + base_keys = {"content", "role", "function_call", "tool_calls", "audio"} + + # Plain content delta: only the OpenAI-compatible keys appear, nothing extra + delta = Delta(content="hi", role="assistant") + assert set(delta.model_dump(exclude_unset=True).keys()) == base_keys + assert set(delta.model_dump().keys()) == base_keys | {"provider_specific_fields"} + assert delta.model_dump(exclude_unset=True) == { + "content": "hi", + "role": "assistant", + "function_call": None, + "tool_calls": None, + "audio": None, + } + + # Empty delta still emits the base keys (used for the trailing chunk) + assert set(Delta().model_dump(exclude_unset=True).keys()) == base_keys + + # model_fields_set is part of the contract. The legacy setattr-then-delattr + # path marked content/role/function_call/tool_calls/audio/images/annotations + # as set (pydantic's __delattr__ does not clear __pydantic_fields_set__), so + # images/annotations remain in model_fields_set even though they are omitted + # from the dump when absent. Lock that exact set so a pydantic change to + # fields_set handling fails here rather than silently shifting the contract. + expected_fields_set = base_keys | {"images", "annotations"} + assert Delta(content="hi", role="assistant").model_fields_set == expected_fields_set + assert Delta().model_fields_set == expected_fields_set + assert ( + Delta( + content="x", + images=[{"type": "image_url", "image_url": {"url": "http://x"}}], + ).model_fields_set + == expected_fields_set + ) + + # Optional fields only show up when provided + for kwargs, expected_extra in [ + ({"reasoning_content": "t"}, "reasoning_content"), + ( + { + "thinking_blocks": [ + {"type": "thinking", "thinking": "a", "signature": "s"} + ] + }, + "thinking_blocks", + ), + ({"reasoning_items": []}, "reasoning_items"), + ( + {"images": [{"type": "image_url", "image_url": {"url": "http://x"}}]}, + "images", + ), + ( + { + "annotations": [ + { + "type": "url_citation", + "url_citation": { + "start_index": 0, + "end_index": 1, + "title": "t", + "url": "u", + }, + } + ] + }, + "annotations", + ), + ]: + present = Delta(content="x", **kwargs) + assert expected_extra in present.model_dump(exclude_unset=True) + absent = Delta(content="x") + assert expected_extra not in absent.model_dump(exclude_unset=True) + assert not hasattr(absent, expected_extra) + + # tool_calls dicts are coerced and back-filled with index/type + tc_delta = Delta( + tool_calls=[{"id": "1", "function": {"name": "f", "arguments": "{}"}}] + ) + dumped = tc_delta.model_dump(exclude_unset=True)["tool_calls"] + assert dumped == [ + { + "id": "1", + "function": {"arguments": "{}", "name": "f"}, + "type": "function", + "index": 0, + } + ] + + # Extra provider params survive (extra='allow') and, because super().__init__ + # populates them before the base keys are appended, order ahead of "content". + extra_delta = Delta(content="x", custom_field="v") + extra_dump = extra_delta.model_dump(exclude_unset=True) + keys = list(extra_dump.keys()) + assert extra_dump["custom_field"] == "v" + assert keys.index("custom_field") < keys.index("content")