From 00a20591746b1d40f9c0ec5407ea06efb5a31234 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 19:51:34 +0000 Subject: [PATCH 01/81] fix(router): resolve realtime session model to routed deployment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 22 +++++++++ tests/test_litellm/test_router.py | 76 +++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index fb2af41dcf2..95509bcdbba 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -29,6 +29,7 @@ import anyio import httpx import openai from openai import AsyncOpenAI +from pydantic import TypeAdapter, ValidationError from typing_extensions import overload import litellm @@ -342,6 +343,26 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) return False +_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) +_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: + """ + Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still + holds the pre-routing model group name, so it has to follow the deployment the router just picked. + + Returns kwargs to merge into the downstream call, empty when there is no session model to resolve. + """ + try: + typed_session: Final = _SESSION_ADAPTER.validate_python(session) + except ValidationError: + return _NO_SESSION_KWARGS + if "model" not in typed_session: + return _NO_SESSION_KWARGS + return MappingProxyType({"session": {**typed_session, "model": model_name}}) + + class RoutingArgs(enum.Enum): ttl = 60 # 1min (RPM/TPM expire key) @@ -4685,6 +4706,7 @@ class Router: "caching": self.cache_responses, **kwargs, "model": model_name, + **_with_router_resolved_session_model(kwargs.get("session"), model_name), } # Only set custom_llm_provider if it's not None if custom_llm_provider is not None: diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bdbf33fb0e1..fbefcc59477 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1298,6 +1298,82 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" +@pytest.mark.asyncio +async def test_ageneric_api_call_resolves_realtime_session_model(): + """ + Regression for #36742: realtime client secret requests carry the model inside `session` too, and the proxy + fills it with the pre-routing model group name. The underlying litellm function reads session.model first, + so it must see the resolved deployment, while a caller's nested transcription model stays untouched. + """ + captured: dict = {} + + async def capture_kwargs(**kwargs): + captured.update(kwargs) + return {"result": "ok"} + + router = litellm.Router( + model_list=[ + { + "model_name": "my-realtime-group", + "litellm_params": { + "model": "openai/gpt-realtime-2.1-mini", + "api_key": "fake-key", + }, + "model_info": {"mode": "realtime"}, + } + ] + ) + + await router._ageneric_api_call_with_fallbacks( + model="my-realtime-group", + original_function=capture_kwargs, + session={ + "type": "realtime", + "model": "my-realtime-group", + "audio": {"input": {"transcription": {"model": "gpt-4o-transcribe"}}}, + }, + ) + + assert captured["model"] == "openai/gpt-realtime-2.1-mini" + assert captured["session"]["model"] == "openai/gpt-realtime-2.1-mini" + assert captured["session"]["audio"]["input"]["transcription"]["model"] == "gpt-4o-transcribe" + + +@pytest.mark.asyncio +async def test_ageneric_api_call_does_not_add_session_model(): + """ + A session that never carried a model must not gain one from routing: the underlying function then falls back + to the resolved `model` kwarg itself, and the outgoing session body keeps the caller's shape. + """ + captured: dict = {} + + async def capture_kwargs(**kwargs): + captured.update(kwargs) + return {"result": "ok"} + + router = litellm.Router( + model_list=[ + { + "model_name": "my-realtime-group", + "litellm_params": { + "model": "openai/gpt-realtime-2.1-mini", + "api_key": "fake-key", + }, + "model_info": {"mode": "realtime"}, + } + ] + ) + + await router._ageneric_api_call_with_fallbacks( + model="my-realtime-group", + original_function=capture_kwargs, + session={"type": "realtime"}, + ) + + assert captured["model"] == "openai/gpt-realtime-2.1-mini" + assert captured["session"] == {"type": "realtime"} + + def test_router_get_model_access_groups_team_only_models(): """ Test that Router.get_model_access_groups returns the correct response for team-only models From e5b54620a4fcf67d13fb5aa2261dab31502be5ab Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 19:58:03 +0000 Subject: [PATCH 02/81] test(router): cover realtime session model resolver directly Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_router.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index fbefcc59477..99e7e9ad6bc 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1374,6 +1374,21 @@ async def test_ageneric_api_call_does_not_add_session_model(): assert captured["session"] == {"type": "realtime"} +@pytest.mark.parametrize( + "session, expected", + [ + ({"type": "realtime", "model": "my-realtime-group"}, {"session": {"type": "realtime", "model": "resolved"}}), + ({"type": "realtime"}, {}), + (None, {}), + ("not-a-session", {}), + ], +) +def test_with_router_resolved_session_model(session, expected): + from litellm.router import _with_router_resolved_session_model + + assert dict(_with_router_resolved_session_model(session, "resolved")) == expected + + def test_router_get_model_access_groups_team_only_models(): """ Test that Router.get_model_access_groups returns the correct response for team-only models From 3275459aec0935b03b60950ba5d625854e2d6c9b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:47:26 -0700 Subject: [PATCH 03/81] fix(mcp): cap tools preview and test-connection at the listing timeout and name the unreachable upstream --- .../mcp_server/rest_endpoints.py | 34 +++++--- .../mcp_server/test_rest_endpoints.py | 77 +++++++++++++++++-- 2 files changed, 96 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3efb6429326..d73277f4417 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal +import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from litellm._logging import verbose_logger +from litellm.constants import MCP_TOOL_LISTING_TIMEOUT from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -68,7 +70,13 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( ) -def _connection_error_message(exc: BaseException) -> str: +def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + if isinstance(exc, TimeoutError): + return ( + f"Failed to connect to MCP server: no response from {url or 'the server'} " + f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " + "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." + ) if isinstance(exc, httpx.LocalProtocolError): return ( "Failed to connect to MCP server: a request header is malformed. " @@ -1136,6 +1144,7 @@ if MCP_AVAILABLE: mcp_auth_header: str | dict[str, str] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + timeout_seconds: float = MCP_TOOL_LISTING_TIMEOUT, ) -> Mapping[str, object]: """ Create a temporary MCP client from *request*, run *operation*, and return the result. @@ -1151,6 +1160,10 @@ if MCP_AVAILABLE: oauth2_headers: Headers extracted from the incoming request (may contain the litellm API key — must NOT be forwarded for M2M servers). raw_headers: Raw request headers forwarded for stdio env construction. + timeout_seconds: Cap on OAuth discovery, connect, handshake, and *operation* + combined. Defaults to ``MCP_TOOL_LISTING_TIMEOUT`` (30s, below common LB + timeouts) so an unreachable upstream yields this endpoint's JSON error + instead of an opaque load-balancer 504 with an empty body. Returns: The dict returned by *operation*, or an error dict on failure. @@ -1240,15 +1253,16 @@ if MCP_AVAILABLE: static_headers=request.static_headers, ) - client: Final = await global_mcp_server_manager._create_mcp_client( - server=server_model, - mcp_auth_header=mcp_auth_header, - extra_headers=merged_headers, - stdio_env=stdio_env, - cred_provider=preview_cred_provider, - ) + with anyio.fail_after(timeout_seconds): + client: Final = await global_mcp_server_manager._create_mcp_client( + server=server_model, + mcp_auth_header=mcp_auth_header, + extra_headers=merged_headers, + stdio_env=stdio_env, + cred_provider=preview_cred_provider, + ) - return await operation(client) + return await operation(client) except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise @@ -1257,7 +1271,7 @@ if MCP_AVAILABLE: return { "status": "error", "error": True, - "message": _connection_error_message(e), + "message": _connection_error_message(e, request.url, timeout_seconds), } async def _preview_openapi_tools(spec_path: str) -> dict: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ef5631218f3..51b946c11b7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,4 +1,5 @@ import asyncio +import inspect import json import sys from datetime import datetime @@ -13,6 +14,7 @@ import pytest from fastapi import HTTPException from starlette.requests import Request +from litellm.constants import MCP_TOOL_LISTING_TIMEOUT from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, @@ -109,6 +111,71 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert "stack_trace" not in result + @pytest.mark.asyncio + async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def hanging_operation(client): + await asyncio.Event().wait() + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, hanging_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + @pytest.mark.asyncio + async def test_timeout_covers_client_creation(self, monkeypatch): + async def hanging_create_client(*args, **kwargs): + await asyncio.Event().wait() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + hanging_create_client, + ) + + async def unreached_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, unreached_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + def test_timeout_defaults_to_tool_listing_timeout(self): + default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default + assert default == MCP_TOOL_LISTING_TIMEOUT + + def test_connection_error_message_timeout_names_url_and_budget(self): + message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) + assert "https://api.example.com/mcp/" in message + assert "30s" in message + @pytest.mark.asyncio async def test_forwards_static_headers(self, monkeypatch): """Ensure static_headers are forwarded to the MCP client during test calls. @@ -2881,17 +2948,17 @@ class TestConnectionErrorMessage: secret = "Bearer sk-super-secret-token" exc = httpx.LocalProtocolError(f"Illegal header value b' {secret}'") - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "header" in message.lower() assert secret not in message def test_connect_error_points_at_reachability(self): - message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed")) + message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0) assert "unreachable" in message.lower() def test_timeout_error_message(self): - message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out")) + message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out"), "https://example.com", 30.0) assert "unreachable" in message.lower() def test_http_status_error_includes_status_code(self): @@ -2901,11 +2968,11 @@ class TestConnectionErrorMessage: request=httpx.Request("POST", "http://x/"), response=response, ) - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message def test_unknown_error_falls_back_to_generic(self): - message = rest_endpoints._connection_error_message(RuntimeError("weird")) + message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message assert "proxy logs" in message.lower() From f4b5449c6a65cf167658fd5bb32695abda9633a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:48:04 -0700 Subject: [PATCH 04/81] fix(openai_like): strip cache_control ttl before forwarding /v1/messages to non-Anthropic providers --- litellm/llms/anthropic/common_utils.py | 33 +++++ litellm/llms/openai_like/README.md | 5 +- .../openai_like/messages/transformation.py | 38 +++++- ..._like_anthropic_messages_transformation.py | 121 ++++++++++++++++++ 4 files changed, 195 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 681a8397f66..695e3a313ef 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1302,6 +1302,39 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format +def _normalized_cache_control(cache_control: dict) -> dict: # mutable-ok: as sibling sanitizers + cache_type: Final = cache_control.get("type") + return {"type": cache_type if isinstance(cache_type, str) else "ephemeral"} # mutable-ok: JSON wire format + + +def _normalize_cache_control_value(value: object) -> object: + if isinstance(value, dict): + return normalize_cache_control_in_anthropic_payload(value) + if isinstance(value, list): + return [_normalize_cache_control_value(item) for item in value] # mutable-ok: JSON wire format + return value + + +def normalize_cache_control_in_anthropic_payload(payload: dict) -> dict: # mutable-ok: as sibling sanitizers + """ + Return a copy of an Anthropic /v1/messages payload with every + ``cache_control`` entry reduced to ``{"type": }``, + recursing through message content blocks, system blocks, and tools. + + Anthropic itself accepts prompt-caching extensions such as ``ttl``, but + strict non-Anthropic implementations of the Messages API validate the field + literally and reject the whole request (``cache_control.ttl: 1h is not + supported``, ``cache_control.type is required``), which 400s clients like + Claude Code that always send cache hints. Non-dict ``cache_control`` values + are dropped entirely. The caller's payload is never mutated. + """ + return { # mutable-ok: JSON wire format, as sibling sanitizers + key: _normalized_cache_control(value) if key == "cache_control" else _normalize_cache_control_value(value) + for key, value in payload.items() + if key != "cache_control" or isinstance(value, dict) + } + + def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "anthropic-ratelimit-requests-limit" in headers: diff --git a/litellm/llms/openai_like/README.md b/litellm/llms/openai_like/README.md index e9aaafe48a1..e1409b81c35 100644 --- a/litellm/llms/openai_like/README.md +++ b/litellm/llms/openai_like/README.md @@ -54,7 +54,10 @@ That's it! The provider will be automatically loaded and available. "constraints": { "temperature_max": 1.0, "temperature_min": 0.0, - "temperature_min_with_n_gt_1": 0.3 + "temperature_min_with_n_gt_1": 0.3, + // /v1/messages providers only: keep Anthropic cache_control extensions + // such as ttl instead of stripping them down to {"type": ...} + "cache_control_ttl": true }, // Optional: Special handling flags diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 11dc236064d..29973fe2101 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -1,11 +1,13 @@ from typing import Any, Final import litellm +from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -19,7 +21,9 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): ``"/v1/messages"``. The inbound Anthropic payload (system, cache_control, thinking, tools, ...) is forwarded essentially unchanged to ``{api_base}/v1/messages``, so Anthropic-only features that the - Anthropic->OpenAI translation would otherwise drop are preserved. Response + Anthropic->OpenAI translation would otherwise drop are preserved. The one + exception is ``cache_control``, whose Anthropic-only extensions (``ttl``) + are stripped unless ``supports_cache_control_ttl`` says otherwise. Response parsing and streaming are inherited from the native Anthropic config. """ @@ -53,6 +57,35 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return False + def supports_cache_control_ttl(self) -> bool: + return False + + def transform_anthropic_messages_request( + self, + model: str, + messages: list[dict], # mutable-ok: matches dict-typed base signature + anthropic_messages_optional_request_params: dict, # mutable-ok: matches dict-typed base signature + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: matches dict-typed base signature + ) -> dict: # mutable-ok: matches dict-typed base signature + """ + Anthropic ignores prompt-caching hints it cannot honor, but strict + non-Anthropic implementations of the Messages API 400 the whole request + on Anthropic-only ``cache_control`` extensions (``cache_control.ttl: 1h + is not supported``), so unless the provider declares ttl support the + hints are reduced to their portable ``{"type": ...}`` core. + """ + request: Final = super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + if self.supports_cache_control_ttl(): + return request + return normalize_cache_control_in_anthropic_payload(request) + def get_complete_url( self, api_base: str | None, @@ -91,6 +124,9 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): def should_strip_billing_metadata(self) -> bool: return True + def supports_cache_control_ttl(self) -> bool: + return bool(self._provider.constraints.get("cache_control_ttl")) + def _resolve_api_key(self, api_key: str | None) -> str | None: return api_key or get_secret_str(self._provider.api_key_env) or litellm.api_key diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 33e677b000e..2cdd969b00f 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -317,3 +317,124 @@ def test_json_provider_messages_config_probes_capabilities_under_provider_slug() ) assert JSONProviderAnthropicMessagesConfig(provider).custom_llm_provider == "exampleprovider" assert OpenAILikeAnthropicMessagesConfig().custom_llm_provider == "anthropic" + + +def _cache_control_request_params() -> tuple[list, dict]: + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "write a regex for a US phone number", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ] + optional_params = { + "max_tokens": 256, + "system": [ + { + "type": "text", + "text": "You are Claude Code.", + "cache_control": {"type": "ephemeral", "ttl": "5m"}, + } + ], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object"}, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + return messages, optional_params + + +def test_request_strips_cache_control_ttl_everywhere(config): + """Regression: Claude Code always sends ``cache_control: {type: ephemeral, + ttl: 1h}``, and strict non-Anthropic /v1/messages validators 400 the whole + request on the ttl extension (``cache_control.ttl: 1h is not supported``).""" + messages, optional_params = _cache_control_request_params() + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_defaults_missing_cache_control_type_and_drops_non_dict(config): + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "a", "cache_control": {"ttl": "1h"}}, + {"type": "text", "text": "b", "cache_control": None}, + ], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + blocks = payload["messages"][0]["content"] + assert blocks[0]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in blocks[1] + + +def test_native_anthropic_config_keeps_cache_control_ttl(): + """Anthropic itself accepts ttl, so the normalization must stay scoped to + the OpenAI-like passthrough and never reach the native Anthropic path.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + messages, optional_params = _cache_control_request_params() + payload = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-sonnet-4-20250514", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} + + +def test_json_provider_constraint_opts_into_cache_control_ttl(): + from litellm.llms.openai_like.json_loader import SimpleProviderConfig + from litellm.llms.openai_like.messages.transformation import ( + JSONProviderAnthropicMessagesConfig, + ) + + base_data = {"base_url": "https://api.example.com/v1", "api_key_env": "EXAMPLE_API_KEY"} + strict = JSONProviderAnthropicMessagesConfig(SimpleProviderConfig(slug="strictprov", data=base_data)) + lenient = JSONProviderAnthropicMessagesConfig( + SimpleProviderConfig(slug="lenientprov", data={**base_data, "constraints": {"cache_control_ttl": True}}) + ) + + def transform(provider_config): + messages, optional_params = _cache_control_request_params() + return provider_config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transform(strict)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert transform(lenient)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} From d3268e4e184f8b7cde6cce4473e14b1acedfc751 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:26:04 -0700 Subject: [PATCH 05/81] test(mcp): wrap over-long connection error message calls --- .../proxy/_experimental/mcp_server/test_rest_endpoints.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 51b946c11b7..e66101f6177 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -2954,11 +2954,15 @@ class TestConnectionErrorMessage: assert secret not in message def test_connect_error_points_at_reachability(self): - message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0) + message = rest_endpoints._connection_error_message( + httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_timeout_error_message(self): - message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out"), "https://example.com", 30.0) + message = rest_endpoints._connection_error_message( + httpx.ConnectTimeout("timed out"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_http_status_error_includes_status_code(self): From 0baf376efd691e139902ca7933b83666ea459a41 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:07:05 -0700 Subject: [PATCH 06/81] fix(openai_like): scope cache_control normalization to Messages API locations Rewrite the sanitizer without recursion (the code-quality gate rejects new recursive functions) and only touch cache_control where the Messages API defines it: the request, system blocks, tools, message content blocks, and tool_result content. Application data such as tool_use.input and tool input_schema is left untouched even when it contains a cache_control key --- litellm/llms/anthropic/common_utils.py | 78 +++++++++++++++---- ..._like_anthropic_messages_transformation.py | 62 +++++++++++++++ 2 files changed, 124 insertions(+), 16 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 695e3a313ef..1ef14362601 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1302,37 +1302,83 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format -def _normalized_cache_control(cache_control: dict) -> dict: # mutable-ok: as sibling sanitizers +def _normalized_cache_control(cache_control: object) -> dict[str, str] | None: # mutable-ok: JSON wire format + if not isinstance(cache_control, Mapping): + return None cache_type: Final = cache_control.get("type") return {"type": cache_type if isinstance(cache_type, str) else "ephemeral"} # mutable-ok: JSON wire format -def _normalize_cache_control_value(value: object) -> object: - if isinstance(value, dict): - return normalize_cache_control_in_anthropic_payload(value) - if isinstance(value, list): - return [_normalize_cache_control_value(item) for item in value] # mutable-ok: JSON wire format - return value +def _with_portable_cache_control(block: Mapping[str, object]) -> dict[str, object]: # mutable-ok: JSON wire format + if "cache_control" not in block: + return dict(block) # mutable-ok: JSON wire format + normalized: Final = _normalized_cache_control(block["cache_control"]) + rest: Final = {key: value for key, value in block.items() if key != "cache_control"} # mutable-ok: JSON wire format + return rest if normalized is None else {**rest, "cache_control": normalized} # mutable-ok: JSON wire format -def normalize_cache_control_in_anthropic_payload(payload: dict) -> dict: # mutable-ok: as sibling sanitizers +def _with_portable_cache_control_in_blocks(blocks: object) -> object: + if isinstance(blocks, str) or not isinstance(blocks, Sequence): + return blocks + return [ # mutable-ok: JSON wire format + _with_portable_cache_control(block) if isinstance(block, Mapping) else block for block in blocks + ] + + +def _with_portable_cache_control_in_content_block(block: object) -> object: + if not isinstance(block, Mapping): + return block + portable: Final = _with_portable_cache_control(block) + if portable.get("type") != "tool_result" or "content" not in portable: + return portable + return { # mutable-ok: JSON wire format + **portable, + "content": _with_portable_cache_control_in_blocks(portable["content"]), + } + + +def _with_portable_cache_control_in_message(message: object) -> object: + if not isinstance(message, Mapping) or "content" not in message: + return message + content: Final = message["content"] + if isinstance(content, str) or not isinstance(content, Sequence): + return message + return { # mutable-ok: JSON wire format + **message, + "content": [_with_portable_cache_control_in_content_block(block) for block in content], + } + + +def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire format + payload: Mapping[str, object], +) -> dict[str, object]: """ Return a copy of an Anthropic /v1/messages payload with every - ``cache_control`` entry reduced to ``{"type": }``, - recursing through message content blocks, system blocks, and tools. + ``cache_control`` entry reduced to ``{"type": }`` + at the places the Messages API defines it: the request itself, system + blocks, tools, message content blocks, and ``tool_result`` content blocks. + Application data such as ``tool_use.input`` and tool ``input_schema`` is + never touched, even when it happens to contain a ``cache_control`` key. Anthropic itself accepts prompt-caching extensions such as ``ttl``, but strict non-Anthropic implementations of the Messages API validate the field literally and reject the whole request (``cache_control.ttl: 1h is not supported``, ``cache_control.type is required``), which 400s clients like - Claude Code that always send cache hints. Non-dict ``cache_control`` values - are dropped entirely. The caller's payload is never mutated. + Claude Code that send cache hints. Non-dict ``cache_control`` values are + dropped entirely. The caller's payload is never mutated. """ - return { # mutable-ok: JSON wire format, as sibling sanitizers - key: _normalized_cache_control(value) if key == "cache_control" else _normalize_cache_control_value(value) - for key, value in payload.items() - if key != "cache_control" or isinstance(value, dict) + portable: Final = _with_portable_cache_control(payload) + scoped: Final = { # mutable-ok: JSON wire format + key: ( + _with_portable_cache_control_in_blocks(value) + if key in ("system", "tools") + else [_with_portable_cache_control_in_message(message) for message in value] + if key == "messages" and isinstance(value, Sequence) and not isinstance(value, str) + else value + ) + for key, value in portable.items() } + return scoped def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 2cdd969b00f..d325492914e 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -438,3 +438,65 @@ def test_json_provider_constraint_opts_into_cache_control_ttl(): assert transform(strict)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} assert transform(lenient)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(config): + """Regression: the sanitizer must only touch ``cache_control`` where the + Messages API defines it (request, system, tools, content blocks, tool_result + content), never application data such as ``tool_use.input`` or a tool's + ``input_schema`` that happens to contain a ``cache_control`` key.""" + tool_input = {"cache_control": {"type": "ephemeral", "ttl": "1h"}, "query": "x"} + input_schema = { + "type": "object", + "properties": {"cache_control": {"type": "string", "ttl": "1h"}}, + } + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": tool_input}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "content": [ + {"type": "text", "text": "result", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], + }, + {"type": "text", "text": "plain string content stays", "cache_control": {"ttl": "1h"}}, + ], + }, + {"role": "user", "content": "a plain string message"}, + ] + optional_params = { + "max_tokens": 64, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "tools": [ + { + "name": "lookup", + "input_schema": input_schema, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["input_schema"] == input_schema + assert payload["messages"][0]["content"][0]["input"] == tool_input + tool_result = payload["messages"][1]["content"][0] + assert tool_result["cache_control"] == {"type": "ephemeral"} + assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][2] == {"role": "user", "content": "a plain string message"} From c32eb41aad3b7b087c7d0023a71876d3dea6511d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:17:13 -0700 Subject: [PATCH 07/81] feat(openai_like): let a passthrough deployment keep cache_control ttl via model_info.cache_control_ttl The supported_endpoints passthrough had no way to keep ttl for an upstream that honors it, so the deployment now opts in with model_info.cache_control_ttl: true, injected into the config the same way the providers.json constraint is for JSON providers --- .../messages/handler.py | 8 +- .../openai_like/messages/transformation.py | 16 +-- ...erimental_pass_through_messages_handler.py | 97 ++++++++++--------- ..._like_anthropic_messages_transformation.py | 17 ++++ 4 files changed, 84 insertions(+), 54 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 69985bcdaa3..b82903d6f87 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -99,6 +99,10 @@ def _deployment_passes_through_anthropic_messages(model_info: object) -> bool: return isinstance(supported_endpoints, (list, tuple)) and "/v1/messages" in supported_endpoints +def _deployment_supports_cache_control_ttl(model_info: object) -> bool: + return isinstance(model_info, dict) and model_info.get("cache_control_ttl") is True + + ####### ENVIRONMENT VARIABLES ################### # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() @@ -568,7 +572,9 @@ def anthropic_messages_handler( OpenAILikeAnthropicMessagesConfig, ) - anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig() + anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig( + cache_control_ttl=_deployment_supports_cache_control_ttl(kwargs.get("model_info")), + ) if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. if _should_route_to_responses_api(custom_llm_provider, original_model, model): diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 29973fe2101..ac99617521c 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -23,10 +23,15 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): ``{api_base}/v1/messages``, so Anthropic-only features that the Anthropic->OpenAI translation would otherwise drop are preserved. The one exception is ``cache_control``, whose Anthropic-only extensions (``ttl``) - are stripped unless ``supports_cache_control_ttl`` says otherwise. Response - parsing and streaming are inherited from the native Anthropic config. + are stripped unless the deployment opts in with + ``model_info.cache_control_ttl: true``. Response parsing and streaming are + inherited from the native Anthropic config. """ + def __init__(self, cache_control_ttl: bool = False) -> None: + super().__init__() + self._cache_control_ttl: Final = cache_control_ttl + def validate_anthropic_messages_environment( self, headers: dict[str, str], @@ -58,7 +63,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): return False def supports_cache_control_ttl(self) -> bool: - return False + return self._cache_control_ttl def transform_anthropic_messages_request( self, @@ -114,7 +119,7 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): """ def __init__(self, provider: SimpleProviderConfig): - super().__init__() + super().__init__(cache_control_ttl=bool(provider.constraints.get("cache_control_ttl"))) self._provider = provider @property @@ -124,9 +129,6 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): def should_strip_billing_metadata(self) -> bool: return True - def supports_cache_control_ttl(self) -> bool: - return bool(self._provider.constraints.get("cache_control_ttl")) - def _resolve_api_key(self, api_key: str | None) -> str | None: return api_key or get_secret_str(self._provider.api_key_env) or litellm.api_key diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index ad4c3d6bfbb..e819433c269 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -296,21 +296,15 @@ async def test_bedrock_converse_budget_tokens_preserved(): mock_acompletion.assert_called_once() call_kwargs = mock_acompletion.call_args.kwargs - print( - "acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str) - ) + print("acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str)) # Verify thinking parameter is passed through with budget_tokens preserved thinking_param = call_kwargs.get("thinking") - assert ( - thinking_param is not None - ), "thinking parameter should be passed to acompletion" - assert ( - thinking_param.get("type") == "enabled" - ), "thinking.type should be 'enabled'" - assert ( - thinking_param.get("budget_tokens") == 1024 - ), f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + assert thinking_param is not None, "thinking parameter should be passed to acompletion" + assert thinking_param.get("type") == "enabled", "thinking.type should be 'enabled'" + assert thinking_param.get("budget_tokens") == 1024, ( + f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + ) def test_openai_model_with_thinking_converts_to_reasoning(): @@ -342,23 +336,18 @@ def test_openai_model_with_thinking_converts_to_reasoning(): call_kwargs = mock_responses.call_args.kwargs # Verify reasoning is set (converted from thinking) - assert ( - "reasoning" in call_kwargs - ), "reasoning should be passed to litellm.responses" + assert "reasoning" in call_kwargs, "reasoning should be passed to litellm.responses" # budget_tokens=1024 -> effort="low" (at the LOW budget threshold) # reasoning_auto_summary is False by default, so no summary key expected_reasoning = {"effort": "low"} assert call_kwargs["reasoning"] == expected_reasoning, ( - f"reasoning should be {expected_reasoning} for budget_tokens=1024, " - f"got {call_kwargs.get('reasoning')}" + f"reasoning should be {expected_reasoning} for budget_tokens=1024, got {call_kwargs.get('reasoning')}" ) assert "summary" not in call_kwargs["reasoning"] # Verify thinking is NOT passed directly to the Responses API - assert ( - "thinking" not in call_kwargs - ), "thinking should NOT be passed directly to litellm.responses" + assert "thinking" not in call_kwargs, "thinking should NOT be passed directly to litellm.responses" class TestThinkingParameterTransformation: @@ -411,9 +400,7 @@ class TestThinkingParameterTransformation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "detailed"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "detailed"}} finally: litellm.reasoning_auto_summary = original @@ -611,9 +598,9 @@ class TestThinkingSummaryPreservation: mock_responses.assert_called_once() call_kwargs = mock_responses.call_args.kwargs reasoning = call_kwargs["reasoning"] - assert ( - reasoning["summary"] == "concise" - ), f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + assert reasoning["summary"] == "concise", ( + f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + ) def test_responses_adapter_preserves_summary(self): """translate_thinking_to_reasoning should include summary when user provides it.""" @@ -622,9 +609,7 @@ class TestThinkingSummaryPreservation: ) thinking = {"type": "enabled", "budget_tokens": 5000, "summary": "concise"} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high", "summary": "concise"} def test_responses_adapter_no_summary_by_default(self): @@ -638,11 +623,7 @@ class TestThinkingSummaryPreservation: try: litellm.reasoning_auto_summary = False thinking = {"type": "enabled", "budget_tokens": 5000} - result = ( - LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high"} assert result is not None and "summary" not in result finally: @@ -659,9 +640,7 @@ class TestThinkingSummaryPreservation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "concise"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "concise"}} def test_translate_thinking_for_model_disabled_stays_plain_string_when_auto_summary_enabled(self): """Disabled thinking must stay a plain string even when reasoning_auto_summary is on.""" @@ -807,9 +786,7 @@ def test_presanitized_flag_not_leaked_to_provider_params(): def fake_base_handler(*args, **kwargs): captured.update(kwargs) - captured["optional"] = kwargs.get( - "anthropic_messages_optional_request_params", {} - ) + captured["optional"] = kwargs.get("anthropic_messages_optional_request_params", {}) return "stub" with patch.object( @@ -974,6 +951,38 @@ def test_gate_passthrough_skipped_when_only_chat_completions_supported(monkeypat assert "config" not in captured +@pytest.mark.parametrize( + "model_info, expected_ttl_support", + [ + ({"supported_endpoints": ["/v1/messages"]}, False), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": True}, True), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": "yes"}, False), + ], +) +def test_gate_passthrough_forwards_cache_control_ttl_only_when_deployment_opts_in( + monkeypatch, model_info, expected_ttl_support +): + """The passthrough config strips cache_control.ttl unless the deployment sets + model_info.cache_control_ttl to exactly true.""" + from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( + anthropic_messages_handler, + ) + + captured, _ = _gate_stubs(monkeypatch) + + result = anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "Hello"}], + model="openai/some-model", + api_key="sk-test", + api_base="https://host/v1", + model_info=model_info, + ) + + assert result == "native-passthrough" + assert captured["config"].supports_cache_control_ttl() is expected_ttl_support + + def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): """Regional and provider-prefixed Claude 4.8+/5 entries carry ``supports_mid_conversation_system``, but the bare first-party keys @@ -987,9 +996,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys import litellm - cost_map_path = os.path.join( - os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json" - ) + cost_map_path = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json") with open(cost_map_path) as f: cost_map = json.load(f) rules = cost_map["fallback_generalizations"]["rules"] @@ -1028,9 +1035,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys ("perplexity/sonar", "sonar", "https://api.perplexity.ai/chat/completions"), ], ) -async def test_messages_strips_provider_prefix_exactly_once( - requested_model, expected_wire_model, expected_url -): +async def test_messages_strips_provider_prefix_exactly_once(requested_model, expected_wire_model, expected_url): """ BerriAI/litellm#37716: only the leading provider segment may be stripped on the way upstream. diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index d325492914e..e33b03afdff 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -414,6 +414,23 @@ def test_native_anthropic_config_keeps_cache_control_ttl(): assert payload["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} +def test_deployment_opt_in_keeps_cache_control_ttl(): + config = OpenAILikeAnthropicMessagesConfig(cache_control_ttl=True) + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 16}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + def test_json_provider_constraint_opts_into_cache_control_ttl(): from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.llms.openai_like.messages.transformation import ( From d4fc54a11d1ac18f10c33741b760b99716faac93 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 08:25:42 +0000 Subject: [PATCH 08/81] chore(techdebt): clear fresh debt from the 2026-08-31 window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 6 +-- litellm/llms/gigachat/authenticator.py | 28 +++++--------- litellm/llms/gigachat/chat/streaming.py | 1 - litellm/llms/gigachat/chat/transformation.py | 38 +++++-------------- .../llms/gigachat/embedding/transformation.py | 14 ++----- .../gigachat/passthrough/transformation.py | 2 - litellm/llms/gigachat/utils.py | 1 - litellm/passthrough/main.py | 3 -- .../llm_passthrough_endpoints.py | 4 +- .../router_strategy/test_complexity_router.py | 6 ++- type-discipline-budget.json | 8 ++-- 11 files changed, 35 insertions(+), 76 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index a07b9352659..d84aacfeaf0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5607 }, "reportMissingTypeArgument": { - "limit": 15310 + "limit": 15308 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38368 + "limit": 38367 }, "reportUnknownParameterType": { "limit": 19633 }, "reportUnknownVariableType": { - "limit": 29908 + "limit": 29906 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index d6b217d5746..73086ba395b 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -8,6 +8,7 @@ Based on official GigaChat SDK authentication flow. import time import uuid from collections.abc import Mapping +from types import MappingProxyType from typing import Final import httpx @@ -32,8 +33,8 @@ GIGACHAT_SCOPE: Final = "GIGACHAT_API_PERS" # Token expiry buffer in milliseconds (refresh token 60s before expiry) TOKEN_EXPIRY_BUFFER_MS: Final = 60000 -# Cache for access tokens _token_cache: Final = InMemoryCache() +_NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) class GigaChatAuthError(BaseLLMException): @@ -80,10 +81,9 @@ def get_access_token( Raises: GigaChatAuthError: If authentication fails """ - if not litellm_params: - litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default + params: Final = litellm_params or _NO_LITELLM_PARAMS - access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") + access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") if access_token: return access_token @@ -94,24 +94,20 @@ def get_access_token( message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope() - effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope() + effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url() - # Check cache cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: _token, _expires_at = cached - # Check if token is still valid (with buffer) if time.time() * 1000 < _expires_at - TOKEN_EXPIRY_BUFFER_MS: verbose_logger.debug("Using cached GigaChat access token") return _token - # Request new token new_token, new_expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str if new_expires_at: - # Cache token ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds) @@ -126,10 +122,9 @@ async def get_access_token_async( litellm_params: Mapping[str, object] | None = None, ) -> str: """Async version of get_access_token.""" - if not litellm_params: - litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default + params: Final = litellm_params or _NO_LITELLM_PARAMS - access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") + access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") if access_token: return access_token @@ -140,10 +135,9 @@ async def get_access_token_async( message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope() - effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope() + effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url() - # Check cache cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: @@ -152,11 +146,9 @@ async def get_access_token_async( verbose_logger.debug("Using cached GigaChat access token") return _token - # Request new token new_token, new_expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str if new_expires_at: - # Cache token ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 2875b30232e..0a4cbd8e520 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -52,7 +52,6 @@ class GigaChatModelResponseIterator: tool_use: ChatCompletionToolCallChunk | None = None # rebind-ok: conditionally assigned on function_call finish_reason: str | None = chunk_finish_reason - # Handle function_call in stream raw_function_call: Final = delta.get("function_call") if chunk_finish_reason == "function_call" and isinstance(raw_function_call, Mapping) and raw_function_call: func_call: Final[Mapping[str, object]] = raw_function_call diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 8f23c5175ec..991a93ccb21 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -111,11 +111,9 @@ class GigaChatConfig(BaseConfig): """ Set up headers with OAuth token. """ - # Get access token credentials: Final = api_key or get_secret_str("GIGACHAT_CREDENTIALS") or get_secret_str("GIGACHAT_API_KEY") access_token: Final = get_access_token(credentials=credentials, litellm_params=litellm_params) - # Store credentials for image uploads self._current_credentials = credentials self._current_api_base = api_base @@ -208,18 +206,16 @@ class GigaChatConfig(BaseConfig): def _convert_tools_to_functions(self, tools: Sequence) -> Sequence[dict]: """Convert OpenAI tools format to GigaChat functions format.""" - functions: Final[list[dict]] = [] # mutable-ok: accumulator for building functions list - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "function": - func = tool.get("function", {}) - functions.append( - { - "name": func.get("name", ""), - "description": func.get("description", ""), - "parameters": func.get("parameters", {}), - } - ) - return functions + return [ + { + "name": function.get("name", ""), + "description": function.get("description", ""), + "parameters": function.get("parameters", {}), + } + for function in ( + tool.get("function", {}) for tool in tools if isinstance(tool, dict) and tool.get("type") == "function" + ) + ] def _map_tool_choice(self, tool_choice: str | Mapping[str, object]) -> str | Mapping[str, object] | None: """ @@ -299,7 +295,6 @@ class GigaChatConfig(BaseConfig): if part.get("type") == "text": texts.append(part.get("text", "")) elif part.get("type") == "image_url": - # Extract image URL and upload to GigaChat image_url: object = part.get("image_url", {}) upload_url: str if isinstance(image_url, str): @@ -322,16 +317,13 @@ class GigaChatConfig(BaseConfig): headers: Mapping[str, object], ) -> dict: # mutable-ok: request payload sent to httpx """Transform OpenAI request to GigaChat format.""" - # Transform messages giga_messages: Final = self._transform_messages(messages) - # Build request request_data: Final[dict[str, object]] = { "model": model.replace("gigachat/", ""), "messages": giga_messages, } - # Add optional params for key in [ "temperature", "top_p", @@ -343,7 +335,6 @@ class GigaChatConfig(BaseConfig): if key in optional_params: request_data[key] = optional_params[key] - # Add functions if present if "functions" in optional_params: request_data["functions"] = optional_params["functions"] if "function_call" in optional_params: @@ -358,10 +349,8 @@ class GigaChatConfig(BaseConfig): for i, msg in enumerate(messages): message = dict(msg) - # Remove unsupported fields message.pop("name", None) - # Transform roles role = message.get("role", "user") if role == "developer": message["role"] = "system" @@ -374,18 +363,15 @@ class GigaChatConfig(BaseConfig): if not isinstance(content, str) or not is_valid_json(content): message["content"] = json.dumps(content, ensure_ascii=False) - # Handle None content if message.get("content") is None: message["content"] = "" - # Handle list content (multimodal) - extract text and images content = message.get("content") if isinstance(content, list): message["content"], attachments = self._transform_list_content(content) if attachments: message["attachments"] = attachments - # Transform tool_calls to function_call tool_calls = message.get("tool_calls") if tool_calls and isinstance(tool_calls, list) and len(tool_calls) > 0: tool_call = tool_calls[0] @@ -436,13 +422,11 @@ class GigaChatConfig(BaseConfig): message_data = choice.get("message", {}) finish_reason = choice.get("finish_reason", "stop") - # Transform function_call to tool_calls or content if finish_reason == "function_call" and message_data.get("function_call"): func_call = message_data["function_call"] args = func_call.get("arguments", {}) if is_structured_output: - # Convert to content for structured output if isinstance(args, dict): content = json.dumps(args, ensure_ascii=False) else: @@ -452,7 +436,6 @@ class GigaChatConfig(BaseConfig): message_data.pop("functions_state_id", None) finish_reason = "stop" else: - # Convert to tool_calls format if isinstance(args, dict): args = json.dumps(args, ensure_ascii=False) message_data["tool_calls"] = [ @@ -468,7 +451,6 @@ class GigaChatConfig(BaseConfig): message_data.pop("function_call", None) finish_reason = "tool_calls" - # Clean up GigaChat-specific fields message_data.pop("functions_state_id", None) choices.append( diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 2ec8324e33c..0db4475be8f 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -112,18 +112,10 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): "input": ["text1", "text2", ...] } """ - # Normalize input to list - if isinstance(input, str): - input_list: list = [input] # rebind-ok: locally scoped conversion - else: - input_list = input - - # Remove gigachat/ prefix from model if present - model = model.removeprefix("gigachat/") # rebind-ok: parameter reassignment for normalization - + normalized_input: Final = [input] if isinstance(input, str) else input # mutable-ok: preserve list API return { - "model": model, - "input": input_list, + "model": model.removeprefix("gigachat/"), + "input": normalized_input, } def transform_embedding_response( diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index a0edc6f5682..e1f73d04275 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -60,7 +60,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): """ Set up headers with OAuth token. """ - # Get access token access_token: Final = get_access_token(credentials=api_key, litellm_params=litellm_params) headers["Authorization"] = f"Bearer {access_token}" # rebind-ok: mutating for OAuth setup @@ -82,7 +81,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager - # cost tracking only for completions and embeddings if "completions" in endpoint: provider_chat_config: Final = ProviderConfigManager.get_provider_chat_config( provider=LlmProviders(custom_llm_provider), diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index cbb35cd1b57..ce7e848ed7f 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -4,7 +4,6 @@ from typing import Final from litellm.secret_managers.main import get_secret_str from litellm.types.utils import PromptTokensDetailsWrapper, Usage -# GigaChat API endpoint GIGACHAT_BASE_URL: Final = "https://gigachat.devices.sberbank.ru/api/v1" diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 9095cee15a9..689c34b7a88 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -113,10 +113,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): ) ) - # Compliant: Save a strong reference to prevent GC self._background_tasks.add(task) - # Remove the task from the set when it finishes to avoid memory leaks task.add_done_callback(self._background_tasks.discard) except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception( @@ -578,7 +576,6 @@ def llm_passthrough_route( else: return response except Exception as e: - # provider_config is guaranteed non-None here due to the earlier guard assert provider_config is not None raise base_llm_http_handler._handle_error( e=e, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 78d8ce296b8..b48b8d81494 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1731,7 +1731,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: def get_vertex_pass_through_handler( - call_type: Literal["discovery", "aiplatform"], # noqa: UP037 + call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here ) -> BaseVertexAIPassThroughHandler: if call_type == "discovery": return VertexAIDiscoveryPassThroughHandler() @@ -2961,7 +2961,6 @@ async def handle_gigachat_passthrough_router_model( """ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - # Detect streaming based on request body is_streaming: Final = request_body.get("stream", False) # pyright: ignore[reportUnknownVariableType] # request_body is dict[Unknown, Unknown] data: dict[str, Any] = await _read_request_body( @@ -2997,7 +2996,6 @@ async def handle_gigachat_passthrough_router_model( data["json"] = request_body data["custom_llm_provider"] = "gigachat" - # Remove sensitive keys from data keys: Final = [ # mutable-ok: list of keys to remove from data "gigachat_auth_url", "gigachat_access_token", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 1ec8be88c9b..93803ce1005 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -10268,7 +10268,8 @@ class TestContextWindowEscalation: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config=_tier_config(session_affinity=True), ) - session_kwargs = lambda: {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} # noqa: E731 + def session_kwargs() -> dict[str, object]: + return {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} first = await router.async_pre_routing_hook( model="test-router", request_kwargs=session_kwargs(), messages=_OVERSIZED_TURNS @@ -10291,7 +10292,8 @@ class TestContextWindowEscalation: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config=_tier_config(session_affinity=True), ) - session_kwargs = lambda: {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} # noqa: E731 + def session_kwargs() -> dict[str, object]: + return {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} pinned = await router.async_pre_routing_hook( model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}] diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 83c49afb538..f65ebd24599 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,12 +1,12 @@ { "LIT001": { - "limit": 22403 + "limit": 22402 }, "LIT002": { "limit": 26780 }, "LIT003": { - "limit": 269 + "limit": 268 }, "LIT004": { "limit": 40 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16512 + "limit": 16511 }, "LIT011": { - "limit": 5537 + "limit": 5535 }, "LIT012": { "limit": 4495 From ab1161344199539bc8dec51161fb59d6d0acf703 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 5 Aug 2026 18:01:29 +0000 Subject: [PATCH 09/81] fix(bedrock): strip client_metadata from converse additionalModelRequestFields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/chat/converse_transformation.py | 1 + .../chat/test_converse_transformation.py | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 395d99a4caa..9e40cb5ee5b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1324,6 +1324,7 @@ class AmazonConverseConfig(BaseConfig): ) additional_request_params.pop("parallel_tool_calls", None) + additional_request_params.pop("client_metadata", None) # Only set the topK value in for models that support it additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params)) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 63f895e1819..bd68857d664 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -979,6 +979,30 @@ def test_config_blocks_do_not_leak_into_inference_config(): assert data["serviceTier"] == {"type": "priority"} +def test_client_metadata_stripped_from_converse_request(): + """``client_metadata`` sent by codex must not reach Bedrock as a passthrough model field. + + Converse forwards ``additionalModelRequestFields`` verbatim to the model, and Anthropic + rejects the request with "client_metadata: Extra inputs are not permitted". + """ + config = AmazonConverseConfig() + + data = config._transform_request_helper( + model="anthropic.claude-opus-4-8", + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "anthropic_beta": ["computer-use-2025-01-24"], + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + fields = data.get("additionalModelRequestFields", {}) + assert "client_metadata" not in fields + assert fields["anthropic_beta"] == ["computer-use-2025-01-24"] + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost From 2063c29f5d95f8dd00eef3fd7dfcbc1df05787b8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 18:45:26 +0000 Subject: [PATCH 10/81] fix(anthropic): upgrade legacy thinking to adaptive on adaptive-only models for chat and Bedrock Converse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/transformation.py | 3 + litellm/llms/anthropic/common_utils.py | 47 +++++++++++- .../messages/transformation.py | 47 +----------- .../bedrock/chat/converse_transformation.py | 3 + .../test_anthropic_chat_transformation.py | 25 ++++++ .../chat/test_converse_transformation.py | 76 +++++++++++++++++++ 6 files changed, 154 insertions(+), 47 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e1387a9068c..319eecfac2c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1544,6 +1544,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params.pop("thinking", None) else: optional_params["thinking"] = value + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider + ) elif param == "reasoning_effort": # Accept both string ("low") and dict ({"effort": "low", # "summary": "concise"}). The Responses->Chat parser keeps the diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 9871001bf66..ede93c6deb2 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -13,7 +13,12 @@ import httpx from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm -from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME +from litellm.constants import ( + DEFAULT_MODEL_CREATED_AT_TIME, + DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, +) from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, ) @@ -490,6 +495,46 @@ class AnthropicModelInfo(BaseLLMModelInfo): ) optional_params.pop("thinking", None) + @staticmethod + def translate_legacy_thinking_for_adaptive_model( + model: str, + optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param, as in maybe_drop_disabled_thinking + custom_llm_provider: str, + ) -> None: + """Translate legacy ``thinking.type=enabled`` to adaptive for the + adaptive-thinking models that reject it (4.7+ and the 5 families). + Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the + legacy shape natively, so it is forwarded verbatim and the caller's + ``budget_tokens`` cap keeps applying. Caller-provided + ``output_config.effort`` is never overridden. + """ + if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): + return + if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): + return + thinking: Final = optional_params.get("thinking") + if not isinstance(thinking, dict) or thinking.get("type") != "enabled": + return + + budget: Final = int(thinking.get("budget_tokens") or 0) + if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and ( + AnthropicModelInfo._supports_model_capability(model, "supports_xhigh_reasoning_effort", custom_llm_provider) + ): + effort = "xhigh" + elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET: + effort = "high" + elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: + effort = "medium" + else: + effort = "low" + + optional_params["thinking"] = {"type": "adaptive"} + existing_output_config = optional_params.get("output_config") + if not isinstance(existing_output_config, dict): + existing_output_config = {} + existing_output_config.setdefault("effort", effort) + optional_params["output_config"] = existing_output_config + def is_effort_used( self, optional_params: dict | None, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 3d62b8b4784..988f81c9eb4 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -3,11 +3,6 @@ from typing import Any, Final import httpx -from litellm.constants import ( - DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, - DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, -) from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import verbose_logger @@ -400,46 +395,6 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): existing_output_config.setdefault("effort", mapped_effort) optional_params["output_config"] = existing_output_config - @staticmethod - def _translate_legacy_thinking_for_adaptive_model( - model: str, optional_params: dict, custom_llm_provider: str - ) -> None: - """Translate legacy ``thinking.type=enabled`` to adaptive for the - adaptive-thinking models that reject it (4.7+ and the 5 families). - Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the - legacy shape natively, so it is forwarded verbatim and the caller's - ``budget_tokens`` cap keeps applying. Caller-provided - ``output_config.effort`` is never overridden. - """ - from litellm.llms.anthropic.chat.transformation import AnthropicConfig - - if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): - return - if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): - return - thinking: Final = optional_params.get("thinking") - if not isinstance(thinking, dict) or thinking.get("type") != "enabled": - return - - budget: Final = int(thinking.get("budget_tokens") or 0) - if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and ( - AnthropicConfig._supports_effort_level(model, "xhigh", custom_llm_provider) - ): - effort = "xhigh" - elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET: - effort = "high" - elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: - effort = "medium" - else: - effort = "low" - - optional_params["thinking"] = {"type": "adaptive"} - existing_output_config = optional_params.get("output_config") - if not isinstance(existing_output_config, dict): - existing_output_config = {} - existing_output_config.setdefault("effort", effort) - optional_params["output_config"] = existing_output_config - @staticmethod def _translate_adaptive_effort_for_non_adaptive_model( model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str @@ -606,7 +561,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): custom_llm_provider=self._resolved_provider, ) - self._translate_legacy_thinking_for_adaptive_model( + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( model=model, optional_params=anthropic_messages_optional_request_params, custom_llm_provider=self._resolved_provider, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 395d99a4caa..0a378dfc11c 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -934,6 +934,9 @@ class AmazonConverseConfig(BaseConfig): litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model) else: optional_params["thinking"] = value + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=model, optional_params=optional_params, custom_llm_provider="bedrock" + ) elif param == "reasoning_effort" and isinstance(value, str): self._handle_reasoning_effort_parameter( model=model, reasoning_effort=value, optional_params=optional_params diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 25e2c3cda80..4f30e7d10f0 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -3127,6 +3127,31 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model( ) +@pytest.mark.parametrize( + "model,budget_tokens,expected", + [ + ("claude-opus-4-8", 4096, ({"type": "adaptive"}, {"effort": "high"})), + ("claude-opus-4-7", 24000, ({"type": "adaptive"}, {"effort": "xhigh"})), + ("claude-opus-4-6", 4096, ({"type": "enabled", "budget_tokens": 4096}, None)), + ("claude-sonnet-4-5-20250929", 4096, ({"type": "enabled", "budget_tokens": 4096}, None)), + ], +) +def test_legacy_thinking_translated_to_adaptive_on_adaptive_only_models(model, budget_tokens, expected): + """Adaptive-only models reject thinking={type: enabled} with a 400, so the + legacy shape must be upgraded to adaptive + output_config.effort on + /chat/completions too, while models that accept it keep the caller's budget.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget_tokens}, "max_tokens": 64000}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert (result["thinking"], result.get("output_config")) == expected + + @pytest.mark.parametrize( "bad_value", [ diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 63f895e1819..b729c9366cc 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -6269,6 +6269,82 @@ def test_adaptive_thinking_passes_through_on_46_plus_converse(model): assert optional_params.get("thinking") == {"type": "adaptive"} +@pytest.mark.parametrize( + "model,budget_tokens,expected_effort", + [ + ("anthropic.claude-opus-4-8", 4096, "high"), + ("us.anthropic.claude-opus-4-8", 2000, "low"), + ("global.anthropic.claude-opus-4-8", 12000, "xhigh"), + ("us.anthropic.claude-opus-4-7", 3000, "medium"), + ("anthropic.claude-fable-5", 4096, "high"), + ], +) +def test_legacy_thinking_translated_to_adaptive_on_adaptive_only_converse(model, budget_tokens, expected_effort): + """Adaptive-only models (4.7+, 5 families) reject thinking={type: enabled} + with a 400 on Bedrock Converse, so the legacy shape from callers like Claude + Code must be upgraded to thinking={type: adaptive} + output_config.effort + derived from budget_tokens, matching the /v1/messages passthrough.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget_tokens}, "max_tokens": 64000}, + optional_params={}, + model=model, + drop_params=False, + ) + request = config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert request["additionalModelRequestFields"]["thinking"] == {"type": "adaptive"} + assert request["additionalModelRequestFields"]["output_config"] == {"effort": expected_effort} + + +def test_legacy_thinking_translation_keeps_caller_output_config_effort_converse(): + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={ + "output_config": {"effort": "low"}, + "thinking": {"type": "enabled", "budget_tokens": 12000}, + "max_tokens": 64000, + }, + optional_params={}, + model="anthropic.claude-opus-4-8", + drop_params=False, + ) + + assert optional_params["thinking"] == {"type": "adaptive"} + assert optional_params["output_config"] == {"effort": "low"} + + +@pytest.mark.parametrize( + "model", + [ + "us.anthropic.claude-opus-4-6", + "anthropic.claude-3-5-sonnet-20241022-v2:0", + ], +) +def test_legacy_thinking_forwarded_verbatim_when_model_accepts_it_converse(model): + """The 4.6 family and pre-adaptive models accept thinking={type: enabled} + natively, so the caller's budget_tokens cap must keep applying.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4096}, "max_tokens": 8192}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert optional_params["thinking"] == {"type": "enabled", "budget_tokens": 4096} + assert "output_config" not in optional_params + + def test_adaptive_thinking_dropped_when_max_tokens_too_small_converse(): """When max_tokens can't fit even the minimum thinking budget, the raw adaptive block must be dropped entirely rather than translated, so the From b8dd27a77fdaaa390761a99bf27737a255c37e3a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:00:46 +0000 Subject: [PATCH 11/81] style(anthropic): keep mutable-ok annotation within ruff format width Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/common_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index ede93c6deb2..f2fcbc6c232 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -498,7 +498,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): @staticmethod def translate_legacy_thinking_for_adaptive_model( model: str, - optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param, as in maybe_drop_disabled_thinking + optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers custom_llm_provider: str, ) -> None: """Translate legacy ``thinking.type=enabled`` to adaptive for the From ed4343a02645e19590657ae257e6ae43b95f8e48 Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 1 Sep 2026 19:04:32 +0000 Subject: [PATCH 12/81] fix(bedrock): scope client_metadata drop to anthropic converse models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/chat/converse_transformation.py | 8 ++++++- .../chat/test_converse_transformation.py | 24 ++++++++++++++++--- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 9e40cb5ee5b..8d1d905a129 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1324,7 +1324,13 @@ class AmazonConverseConfig(BaseConfig): ) additional_request_params.pop("parallel_tool_calls", None) - additional_request_params.pop("client_metadata", None) + + if base_model.startswith("anthropic") and additional_request_params.pop("client_metadata", None) is not None: + litellm.verbose_logger.debug( + "Bedrock Converse: dropping `client_metadata` for model=%s, Anthropic rejects it with " + "'client_metadata: Extra inputs are not permitted'", + model, + ) # Only set the topK value in for models that support it additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params)) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index bd68857d664..3600e366b5a 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -979,8 +979,9 @@ def test_config_blocks_do_not_leak_into_inference_config(): assert data["serviceTier"] == {"type": "priority"} -def test_client_metadata_stripped_from_converse_request(): - """``client_metadata`` sent by codex must not reach Bedrock as a passthrough model field. +@pytest.mark.parametrize("model", ["anthropic.claude-opus-4-8", "us.anthropic.claude-opus-4-8"]) +def test_client_metadata_stripped_for_anthropic_converse_request(model): + """``client_metadata`` sent by codex must not reach Anthropic as a passthrough model field. Converse forwards ``additionalModelRequestFields`` verbatim to the model, and Anthropic rejects the request with "client_metadata: Extra inputs are not permitted". @@ -988,7 +989,7 @@ def test_client_metadata_stripped_from_converse_request(): config = AmazonConverseConfig() data = config._transform_request_helper( - model="anthropic.claude-opus-4-8", + model=model, system_content_blocks=[], optional_params={ "maxTokens": 16, @@ -1003,6 +1004,23 @@ def test_client_metadata_stripped_from_converse_request(): assert fields["anthropic_beta"] == ["computer-use-2025-01-24"] +def test_client_metadata_kept_for_non_anthropic_converse_request(): + """Only Anthropic is known to reject ``client_metadata``, so other families keep the passthrough.""" + config = AmazonConverseConfig() + + data = config._transform_request_helper( + model="amazon.nova-pro-v1:0", + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost From b96121efb173965405aa521ffbbba026ce73d3a4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:10:22 +0000 Subject: [PATCH 13/81] refactor(anthropic): build adaptive output_config in one shot in legacy thinking helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/common_utils.py | 37 +++++++++++++++----------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index f2fcbc6c232..17fa8022388 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -516,24 +516,29 @@ class AnthropicModelInfo(BaseLLMModelInfo): if not isinstance(thinking, dict) or thinking.get("type") != "enabled": return - budget: Final = int(thinking.get("budget_tokens") or 0) - if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and ( + effort: Final = AnthropicModelInfo._legacy_budget_to_effort( + model=model, + budget_tokens=int(thinking.get("budget_tokens") or 0), + custom_llm_provider=custom_llm_provider, + ) + existing_output_config: Final = optional_params.get("output_config") + optional_params["thinking"] = {"type": "adaptive"} + optional_params["output_config"] = { + "effort": effort, + **(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})), + } + + @staticmethod + def _legacy_budget_to_effort(model: str, budget_tokens: int, custom_llm_provider: str) -> str: + if budget_tokens >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and ( AnthropicModelInfo._supports_model_capability(model, "supports_xhigh_reasoning_effort", custom_llm_provider) ): - effort = "xhigh" - elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET: - effort = "high" - elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: - effort = "medium" - else: - effort = "low" - - optional_params["thinking"] = {"type": "adaptive"} - existing_output_config = optional_params.get("output_config") - if not isinstance(existing_output_config, dict): - existing_output_config = {} - existing_output_config.setdefault("effort", effort) - optional_params["output_config"] = existing_output_config + return "xhigh" + if budget_tokens >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET: + return "high" + if budget_tokens >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: + return "medium" + return "low" def is_effort_used( self, From bba951c5ebe20871caab9848c474585c91fa3535 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 1 Sep 2026 19:26:19 +0000 Subject: [PATCH 14/81] fix(bedrock): drop client_metadata for ARNs that hide the model family Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/chat/converse_transformation.py | 4 +- litellm/llms/bedrock/common_utils.py | 9 ++++ .../chat/test_converse_transformation.py | 43 +++++++++++++++++++ 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 8d1d905a129..99d45bfc94b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -86,6 +86,7 @@ from litellm.utils import ( from ..common_utils import ( BedrockError, BedrockModelInfo, + bedrock_arn_hides_model_family, bedrock_converse_supports_parallel_tool_use_config, get_anthropic_beta_from_headers, get_bedrock_tool_name, @@ -1325,7 +1326,8 @@ class AmazonConverseConfig(BaseConfig): additional_request_params.pop("parallel_tool_calls", None) - if base_model.startswith("anthropic") and additional_request_params.pop("client_metadata", None) is not None: + drops_client_metadata: Final = base_model.startswith("anthropic") or bedrock_arn_hides_model_family(model) + if drops_client_metadata and additional_request_params.pop("client_metadata", None) is not None: litellm.verbose_logger.debug( "Bedrock Converse: dropping `client_metadata` for model=%s, Anthropic rejects it with " "'client_metadata: Extra inputs are not permitted'", diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9cbceb4880c..3b82d98ceee 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -720,6 +720,15 @@ def get_bedrock_base_model(model: str) -> str: return model +def bedrock_arn_hides_model_family(model: str) -> bool: + """ + True for an ARN-addressed model whose base name carries no ``provider.model`` + id, such as an application inference profile or a provisioned throughput ARN. + Callers that gate behavior on the model family cannot resolve one here. + """ + return "arn:" in model.lower() and "." not in get_bedrock_base_model(model) + + def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: return any( (litellm.model_cost.get(candidate) or {}).get("supports_parallel_tool_use_config") is True diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 3600e366b5a..423e3a5a7c3 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -1021,6 +1021,49 @@ def test_client_metadata_kept_for_non_anthropic_converse_request(): assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} +@pytest.mark.parametrize( + "model", + [ + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456", + "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/abcdef123456", + ], +) +def test_client_metadata_stripped_for_arn_models_converse(model): + """An ARN hides which family serves the request, and pointing one at Claude is how + teams route codex traffic, so the field has to go there too or the 400 comes back.""" + config = AmazonConverseConfig() + + data = config._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + assert "client_metadata" not in data.get("additionalModelRequestFields", {}) + + +def test_client_metadata_kept_for_arn_naming_a_non_anthropic_family(): + """An inference profile ARN that still spells out the family is resolvable, so a + non-Anthropic one keeps its passthrough.""" + config = AmazonConverseConfig() + + data = config._transform_request_helper( + model="arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-pro-v1:0", + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost From 1eea8e283157d3c92be36749b2713607aebc9786 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:54:34 -0700 Subject: [PATCH 15/81] fix(deps): raise the tornado floor to 6.5.8 for GHSA-8423-8fgw-73vq and GHSA-wwv5-g3v4-889x --- pyproject.toml | 2 +- uv.lock | 26 +++++++++++++------------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2866e27e84c..d0f14722acd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -292,7 +292,7 @@ exclude = [ [tool.uv] constraint-dependencies = [ - "tornado>=6.5.6", + "tornado>=6.5.8", "aiohttp>=3.14.2,<4.0", "packaging>=24.0", "soupsieve>=2.8.4", diff --git a/uv.lock b/uv.lock index 27be919eea1..8bac024d49e 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-29T17:58:57.633306Z" +exclude-newer = "2026-08-29T20:52:40.322465Z" exclude-newer-span = "P3D" [manifest] @@ -25,7 +25,7 @@ constraints = [ { name = "packaging", specifier = ">=24.0" }, { name = "setuptools", specifier = ">=83.0.0" }, { name = "soupsieve", specifier = ">=2.8.4" }, - { name = "tornado", specifier = ">=6.5.6" }, + { name = "tornado", specifier = ">=6.5.8" }, ] overrides = [ { name = "cryptography", specifier = ">=50.0.0,<51.0" }, @@ -9441,19 +9441,19 @@ wheels = [ [[package]] name = "tornado" -version = "6.5.7" +version = "6.5.8" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } +sdist = { url = "https://files.pythonhosted.org/packages/10/d3/343e5bb989d6515b1646cf3d40135d73f3d5e45339bded401b56cdac24dd/tornado-6.5.8.tar.gz", hash = "sha256:9452e1b208a8bd771e2cb1f2ff564985b9b214bdebbe622793e1799e0a6bd23f", size = 520493, upload-time = "2026-08-07T02:12:42.971Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, - { url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" }, - { url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" }, - { url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" }, - { url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" }, - { url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" }, - { url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" }, - { url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" }, - { url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" }, + { url = "https://files.pythonhosted.org/packages/f2/d5/007086fd8df5489338e204f65adce33fd4f21a4999dbb2b9cff2f897b5f4/tornado-6.5.8-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:cc6aa787d7cfab7c3d35189dc7a56fbd2399a569624c730c6b55b3d6531d0403", size = 449487, upload-time = "2026-08-07T02:12:28.682Z" }, + { url = "https://files.pythonhosted.org/packages/70/c8/5a24a99495903f594f6a199dd7beead1cbc0a13e2cb9102727bcaaf2a997/tornado-6.5.8-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9715b5eb79735b2bcd454ce216a9275b7c0470e64ea1bf5742f78b2f72b26eeb", size = 447649, upload-time = "2026-08-07T02:12:30.306Z" }, + { url = "https://files.pythonhosted.org/packages/6e/de/f2e733f386b85962d1b1dc82cd63d169b5b4580062b35397eac9244a41fe/tornado-6.5.8-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:547d63f450d570c14fe0e8db2cfb14c9bbd1c2503b4a6612586267955aa47b58", size = 450707, upload-time = "2026-08-07T02:12:31.95Z" }, + { url = "https://files.pythonhosted.org/packages/0b/94/20efeee9a01c141e9ac47c397f81679dfda24b32768fc4fff24e76d36c2c/tornado-6.5.8-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e2360a0ffbe145eca8af0b19cb7203d79b1a98dd4cccdd6b368f6f49c2e3808", size = 451677, upload-time = "2026-08-07T02:12:33.512Z" }, + { url = "https://files.pythonhosted.org/packages/42/ec/a96ccb8ccf0de2b7bc2c5fa1608a4803735018242e90c4882365a9fd418f/tornado-6.5.8-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5d242290bdf7ab3151bc1065fdd75c0dcc21cbc7b49f22a4c56329c2d6566d22", size = 451510, upload-time = "2026-08-07T02:12:35.346Z" }, + { url = "https://files.pythonhosted.org/packages/29/b5/93185859245ad3f00e62175f29607346788b696369347f0146e0421286bb/tornado-6.5.8-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7b94ff0e128fe0542f3bd331fb44d06260fc4ac16881545159f34ef08aad4195", size = 450917, upload-time = "2026-08-07T02:12:36.963Z" }, + { url = "https://files.pythonhosted.org/packages/97/cf/fe33cf062834487d34d1559746a4a12521033c22645b6d74d4bca702e018/tornado-6.5.8-cp39-abi3-win32.whl", hash = "sha256:67832909c4779c64942380cb5f044a5c6163d00831472d80e25e115de9917836", size = 451952, upload-time = "2026-08-07T02:12:38.512Z" }, + { url = "https://files.pythonhosted.org/packages/cb/e1/468ad54333e92ccb62627e62cb88e5fc14a2171daa67ed47b1b8542d5b86/tornado-6.5.8-cp39-abi3-win_amd64.whl", hash = "sha256:11881db6b7c168494be2c2d12e65931451bdf7ee718535418ae1d8855dd5a0ee", size = 452391, upload-time = "2026-08-07T02:12:39.971Z" }, + { url = "https://files.pythonhosted.org/packages/ad/3e/cd5e4f06e34cde33b8ef66cf36aa2b5ad46354cc1af7d2136bbe365fee1d/tornado-6.5.8-cp39-abi3-win_arm64.whl", hash = "sha256:68a7468c7e289f8514d7d664101753903217eff1bb6822c6b5994a0b5f5bcb26", size = 451411, upload-time = "2026-08-07T02:12:41.469Z" }, ] [[package]] From 4a68abfd49eea3b7d4fdfeaadd55598369848ef2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 23:26:23 +0000 Subject: [PATCH 16/81] fix(proxy): route container create and list through model_list deployments Container create and list requests had no container ID to decode, so the router called the provider handler directly and the OpenAI transformation fell back to the global OPENAI_API_KEY. Proxies configured only with model_list credentials sent Authorization: Bearer None. Route through _ageneric_api_call_with_fallbacks when the caller passes a model, expose the list endpoint's model query param to the router, and encode the managed container ID on the async create path so follow-up calls route to the same deployment. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/containers/main.py | 25 +++++--- .../proxy/container_endpoints/endpoints.py | 2 +- litellm/router.py | 11 +++- .../test_router_endpoints.py | 59 +++++++++++++++++++ .../containers/test_container_api.py | 36 +++++++++++ 5 files changed, 122 insertions(+), 11 deletions(-) diff --git a/litellm/containers/main.py b/litellm/containers/main.py index 97ca11872c1..90d59af009f 100644 --- a/litellm/containers/main.py +++ b/litellm/containers/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars import json -from collections.abc import Coroutine, Mapping +from collections.abc import Callable, Coroutine, Mapping from functools import partial from typing import Final, Literal, overload @@ -47,6 +47,13 @@ __all__ = [ ##### Container Create ####################### +async def _encode_created_container_id( + pending: Coroutine[object, object, ContainerObject], + encode: Callable[[ContainerObject], ContainerObject], +) -> ContainerObject: + return encode(await pending) + + @client async def acreate_container( name: str, @@ -256,16 +263,16 @@ def create_container( _is_async=_is_async, ) - # Encode container_id with provider/model metadata for routing + encode: Final = partial( + ContainerRequestUtils.encode_container_id_in_response, + custom_llm_provider=custom_llm_provider, + litellm_metadata=kwargs.get("litellm_metadata"), + extra_body=extra_body, + ) if isinstance(container_obj, ContainerObject): - container_obj = ContainerRequestUtils.encode_container_id_in_response( - response_obj=container_obj, - custom_llm_provider=custom_llm_provider, - litellm_metadata=kwargs.get("litellm_metadata"), - extra_body=extra_body, - ) + return encode(container_obj) - return container_obj + return _encode_created_container_id(pending=container_obj, encode=encode) except Exception as e: raise litellm.exception_type( diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 4a088140725..eaa3db336a9 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -208,7 +208,7 @@ async def list_containers( # Read query parameters query_params: Final = dict(request.query_params) - data: Final[dict[str, Any]] = {"query_params": query_params} + data: Final[dict[str, Any]] = {"query_params": query_params, "model": query_params.get("model")} # Extract custom_llm_provider using priority chain custom_llm_provider: Final = ( diff --git a/litellm/router.py b/litellm/router.py index 471a1116f44..52e858a0b6c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6589,7 +6589,9 @@ class Router: metadata. When present, decode the ID, replace ``container_id`` with the upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so deployment credentials (e.g. regional ``api_base`` for Azure) match - :meth:`_init_responses_api_endpoints`. Otherwise call the handler directly. + :meth:`_init_responses_api_endpoints`. Create/list calls carry no container ID, so + they route through the deployment named by ``model`` when the caller passes one. + Otherwise call the handler directly with global provider credentials. """ if custom_llm_provider and "custom_llm_provider" not in kwargs: kwargs["custom_llm_provider"] = custom_llm_provider @@ -6621,6 +6623,13 @@ class Router: **kwargs, ) + requested_model: Final = kwargs.get("model") + if isinstance(requested_model, str) and requested_model.strip(): + return await self._ageneric_api_call_with_fallbacks( + original_function=original_function, + **kwargs, + ) + return await original_function(**kwargs) async def _init_responses_api_endpoints( diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index d37af5b456a..6fcbd77e054 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -1324,6 +1324,65 @@ async def test_init_containers_api_endpoints_managed_id_without_model_id_applies assert call_kw["custom_llm_provider"] == "azure" +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_create_with_model_uses_deployment_credentials(monkeypatch): + """ + ``POST /v1/containers`` carries no container ID, so a ``model`` in the request + body is the only way to pick a deployment. The upstream call must receive that + deployment's ``api_key``/``api_base`` instead of falling back to the global + ``OPENAI_API_KEY`` (which may be unset on the proxy). + """ + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + router = Router( + model_list=[ + { + "model_name": "gpt-5.4", + "litellm_params": { + "model": "openai/gpt-5.4", + "api_key": "sk-model-list-key", + "api_base": "https://custom.openai.example/v1", + }, + } + ] + ) + mock_original_function = AsyncMock(return_value={"id": "cntr_test", "name": "Test Container"}) + + await router._init_containers_api_endpoints( + original_function=mock_original_function, + custom_llm_provider="openai", + name="Test Container", + model="gpt-5.4", + ) + + mock_original_function.assert_called_once() + call_kw = mock_original_function.call_args.kwargs + assert call_kw["api_key"] == "sk-model-list-key" + assert call_kw["api_base"] == "https://custom.openai.example/v1" + assert call_kw["model"] == "openai/gpt-5.4" + assert call_kw["name"] == "Test Container" + + +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_create_without_model_calls_directly(): + """ + Without ``model`` (or with ``model=None`` as the proxy forwards it), create/list + must keep calling the handler directly with global provider credentials. + """ + router = Router(model_list=[]) + router._ageneric_api_call_with_fallbacks = AsyncMock() + mock_original_function = AsyncMock(return_value={"id": "cntr_test"}) + + await router._init_containers_api_endpoints( + original_function=mock_original_function, + custom_llm_provider="openai", + name="Test Container", + model=None, + ) + + router._ageneric_api_call_with_fallbacks.assert_not_called() + mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None) + + def test_router_model_group_encrypted_content_affinity_callback_registration(): from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, diff --git a/tests/test_litellm/containers/test_container_api.py b/tests/test_litellm/containers/test_container_api.py index 885c4cd294a..8f1d87f9603 100644 --- a/tests/test_litellm/containers/test_container_api.py +++ b/tests/test_litellm/containers/test_container_api.py @@ -152,6 +152,42 @@ class TestContainerAPI: assert response.id == "cntr_async_123" assert response.name == "Async Test Container" + @pytest.mark.asyncio + async def test_acreate_container_encodes_router_model_id(self): + """ + The async handler returns a coroutine, so the managed-ID encoding must run + after it resolves. Otherwise follow-up calls (retrieve/delete/files) lose the + deployment and fall back to global provider credentials. + """ + upstream_response = ContainerObject( + id="cntr_upstream_123", + object="container", + created_at=1747857508, + status="running", + expires_after={"anchor": "last_active_at", "minutes": 20}, + last_active_at=1747857508, + name="Routed Container", + ) + + async def _resolve_upstream(): + return upstream_response + + with patch.object( # test-quality-ok: create_container does not forward a client, so the handler is the only seam + base_llm_http_handler, + "container_create_handler", + side_effect=lambda **kwargs: _resolve_upstream() if kwargs["_is_async"] else upstream_response, + ): + response = await acreate_container( + name="Routed Container", + custom_llm_provider="openai", + litellm_metadata={"model_info": {"id": "deployment-abc"}}, + ) + + decoded = ResponsesAPIRequestUtils._decode_container_id(response.id) + assert decoded["model_id"] == "deployment-abc" + assert decoded["custom_llm_provider"] == "openai" + assert decoded["response_id"] == "cntr_upstream_123" + @pytest.mark.asyncio async def test_alist_containers_basic(self): """Test basic async container listing functionality.""" From acc65b27d210158a18db74d6d0cd48d12aeb3a40 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 23:39:42 +0000 Subject: [PATCH 17/81] fix(router): pass container create/list through when model names no deployment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 6 ++-- .../test_router_endpoints.py | 33 +++++++++++++++++++ .../containers/test_container_api.py | 2 +- 3 files changed, 38 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 52e858a0b6c..7b471a69ecf 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6590,8 +6590,9 @@ class Router: upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so deployment credentials (e.g. regional ``api_base`` for Azure) match :meth:`_init_responses_api_endpoints`. Create/list calls carry no container ID, so - they route through the deployment named by ``model`` when the caller passes one. - Otherwise call the handler directly with global provider credentials. + they route through the deployment named by ``model`` when the caller passes one, + falling back to the direct call when no deployment matches. Otherwise call the + handler directly with global provider credentials. """ if custom_llm_provider and "custom_llm_provider" not in kwargs: kwargs["custom_llm_provider"] = custom_llm_provider @@ -6627,6 +6628,7 @@ class Router: if isinstance(requested_model, str) and requested_model.strip(): return await self._ageneric_api_call_with_fallbacks( original_function=original_function, + passthrough_on_no_deployment=True, **kwargs, ) diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index 6fcbd77e054..a06aaa363ad 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -1383,6 +1383,39 @@ async def test_init_containers_api_endpoints_create_without_model_calls_directly mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None) +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_create_with_unknown_model_passes_through(monkeypatch): + """ + A ``model`` that names no configured deployment must not turn into a 400. The call + falls through to the handler with the caller's model and no injected deployment + credentials, matching the behaviour before model-based routing existed. + """ + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + router = Router( + model_list=[ + { + "model_name": "gpt-5.4", + "litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"}, + } + ] + ) + mock_original_function = AsyncMock(return_value={"id": "cntr_test"}) + + await router._init_containers_api_endpoints( + original_function=mock_original_function, + custom_llm_provider="openai", + name="Test Container", + model="does-not-exist", + ) + + mock_original_function.assert_called_once() + call_kw = mock_original_function.call_args.kwargs + assert call_kw["model"] == "does-not-exist" + assert call_kw["name"] == "Test Container" + assert "api_key" not in call_kw + assert "api_base" not in call_kw + + def test_router_model_group_encrypted_content_affinity_callback_registration(): from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, diff --git a/tests/test_litellm/containers/test_container_api.py b/tests/test_litellm/containers/test_container_api.py index 8f1d87f9603..16a3844431e 100644 --- a/tests/test_litellm/containers/test_container_api.py +++ b/tests/test_litellm/containers/test_container_api.py @@ -172,7 +172,7 @@ class TestContainerAPI: async def _resolve_upstream(): return upstream_response - with patch.object( # test-quality-ok: create_container does not forward a client, so the handler is the only seam + with patch.object( # test-quality-ok: create_container exposes no client seam, only the handler base_llm_http_handler, "container_create_handler", side_effect=lambda **kwargs: _resolve_upstream() if kwargs["_is_async"] else upstream_response, From ac19d0dbdf61a3e2707b03d2deed225ba9d68389 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:44:29 -0700 Subject: [PATCH 18/81] fix(spend): keep every-deployment scope on gateway cache-injection marks The caching-savings marker litellm_gateway_injected_cache credits gateway-earned prompt-caching savings to the deployment it names, or to every deployment via the empty-string sentinel. Two paths lost that scope: - the router prompt-management factory stamps a provisional deployment's model_info into kwargs before the prompt pass runs, so an injection recorded there named that provisional pick and a differently-billed deployment lost the credit - record_gateway_injection overwrote on every positive delta, so a per-leg stamp (the Bedrock converse tool_config one included) downgraded an existing every-deployment mark and the leg billed after a failover lost the credit record_gateway_injection now takes injected_for_every_deployment, the two pre-choice callers declare it, and an every-deployment mark is never narrowed by a later per-leg stamp. Per-leg marks still overwrite each other. Spend amounts are untouched; only the savings attribution is affected. Also unblocks make lint at the staging tip: tests/e2e/test_junit_properties.py landed three basedpyright reds via an e2e-only PR whose lint job skipped, now suppressed as the deliberate duck-typed double they are. --- .../anthropic_cache_control_hook.py | 32 +++++++++--- litellm/litellm_core_utils/litellm_logging.py | 4 ++ litellm/proxy/utils.py | 1 + litellm/router.py | 1 + tests/e2e/test_junit_properties.py | 6 +-- .../test_anthropic_cache_control_hook.py | 21 ++++++++ .../test_litellm_logging.py | 26 +++++++++- tests/test_litellm/test_router.py | 52 +++++++++++++++++++ 8 files changed, 131 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 545b0f40018..3519240dda9 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -755,6 +755,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): def record_gateway_injection( request_kwargs: Mapping[str, object], added: int, + injected_for_every_deployment: bool = False, ) -> None: """Name the deployment whose payload the gateway, not the client, put breakpoints on. @@ -771,7 +772,16 @@ class AnthropicCacheControlHook(CustomPromptManagement): A pass that runs before a deployment is chosen, which is what the proxy does for prompt templates, injects into the payload every leg goes on to send, so it marks - the request for all of them rather than for one. + the request for all of them rather than for one. Such a pass says so with + ``injected_for_every_deployment`` instead of relying on the shape of + ``request_kwargs``: the router's prompt-management factory stamps a provisional + deployment's ``model_info`` into kwargs before the prompt pass runs, and billing + the request through any other deployment would silently drop the credit. An + every-deployment mark, once written, also never narrows: a later per-leg stamp + (the Bedrock converse tool_config one included) describes one leg of a payload + every leg sends, so narrowing to it would uncredit whichever leg gets billed + after a failover. Both losses are fail-closed under-crediting, which is why the + guard only protects the sentinel and per-leg marks still overwrite each other. Only what this pass actually placed counts. A ``tool_config`` point is placed by the Bedrock converse transform, and only when the request carries tools, so the @@ -801,13 +811,19 @@ class AnthropicCacheControlHook(CustomPromptManagement): ), None, ) - if bucket is not None: - model_info: Final = request_kwargs.get("model_info") - bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = ( - model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT) - if isinstance(model_info, dict) - else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT - ) + if bucket is None: + return + if bucket.get(GATEWAY_INJECTED_CACHE_METADATA_KEY) == GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT: + return + if injected_for_every_deployment: + bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT + return + model_info: Final = request_kwargs.get("model_info") + bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = ( + model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT) + if isinstance(model_info, dict) + else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT + ) @staticmethod def maybe_inject_cache_control( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9a6fb11f978..f94e86b4460 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -901,6 +901,7 @@ class Logging(LiteLLMLoggingBaseClass): prompt_label: str | None = None, prompt_version: int | None = None, request_kwargs: dict[str, object] | None = None, # mutable-ok: marker stamped into live request kwargs + injected_for_every_deployment: bool = False, ) -> tuple[str, list[AllMessageValues], dict]: from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook @@ -933,6 +934,7 @@ class Logging(LiteLLMLoggingBaseClass): AnthropicCacheControlHook.record_gateway_injection( request_kwargs, AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before, + injected_for_every_deployment=injected_for_every_deployment, ) self.messages = messages return model, messages, non_default_params @@ -950,6 +952,7 @@ class Logging(LiteLLMLoggingBaseClass): prompt_label: str | None = None, prompt_version: int | None = None, request_kwargs: dict[str, object] | None = None, # mutable-ok: marker stamped into live request kwargs + injected_for_every_deployment: bool = False, ) -> tuple[str, list[AllMessageValues], dict]: from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook @@ -985,6 +988,7 @@ class Logging(LiteLLMLoggingBaseClass): AnthropicCacheControlHook.record_gateway_injection( request_kwargs, AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before, + injected_for_every_deployment=injected_for_every_deployment, ) self.messages = messages return model, messages, non_default_params diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 051d36c4d0f..cab2bd6d9db 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1524,6 +1524,7 @@ class ProxyLogging: prompt_label=data.pop("prompt_label", None) or {}, prompt_version=data.pop("prompt_version", None) or {}, request_kwargs=data, + injected_for_every_deployment=True, ) data.update(optional_params) diff --git a/litellm/router.py b/litellm/router.py index 6e4405ebfef..462d5414456 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4006,6 +4006,7 @@ class Router: prompt_variables=prompt_variables, prompt_label=prompt_label, request_kwargs=kwargs, + injected_for_every_deployment=True, ) # Filter out prompt management specific parameters from data before merging diff --git a/tests/e2e/test_junit_properties.py b/tests/e2e/test_junit_properties.py index c0596177cc1..f7d1f70c5ec 100644 --- a/tests/e2e/test_junit_properties.py +++ b/tests/e2e/test_junit_properties.py @@ -115,7 +115,7 @@ class TestResultProperties: ("logging/test_x.py", 40, "TestFoo.test_bar"), (FakeMarker("covers", "LOG-1", "LOG-2"),), ) - assert result_properties(item) == ( + assert result_properties(item) == ( # pyright: ignore[reportArgumentType] # duck-typed Item double ("package", "logging"), ("covers", "LOG-1,LOG-2"), ("source", "tests/e2e/logging/test_x.py:41"), @@ -125,8 +125,8 @@ class TestResultProperties: """Collection can run the hook more than once; a second pass must not double the entries in the report.""" item = FakeItem("logging/test_x.py::test_bar", ("logging/test_x.py", 40, "test_bar")) - attach_result_properties(item) - attach_result_properties(item) + attach_result_properties(item) # pyright: ignore[reportArgumentType] # duck-typed Item double + attach_result_properties(item) # pyright: ignore[reportArgumentType] # duck-typed Item double assert [name for name, _ in item.user_properties] == ["package", "covers", "source"] 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 e995cbae782..de8b654987b 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -2858,6 +2858,27 @@ class TestRecordGatewayInjection: AnthropicCacheControlHook.record_gateway_injection(kwargs, 0) assert kwargs["litellm_metadata"][self.KEY] == self.DEPLOYMENT + def test_an_every_deployment_mark_survives_a_later_per_deployment_stamp(self): + """A per-leg stamp like the Bedrock converse tool_config one describes one leg of + a payload every leg sends, so narrowing an every-deployment mark to that leg's + deployment would uncredit whichever leg gets billed after a failover.""" + kwargs: dict = {"litellm_metadata": {self.KEY: ""}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1) + assert kwargs["litellm_metadata"][self.KEY] == "" + + def test_a_pre_choice_pass_stamps_the_sentinel_over_a_provisional_deployment(self): + """The router's prompt-management factory stamps a provisional deployment's + model_info into kwargs before the prompt pass runs, and any other deployment can + end up billed, so the pass declares every-deployment scope explicitly.""" + kwargs: dict = {"litellm_metadata": {}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1, injected_for_every_deployment=True) + assert kwargs["litellm_metadata"][self.KEY] == "" + + def test_a_per_deployment_mark_still_follows_the_latest_leg(self): + kwargs: dict = {"litellm_metadata": {self.KEY: "dep-old"}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1) + assert kwargs["litellm_metadata"][self.KEY] == self.DEPLOYMENT + def test_v1_messages_auto_injection_stamps_the_marker(self, monkeypatch): monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) kwargs: dict = {"litellm_metadata": {}, "model_info": {"id": self.DEPLOYMENT}} diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 366f61ded49..f1de7390b5b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6002,7 +6002,9 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o """The savings gate reads litellm_gateway_injected_cache from the request's metadata bucket. Recording lives in the shared prompt-hook wrappers, so chat, /v1/responses, router prompt deployments, and proxy prompt templates all mark - injected requests the same way; a hook that injects nothing leaves no marker.""" + injected requests the same way; a hook that injects nothing leaves no marker. + A pass that runs before deployment choice declares it and gets the every-deployment + sentinel, which a later per-deployment pass never narrows.""" from litellm.integrations.custom_prompt_management import CustomPromptManagement class _InjectingHook(CustomPromptManagement): @@ -6086,6 +6088,28 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o ) assert "litellm_gateway_injected_cache" not in untouched["metadata"] + pre_choice = {"metadata": {}, "model_info": {"id": "dep-of-this-attempt"}} + logging_obj.get_chat_completion_prompt( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + non_default_params={}, + prompt_variables=None, + prompt_management_logger=_InjectingHook(), + request_kwargs=pre_choice, + injected_for_every_deployment=True, + ) + assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" + + await logging_obj.async_get_chat_completion_prompt( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "a fresh turn"}], + non_default_params={}, + prompt_variables=None, + prompt_management_logger=_InjectingHook(), + request_kwargs=pre_choice, + ) + assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" + def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 84f6344be35..058ed5bb3a7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11730,3 +11730,55 @@ class TestPreRoutingTierDrivesFallbacks: response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-b" + + +@pytest.mark.asyncio +async def test_prompt_management_factory_marks_injection_for_every_deployment(monkeypatch): + """The factory stamps a provisional deployment's model_info into kwargs before the + prompt pass runs, then routes on the returned model, so any deployment can end up + billed. An injection recorded there must carry the every-deployment sentinel, never + the provisional deployment's id, or a differently-billed deployment loses the credit.""" + import time + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + router = litellm.Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + "model_info": {"id": "provisional-dep"}, + } + ] + ) + captured: dict = {} + + async def _capture_acompletion(**kwargs): + captured.update(kwargs) + return litellm.ModelResponse() + + monkeypatch.setattr(litellm, "acompletion", _capture_acompletion) + logging_obj = LiteLLMLogging( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="lit-6445", + function_id="f", + ) + await router.acompletion( + model="cached-claude", + messages=[ + {"role": "system", "content": "a static system prompt"}, + {"role": "user", "content": "hi"}, + ], + cache_control_injection_points=[{"location": "message", "role": "system"}], + litellm_logging_obj=logging_obj, + ) + bucket = captured.get("litellm_metadata") or captured["metadata"] + assert captured["model_info"]["id"] == "provisional-dep" + assert bucket["litellm_gateway_injected_cache"] == "" From c18511be7d76f0a1dfd18aa07feaa80784afdb9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 22:28:38 -0700 Subject: [PATCH 19/81] fix(guardrails): track and tear down presidio sibling callbacks initialize_presidio registers up to three callbacks per guardrail but the registry only kept the first, so deleting or re-syncing the guardrail left the post_call siblings serving the old config. The initializer now returns every callback it registered, the registry tracks primary and siblings per guardrail id, delete purges all of them from every callback list, and update pushes the new params into each while siblings keep their stage. --- .../guardrails/guardrail_initializers.py | 39 ++--- .../proxy/guardrails/guardrail_registry.py | 159 +++++++++++------- .../guardrail_hooks/test_presidio.py | 27 ++- .../guardrails/test_guardrail_registry.py | 129 ++++++++++++++ 4 files changed, 267 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 76dea1b7784..16369abbfb0 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -2,6 +2,7 @@ from typing import Any, Final import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import CommonProxyErrors from litellm.types.guardrails import * @@ -85,7 +86,7 @@ def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail): return _lakera_v2_callback -def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): +def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) -> tuple[CustomGuardrail, ...]: from litellm.proxy.guardrails.guardrail_hooks.presidio import ( _OPTIONAL_PresidioPIIMasking, ) @@ -94,7 +95,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): run_input: Final = filter_scope in ("input", "both") run_output: Final = filter_scope in ("output", "both") - def _make_presidio_callback(**overrides): + def _make_presidio_callback(**overrides) -> CustomGuardrail: params: Final = dict( guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, @@ -120,27 +121,27 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): litellm.logging_callback_manager.add_litellm_callback(callback) return callback - primary_callback = None - - if run_input: - primary_callback = _make_presidio_callback() - - if litellm_params.output_parse_pii: - _make_presidio_callback( - output_parse_pii=True, - event_hook=GuardrailEventHooks.post_call.value, - ) - - if run_output: - output_callback: Final = _make_presidio_callback( + input_callback: Final = _make_presidio_callback() if run_input else None + unmask_output_callback: Final = ( + _make_presidio_callback( + output_parse_pii=True, + event_hook=GuardrailEventHooks.post_call.value, + ) + if run_input and litellm_params.output_parse_pii + else None + ) + mask_output_callback: Final = ( + _make_presidio_callback( apply_to_output=True, event_hook=GuardrailEventHooks.post_call.value, output_parse_pii=False, ) - if primary_callback is None: - primary_callback = output_callback - - return primary_callback + if run_output + else None + ) + return tuple( + callback for callback in (input_callback, unmask_output_callback, mask_output_callback) if callback is not None + ) def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail): diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index dc13c09dd38..bd35782444b 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -3,10 +3,10 @@ import asyncio import importlib import os -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime, timezone from itertools import chain, count -from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeAlias, cast from pydantic import ValidationError @@ -90,6 +90,8 @@ guardrail_initializer_registry: Final = { CONFIG_GUARDRAIL_ID_NAMESPACE: Final = uuid.UUID("625f63f4-935a-50e5-98b5-fbe77babc74a") +GuardrailCallbacks: TypeAlias = tuple[CustomGuardrail, ...] + guardrail_class_registry: Final[dict[str, type[CustomGuardrail]]] = { SupportedGuardrailIntegrations.BEDROCK.value: BedrockGuardrail, SupportedGuardrailIntegrations.GRAYSWAN.value: GraySwanGuardrail, @@ -424,6 +426,41 @@ def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params: instance.scan_raw_request = bool(litellm_params.scan_raw_request) +def _as_callback_tuple( + initialized: CustomGuardrail | Sequence[CustomGuardrail] | None, +) -> GuardrailCallbacks: + if initialized is None: + return () + if isinstance(initialized, (list, tuple)): + return tuple(initialized) + return (initialized,) + + +def _configure_callback_scoping( + custom_guardrail_callback: CustomGuardrail, guardrail_name: str, litellm_params: LitellmParams +) -> None: + for scoping_param in ( + "skip_system_message_in_guardrail", + "skip_tool_message_in_guardrail", + "scan_only_tool_results", + ): + setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None)) + scan_only_tool_results_enabled: Final = effective_scan_only_tool_results_for_guardrail(custom_guardrail_callback) + if scan_only_tool_results_enabled and not custom_guardrail_callback.supports_scan_only_tool_results(): + raise ValueError( + f"Guardrail {guardrail_name}: scan_only_tool_results is enabled, but this " + "guardrail's role filtering never scans tool results, so no request content would ever " + "be scanned. Remove scan_only_tool_results or the guardrail's role-filtering option." + ) + if scan_only_tool_results_enabled and effective_skip_tool_message_for_guardrail(custom_guardrail_callback): + raise ValueError( + f"Guardrail {guardrail_name}: scan_only_tool_results and " + "skip_tool_message_in_guardrail are enabled together, which excludes every message from " + "scanning, so no request content would ever be scanned. Remove one of the two." + ) + _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params) + + class InMemoryGuardrailHandler: """ Class that handles initializing guardrails and adding them to the CallbackManager @@ -440,6 +477,8 @@ class InMemoryGuardrailHandler: Guardrail id to CustomGuardrail object mapping """ + self.guardrail_id_to_sibling_callbacks: dict[str, GuardrailCallbacks] = {} # mutable-ok: per-id registry + self._sources: dict[str, Literal["db", "config"]] = {} """ Guardrail id to provenance marker. "db" entries are reconciled against @@ -474,7 +513,6 @@ class InMemoryGuardrailHandler: self._sources[guardrail_id] = source return self.IN_MEMORY_GUARDRAILS[guardrail_id] - custom_guardrail_callback: CustomGuardrail | None = None litellm_params_data: Final = guardrail["litellm_params"] verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data) @@ -498,54 +536,15 @@ class InMemoryGuardrailHandler: if guardrail_type is None: raise ValueError("guardrail_type is required") - initializer: Final = guardrail_initializer_registry.get(guardrail_type) - - if initializer: - # Try to call with llm_router first, fall back to without if it fails - import inspect - - sig: Final = inspect.signature(initializer) - if "llm_router" in sig.parameters: - custom_guardrail_callback = initializer( - litellm_params, - guardrail, - llm_router, - ) - else: - custom_guardrail_callback = initializer(litellm_params, guardrail) - elif isinstance(guardrail_type, str) and "." in guardrail_type: - custom_guardrail_callback = self.initialize_custom_guardrail( - guardrail=guardrail, - guardrail_type=guardrail_type, - litellm_params=litellm_params, - config_file_path=config_file_path, - ) - else: - raise ValueError(f"Unsupported guardrail: {guardrail_type}") - - if custom_guardrail_callback is not None: - for scoping_param in ( - "skip_system_message_in_guardrail", - "skip_tool_message_in_guardrail", - "scan_only_tool_results", - ): - setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None)) - scan_only_tool_results_enabled: Final = effective_scan_only_tool_results_for_guardrail( - custom_guardrail_callback - ) - if scan_only_tool_results_enabled and not custom_guardrail_callback.supports_scan_only_tool_results(): - raise ValueError( - f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results is enabled, but this " - "guardrail's role filtering never scans tool results, so no request content would ever " - "be scanned. Remove scan_only_tool_results or the guardrail's role-filtering option." - ) - if scan_only_tool_results_enabled and effective_skip_tool_message_for_guardrail(custom_guardrail_callback): - raise ValueError( - f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results and " - "skip_tool_message_in_guardrail are enabled together, which excludes every message from " - "scanning, so no request content would ever be scanned. Remove one of the two." - ) - _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params) + created_callbacks: Final = self._create_callbacks( + guardrail=guardrail, + guardrail_type=guardrail_type, + litellm_params=litellm_params, + config_file_path=config_file_path, + llm_router=llm_router, + ) + for custom_guardrail_callback in created_callbacks: + _configure_callback_scoping(custom_guardrail_callback, guardrail["guardrail_name"], litellm_params) parsed_guardrail: Final = Guardrail( guardrail_id=guardrail.get("guardrail_id"), @@ -556,11 +555,44 @@ class InMemoryGuardrailHandler: # store references to the guardrail in memory self.IN_MEMORY_GUARDRAILS[guardrail_id] = parsed_guardrail - self.guardrail_id_to_custom_guardrail[guardrail_id] = custom_guardrail_callback + self.guardrail_id_to_custom_guardrail[guardrail_id] = created_callbacks[0] if created_callbacks else None + self.guardrail_id_to_sibling_callbacks[guardrail_id] = created_callbacks[1:] self._sources[guardrail_id] = source return parsed_guardrail + def _create_callbacks( + self, + guardrail: Guardrail, + guardrail_type: str, + litellm_params: LitellmParams, + config_file_path: str | None, + llm_router: Optional["Router"], + ) -> GuardrailCallbacks: + initializer: Final = guardrail_initializer_registry.get(guardrail_type) + if initializer: + import inspect + + sig: Final = inspect.signature(initializer) + if "llm_router" in sig.parameters: + return _as_callback_tuple(initializer(litellm_params, guardrail, llm_router)) + return _as_callback_tuple(initializer(litellm_params, guardrail)) + if isinstance(guardrail_type, str) and "." in guardrail_type: + return _as_callback_tuple( + self.initialize_custom_guardrail( + guardrail=guardrail, + guardrail_type=guardrail_type, + litellm_params=litellm_params, + config_file_path=config_file_path, + ) + ) + raise ValueError(f"Unsupported guardrail: {guardrail_type}") + + def _tracked_callbacks(self, guardrail_id: str) -> GuardrailCallbacks: + primary: Final = self.guardrail_id_to_custom_guardrail.get(guardrail_id) + siblings: Final = self.guardrail_id_to_sibling_callbacks.get(guardrail_id, ()) + return (() if primary is None else (primary,)) + siblings + def initialize_custom_guardrail( self, guardrail: Guardrail, @@ -630,10 +662,15 @@ class InMemoryGuardrailHandler: self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail self._sources[guardrail_id] = source - custom_guardrail_callback: Final = self.guardrail_id_to_custom_guardrail.get(guardrail_id) - if custom_guardrail_callback: - updated_litellm_params: Final = cast(LitellmParams, guardrail.get("litellm_params", {})) - custom_guardrail_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params) + tracked_callbacks: Final = self._tracked_callbacks(guardrail_id) + if not tracked_callbacks: + return + updated_litellm_params: Final = cast(LitellmParams, guardrail.get("litellm_params", {})) + tracked_callbacks[0].update_in_memory_litellm_params(litellm_params=updated_litellm_params) + for sibling_callback in tracked_callbacks[1:]: + sibling_stage = sibling_callback.event_hook + sibling_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params) + sibling_callback.event_hook = sibling_stage def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ @@ -648,11 +685,11 @@ class InMemoryGuardrailHandler: self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None) self._sources.pop(guardrail_id, None) - custom_guardrail_callback: Final = self.guardrail_id_to_custom_guardrail.pop(guardrail_id, None) - if custom_guardrail_callback is None: - return - - litellm.logging_callback_manager.remove_callback_from_all_lists(custom_guardrail_callback) + tracked_callbacks: Final = self._tracked_callbacks(guardrail_id) + self.guardrail_id_to_custom_guardrail.pop(guardrail_id, None) + self.guardrail_id_to_sibling_callbacks.pop(guardrail_id, None) + for custom_guardrail_callback in tracked_callbacks: + litellm.logging_callback_manager.remove_callback_from_all_lists(custom_guardrail_callback) def list_in_memory_guardrails(self) -> list[Guardrail]: """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 4ee6741ee02..fcf940afd0d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -842,24 +842,37 @@ async def test_presidio_filter_scope_initializer(monkeypatch): params_input = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="input") guardrail_dict = {"guardrail_name": "g1"} - cb = initialize_presidio(params_input, guardrail_dict) - assert cb is created[0] + callbacks = initialize_presidio(params_input, guardrail_dict) + assert callbacks == (created[0],) assert created[0].apply_to_output is False # output-only created.clear() params_output = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="output") - cb = initialize_presidio(params_output, guardrail_dict) + callbacks = initialize_presidio(params_output, guardrail_dict) assert len(created) == 1 + assert callbacks == (created[0],) assert created[0].apply_to_output is True - # both -> expect two callbacks (input + output) + # both -> expect two callbacks (input + output), both returned, input first created.clear() params_both = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="both") - cb = initialize_presidio(params_both, guardrail_dict) + callbacks = initialize_presidio(params_both, guardrail_dict) assert len(created) == 2 - assert any(not c.apply_to_output for c in created) - assert any(c.apply_to_output for c in created) + assert callbacks == tuple(created) + assert callbacks[0].apply_to_output is False + assert callbacks[1].apply_to_output is True + + # both + output_parse_pii -> three callbacks, all returned, input first + created.clear() + params_all = LitellmParams( + guardrail="presidio", mode="pre_call", presidio_filter_scope="both", output_parse_pii=True + ) + callbacks = initialize_presidio(params_all, guardrail_dict) + assert len(created) == 3 + assert callbacks == tuple(created) + assert callbacks[0].apply_to_output is False + assert mgr.added[-3:] == list(created) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 2c0735970d3..b8a58f5e3da 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -491,6 +491,135 @@ def test_repeated_db_sync_does_not_accumulate_runner_instances(): cb_list[:] = snapshot +PRESIDIO_SIBLINGS_GID = "55555555-5555-5555-5555-555555555555" +PRESIDIO_SIBLINGS_NAME = "presidio-siblings" + + +def _presidio_db_guardrail(pii_entities_config: dict) -> Guardrail: + return Guardrail( + guardrail_id=PRESIDIO_SIBLINGS_GID, + guardrail_name=PRESIDIO_SIBLINGS_NAME, + litellm_params={ + "guardrail": "presidio", + "mode": "pre_call", + "default_on": True, + "output_parse_pii": True, + "presidio_filter_scope": "both", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + "pii_entities_config": pii_entities_config, + }, + ) + + +def _presidio_callbacks_in(cb_list) -> list: + return [ + callback + for callback in cb_list + if isinstance(callback, CustomGuardrail) and getattr(callback, "guardrail_name", None) == PRESIDIO_SIBLINGS_NAME + ] + + +def test_presidio_siblings_are_tracked_and_deleted_together(): + """ + A presidio guardrail scoped to both stages registers the pre_call primary plus + the post_call unmask and mask-output siblings. Deleting the guardrail must remove + all three from every callback list, not just the primary. + """ + import litellm + + handler = InMemoryGuardrailHandler() + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.initialize_guardrail(_presidio_db_guardrail({"EMAIL_ADDRESS": "MASK"})) + + registered = _presidio_callbacks_in(litellm.callbacks) + assert len(registered) == 3 + primary = handler.guardrail_id_to_custom_guardrail[PRESIDIO_SIBLINGS_GID] + siblings = handler.guardrail_id_to_sibling_callbacks[PRESIDIO_SIBLINGS_GID] + assert primary is registered[0] + assert siblings == tuple(registered[1:]) + assert [sibling.event_hook for sibling in siblings] == [GuardrailEventHooks.post_call] * 2 + + for cb_list in lists[1:]: + cb_list.extend(registered) + + handler.delete_in_memory_guardrail(PRESIDIO_SIBLINGS_GID) + + for cb_list in lists: + assert _presidio_callbacks_in(cb_list) == [] + assert PRESIDIO_SIBLINGS_GID not in handler.guardrail_id_to_custom_guardrail + assert PRESIDIO_SIBLINGS_GID not in handler.guardrail_id_to_sibling_callbacks + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_stage(): + import litellm + + handler = InMemoryGuardrailHandler() + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.initialize_guardrail(_presidio_db_guardrail({"EMAIL_ADDRESS": "MASK", "IP_ADDRESS": "MASK"})) + tracked = _presidio_callbacks_in(litellm.callbacks) + roles_before = [(callback.apply_to_output, callback.event_hook) for callback in tracked] + + updated = Guardrail( + guardrail_id=PRESIDIO_SIBLINGS_GID, + guardrail_name=PRESIDIO_SIBLINGS_NAME, + litellm_params=LitellmParams( + guardrail="presidio", + mode="pre_call", + default_on=True, + output_parse_pii=True, + presidio_filter_scope="both", + presidio_analyzer_api_base="https://fakelink.com/v1/presidio/analyze", + presidio_anonymizer_api_base="https://fakelink.com/v1/presidio/anonymize", + pii_entities_config={"EMAIL_ADDRESS": "MASK"}, + ), + ) + handler.update_in_memory_guardrail(guardrail_id=PRESIDIO_SIBLINGS_GID, guardrail=updated) + + assert [callback.pii_entities_config for callback in tracked] == [{"EMAIL_ADDRESS": "MASK"}] * 3 + assert [(callback.apply_to_output, callback.event_hook) for callback in tracked] == roles_before + assert _presidio_callbacks_in(litellm.callbacks) == tracked + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_repeated_db_sync_replaces_presidio_siblings_instead_of_leaking_stale_ones(): + """ + The callback manager dedupes custom loggers by their scalar attributes, so a + leaked post_call sibling blocks the re-initialized sibling from registering and + keeps serving the previous entity config. After every DB re-sync, each callback + list must hold exactly the three current instances, all on the latest config. + """ + import litellm + + handler = InMemoryGuardrailHandler() + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + entity_configs = [{"EMAIL_ADDRESS": "MASK"}, {"EMAIL_ADDRESS": "MASK", "IP_ADDRESS": "MASK"}] + for cycle in range(4): + latest = entity_configs[cycle % 2] + handler.sync_guardrail_from_db(_presidio_db_guardrail(latest)) + for cb_list in lists[1:]: + cb_list.extend(_presidio_callbacks_in(litellm.callbacks)) + + for cb_list in lists: + current = _presidio_callbacks_in(cb_list) + assert len({id(callback) for callback in current}) == 3 + assert all(callback.pii_entities_config == latest for callback in current) + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + def _judge_guardrail(guardrail_id: str) -> Guardrail: return Guardrail( guardrail_id=guardrail_id, From 7cde2cd77f9c39306dddd8c614769e49a507b84b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 22:39:17 -0700 Subject: [PATCH 20/81] test(guardrails): type the presidio sibling test helpers precisely --- .../test_litellm/proxy/guardrails/test_guardrail_registry.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index b8a58f5e3da..5cbdef5f92f 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -1,3 +1,4 @@ +from collections.abc import Iterable from unittest.mock import AsyncMock, MagicMock import pytest @@ -495,7 +496,7 @@ PRESIDIO_SIBLINGS_GID = "55555555-5555-5555-5555-555555555555" PRESIDIO_SIBLINGS_NAME = "presidio-siblings" -def _presidio_db_guardrail(pii_entities_config: dict) -> Guardrail: +def _presidio_db_guardrail(pii_entities_config: dict[str, str]) -> Guardrail: return Guardrail( guardrail_id=PRESIDIO_SIBLINGS_GID, guardrail_name=PRESIDIO_SIBLINGS_NAME, @@ -512,7 +513,7 @@ def _presidio_db_guardrail(pii_entities_config: dict) -> Guardrail: ) -def _presidio_callbacks_in(cb_list) -> list: +def _presidio_callbacks_in(cb_list: Iterable[object]) -> list[CustomGuardrail]: return [ callback for callback in cb_list From 1400070d711f645290fb382564e1f21200c5e610 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 07:58:09 +0000 Subject: [PATCH 21/81] chore(techdebt): clear fresh debt from the 2026-09-01 window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 6 +++--- .../integrations/SlackAlerting/slack_alerting.py | 9 ++++----- .../websearch_interception/handler.py | 1 - .../_experimental/mcp_server/rest_endpoints.py | 13 +++++++++---- .../guardrails/guardrail_hooks/alice/alice.py | 16 ++++++++-------- type-discipline-budget.json | 6 +++--- 6 files changed, 27 insertions(+), 24 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 2d4adc02234..f14f8e002dd 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15300 + "limit": 15298 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38347 + "limit": 38344 }, "reportUnknownParameterType": { "limit": 19626 }, "reportUnknownVariableType": { - "limit": 29884 + "limit": 29880 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 748ef938cea..dc41c7dadc8 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1955,11 +1955,10 @@ Model Info: if not thresholds_enabled and not anomalies_enabled: return - if prisma_client is None: - from litellm.proxy.proxy_server import prisma_client as global_prisma_client + from litellm.proxy.proxy_server import prisma_client as global_prisma_client - prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client - if prisma_client is None: + client: Final = prisma_client if prisma_client is not None else global_prisma_client + if client is None: return from litellm.integrations.SlackAlerting.user_spend_alerts import ( @@ -1970,7 +1969,7 @@ Model Info: try: today: Final = datetime.datetime.now(datetime.timezone.utc).date() rows: Final = await fetch_user_spend_rows( - prisma_client=prisma_client, + prisma_client=client, today=today, baseline_days=self.alerting_args.spend_anomaly_baseline_days, ) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index dc61ee38a8c..2d737bc34e7 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -419,7 +419,6 @@ class WebSearchInterceptionLogger(CustomLogger): if call_type in (CallTypes.responses, CallTypes.aresponses): return self._convert_responses_tools(kwargs=kwargs, tools=tools) - # Check if any tool is a web search tool (native or already LiteLLM standard) has_websearch: Final = any(is_web_search_tool(t) for t in tools) if not has_websearch: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d1ef73a15cd..90474bfc5e6 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -92,6 +92,9 @@ def _connection_error_message(exc: BaseException) -> str: if MCP_AVAILABLE: + from mcp.types import Tool as MCPTool + + from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -876,7 +879,6 @@ if MCP_AVAILABLE: return (), classify_list_exception(e) return tools_result, ServerListOk(tool_count=len(tools_result)) - # Query all servers the user has access to queried_servers: Final = tuple( server for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids) @@ -1141,6 +1143,11 @@ if MCP_AVAILABLE: scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None return client_id, client_secret, scopes + async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None: + with anyio.move_on_after(deadline): + return await client.list_tools(raise_on_error=True) + return None + async def _execute_with_mcp_client( request: NewMCPServerRequest, operation: Callable[..., Awaitable[Mapping[str, object]]], @@ -1422,9 +1429,7 @@ if MCP_AVAILABLE: getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT, ) - list_tools_result = None # rebind-ok: set inside the timeout scope below - with anyio.move_on_after(listing_deadline): - list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above + list_tools_result: Final = await _list_tools_within(client, listing_deadline) if list_tools_result is None: verbose_logger.warning( "MCP tools/list preview timed out after %s seconds while paginating upstream tools", diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py index 27018769909..9cabac2d0fa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py +++ b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py @@ -8,6 +8,7 @@ import json import os from collections.abc import Mapping +from itertools import islice from typing import ( TYPE_CHECKING, Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__; see ruff-strict.toml @@ -341,19 +342,18 @@ def _json_safe( if depth >= _MAX_DEPTH or id(value) in seen: return None - nested: Final = seen | {id(value)} # mutable-ok: one-shot set literal, unioned into a frozenset immediately + nested: Final = seen | frozenset((id(value),)) if isinstance(value, dict): - out: dict[str, object] = {} # mutable-ok: bounded accumulator local to this call, never escapes as-is - for key, item in list(value.items())[:_MAX_ITEMS]: # mutable-ok: list() only to slice an unordered view - if isinstance(key, str) and key not in strip_keys: - out[key] = _json_safe(item, depth + 1, nested, strip_keys) - return out + return { + key: _json_safe(item, depth + 1, nested, strip_keys) + for key, item in islice(value.items(), _MAX_ITEMS) + if isinstance(key, str) and key not in strip_keys + } if isinstance(value, (list, tuple, set, frozenset)): return [ # mutable-ok: return value is a one-shot list, discarded by the caller after use - _json_safe(item, depth + 1, nested, strip_keys) - for item in list(value)[:_MAX_ITEMS] # mutable-ok: list() only to slice an unordered view + _json_safe(item, depth + 1, nested, strip_keys) for item in islice(value, _MAX_ITEMS) ] dump: Final = getattr(value, "model_dump", None) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index fbf998533f8..0f39b32670a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -6,7 +6,7 @@ "limit": 26777 }, "LIT003": { - "limit": 266 + "limit": 265 }, "LIT004": { "limit": 40 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16504 + "limit": 16502 }, "LIT011": { - "limit": 5531 + "limit": 5529 }, "LIT012": { "limit": 4495 From dba190842cea106ab4b03880861038ddbe6aae42 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:26:59 -0700 Subject: [PATCH 22/81] fix(anthropic): keep the cache_control normalizer inside the type-discipline budget --- litellm/llms/anthropic/common_utils.py | 38 +++++++++++++++++--------- 1 file changed, 25 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index b5bfb32c0c6..19d3d6d7043 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1404,13 +1404,33 @@ def _with_portable_cache_control_in_message(message: object) -> object: return message return { # mutable-ok: JSON wire format **message, - "content": [_with_portable_cache_control_in_content_block(block) for block in content], + "content": [ # mutable-ok: JSON wire format + _with_portable_cache_control_in_content_block(block) for block in content + ], } -def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire format +def _with_portable_cache_control_in_messages(messages: object) -> object: + if isinstance(messages, str) or not isinstance(messages, Sequence): + return messages + return [ # mutable-ok: JSON wire format + _with_portable_cache_control_in_message(message) for message in messages + ] + + +def _with_portable_cache_control_in_scoped_value(key: str, value: object) -> object: + match key: + case "system" | "tools": + return _with_portable_cache_control_in_blocks(value) + case "messages": + return _with_portable_cache_control_in_messages(value) + case _: + return value + + +def normalize_cache_control_in_anthropic_payload( payload: Mapping[str, object], -) -> dict[str, object]: +) -> dict[str, object]: # mutable-ok: JSON wire format """ Return a copy of an Anthropic /v1/messages payload with every ``cache_control`` entry reduced to ``{"type": }`` @@ -1427,17 +1447,9 @@ def normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire forma dropped entirely. The caller's payload is never mutated. """ portable: Final = _with_portable_cache_control(payload) - scoped: Final = { # mutable-ok: JSON wire format - key: ( - _with_portable_cache_control_in_blocks(value) - if key in ("system", "tools") - else [_with_portable_cache_control_in_message(message) for message in value] - if key == "messages" and isinstance(value, Sequence) and not isinstance(value, str) - else value - ) - for key, value in portable.items() + return { # mutable-ok: JSON wire format + key: _with_portable_cache_control_in_scoped_value(key, value) for key, value in portable.items() } - return scoped def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: From 53da9bca8e45af86507cb6b5c83736290913ba71 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:29:44 -0700 Subject: [PATCH 23/81] fix(bedrock): drop client_metadata for every converse model --- .../bedrock/chat/converse_transformation.py | 10 +-- litellm/llms/bedrock/common_utils.py | 9 -- .../chat/test_converse_transformation.py | 85 +++---------------- 3 files changed, 15 insertions(+), 89 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 38b9569856a..df52b78f6b5 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -86,7 +86,6 @@ from litellm.utils import ( from ..common_utils import ( BedrockError, BedrockModelInfo, - bedrock_arn_hides_model_family, bedrock_converse_supports_parallel_tool_use_config, bedrock_model_accepts_cache_points, get_anthropic_beta_from_headers, @@ -1335,14 +1334,7 @@ class AmazonConverseConfig(BaseConfig): ) additional_request_params.pop("parallel_tool_calls", None) - - drops_client_metadata: Final = base_model.startswith("anthropic") or bedrock_arn_hides_model_family(model) - if drops_client_metadata and additional_request_params.pop("client_metadata", None) is not None: - litellm.verbose_logger.debug( - "Bedrock Converse: dropping `client_metadata` for model=%s, Anthropic rejects it with " - "'client_metadata: Extra inputs are not permitted'", - model, - ) + additional_request_params.pop("client_metadata", None) # Only set the topK value in for models that support it additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params)) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index cf18e3e3ec8..66ee5f10679 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -809,15 +809,6 @@ def get_bedrock_base_model(model: str) -> str: return model -def bedrock_arn_hides_model_family(model: str) -> bool: - """ - True for an ARN-addressed model whose base name carries no ``provider.model`` - id, such as an application inference profile or a provisioned throughput ARN. - Callers that gate behavior on the model family cannot resolve one here. - """ - return "arn:" in model.lower() and "." not in get_bedrock_base_model(model) - - def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: return any( (litellm.model_cost.get(candidate) or {}).get("supports_parallel_tool_use_config") is True diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 7556e1624be..fb165c38ef9 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -979,16 +979,19 @@ def test_config_blocks_do_not_leak_into_inference_config(): assert data["serviceTier"] == {"type": "priority"} -@pytest.mark.parametrize("model", ["anthropic.claude-opus-4-8", "us.anthropic.claude-opus-4-8"]) -def test_client_metadata_stripped_for_anthropic_converse_request(model): - """``client_metadata`` sent by codex must not reach Anthropic as a passthrough model field. - - Converse forwards ``additionalModelRequestFields`` verbatim to the model, and Anthropic - rejects the request with "client_metadata: Extra inputs are not permitted". - """ - config = AmazonConverseConfig() - - data = config._transform_request_helper( +@pytest.mark.parametrize( + "model", + [ + "anthropic.claude-opus-4-8", + "us.anthropic.claude-opus-4-8", + "amazon.nova-pro-v1:0", + "us.meta.llama4-maverick-17b-instruct-v1:0", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456", + "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-pro-v1:0", + ], +) +def test_client_metadata_stripped_from_converse_request(model): + data = AmazonConverseConfig()._transform_request_helper( model=model, system_content_blocks=[], optional_params={ @@ -999,71 +1002,11 @@ def test_client_metadata_stripped_for_anthropic_converse_request(model): messages=None, ) - fields = data.get("additionalModelRequestFields", {}) + fields = data["additionalModelRequestFields"] assert "client_metadata" not in fields assert fields["anthropic_beta"] == ["computer-use-2025-01-24"] -def test_client_metadata_kept_for_non_anthropic_converse_request(): - """Only Anthropic is known to reject ``client_metadata``, so other families keep the passthrough.""" - config = AmazonConverseConfig() - - data = config._transform_request_helper( - model="amazon.nova-pro-v1:0", - system_content_blocks=[], - optional_params={ - "maxTokens": 16, - "client_metadata": {"originator": "codex_cli_rs"}, - }, - messages=None, - ) - - assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} - - -@pytest.mark.parametrize( - "model", - [ - "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456", - "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/abcdef123456", - ], -) -def test_client_metadata_stripped_for_arn_models_converse(model): - """An ARN hides which family serves the request, and pointing one at Claude is how - teams route codex traffic, so the field has to go there too or the 400 comes back.""" - config = AmazonConverseConfig() - - data = config._transform_request_helper( - model=model, - system_content_blocks=[], - optional_params={ - "maxTokens": 16, - "client_metadata": {"originator": "codex_cli_rs"}, - }, - messages=None, - ) - - assert "client_metadata" not in data.get("additionalModelRequestFields", {}) - - -def test_client_metadata_kept_for_arn_naming_a_non_anthropic_family(): - """An inference profile ARN that still spells out the family is resolvable, so a - non-Anthropic one keeps its passthrough.""" - config = AmazonConverseConfig() - - data = config._transform_request_helper( - model="arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-pro-v1:0", - system_content_blocks=[], - optional_params={ - "maxTokens": 16, - "client_metadata": {"originator": "codex_cli_rs"}, - }, - messages=None, - ) - - assert data["additionalModelRequestFields"]["client_metadata"] == {"originator": "codex_cli_rs"} - - def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost From 0c539614451869b790897c0576d202820d85ad59 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:44:51 -0700 Subject: [PATCH 24/81] fix(proxy-extras): give prisma migrate deploy its own timeout budget --- .../litellm_proxy_extras/prisma_toolchain.py | 25 +++++- .../litellm_proxy_extras/utils.py | 24 ++++-- .../test_prisma_toolchain.py | 78 ++++++++++++++++++- 3 files changed, 114 insertions(+), 13 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index f3b55fd4d96..5feb7a953b4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -14,9 +14,19 @@ then fails on a Node binary that was never written. Deleting a cache directory that exists without a Node binary is what turns a killed bootstrap back into a recoverable one. -Both budgets are overridable so an operator can widen them without a release: -``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install and -``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every individual Prisma command. +``prisma migrate deploy`` is the other command whose runtime is not a +constant: it grows with the number of pending migrations, so a fresh database +that has to replay every migration this package ships overruns a per-command +budget sized for the short bookkeeping commands, on a laptop as much as on a +slow CI runner. The Python ``prisma`` wrapper spawns Node and the schema engine +as separate children, so killing the wrapper on timeout leaves them running: +the retry then contends with that orphan for Prisma's advisory lock and cannot +finish any sooner. Migrate deploy therefore runs under its own budget. + +All three budgets are overridable so an operator can widen them without a +release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install, +``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and +``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. """ import math @@ -36,10 +46,12 @@ except ImportError: PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT" PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT" +PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT" NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR" DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0 DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0 +DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0 BOOTSTRAP_ARG = "--version" @@ -88,6 +100,13 @@ def prisma_bootstrap_timeout() -> float: ) +def prisma_migrate_deploy_timeout() -> float: + """Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending.""" + return _timeout_from_env( + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT + ) + + def nodeenv_cache_dir() -> Optional[Path]: """Where Prisma installs its private Node runtime, or None if unknowable.""" override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..ab9ec1e8a3a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -15,8 +15,10 @@ from litellm_proxy_extras.replica_identity import ( apply_replica_identity_full, ) from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ensure_prisma_toolchain, prisma_command_timeout, + prisma_migrate_deploy_timeout, ) @@ -698,12 +700,13 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) + deploy_timeout = prisma_migrate_deploy_timeout() try: for attempt in range(4): try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], - timeout=prisma_command_timeout(), + timeout=deploy_timeout, check=True, capture_output=True, text=True, @@ -713,8 +716,12 @@ class ProxyExtrasDBManager: return True except subprocess.TimeoutExpired: - logger.info( - f"prisma migrate deploy attempt {attempt + 1} timed out, retrying" + logger.warning( + "prisma migrate deploy attempt %s timed out after %ss, retrying. " + "Raise %s if this database needs longer to apply its pending migrations.", + attempt + 1, + deploy_timeout, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ) time.sleep(random.randrange(5, 15)) continue @@ -823,7 +830,8 @@ class ProxyExtrasDBManager: "Database migration failed after 4 attempts (retry loop " "exhausted by timeouts or repeated idempotent-recovery " "continues). Check database connectivity, load, and " - "_prisma_migrations ledger state." + "_prisma_migrations ledger state, and raise " + f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out." ) finally: os.chdir(original_dir) @@ -908,7 +916,7 @@ class ProxyExtrasDBManager: # Set migrations directory for Prisma result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], - timeout=prisma_command_timeout(), + timeout=prisma_migrate_deploy_timeout(), check=True, capture_output=True, text=True, @@ -1126,7 +1134,11 @@ class ProxyExtrasDBManager: ) return True except subprocess.TimeoutExpired: - logger.info(f"Attempt {attempt + 1} timed out") + logger.warning( + "Attempt %s timed out. Raise %s if this database needs longer to apply its pending migrations.", + attempt + 1, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, + ) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index d12a2c4dd4e..1bd440d37b4 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -7,6 +7,11 @@ attempt fails identically. These tests pin the two behaviours that keep a container recoverable: an incomplete cache is deleted before Prisma is invoked, and the install gets a budget of its own rather than sharing the one that bounds each migration command. + +``prisma migrate deploy`` gets a budget of its own for the same reason: its +runtime grows with the number of pending migrations, so a fresh database that +replays every migration overran the per-command budget on slow machines and +the proxy gave up after four identical timeouts. """ import ast @@ -14,19 +19,23 @@ import json import os import sys import time +from collections.abc import Callable from pathlib import Path import pytest from litellm_proxy_extras.prisma_toolchain import ( DEFAULT_PRISMA_COMMAND_TIMEOUT, + DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, PRISMA_COMMAND_TIMEOUT_ENV_VAR, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ensure_prisma_toolchain, heal_incomplete_nodeenv_cache, node_binary_path, prisma_bootstrap_timeout, prisma_command_timeout, + prisma_migrate_deploy_timeout, ) from litellm_proxy_extras.utils import ProxyExtrasDBManager @@ -42,13 +51,24 @@ import time args = sys.argv[1:] cache_dir = os.environ["PRISMA_NODEENV_CACHE_DIR"] -with pathlib.Path(os.environ["FAKE_PRISMA_LOG"]).open("a") as log: +log_path = pathlib.Path(os.environ["FAKE_PRISMA_LOG"]) +earlier_deploys = sum( + 1 + for line in (log_path.read_text().splitlines() if log_path.exists() else []) + if json.loads(line)["args"][:2] == ["migrate", "deploy"] +) +with log_path.open("a") as log: log.write( json.dumps({{"args": args, "cache_dir_present": os.path.isdir(cache_dir)}}) + "\\n" ) time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0"))) if args[:2] == ["migrate", "deploy"]: + if earlier_deploys == 0: + time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0"))) + elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"): + print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr) + sys.exit(1) print("No pending migrations to apply") sys.exit(0) """ @@ -80,9 +100,14 @@ def toolchain_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") monkeypatch.delenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raising=False) monkeypatch.delenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, raising=False) + monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False) return cache_dir, log_path +def _deploy_calls(log_path: Path) -> list[list[str]]: + return [call["args"] for call in _fake_prisma_calls(log_path) if call["args"][:2] == ["migrate", "deploy"]] + + def _make_incomplete_cache(cache_dir: Path) -> None: (cache_dir / "lib").mkdir(parents=True) (cache_dir / "bin").mkdir() @@ -209,25 +234,70 @@ def test_setup_database_prepares_the_toolchain_before_migrating( assert calls[0]["cache_dir_present"] is False +@pytest.mark.parametrize("use_v2_resolver", [False, True], ids=["v1", "v2"]) +def test_migrate_deploy_is_not_bounded_by_the_per_command_timeout( + toolchain_env: tuple[Path, Path], + monkeypatch: pytest.MonkeyPatch, + use_v2_resolver: bool, +) -> None: + """A fresh database replays every migration, which takes longer than any bookkeeping command.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "3") + + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + assert _deploy_calls(log_path) == [["migrate", "deploy"]] + + +def test_migrate_deploy_stops_at_its_own_timeout( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + """The deploy budget still bounds a deploy that hangs, so boot cannot wait forever.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") + monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") + + started = time.monotonic() + with pytest.raises(RuntimeError, match="insufficient permissions"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + elapsed = time.monotonic() - started + + assert len(_deploy_calls(log_path)) == 2 + assert elapsed < 30 + + @pytest.mark.parametrize( "raw", ["", "0", "-5", "not-a-number", "nan", "inf", "-inf", "1e400"], ) +@pytest.mark.parametrize( + ("env_var", "read_timeout", "default"), + [ + (PRISMA_COMMAND_TIMEOUT_ENV_VAR, prisma_command_timeout, DEFAULT_PRISMA_COMMAND_TIMEOUT), + (PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, prisma_migrate_deploy_timeout, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT), + ], + ids=["command", "migrate_deploy"], +) def test_unusable_timeout_override_falls_back_to_the_default( - raw: str, monkeypatch: pytest.MonkeyPatch + raw: str, env_var: str, read_timeout: Callable[[], float], default: float, monkeypatch: pytest.MonkeyPatch ) -> None: """A non-finite override would silently disable the timeout it configures.""" - monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raw) + monkeypatch.setenv(env_var, raw) - assert prisma_command_timeout() == DEFAULT_PRISMA_COMMAND_TIMEOUT + assert read_timeout() == default def test_timeout_overrides_are_independent(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "12") monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, "900") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1200") assert prisma_command_timeout() == 12 assert prisma_bootstrap_timeout() == 900 + assert prisma_migrate_deploy_timeout() == 1200 @pytest.mark.parametrize("module", ["utils.py", "replica_identity.py"]) From e0be9a35e6838505b5a2ae2ecf93c01579a02f56 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:02:43 -0700 Subject: [PATCH 25/81] fix(deps): raise the pypdf floor to 6.16.1 for three new advisories GHSA-jp53-mhqp-8xcg (fixed in 6.16.0), GHSA-23w6-3w8w-8484 and GHSA-763m-79hh-57f2 (fixed in 6.16.1) flag pypdf 6.15.0 in uv.lock and keep osv-scan red alongside the tornado advisories. The proxy-runtime extra now requires pypdf>=6.16.1 and the lock resolves 6.16.2. --- pyproject.toml | 2 +- uv.lock | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d0f14722acd..60162544612 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -161,7 +161,7 @@ proxy-runtime = [ "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", "azure-storage-file-datalake>=12.20.0,<13.0", - "pypdf>=6.12.0,<7.0", + "pypdf>=6.16.1,<7.0", "llm-sandbox>=0.3.39,<1.0", "detect-secrets>=1.5.0,<2.0", ] diff --git a/uv.lock b/uv.lock index 8bac024d49e..aa59ff7b229 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-29T20:52:40.322465Z" +exclude-newer = "2026-08-30T17:51:25.171404Z" exclude-newer-span = "P3D" [manifest] @@ -4552,7 +4552,7 @@ requires-dist = [ { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, - { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.12.0,<7.0" }, + { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.16.1,<7.0" }, { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0" }, { name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" }, @@ -7564,14 +7564,14 @@ wheels = [ [[package]] name = "pypdf" -version = "6.15.0" +version = "6.16.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/17/17/ee75a92718ec7212de831e71454d702225aa5e474a805cce169806044453/pypdf-6.15.0.tar.gz", hash = "sha256:d39c4d955a76409284a905e2d65b40076d77ab76129e0faaeeb6612403ecfc79", size = 6993794, upload-time = "2026-08-06T13:06:49.929Z" } +sdist = { url = "https://files.pythonhosted.org/packages/44/66/54212e75406afd9f3e933d0dda23072f6aecc55c5a273077dc2e0b028b23/pypdf-6.16.2.tar.gz", hash = "sha256:595647f6191de6f402cfde1d0c455d6cbccbd509aac32b34783009c032de5d6e", size = 7008996, upload-time = "2026-08-23T13:50:07.135Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/af/72/ce3067ac31e214a66388159f8462ddb8c13dd00170f24d555a1f1ae8ee91/pypdf-6.15.0-py3-none-any.whl", hash = "sha256:14e001d6504822cb1ca9c7ed9a69bccb320f59b320730f55af804361abe4d5ee", size = 378123, upload-time = "2026-08-06T13:06:47.709Z" }, + { url = "https://files.pythonhosted.org/packages/13/f1/a2da3b55acd4ab737bf728c97edaaed5ec1d3c1236acb639dcdfa97e42c7/pypdf-6.16.2-py3-none-any.whl", hash = "sha256:c8b09a59399062fb45a1b8156c18a787a10a3dae03ac9674397a226712c94604", size = 385060, upload-time = "2026-08-23T13:50:05.349Z" }, ] [[package]] From dbc126cfc97734e47066ab988c3ac1090e7f830a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:07:30 -0700 Subject: [PATCH 26/81] fix(hosted_vllm): forward truncate_prompt_tokens on rerank requests --- .../llms/hosted_vllm/rerank/transformation.py | 21 ++- litellm/types/rerank.py | 21 ++- .../test_hosted_vllm_rerank_transformation.py | 122 +++++++++++++++++- 3 files changed, 154 insertions(+), 10 deletions(-) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 0e8fa294f5d..265eb350fc6 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -3,6 +3,7 @@ Transformation logic for Hosted VLLM rerank """ from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final import httpx @@ -13,6 +14,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( + HostedVLLMRerankTruncationParams, OptionalRerankParams, RerankBilledUnits, RerankRequest, @@ -62,7 +64,11 @@ class HostedVLLMRerankConfig(BaseRerankConfig): "top_n", "rank_fields", "return_documents", + "max_tokens_per_doc", "instruction", + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", ] def map_cohere_rerank_params( @@ -100,7 +106,15 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if instruction is not None: mapped_params["instruction"] = instruction - return dict(mapped_params) + truncation: Final = HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + forwarded: Final[OptionalRerankParams] = { + **mapped_params, + "max_tokens_per_doc": max_tokens_per_doc, + "truncate_prompt_tokens": truncation.truncate_prompt_tokens, + "truncation_side": truncation.truncation_side, + "max_tokens_per_query": truncation.max_tokens_per_query, + } + return dict(forwarded) def validate_environment( self, @@ -138,6 +152,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if "documents" not in optional_rerank_params: raise ValueError("documents is required for Hosted VLLM rerank") + truncation: Final = HostedVLLMRerankTruncationParams.model_validate(optional_rerank_params) rerank_request: Final = RerankRequest( model=model, query=optional_rerank_params["query"], @@ -146,6 +161,10 @@ class HostedVLLMRerankConfig(BaseRerankConfig): rank_fields=optional_rerank_params.get("rank_fields", None), return_documents=optional_rerank_params.get("return_documents", None), instruction=optional_rerank_params.get("instruction", None), + max_tokens_per_doc=truncation.max_tokens_per_doc, + truncate_prompt_tokens=truncation.truncate_prompt_tokens, + truncation_side=truncation.truncation_side, + max_tokens_per_query=truncation.max_tokens_per_query, ) return rerank_request.model_dump(exclude_none=True) diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 903781b2ccd..a76e6cf1187 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,8 +4,10 @@ https://docs.cohere.com/reference/rerank """ -from pydantic import BaseModel, PrivateAttr -from typing_extensions import Required, TypedDict +from typing import Literal + +from pydantic import BaseModel, ConfigDict, PrivateAttr +from typing_extensions import ReadOnly, Required, TypedDict class RerankRequest(BaseModel): @@ -21,6 +23,18 @@ class RerankRequest(BaseModel): # (e.g. hosted vLLM / Qwen3-Reranker, DeepInfra). Omitted from the outgoing # request when None, so this is fully backward-compatible. instruction: str | None = None + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + + +class HostedVLLMRerankTruncationParams(BaseModel): + model_config = ConfigDict(frozen=True) + + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + max_tokens_per_doc: int | None = None class OptionalRerankParams(TypedDict, total=False): @@ -32,6 +46,9 @@ class OptionalRerankParams(TypedDict, total=False): max_chunks_per_doc: int | None max_tokens_per_doc: int | None instruction: str | None + truncate_prompt_tokens: ReadOnly[int | None] + truncation_side: ReadOnly[Literal["left", "right"] | None] + max_tokens_per_query: ReadOnly[int | None] class RerankBilledUnits(TypedDict, total=False): diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index e6e6aa946d5..da27ea1ac58 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,8 +1,13 @@ +import json import os import sys +from unittest.mock import MagicMock, patch +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.types.rerank import ( @@ -87,9 +92,7 @@ class TestHostedVLLMRerankTransform: assert "instruction" not in body def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self): - with pytest.raises( - ValueError, match="Hosted VLLM does not support max_chunks_per_doc" - ): + with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"): self.config.map_cohere_rerank_params( non_default_params=None, model=self.model, @@ -104,12 +107,10 @@ class TestHostedVLLMRerankTransform: url = self.config.get_complete_url(base, self.model) assert url == "https://api.example.com/rerank" # Already ends with /rerank - url2 = self.config.get_complete_url( - "https://api.example.com/rerank", self.model - ) + url2 = self.config.get_complete_url("https://api.example.com/rerank", self.model) assert url2 == "https://api.example.com/rerank" # Raises if api_base is None - with pytest.raises(ValueError, match='api_base must be provided for Hosted VLLM rerank'): + with pytest.raises(ValueError, match="api_base must be provided for Hosted VLLM rerank"): self.config.get_complete_url(None, self.model) def test_transform_response(self): @@ -173,3 +174,110 @@ class TestGetOptionalRerankParamsInstruction: documents=["doc1", "doc2"], ) assert "instruction" not in params + + +class TestHostedVLLMRerankTruncationParams: + def setup_method(self): + self.config = HostedVLLMRerankConfig() + self.model = "hosted-vllm-model" + + def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self): + params = self.config.map_cohere_rerank_params( + non_default_params={ + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "metadata": {"user_api_key": "sk-test"}, + }, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + max_tokens_per_doc=128, + ) + assert params["truncate_prompt_tokens"] == 512 + assert params["truncation_side"] == "left" + assert params["max_tokens_per_query"] == 64 + assert params["max_tokens_per_doc"] == 128 + assert "metadata" not in params + + def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): + params = self.config.map_cohere_rerank_params( + non_default_params={"metadata": {"user_api_key": "sk-test"}}, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + body = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) + truncation_keys = {"truncate_prompt_tokens", "truncation_side", "max_tokens_per_query", "max_tokens_per_doc"} + assert not truncation_keys & body.keys() + assert body == { + "model": self.model, + "query": "test query", + "documents": ["doc1", "doc2"], + "return_documents": True, + } + + def test_map_cohere_rerank_params_rejects_invalid_truncation_side(self): + with pytest.raises(ValueError, match="truncation_side"): + self.config.map_cohere_rerank_params( + non_default_params={"truncation_side": "middle"}, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + + def test_transform_request_forwards_truncation_params(self): + body = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={ + "query": "test query", + "documents": ["doc1", "doc2"], + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "max_tokens_per_doc": 128, + }, + headers={}, + ) + assert body["truncate_prompt_tokens"] == 512 + assert body["truncation_side"] == "left" + assert body["max_tokens_per_query"] == 64 + assert body["max_tokens_per_doc"] == 128 + + def test_transform_request_omits_truncation_params_when_absent(self): + body = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]}, + headers={}, + ) + assert "truncate_prompt_tokens" not in body + assert "truncation_side" not in body + assert "max_tokens_per_query" not in body + assert "max_tokens_per_doc" not in body + + def test_rerank_sends_truncate_prompt_tokens_to_vllm(self): + client = HTTPHandler() + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "score-1", + "results": [{"index": 0, "relevance_score": 0.5}], + "usage": {"total_tokens": 512}, + } + with patch.object(client, "post", return_value=mock_response) as mock_post: + litellm.rerank( + model="hosted_vllm/BAAI/bge-reranker-base", + api_base="http://vllm.local:8000", + query="List all the unique case ids", + documents=["a document longer than the reranker context window"], + truncate_prompt_tokens=512, + truncation_side="left", + client=client, + ) + sent_body = json.loads(mock_post.call_args.kwargs["data"]) + assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank" + assert sent_body["truncate_prompt_tokens"] == 512 + assert sent_body["truncation_side"] == "left" From dfaf23523453de12278e3d30801075c4da6ee903 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:08:45 -0700 Subject: [PATCH 27/81] fix(bedrock): honor BEDROCK_MANTLE_API_BASE on bedrock/mantle messages and chat URLs --- litellm/llms/bedrock/common_utils.py | 7 ++-- .../test_litellm/llms/bedrock/test_mantle.py | 42 +++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 66ee5f10679..048d023a1bd 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -758,12 +758,13 @@ def build_mantle_messages_url( """Build the bedrock-mantle Anthropic /messages URL. Honors an explicit endpoint override (``api_base``, then - ``aws_bedrock_runtime_endpoint``) so private VPC / VPCE / GovCloud Mantle - endpoints are reachable; otherwise falls back to the public regional host. + ``aws_bedrock_runtime_endpoint``, then ``BEDROCK_MANTLE_API_BASE``) so + private VPC / VPCE / GovCloud Mantle endpoints are reachable; otherwise + falls back to the public regional host. The mantle messages path is appended unless the override already carries it, so callers can pass either the host or the full messages URL. """ - override: Final = api_base or aws_bedrock_runtime_endpoint + override: Final = api_base or aws_bedrock_runtime_endpoint or get_secret_str("BEDROCK_MANTLE_API_BASE") if override: base: Final = override.rstrip("/") if base.endswith(MANTLE_MESSAGES_PATH): diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index d34517f61f6..d1d1ba447fb 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -128,6 +128,12 @@ def test_mantle_messages_url_construction(): _VPC_ENDPOINT = "https://vpce-0a1b2c3d.bedrock-mantle.us-gov-west-1.vpce.amazonaws.com" +@pytest.fixture(autouse=True) +def no_ambient_mantle_api_base(monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + + + def test_mantle_chat_url_honors_api_base_host(): config = AmazonMantleConfig() url = config.get_complete_url( @@ -193,6 +199,42 @@ def test_mantle_messages_url_honors_aws_bedrock_runtime_endpoint(): assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" +_ENV_ENDPOINT = "https://bedrock-mantle.us-east-1.api.aws.internal.example.com" + + +@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) +def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) + url = config_cls().get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + ) + assert url == f"{_ENV_ENDPOINT}/anthropic/v1/messages" + + +@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) +@pytest.mark.parametrize( + ("api_base", "optional_params"), + [ + (_VPC_ENDPOINT, {"aws_region_name": "us-gov-west-1"}), + (None, {"aws_region_name": "us-gov-west-1", "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT}), + ], +) +def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env(monkeypatch, config_cls, api_base, optional_params): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) + url = config_cls().get_complete_url( + api_base=api_base, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params=optional_params, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + def test_mantle_transform_request_strips_prefix_and_adds_model(): config = AmazonMantleConfig() request = config.transform_request( From 8d00220acef335234797b00cf3b1f1ee73db2b3a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:19:20 -0700 Subject: [PATCH 28/81] test(hosted_vllm): annotate rerank truncation test locals as Final --- .../test_hosted_vllm_rerank_transformation.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index da27ea1ac58..49d58cc28c4 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,6 +1,7 @@ import json import os import sys +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -182,7 +183,7 @@ class TestHostedVLLMRerankTruncationParams: self.model = "hosted-vllm-model" def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self): - params = self.config.map_cohere_rerank_params( + params: Final = self.config.map_cohere_rerank_params( non_default_params={ "truncate_prompt_tokens": 512, "truncation_side": "left", @@ -202,15 +203,20 @@ class TestHostedVLLMRerankTruncationParams: assert "metadata" not in params def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): - params = self.config.map_cohere_rerank_params( + params: Final = self.config.map_cohere_rerank_params( non_default_params={"metadata": {"user_api_key": "sk-test"}}, model=self.model, drop_params=False, query="test query", documents=["doc1", "doc2"], ) - body = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) - truncation_keys = {"truncate_prompt_tokens", "truncation_side", "max_tokens_per_query", "max_tokens_per_doc"} + body: Final = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) + truncation_keys: Final = { + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", + "max_tokens_per_doc", + } assert not truncation_keys & body.keys() assert body == { "model": self.model, @@ -230,7 +236,7 @@ class TestHostedVLLMRerankTruncationParams: ) def test_transform_request_forwards_truncation_params(self): - body = self.config.transform_rerank_request( + body: Final = self.config.transform_rerank_request( model=self.model, optional_rerank_params={ "query": "test query", @@ -248,7 +254,7 @@ class TestHostedVLLMRerankTruncationParams: assert body["max_tokens_per_doc"] == 128 def test_transform_request_omits_truncation_params_when_absent(self): - body = self.config.transform_rerank_request( + body: Final = self.config.transform_rerank_request( model=self.model, optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]}, headers={}, @@ -259,8 +265,8 @@ class TestHostedVLLMRerankTruncationParams: assert "max_tokens_per_doc" not in body def test_rerank_sends_truncate_prompt_tokens_to_vllm(self): - client = HTTPHandler() - mock_response = MagicMock(spec=httpx.Response) + client: Final = HTTPHandler() + mock_response: Final = MagicMock(spec=httpx.Response) mock_response.status_code = 200 mock_response.json.return_value = { "id": "score-1", @@ -277,7 +283,7 @@ class TestHostedVLLMRerankTruncationParams: truncation_side="left", client=client, ) - sent_body = json.loads(mock_post.call_args.kwargs["data"]) + sent_body: Final = json.loads(mock_post.call_args.kwargs["data"]) assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank" assert sent_body["truncate_prompt_tokens"] == 512 assert sent_body["truncation_side"] == "left" From ef14bed0296b4a6c48777b8f9df8018b11179c4b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:31:19 -0700 Subject: [PATCH 29/81] fix(hosted_vllm): reject invalid rerank truncation params with a 400 --- .../llms/hosted_vllm/rerank/transformation.py | 11 +++++++- .../test_hosted_vllm_rerank_transformation.py | 26 ++++++++++++------- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 265eb350fc6..764d80c6f82 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -7,8 +7,10 @@ from types import MappingProxyType from typing import Any, Final import httpx +from pydantic import ValidationError from litellm._uuid import uuid +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -36,6 +38,13 @@ class HostedVLLMRerankError(BaseLLMException): super().__init__(status_code=status_code, message=message, headers=headers) +def validated_truncation_params(non_default_params: Mapping[str, object] | None) -> HostedVLLMRerankTruncationParams: + try: + return HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + except ValidationError as error: + raise UnsupportedParamsError(status_code=400, message=f"hosted_vllm rerank: {error}") from error + + class HostedVLLMRerankConfig(BaseRerankConfig): def __init__(self) -> None: pass @@ -106,7 +115,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if instruction is not None: mapped_params["instruction"] = instruction - truncation: Final = HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({})) + truncation: Final = validated_truncation_params(non_default_params) forwarded: Final[OptionalRerankParams] = { **mapped_params, "max_tokens_per_doc": max_tokens_per_doc, diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 49d58cc28c4..9a62fcf6f0f 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -202,6 +202,22 @@ class TestHostedVLLMRerankTruncationParams: assert params["max_tokens_per_doc"] == 128 assert "metadata" not in params + @pytest.mark.parametrize( + "bad_params", + [{"truncation_side": "middle"}, {"truncate_prompt_tokens": "lots"}, {"max_tokens_per_query": -1.5}], + ) + def test_map_cohere_rerank_params_rejects_invalid_truncation_params_as_400(self, bad_params: dict[str, object]): + with pytest.raises(litellm.UnsupportedParamsError) as raised: + self.config.map_cohere_rerank_params( + non_default_params=dict(bad_params), + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + assert raised.value.status_code == 400 + assert next(iter(bad_params)) in str(raised.value) + def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): params: Final = self.config.map_cohere_rerank_params( non_default_params={"metadata": {"user_api_key": "sk-test"}}, @@ -225,16 +241,6 @@ class TestHostedVLLMRerankTruncationParams: "return_documents": True, } - def test_map_cohere_rerank_params_rejects_invalid_truncation_side(self): - with pytest.raises(ValueError, match="truncation_side"): - self.config.map_cohere_rerank_params( - non_default_params={"truncation_side": "middle"}, - model=self.model, - drop_params=False, - query="test query", - documents=["doc1", "doc2"], - ) - def test_transform_request_forwards_truncation_params(self): body: Final = self.config.transform_rerank_request( model=self.model, From 49c69c46b25af2dd962b322bd6c8e6c5668546c6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:51:27 -0700 Subject: [PATCH 30/81] fix(bedrock): drop the OpenAI base suffix from BEDROCK_MANTLE_API_BASE before the mantle messages path --- litellm/llms/bedrock/common_utils.py | 15 +++++++++++++-- tests/test_litellm/llms/bedrock/test_mantle.py | 12 +++++++++--- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 048d023a1bd..1e5329c90dd 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -748,6 +748,15 @@ def strip_bedrock_throughput_suffix(model: str) -> str: MANTLE_MESSAGES_PATH: Final = "/anthropic/v1/messages" +_MANTLE_OPENAI_BASE_SUFFIXES: Final = ("/openai/v1", "/v1") + + +def _mantle_api_base_from_env() -> str | None: + env_base: Final = get_secret_str("BEDROCK_MANTLE_API_BASE") + if env_base is None: + return None + base: Final = env_base.rstrip("/") + return next((base[: -len(suffix)] for suffix in _MANTLE_OPENAI_BASE_SUFFIXES if base.endswith(suffix)), base) def build_mantle_messages_url( @@ -762,9 +771,11 @@ def build_mantle_messages_url( private VPC / VPCE / GovCloud Mantle endpoints are reachable; otherwise falls back to the public regional host. The mantle messages path is appended unless the override already carries it, - so callers can pass either the host or the full messages URL. + so callers can pass either the host or the full messages URL. The env var is + shared with the OpenAI-surface ``bedrock_mantle/*`` routes, which need it to + carry their ``/v1`` or ``/openai/v1`` base, so that suffix is dropped first. """ - override: Final = api_base or aws_bedrock_runtime_endpoint or get_secret_str("BEDROCK_MANTLE_API_BASE") + override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env() if override: base: Final = override.rstrip("/") if base.endswith(MANTLE_MESSAGES_PATH): diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index d1d1ba447fb..09be2118001 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -203,8 +203,12 @@ _ENV_ENDPOINT = "https://bedrock-mantle.us-east-1.api.aws.internal.example.com" @pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) -def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls): - monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) +@pytest.mark.parametrize( + "env_value", + [_ENV_ENDPOINT, f"{_ENV_ENDPOINT}/", f"{_ENV_ENDPOINT}/v1", f"{_ENV_ENDPOINT}/openai/v1"], +) +def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls, env_value): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", env_value) url = config_cls().get_complete_url( api_base=None, api_key=None, @@ -223,7 +227,9 @@ def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls): (None, {"aws_region_name": "us-gov-west-1", "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT}), ], ) -def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env(monkeypatch, config_cls, api_base, optional_params): +def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env( + monkeypatch, config_cls, api_base, optional_params +): monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) url = config_cls().get_complete_url( api_base=api_base, From cc2cbb36f326a87cd72421a4448349734f872201 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:57:45 -0700 Subject: [PATCH 31/81] fix(otel): stamp Langfuse root observation input and output from the request task --- litellm/integrations/otel/langfuse_logger.py | 69 ++++ litellm/integrations/otel/logger.py | 32 +- litellm/integrations/otel/mappers/langfuse.py | 8 +- litellm/integrations/otel/model/request_io.py | 90 +++++ litellm/litellm_core_utils/litellm_logging.py | 14 +- .../integrations/otel/test_langfuse_logger.py | 317 ++++++++++++++++++ 6 files changed, 519 insertions(+), 11 deletions(-) create mode 100644 litellm/integrations/otel/langfuse_logger.py create mode 100644 litellm/integrations/otel/model/request_io.py create mode 100644 tests/test_litellm/integrations/otel/test_langfuse_logger.py diff --git a/litellm/integrations/otel/langfuse_logger.py b/litellm/integrations/otel/langfuse_logger.py new file mode 100644 index 00000000000..9986eae4d0a --- /dev/null +++ b/litellm/integrations/otel/langfuse_logger.py @@ -0,0 +1,69 @@ +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_logger +from litellm.integrations.otel.logger import OpenTelemetryV2 +from litellm.integrations.otel.mappers.langfuse import LANGFUSE_OBSERVATION_INPUT, LANGFUSE_OBSERVATION_OUTPUT +from litellm.integrations.otel.model.request_io import request_input, response_output, stream_output +from litellm.integrations.otel.plumbing.context import request_root_span + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import CallTypesLiteral, ModelResponseStream + +ROOT_OBSERVATION_IO_CALL_TYPES: Final = frozenset( + {"completion", "acompletion", "responses", "aresponses", "anthropic_messages", "aanthropic_messages"} +) + + +class LangfuseOpenTelemetryV2(OpenTelemetryV2): + """Stamps the request's input and output on the root observation while it is still recording. + + Langfuse shows a trace's input and output from its root observation. The proxy's root span ends + when the response is sent, before the success callback runs, so the stamps have to come from the + request-task hooks: input at pre-call, output at post-call success or at the end of the stream. + """ + + async def async_pre_call_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + cache: "DualCache", + data: Mapping[str, object], + call_type: "CallTypesLiteral", + ) -> None: + await super().async_pre_call_hook(user_api_key_dict, cache, data, call_type) + if call_type in ROOT_OBSERVATION_IO_CALL_TYPES: + self._stamp_root(LANGFUSE_OBSERVATION_INPUT, lambda: request_input(data)) + + async def async_post_call_success_hook( + self, + data: Mapping[str, object], + user_api_key_dict: "UserAPIKeyAuth", + response: object, + ) -> None: + self._stamp_root(LANGFUSE_OBSERVATION_OUTPUT, lambda: response_output(response)) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + response: "AsyncIterator[ModelResponseStream]", + request_data: Mapping[str, object], + ) -> "AsyncGenerator[ModelResponseStream, None]": + relayed: Final[list[ModelResponseStream]] = [] # mutable-ok: relayed as they arrive, assembled at end of stream + async for chunk in response: + relayed.append(chunk) + yield chunk + self._stamp_root(LANGFUSE_OBSERVATION_OUTPUT, lambda: stream_output(tuple(relayed), request_data)) + + def _stamp_root(self, key: str, render: Callable[[], str | None]) -> None: + root: Final = request_root_span() + if root is None or not root.is_recording(): + return + try: + value: Final = render() + except Exception: # noqa: BLE001 # telemetry must never fail the request it describes + verbose_logger.debug("otel v2 langfuse: could not render %s for the root observation", key, exc_info=True) + return + if value is not None: + root.set_attribute(key, value) diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index d2a32ef73b6..4ab1c738488 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -4,6 +4,7 @@ from collections import OrderedDict from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast from opentelemetry.context import Context, attach, get_current @@ -722,14 +723,13 @@ class OpenTelemetryV2(CustomLogger): self, user_api_key_dict: "UserAPIKeyAuth", cache: "DualCache", - data: dict, + data: Mapping[str, object], call_type: "CallTypesLiteral", - ) -> dict: + ) -> None: self.seed_request_identity( user_api_key_dict, model=model_from_request_data(data), ) - return data def record_error_attributes_on_span( self, @@ -909,3 +909,29 @@ def phase_span(name: str) -> "Iterator[Span | None]": return with logger.start_phase_span(name) as span: yield span + + +def build_otel_v2_logger( + config: OpenTelemetryV2Config, + callback_name: str | None = None, + tracer_provider: TracerProvider | None = None, + logger_provider: LoggerProvider | None = None, + meter_provider: "MeterProvider | None" = None, + settings: Mapping[str, object] = MappingProxyType({}), +) -> OpenTelemetryV2: + return _logger_class(config)( + config=config, + callback_name=callback_name, + tracer_provider=tracer_provider, + logger_provider=logger_provider, + meter_provider=meter_provider, + **settings, + ) + + +def _logger_class(config: OpenTelemetryV2Config) -> type[OpenTelemetryV2]: + if "langfuse" not in config.mapper_names or not config.capture_span_content: + return OpenTelemetryV2 + from litellm.integrations.otel.langfuse_logger import LangfuseOpenTelemetryV2 + + return LangfuseOpenTelemetryV2 diff --git a/litellm/integrations/otel/mappers/langfuse.py b/litellm/integrations/otel/mappers/langfuse.py index 6d4f1b4fd0a..01063d85355 100644 --- a/litellm/integrations/otel/mappers/langfuse.py +++ b/litellm/integrations/otel/mappers/langfuse.py @@ -11,6 +11,7 @@ the JSON-serialized payloads. ``_llm_call`` just applies both tables. import json from collections.abc import Callable +from typing import Final from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData from litellm.integrations.otel.mappers.utils import ( @@ -25,6 +26,9 @@ from litellm.integrations.otel.model.payloads import ( LLMUsage, ) +LANGFUSE_OBSERVATION_INPUT: Final = "langfuse.observation.input" +LANGFUSE_OBSERVATION_OUTPUT: Final = "langfuse.observation.output" + class LangfuseMapper: _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { @@ -56,8 +60,8 @@ class LangfuseMapper: "langfuse.observation.model.parameters": lambda d: json_if( collect(LangfuseMapper._MODEL_PARAMS, d.request_params) ), - "langfuse.observation.input": lambda d: serialize_messages(d.messages_in), - "langfuse.observation.output": lambda d: serialize_messages(output_messages(d)), + LANGFUSE_OBSERVATION_INPUT: lambda d: serialize_messages(d.messages_in), + LANGFUSE_OBSERVATION_OUTPUT: lambda d: serialize_messages(output_messages(d)), "langfuse.observation.usage_details": lambda d: json_if(collect(LangfuseMapper._USAGE_FIELDS, d.usage)), "langfuse.observation.cost_details": lambda d: ( json.dumps({"total": d.response_cost}) if d.response_cost is not None else None diff --git a/litellm/integrations/otel/model/request_io.py b/litellm/integrations/otel/model/request_io.py new file mode 100644 index 00000000000..4e80fb91993 --- /dev/null +++ b/litellm/integrations/otel/model/request_io.py @@ -0,0 +1,90 @@ +from collections.abc import Mapping, Sequence +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm.integrations.otel.mappers.utils import json_or_none +from litellm.proxy.guardrails.anthropic_sse import assemble_anthropic_sse_stream, is_raw_sse_stream +from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.utils import ModelResponse, ModelResponseStream + +_SYSTEM_KEYS: Final = ("system", "instructions") +_TURNS: Final = TypeAdapter(tuple[object, ...]) +_MESSAGES: Final = TypeAdapter(list[object] | None) + + +class _Turn(TypedDict): + role: ReadOnly[str] + content: ReadOnly[object] + + +class _AnthropicMessage(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["message"] = Field(exclude=True) + role: str = "assistant" + content: object = None + + +def request_input(data: Mapping[str, object]) -> str | None: + turns: Final = data.get("messages", data.get("input")) + if turns is None: + return None + return json_or_none((*_system_turns(data), *_user_turns(turns))) + + +def _system_turns(data: Mapping[str, object]) -> tuple[_Turn, ...]: + return tuple(_Turn(role="system", content=data[key]) for key in _SYSTEM_KEYS if data.get(key) is not None) + + +def _user_turns(turns: object) -> tuple[object, ...]: + if isinstance(turns, str): + return (_Turn(role="user", content=turns),) + try: + return _TURNS.validate_python(turns) + except ValidationError: + return (_Turn(role="user", content=turns),) + + +def response_output(response: object) -> str | None: + match response: + case ModelResponse(): + return json_or_none(tuple(choice.message.model_dump(exclude_none=True) for choice in response.choices)) + case ResponsesAPIResponse(): + return json_or_none(response.model_dump(exclude_none=True).get("output")) + case _: + return _anthropic_message_output(response) + + +def _anthropic_message_output(message: object) -> str | None: + try: + parsed: Final = _AnthropicMessage.model_validate(message) + except ValidationError: + return None + return json_or_none((parsed.model_dump(),)) + + +def stream_output(chunks: Sequence[object], data: Mapping[str, object]) -> str | None: + if not chunks: + return None + if is_raw_sse_stream(chunks): + return response_output(assemble_anthropic_sse_stream(chunks)) + if all(isinstance(chunk, ModelResponseStream) for chunk in chunks): + return response_output(_assembled_chat_stream(chunks, data)) + return response_output(_completed_response(chunks)) + + +def _assembled_chat_stream(chunks: Sequence[object], data: Mapping[str, object]) -> object: + try: + return litellm.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType] # upstream types chunks as a bare list + chunks=list(chunks), # mutable-ok: stream_chunk_builder takes a list + messages=_MESSAGES.validate_python(data.get("messages")), + ) + except (litellm.APIError, ValidationError): + return None + + +def _completed_response(chunks: Sequence[object]) -> ResponsesAPIResponse | None: + return next((chunk.response for chunk in reversed(chunks) if isinstance(chunk, ResponseCompletedEvent)), None) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9a6fb11f978..8e8575eff9a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4390,13 +4390,15 @@ def _init_custom_logger_compatible_class( from litellm.integrations.otel.model.config import is_otel_v2_enabled if is_otel_v2_enabled(): - from litellm.integrations.otel.logger import OpenTelemetryV2 + from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger + from litellm.integrations.otel.model.config import OpenTelemetryV2Config for callback in _in_memory_loggers: - if type(callback) is OpenTelemetryV2: + if isinstance(callback, OpenTelemetryV2): return callback - otel_logger_v2: Final = OpenTelemetryV2( - **_get_custom_logger_settings_from_proxy_server(callback_name=logging_integration) + otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration) + otel_logger_v2: Final = build_otel_v2_logger( + config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings ) _in_memory_loggers.append(otel_logger_v2) _maybe_auto_initialize_arize_phoenix(_in_memory_loggers) @@ -4759,7 +4761,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom if not is_otel_v2_enabled(): return None - from litellm.integrations.otel.logger import OpenTelemetryV2 + from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger from litellm.integrations.otel.presets import PRESET_BY_CALLBACK preset_fn: Final = PRESET_BY_CALLBACK.get(callback_name) @@ -4774,7 +4776,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom # If env vars are missing or the preset raises, defer to the legacy path # so customers get the same error story they had before V2 landed. return None - v2_logger: Final = OpenTelemetryV2(config=config, callback_name=callback_name) + v2_logger: Final = build_otel_v2_logger(config=config, callback_name=callback_name) _in_memory_loggers.append(v2_logger) return v2_logger diff --git a/tests/test_litellm/integrations/otel/test_langfuse_logger.py b/tests/test_litellm/integrations/otel/test_langfuse_logger.py new file mode 100644 index 00000000000..af0597517dc --- /dev/null +++ b/tests/test_litellm/integrations/otel/test_langfuse_logger.py @@ -0,0 +1,317 @@ +"""Tests for ``LangfuseOpenTelemetryV2``: the root observation's input and output are stamped from the +request-task hooks, while the root span is still recording, so Langfuse can show them on the trace.""" + +import asyncio +import json +from collections.abc import AsyncIterator, Sequence +from typing import Final + +import pytest + +pytest.importorskip("opentelemetry") + +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter # noqa: E402 + +from litellm.caching.dual_cache import DualCache # noqa: E402 +from litellm.integrations.otel.langfuse_logger import LangfuseOpenTelemetryV2 # noqa: E402 +from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger # noqa: E402 +from litellm.integrations.otel.model.config import OpenTelemetryV2Config, is_otel_v2_enabled # noqa: E402 +from litellm.integrations.otel.model.spans import LITELLM_PROXY_REQUEST_SPAN_NAME, SpanRole # noqa: E402 +from litellm.integrations.otel.plumbing import context as otel_context # noqa: E402 +from litellm.integrations.otel.plumbing import providers # noqa: E402 +from litellm.integrations.otel.plumbing.context import set_request_root_span # noqa: E402 +from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 # noqa: E402 +from litellm.proxy._types import UserAPIKeyAuth # noqa: E402 +from litellm.types.llms.openai import ( # noqa: E402 + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) +from litellm.types.utils import ( # noqa: E402 + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) + +INPUT_ATTR: Final = "langfuse.observation.input" +OUTPUT_ATTR: Final = "langfuse.observation.output" +CHAT_DATA: Final = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "ping"}]} + + +@pytest.fixture(autouse=True) +def _reset_request_root_span(): + otel_context._request_root_span.set(None) + yield + otel_context._request_root_span.set(None) + + +def _logger(*, capture: str = "span_only", mappers: Sequence[str] = ("genai", "langfuse")): + cfg = OpenTelemetryV2Config(exporter="in_memory", mapper_names=list(mappers), capture_message_content=capture) + exporter = InMemorySpanExporter() + tracer_provider = providers.build_tracer_provider(cfg, exporter=exporter) + return build_otel_v2_logger(config=cfg, tracer_provider=tracer_provider), exporter + + +def _start_root(logger): + root = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(root) + return root + + +def _root_attrs(exporter): + by_name = {span.name: span for span in exporter.get_finished_spans()} + return dict(by_name[LITELLM_PROXY_REQUEST_SPAN_NAME].attributes or {}) + + +def _run_request(logger, data: dict, call_type: str, response: object): + root = _start_root(logger) + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, call_type)) + asyncio.run(logger.async_post_call_success_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), response=response)) + root.end() + + +async def _relay(logger, chunks: Sequence[object], data: dict) -> list[object]: + async def source() -> AsyncIterator[object]: + for chunk in chunks: + yield chunk + + return [chunk async for chunk in logger.async_post_call_streaming_iterator_hook(UserAPIKeyAuth(), source(), data)] + + +def _run_stream(logger, data: dict, chunks: Sequence[object]) -> list[object]: + root = _start_root(logger) + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion")) + relayed = asyncio.run(_relay(logger, chunks, data)) + root.end() + return relayed + + +def _chat_chunk(content: str | None, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + created=1, + model="gpt-5.4-mini", + choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + +def _responses_api_response() -> ResponsesAPIResponse: + return ResponsesAPIResponse( + id="resp_1", + created_at=1, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + } + ], + ) + + +def _anthropic_sse_frames() -> tuple[bytes, ...]: + events = ( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "po"}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ng"}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}}, + {"type": "message_stop"}, + ) + return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) + + +def test_chat_request_stamps_root_observation_input_and_output(): + logger, exporter = _logger() + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) + + _run_request(logger, CHAT_DATA, "acompletion", response) + + attrs = _root_attrs(exporter) + assert json.loads(attrs[INPUT_ATTR]) == [{"role": "user", "content": "ping"}] + output = json.loads(attrs[OUTPUT_ATTR]) + assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")] + + +def test_responses_request_folds_instructions_into_input_and_stamps_output_items(): + logger, exporter = _logger() + data = {"model": "gpt-5.4-mini", "instructions": "be terse", "input": "ping"} + + _run_request(logger, data, "aresponses", _responses_api_response()) + + attrs = _root_attrs(exporter) + assert json.loads(attrs[INPUT_ATTR]) == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "ping"}, + ] + output = json.loads(attrs[OUTPUT_ATTR]) + assert output[0]["role"] == "assistant" + assert output[0]["content"][0]["text"] == "pong" + + +def test_anthropic_messages_request_folds_system_into_input_and_stamps_content_blocks(): + logger, exporter = _logger() + data = {"model": "claude-sonnet-4-5", "system": "be terse", "messages": [{"role": "user", "content": "ping"}]} + response = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "pong"}]} + + _run_request(logger, data, "aanthropic_messages", response) + + attrs = _root_attrs(exporter) + assert json.loads(attrs[INPUT_ATTR]) == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "ping"}, + ] + assert json.loads(attrs[OUTPUT_ATTR]) == [{"role": "assistant", "content": [{"type": "text", "text": "pong"}]}] + + +def test_chat_stream_relays_chunks_untouched_and_stamps_assembled_output(): + logger, exporter = _logger() + chunks = (_chat_chunk("po"), _chat_chunk("ng"), _chat_chunk(None, finish_reason="stop")) + + relayed = _run_stream(logger, CHAT_DATA, chunks) + + assert [id(chunk) for chunk in relayed] == [id(chunk) for chunk in chunks] + output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR]) + assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")] + + +def test_responses_stream_stamps_output_from_the_completed_event(): + logger, exporter = _logger() + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_responses_api_response() + ) + chunks = ({"type": "response.created"}, {"type": "response.output_text.delta", "delta": "pong"}, completed) + + relayed = _run_stream(logger, {"model": "gpt-5.4-mini", "input": "ping"}, chunks) + + assert relayed == list(chunks) + output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR]) + assert output[0]["content"][0]["text"] == "pong" + + +def test_anthropic_sse_stream_stamps_output_from_the_assembled_frames(): + logger, exporter = _logger() + frames = _anthropic_sse_frames() + + relayed = _run_stream( + logger, {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "ping"}]}, frames + ) + + assert relayed == list(frames) + output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR]) + assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")] + + +def test_root_observation_io_survives_the_root_ending_before_the_success_callback(): + logger, exporter = _logger() + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) + root = _start_root(logger) + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion")) + logger.log_pre_api_call( + model="gpt-5.4-mini", + messages=[], + kwargs={"litellm_call_id": "call_1", "litellm_params": {"metadata": {}}}, + ) + asyncio.run( + logger.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response) + ) + root.end() + + payload = { + "call_type": "acompletion", + "custom_llm_provider": "openai", + "model": "gpt-5.4-mini", + "messages": CHAT_DATA["messages"], + "response": response.model_dump(), + "status": "success", + "litellm_call_id": "call_1", + "metadata": {}, + "hidden_params": {}, + } + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": payload, "litellm_params": {"metadata": {}}}, response, None, None + ) + ) + + attrs = _root_attrs(exporter) + assert INPUT_ATTR in attrs and OUTPUT_ATTR in attrs + generation = next(span for span in exporter.get_finished_spans() if span.name != LITELLM_PROXY_REQUEST_SPAN_NAME) + assert OUTPUT_ATTR in dict(generation.attributes or {}) + + +def test_root_already_ended_is_left_alone(): + logger, exporter = _logger() + root = _start_root(logger) + root.end() + + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion")) + + assert INPUT_ATTR not in _root_attrs(exporter) + + +def test_non_chat_call_types_do_not_stamp_input(): + logger, exporter = _logger() + root = _start_root(logger) + + asyncio.run( + logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), {"model": "e", "input": "ping"}, "aembedding") + ) + root.end() + + assert INPUT_ATTR not in _root_attrs(exporter) + + +def test_unrenderable_output_never_raises_into_the_request(): + logger, exporter = _logger() + + _run_request(logger, CHAT_DATA, "acompletion", object()) + + assert OUTPUT_ATTR not in _root_attrs(exporter) + + +@pytest.mark.parametrize( + ("capture", "mappers"), + [("no_content", ("genai", "langfuse")), ("span_only", ("genai",))], +) +def test_factory_keeps_the_base_logger_unless_langfuse_content_capture_is_on(capture, mappers): + logger, exporter = _logger(capture=capture, mappers=mappers) + + assert type(logger) is OpenTelemetryV2 + _run_request(logger, CHAT_DATA, "acompletion", ModelResponse()) + attrs = _root_attrs(exporter) + assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs + + +def test_langfuse_otel_preset_builds_the_langfuse_logger(monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk") + monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com") + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "span_only") + is_otel_v2_enabled.cache_clear() + + loggers: list = [] + try: + built = _maybe_construct_otel_v2("langfuse_otel", loggers) + assert isinstance(built, LangfuseOpenTelemetryV2) + assert _maybe_construct_otel_v2("langfuse_otel", loggers) is built + finally: + is_otel_v2_enabled.cache_clear() From a677242d6f07af683b9c146133287f07b4e1459c Mon Sep 17 00:00:00 2001 From: Ali Ahmed <128928915+QuantumBreakz@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:07:13 +0500 Subject: [PATCH 32/81] fix(headroom): stop re-compressing retrieved CCR content in client tool loops (#38591) When the headroom_retrieve tool is exposed to a client that runs its own tool-execution loop (the LiteLLM MCP gateway path), the client executes the retrieve call and sends the recovered original content back as a tool result on the next turn. The guardrail then compressed that row again, and because CCR is content-addressed it collapsed back to the exact same hash it was just retrieved from. The model never saw the expansion and the agent looped. Hold tool-result rows that carry headroom_retrieve output back from the compression service, the same way the live turn and trailing tool exchange are already protected, so the expansion survives. Retrieve calls are matched by the direct headroom_retrieve name and the mcp____headroom_retrieve gateway name. Because a long gateway name is truncated past 64 chars in the OpenAI-translated view the guardrail scans, the pairing also falls back to the tool-call id read from the request's own untranslated messages, which is never truncated. Fixes #38558 --- .../guardrail_hooks/headroom/headroom.py | 120 ++++++++++++- .../guardrail_hooks/test_headroom.py | 158 ++++++++++++++++++ 2 files changed, 274 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index d8c8c2f4974..fc881a60f43 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard import httpx from fastapi import HTTPException from httpx import Response as HttpxResponse +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -52,6 +53,10 @@ BYPASS_HEADER: Final = "x-headroom-bypass" HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" _HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 +# Narrows the base class's bare-dict ``request_data`` at the boundary so its +# untranslated messages can be read with concrete types (values pass through by +# reference, so this is a shallow top-level reconstruction). +_REQUEST_DATA_ADAPTER: Final = TypeAdapter(dict[str, object]) def _is_str_object_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip @@ -116,16 +121,119 @@ def _restore_content_shapes( return restored -def _protected_indices(messages: Sequence[Mapping[str, object]]) -> frozenset[int]: +def _tool_call_name(tool_call: Mapping[str, object]) -> str | None: + function: Final = tool_call.get("function") + if not _is_str_object_dict(function): + return None + name: Final = function.get("name") + return name if isinstance(name, str) else None + + +def _is_retrieve_tool_name(name: str | None) -> bool: + """Match the retrieve tool whether called directly or via the MCP gateway. + + Server-side the tool is ``headroom_retrieve``; exposed through LiteLLM's MCP + gateway a client calls it as ``mcp____headroom_retrieve``. + """ + return name is not None and ( + name == HEADROOM_RETRIEVE_TOOL_NAME or name.endswith(f"__{HEADROOM_RETRIEVE_TOOL_NAME}") + ) + + +def _retrieve_call_ids_in_message(message: Mapping[str, object]) -> frozenset[str]: + if message.get("role") != "assistant": + return frozenset() + tool_calls: Final = message.get("tool_calls") + if not _is_object_list(tool_calls): + return frozenset() + return frozenset( + str(tool_call["id"]) + for tool_call in tool_calls + if _is_str_object_dict(tool_call) and tool_call.get("id") and _is_retrieve_tool_name(_tool_call_name(tool_call)) + ) + + +def _anthropic_tool_use_retrieve_id(block: object) -> str | None: + if not _is_str_object_dict(block) or block.get("type") != "tool_use": + return None + name: Final = block.get("name") + call_id: Final = block.get("id") + if isinstance(name, str) and call_id is not None and _is_retrieve_tool_name(name): + return str(call_id) + return None + + +def _anthropic_retrieve_ids_in_message(message: Mapping[str, object]) -> frozenset[str]: + content: Final = message.get("content") + if not _is_object_list(content): + return frozenset() + return frozenset(call_id for block in content if (call_id := _anthropic_tool_use_retrieve_id(block)) is not None) + + +def _raw_retrieve_call_ids(messages: object) -> frozenset[str]: + """Retrieve-tool call ids read from the request's own, untranslated messages. + + The guardrail otherwise scans an OpenAI-translated view where a tool name + over 64 chars is truncated to ``{prefix}_{hash}``, which drops the + ``__headroom_retrieve`` suffix a long ``mcp____`` prefix pushes past + the limit. Tool-call ids are never truncated, so pairing the tool result to + an id read from the original request keeps the match intact. Both wire + shapes are handled: OpenAI ``tool_calls`` and Anthropic ``tool_use`` blocks. + """ + if not _is_object_list(messages): + return frozenset() + return frozenset( + call_id + for message in messages + if _is_str_object_dict(message) + for call_id in _retrieve_call_ids_in_message(message) | _anthropic_retrieve_ids_in_message(message) + ) + + +def _retrieval_result_indices( + messages: Sequence[Mapping[str, object]], extra_retrieve_call_ids: frozenset[str] = frozenset() +) -> frozenset[int]: + """Indices of tool-result rows that carry ``headroom_retrieve`` output. + + When the retrieve tool is exposed to a client that runs its own tool loop + (the LiteLLM MCP gateway path), the client executes the call and sends the + recovered original content back as a tool result on the next turn. That + content is exactly what a prior compression stubbed, so compressing it again + re-derives the identical content hash: a no-op that strands the model on the + marker and loops the agent. Hold those rows back so the expansion survives. + + ``extra_retrieve_call_ids`` carries ids recovered from the untruncated + request so the pairing survives tool-name truncation (see + ``_raw_retrieve_call_ids``). + """ + retrieve_call_ids: Final = extra_retrieve_call_ids | frozenset( + call_id for message in messages for call_id in _retrieve_call_ids_in_message(message) + ) + if not retrieve_call_ids: + return frozenset() + return frozenset( + index + for index, message in enumerate(messages) + if message.get("role") in ("tool", "function") and str(message.get("tool_call_id")) in retrieve_call_ids + ) + + +def _protected_indices( + messages: Sequence[Mapping[str, object]], extra_retrieve_call_ids: frozenset[str] = frozenset() +) -> frozenset[int]: """Indices headroom must not send to the compression service. ``get_protected_indices`` is litellm's own compression policy: the system - rows, the last user row, the last assistant row. It is expanded over whole + rows, the last user row, the last assistant row. Rows carrying just-retrieved + ``headroom_retrieve`` output are added so re-compression can't collapse them + back to the marker they were expanded from. The union is expanded over whole tool exchanges the way ``compress()`` expands it, so a protected assistant tool call cannot end up answered by a marker standing in for the result the model just asked for. """ - protected: Final = frozenset(get_protected_indices(messages)) + protected: Final = frozenset(get_protected_indices(messages)) | _retrieval_result_indices( + messages, extra_retrieve_call_ids + ) return protected | frozenset( index for group in group_tool_exchanges(messages) @@ -634,7 +742,11 @@ class HeadroomGuardrail(CustomGuardrail): # /v1/compress grows a field for sending the live turn as the retrieval # query without compressing it: query-aware compression reads the newest # user message, so it is withheld here at some cost to history ranking. - protected_indices: Final = _protected_indices(messages) + # request_data is a bare dict on the base signature; narrow it before + # reading the untranslated messages so long tool names can be recovered. + raw_messages: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data).get("messages") + raw_retrieve_call_ids: Final = _raw_retrieve_call_ids(raw_messages) + protected_indices: Final = _protected_indices(messages, raw_retrieve_call_ids) compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices] if not compressible: return inputs diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 1fbc975e40a..c04fb7b30ec 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -2199,6 +2199,164 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): assert has_headroom_retrieve_tool(result.get("tools") or []) +# --------------------------------------------------------------------------- +# #38558: a client that runs its own tool loop (e.g. Claude Code via the MCP +# gateway) executes headroom_retrieve and echoes the recovered original content +# back as a tool result. Compressing that row re-derives the same content hash +# it was just retrieved from -- the marker returns and the agent loops. The +# retrieved row must be held back from the compression service. +# --------------------------------------------------------------------------- + +RETRIEVE_ECHO_MESSAGES = [ + {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "Expanding the marker.", + "tool_calls": [ + { + "id": "hr_1", + "type": "function", + "function": { + "name": "mcp__headroom__headroom_retrieve", + "arguments": '{"hash": "b573993006976af767214fac"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "hr_1", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older answer. " + "O" * 5000}, + {"role": "user", "content": "now summarize the description"}, +] + + +@pytest.mark.asyncio +async def test_retrieved_content_is_never_recompressed(guardrail: HeadroomGuardrail): + """The tool result carrying headroom_retrieve output is held back, so it can + never collapse back to the hash it was just retrieved from.""" + wire, result = await _wire_and_result(guardrail, RETRIEVE_ECHO_MESSAGES) + + assert not any(row.get("tool_call_id") == "hr_1" for row in wire) + assert not any("RETRIEVED BODY" in json.dumps(row) for row in wire) + # Reaches the model byte-identical, so no marker stands in for the expansion. + assert result["structured_messages"][3] == RETRIEVE_ECHO_MESSAGES[3] + # Negative control: unrelated history is still compressed, not a no-op. + assert any(row.get("content") == "H" * 5000 for row in wire) + + +@pytest.mark.asyncio +async def test_retrieved_content_guard_matches_direct_tool_name(guardrail: HeadroomGuardrail): + """Server-side the tool is named headroom_retrieve (no MCP prefix); its + result must be protected the same way.""" + messages = [ + {"role": "system", "content": "sys " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "hr_direct", + "type": "function", + "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME, "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "hr_direct", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older. " + "O" * 5000}, + {"role": "user", "content": "summarize"}, + ] + wire, result = await _wire_and_result(guardrail, messages) + + assert not any(row.get("tool_call_id") == "hr_direct" for row in wire) + assert result["structured_messages"][3] == messages[3] + + +@pytest.mark.asyncio +async def test_retrieved_content_protected_when_mcp_tool_name_is_truncated(guardrail: HeadroomGuardrail): + """A long mcp____headroom_retrieve name is truncated past 64 chars in + the OpenAI-translated view the guardrail scans, dropping the suffix. The call + id read from the request's own Anthropic tool_use (never truncated) still + pairs the retrieved row so it is held back.""" + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + truncate_tool_name, + ) + + long_name = "mcp__" + "s" * 45 + "__" + HEADROOM_RETRIEVE_TOOL_NAME + assert len(long_name) > 64 + truncated = truncate_tool_name(long_name) + assert not truncated.endswith(HEADROOM_RETRIEVE_TOOL_NAME) + + # What the guardrail scans: OpenAI-translated messages with the truncated name. + structured = [ + {"role": "system", "content": "sys " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "hr_long", "type": "function", "function": {"name": truncated, "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "hr_long", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older. " + "O" * 5000}, + {"role": "user", "content": "summarize"}, + ] + # The request's own messages, untranslated: Anthropic tool_use carries the full name. + raw_messages = [ + {"role": "assistant", "content": [{"type": "tool_use", "id": "hr_long", "name": long_name, "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "hr_long", "content": "RETRIEVED BODY"}]}, + ] + + inputs = GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(structured))) + sent: dict = {} + + def _echo(**kwargs): + sent["messages"] = kwargs["json"]["messages"] + return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"]))) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-sonnet-4-5-20250929", "messages": raw_messages}, + input_type="request", + ) + + assert not any(row.get("tool_call_id") == "hr_long" for row in sent["messages"]) + assert result["structured_messages"][3] == structured[3] + assert any(row.get("content") == "H" * 5000 for row in sent["messages"]) + + +def test_raw_retrieve_call_ids_covers_both_shapes_and_ignores_others(): + """Retrieve ids are read from OpenAI tool_calls and Anthropic tool_use blocks; + non-retrieve calls, non-tool_use blocks, string content, and non-list inputs + yield nothing.""" + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import _raw_retrieve_call_ids + + messages = [ + { + "role": "assistant", + "tool_calls": [ + {"id": "oa1", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}, + {"id": "other", "function": {"name": "get_weather"}}, + {"id": "malformed", "function": {"name": 123}}, + {"id": "nofunc"}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "an1", "name": "mcp__hr__headroom_retrieve", "input": {}}, + {"type": "tool_use", "id": "an2", "name": "jira_get_issue", "input": {}}, + {"type": "text", "text": "noise"}, + ], + }, + {"role": "user", "content": "plain string content, not a list"}, + ] + + assert _raw_retrieve_call_ids(messages) == frozenset({"oa1", "an1"}) + assert _raw_retrieve_call_ids("not a list") == frozenset() + assert _raw_retrieve_call_ids(None) == frozenset() + + @pytest.mark.asyncio async def test_nothing_compressible_returns_inputs_untouched(guardrail: HeadroomGuardrail): """A single-turn request is all protected, so there is nothing to send and From c83b4a1d1985e79253b7d32f1454cfc96f48e01c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:10:54 -0700 Subject: [PATCH 33/81] fix(proxy-extras): honor a raised command timeout for migrate deploy and name the right knob on db push timeouts --- .../litellm_proxy_extras/prisma_toolchain.py | 13 +++-- .../litellm_proxy_extras/utils.py | 5 +- .../test_prisma_toolchain.py | 51 +++++++++++++++++-- 3 files changed, 60 insertions(+), 9 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index 5feb7a953b4..2283814ab35 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -26,7 +26,10 @@ finish any sooner. Migrate deploy therefore runs under its own budget. All three budgets are overridable so an operator can widen them without a release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install, ``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and -``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. +``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. The +per-command budget used to bound migrate deploy as well, so a deployment that +raised it above the deploy default keeps that larger budget for deploy unless +the deploy override says otherwise. """ import math @@ -102,9 +105,11 @@ def prisma_bootstrap_timeout() -> float: def prisma_migrate_deploy_timeout() -> float: """Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending.""" - return _timeout_from_env( - PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT - ) + if os.getenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR) is not None: + return _timeout_from_env( + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT + ) + return max(DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, prisma_command_timeout()) def nodeenv_cache_dir() -> Optional[Path]: diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index ab9ec1e8a3a..f6647268624 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -15,6 +15,7 @@ from litellm_proxy_extras.replica_identity import ( apply_replica_identity_full, ) from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ensure_prisma_toolchain, prisma_command_timeout, @@ -1135,9 +1136,9 @@ class ProxyExtrasDBManager: return True except subprocess.TimeoutExpired: logger.warning( - "Attempt %s timed out. Raise %s if this database needs longer to apply its pending migrations.", + "Attempt %s timed out. Raise %s if this database needs longer to apply its schema.", attempt + 1, - PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR, ) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 1bd440d37b4..4c258c4d007 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -16,11 +16,13 @@ the proxy gave up after four identical timeouts. import ast import json +import logging import os import sys import time from collections.abc import Callable from pathlib import Path +from typing import Optional import pytest @@ -52,10 +54,10 @@ import time args = sys.argv[1:] cache_dir = os.environ["PRISMA_NODEENV_CACHE_DIR"] log_path = pathlib.Path(os.environ["FAKE_PRISMA_LOG"]) -earlier_deploys = sum( +earlier_same_command = sum( 1 for line in (log_path.read_text().splitlines() if log_path.exists() else []) - if json.loads(line)["args"][:2] == ["migrate", "deploy"] + if json.loads(line)["args"][:2] == args[:2] ) with log_path.open("a") as log: log.write( @@ -64,12 +66,14 @@ with log_path.open("a") as log: ) time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0"))) if args[:2] == ["migrate", "deploy"]: - if earlier_deploys == 0: + if earlier_same_command == 0: time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0"))) elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"): print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr) sys.exit(1) print("No pending migrations to apply") +if args[:2] == ["db", "push"] and earlier_same_command == 0: + time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_PUSH_SLEEP", "0"))) sys.exit(0) """ @@ -269,6 +273,47 @@ def test_migrate_deploy_stops_at_its_own_timeout( assert elapsed < 30 +def test_db_push_timeout_hint_names_the_per_command_budget( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """``db push`` keeps the per-command budget, so its timeout hint has to name that variable.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_PUSH_SLEEP", "3") + + with caplog.at_level(logging.WARNING, logger="litellm_proxy_extras"): + assert ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=False) is True + + assert [call["args"][:2] for call in _fake_prisma_calls(log_path)].count(["db", "push"]) == 2 + assert [record.getMessage() for record in caplog.records if "timed out" in record.getMessage()] == [ + f"Attempt 1 timed out. Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer to apply its schema." + ] + + +@pytest.mark.parametrize( + ("command_timeout", "deploy_timeout", "expected"), + [ + ("900", None, 900.0), + ("12", None, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT), + ("900", "1200", 1200.0), + ("900", "300", 300.0), + ], + ids=["raised_command_budget_carries_over", "lowered_command_budget_does_not", "override_wins_upward", "override_wins_downward"], +) +def test_migrate_deploy_budget_keeps_a_raised_command_budget( + command_timeout: str, deploy_timeout: Optional[str], expected: float, monkeypatch: pytest.MonkeyPatch +) -> None: + """Deployments that raised the per-command budget to survive a long deploy keep that budget for deploy.""" + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, command_timeout) + if deploy_timeout is None: + monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False) + else: + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, deploy_timeout) + + assert prisma_migrate_deploy_timeout() == expected + + @pytest.mark.parametrize( "raw", ["", "0", "-5", "not-a-number", "nan", "inf", "-inf", "1e400"], From 25987cb961567541922cef80e1433137f42c7b77 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 14:42:48 -0700 Subject: [PATCH 34/81] test(build): validate release wheel contracts --- .github/scripts/smoke_test_native_wheel.py | 68 ++++++++ .github/scripts/verify_linux_native_wheel.py | 161 +++++++++++++++++++ .github/workflows/test-rust.yml | 111 +++++++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 8 + 5 files changed, 349 insertions(+) create mode 100644 .github/scripts/smoke_test_native_wheel.py create mode 100644 .github/scripts/verify_linux_native_wheel.py diff --git a/.github/scripts/smoke_test_native_wheel.py b/.github/scripts/smoke_test_native_wheel.py new file mode 100644 index 00000000000..577bb32fcf0 --- /dev/null +++ b/.github/scripts/smoke_test_native_wheel.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import subprocess +import sys +import tempfile +import zipfile +from pathlib import Path +from typing import Final + +CHILD_SCRIPT: Final = """ +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path +import sys + +native_path = Path(sys.argv[1]) +spec = spec_from_file_location("litellm.rust_bridge._native", native_path) +if spec is None or spec.loader is None: + raise RuntimeError("cannot create native extension import specification") +module = module_from_spec(spec) +spec.loader.exec_module(module) + +before = module.gil_stats() +if not isinstance(before.get("releases"), int): + raise AssertionError(f"unexpected gil_stats result: {before!r}") + +try: + module._panic_for_test() +except BaseException as error: + if type(error).__name__ != "PanicException": + raise AssertionError(f"expected PanicException, got {type(error).__name__}") from error +else: + raise AssertionError("Rust panic returned without raising") + +after = module.gil_stats() +if not isinstance(after.get("releases"), int): + raise AssertionError(f"native module unusable after panic: {after!r}") +""" + + +def main() -> int: + if len(sys.argv) != 2: + sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") + return 2 + + wheel: Final = Path(sys.argv[1]) + with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive: + native_members: Final = tuple( + member + for member in archive.infolist() + if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so") + ) + if len(native_members) != 1: + sys.stderr.write(f"expected one native extension, found {len(native_members)}\n") + return 1 + + native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name + native_path.write_bytes(archive.read(native_members[0])) + result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False) + + if result.returncode != 0: + sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n") + return 1 + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py new file mode 100644 index 00000000000..1de98ae0168 --- /dev/null +++ b/.github/scripts/verify_linux_native_wheel.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import importlib.util +import os +import re +import subprocess +import sys +import zipfile +from pathlib import Path, PurePosixPath +from typing import Final + + +def _loads_native_module(native_path: Path) -> bool: + module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path) + if module_spec is None or module_spec.loader is None: + return False + try: + native_module: Final = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(native_module) + except Exception as error: + sys.stderr.write(f"native module load failed: {error}\n") + return False + return True + + +def main() -> int: + if len(sys.argv) != 2: + sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") + return 2 + + wheel: Final = Path(sys.argv[1]) + with zipfile.ZipFile(wheel) as archive: + wheel_members: Final = archive.infolist() + native_members: Final = tuple( + member + for member in wheel_members + if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so") + ) + if len(native_members) != 1: + sys.stderr.write(f"expected one native extension, found {len(native_members)}\n") + return 1 + + unexpected_members: Final = tuple( + member.filename + for member in wheel_members + if member.filename.endswith((".pdb", ".dwp", ".rlib", ".rmeta", "Cargo.toml", "Cargo.lock")) + or any(part.endswith(".dSYM") for part in PurePosixPath(member.filename).parts) + ) + native_member: Final = native_members[0] + uncompressed_wheel_size: Final = sum(member.file_size for member in wheel_members) + native_path: Final = wheel.parent / "native" / Path(native_member.filename).name + native_path.parent.mkdir(parents=True, exist_ok=True) + native_path.write_bytes(archive.read(native_member)) + + wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3) + if len(wheel_tags) != 4: + sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n") + return 1 + + python_tag: Final = wheel_tags[1] + abi_tag: Final = wheel_tags[2] + platform_tag: Final = wheel_tags[3] + commit_sha: Final = os.environ.get("RELEASE_WHEEL_COMMIT_SHA", os.environ.get("GITHUB_SHA", "unknown")) + rustc_version: Final = subprocess.run( + ("rustc", "--version"), + check=True, + capture_output=True, + text=True, + ).stdout.strip() + pyproject: Final = (Path(__file__).parents[2] / "pyproject.toml").read_text() + maturin_match: Final = re.search(r'"maturin==([^";]+)', pyproject) + if maturin_match is None: + sys.stderr.write("build-system does not pin an exact Maturin version\n") + return 1 + + maturin_version: Final = maturin_match.group(1) + native_percentage: Final = native_member.file_size / uncompressed_wheel_size * 100 + size_report: Final = "\n".join( + ( + "## Native wheel build report", + "", + "| Build | Value |", + "| --- | --- |", + f"| Commit | `{commit_sha}` |", + f"| Platform | `{platform_tag}` |", + f"| Python ABI | `{python_tag}-{abi_tag}` |", + f"| Rust compiler | `{rustc_version}` |", + f"| Maturin | `{maturin_version}` |", + "| Cargo profile | `release` |", + "", + "| Artifact | Size |", + "| --- | ---: |", + f"| Compressed wheel | {wheel.stat().st_size / 1_000_000:.2f} MB |", + f"| Uncompressed wheel | {uncompressed_wheel_size / 1_000_000:.2f} MB |", + f"| Native extension | {native_member.file_size / 1_000_000:.2f} MB |", + f"| Native share | {native_percentage:.2f}% |", + "", + ) + ) + summary_path: Final = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_path is None: + sys.stdout.write(size_report) + else: + Path(summary_path).write_text(size_report) + + sections: Final = subprocess.run( + ("readelf", "--sections", "--wide", native_path), + check=True, + capture_output=True, + text=True, + ).stdout + debug_sections: Final = tuple(section for section in (".debug_", ".zdebug_") if section in sections) + debug_sections_absent: Final = not debug_sections + static_symbol_table_absent: Final = ".symtab" not in sections + + dynamic_symbols: Final = subprocess.run( + ("readelf", "--dyn-syms", "--wide", native_path), + check=True, + capture_output=True, + text=True, + ).stdout + extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols + native_module_loads: Final = _loads_native_module(native_path) + native_size_limit: Final = 20_000_000 + native_size_within_limit: Final = native_member.file_size <= native_size_limit + validations: Final = ( + ("Debug sections are absent", debug_sections_absent), + ("Static symbol table is absent", static_symbol_table_absent), + ("Python extension entry point is present", extension_entry_point_present), + ("Native module loads", native_module_loads), + ("Native extension does not exceed 20 MB", native_size_within_limit), + ("Wheel contents are valid", not unexpected_members), + ) + + verified_report: Final = size_report + "\n".join( + ("", "| Validation | Expected | Result |", "| --- | --- | :---: |") + + tuple(f"| {label} | Yes | {'O' if passed else 'X'} |" for label, passed in validations) + + ("",) + ) + report_path: Final = os.environ.get("RELEASE_WHEEL_REPORT") + if report_path is not None: + Path(report_path).write_text(size_report) + if summary_path is not None: + Path(summary_path).write_text(verified_report) + + if debug_sections: + sys.stderr.write(f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}\n") + if not static_symbol_table_absent: + sys.stderr.write(f"{native_member.filename} contains a static symbol table\n") + if not extension_entry_point_present: + sys.stderr.write("native extension does not export PyInit__native\n") + if not native_size_within_limit: + sys.stderr.write(f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB\n") + if unexpected_members: + sys.stderr.write(f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}\n") + + return 0 if all(passed for _, passed in validations) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 21e1bcb90c6..b05ddb81740 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,6 +4,10 @@ on: push: paths: - "litellm-rust/**" + - ".cargo/**" + - "pyproject.toml" + - ".github/scripts/smoke_test_native_wheel.py" + - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -13,6 +17,10 @@ on: - "litellm_**" paths: - "litellm-rust/**" + - ".cargo/**" + - "pyproject.toml" + - ".github/scripts/smoke_test_native_wheel.py" + - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" permissions: @@ -69,3 +77,106 @@ jobs: - name: Run core tests with Bedrock auth run: cargo test -p litellm-core --features bedrock-auth --locked + + release-wheel: + name: release wheel + runs-on: ubuntu-latest + timeout-minutes: 20 + permissions: + contents: read + pull-requests: write + env: + CARGO_TERM_COLOR: always + + 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: Set up Rust + run: | + rustup toolchain install stable --profile minimal + rustup default stable + + - name: Cache release build + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo-release- + + - name: Build release wheel + run: uv build --wheel --out-dir dist + + - name: Verify stripped native extension + env: + RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + RELEASE_WHEEL_REPORT: dist/release-wheel-report.md + run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl + + - name: Build panic contract wheel + run: >- + uv build --wheel --out-dir panic-dist + --config-setting "maturin.build-args=--features panic-test" + + - name: Smoke-test native panic unwinding + run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl + + - name: Report release wheel size on PR + if: >- + github.event_name == 'pull_request' && + github.event.pull_request.head.repo.full_name == github.repository + uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 + env: + COMMENT_MARKER: "" + REPORT_PATH: dist/release-wheel-report.md + with: + script: | + const fs = require("fs"); + const marker = process.env.COMMENT_MARKER; + const report = fs.readFileSync(process.env.REPORT_PATH, "utf8"); + const body = `${marker}\n${report}`; + const comments = await github.paginate(github.rest.issues.listComments, { + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + per_page: 100, + }); + const existing = comments.find((comment) => comment.body?.includes(marker)); + if (existing) { + await github.rest.issues.updateComment({ + owner: context.repo.owner, + repo: context.repo.repo, + comment_id: existing.id, + body, + }); + } else { + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + body, + }); + } + + - name: Upload release wheel + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: litellm-release-wheel-linux-x86_64 + path: dist/*.whl + if-no-files-found: error diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index d461a483ae0..b1fdfd7677a 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -13,6 +13,7 @@ crate-type = ["cdylib"] default = ["abi3"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] +panic-test = [] [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f9e75f45f75..18e0b05cbb5 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -614,6 +614,12 @@ fn gil_stats(py: Python<'_>) -> PyResult> { Ok(stats.into_any().unbind()) } +#[cfg(feature = "panic-test")] +#[pyfunction] +fn _panic_for_test() { + panic!("intentional PyO3 panic smoke test"); +} + #[pymodule] fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { let py = module.py(); @@ -630,5 +636,7 @@ fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(achat_completions, module)?)?; module.add_class::()?; module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + #[cfg(feature = "panic-test")] + module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; Ok(()) } From 2f362cfec2d122d4d2c14b73a988a20fe61e6629 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 14:59:03 -0700 Subject: [PATCH 35/81] fix(ci): preserve release wheel contract parity --- .github/scripts/verify_linux_native_wheel.py | 2 +- .github/workflows/test-rust.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 1de98ae0168..b0ab34df98d 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -139,7 +139,7 @@ def main() -> int: ) report_path: Final = os.environ.get("RELEASE_WHEEL_REPORT") if report_path is not None: - Path(report_path).write_text(size_report) + Path(report_path).write_text(verified_report) if summary_path is not None: Path(summary_path).write_text(verified_report) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index b05ddb81740..1a9f36f24d2 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -132,7 +132,7 @@ jobs: - name: Build panic contract wheel run: >- uv build --wheel --out-dir panic-dist - --config-setting "maturin.build-args=--features panic-test" + --config-setting "maturin.build-args=--features panic-test,extension-module" - name: Smoke-test native panic unwinding run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl From 9dc9cd325cf5adca79d910a30c420381a28d1923 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 20:21:37 -0700 Subject: [PATCH 36/81] fix(ci): isolate release wheel reporting permissions --- .github/scripts/verify_linux_native_wheel.py | 3 - .../workflows/report-rust-release-wheel.yml | 95 +++++++++++++++++++ .github/workflows/test-rust.yml | 39 -------- 3 files changed, 95 insertions(+), 42 deletions(-) create mode 100644 .github/workflows/report-rust-release-wheel.yml diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index b0ab34df98d..96264beb632 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -137,9 +137,6 @@ def main() -> int: + tuple(f"| {label} | Yes | {'O' if passed else 'X'} |" for label, passed in validations) + ("",) ) - report_path: Final = os.environ.get("RELEASE_WHEEL_REPORT") - if report_path is not None: - Path(report_path).write_text(verified_report) if summary_path is not None: Path(summary_path).write_text(verified_report) diff --git a/.github/workflows/report-rust-release-wheel.yml b/.github/workflows/report-rust-release-wheel.yml new file mode 100644 index 00000000000..76aefd70ef6 --- /dev/null +++ b/.github/workflows/report-rust-release-wheel.yml @@ -0,0 +1,95 @@ +name: Report LiteLLM Rust release wheel + +on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs + workflow_run: + workflows: + - LiteLLM Rust + types: + - completed + +permissions: {} + +jobs: + report-release-wheel: + name: report release wheel + if: >- + github.event.workflow_run.conclusion == 'success' && + github.event.workflow_run.event == 'pull_request' && + github.event.workflow_run.path == '.github/workflows/test-rust.yml' && + github.event.workflow_run.head_repository.full_name == github.repository && + github.event.workflow_run.pull_requests[0].number != null + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + issues: write + + steps: + - name: Link release wheel report on PR + uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 + env: + COMMENT_MARKER: "" + with: + script: | + const marker = process.env.COMMENT_MARKER; + const workflowRun = context.payload.workflow_run; + if ( + workflowRun.conclusion !== "success" || + workflowRun.event !== "pull_request" || + workflowRun.path !== ".github/workflows/test-rust.yml" || + workflowRun.head_repository?.full_name !== + `${context.repo.owner}/${context.repo.repo}` || + workflowRun.pull_requests?.length !== 1 + ) { + throw new Error("unexpected source workflow"); + } + const pullRequest = workflowRun.pull_requests[0]; + const pullRequestNumber = pullRequest.number; + const headSha = workflowRun.head_sha; + const runId = workflowRun.id; + if ( + !Number.isSafeInteger(pullRequestNumber) || + pullRequestNumber <= 0 || + !Number.isSafeInteger(runId) || + runId <= 0 || + !/^[0-9a-f]{40}$/.test(headSha) || + pullRequest.head?.sha !== headSha + ) { + throw new Error("invalid source workflow metadata"); + } + const runUrl = + `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` + + `/actions/runs/${runId}`; + const body = [ + marker, + "## LiteLLM Rust workflow", + "", + `Workflow completed successfully for \`${headSha}\``, + "", + `[View workflow run](${runUrl})`, + ].join("\n"); + const comments = await github.paginate(github.rest.issues.listComments, { + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pullRequestNumber, + per_page: 100, + }); + const existing = comments.find( + (comment) => + comment.user?.login === "github-actions[bot]" && + comment.body?.startsWith(marker), + ); + if (existing) { + await github.rest.issues.updateComment({ + owner: context.repo.owner, + repo: context.repo.repo, + comment_id: existing.id, + body, + }); + } else { + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pullRequestNumber, + body, + }); + } diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 1a9f36f24d2..b01e7fabe4a 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -84,7 +84,6 @@ jobs: timeout-minutes: 20 permissions: contents: read - pull-requests: write env: CARGO_TERM_COLOR: always @@ -126,7 +125,6 @@ jobs: - name: Verify stripped native extension env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - RELEASE_WHEEL_REPORT: dist/release-wheel-report.md run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl - name: Build panic contract wheel @@ -137,43 +135,6 @@ jobs: - name: Smoke-test native panic unwinding run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - - name: Report release wheel size on PR - if: >- - github.event_name == 'pull_request' && - github.event.pull_request.head.repo.full_name == github.repository - uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 - env: - COMMENT_MARKER: "" - REPORT_PATH: dist/release-wheel-report.md - with: - script: | - const fs = require("fs"); - const marker = process.env.COMMENT_MARKER; - const report = fs.readFileSync(process.env.REPORT_PATH, "utf8"); - const body = `${marker}\n${report}`; - const comments = await github.paginate(github.rest.issues.listComments, { - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: context.issue.number, - per_page: 100, - }); - const existing = comments.find((comment) => comment.body?.includes(marker)); - if (existing) { - await github.rest.issues.updateComment({ - owner: context.repo.owner, - repo: context.repo.repo, - comment_id: existing.id, - body, - }); - } else { - await github.rest.issues.createComment({ - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: context.issue.number, - body, - }); - } - - name: Upload release wheel uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 with: From e6a317e0790c575d6452f59c90e861f110b54a0c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 20:53:17 -0700 Subject: [PATCH 37/81] fix(ci): enforce release wheel metadata contract --- .github/scripts/verify_linux_native_wheel.py | 110 ++++++++++-- .../test_verify_linux_native_wheel.py | 170 ++++++++++++++++++ 2 files changed, 266 insertions(+), 14 deletions(-) create mode 100644 tests/test_litellm/test_verify_linux_native_wheel.py diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 96264beb632..61f176ab183 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -6,21 +6,44 @@ import re import subprocess import sys import zipfile +from email import policy +from email.parser import BytesParser +from itertools import product from pathlib import Path, PurePosixPath -from typing import Final +from types import ModuleType +from typing import Final, cast + +EXPECTED_PYTHON_TAG: Final = "cp310" +EXPECTED_ABI_TAG: Final = "abi3" +EXPECTED_PLATFORM_TAG: Final = "linux_x86_64" -def _loads_native_module(native_path: Path) -> bool: +def _dist_info_directory(member: zipfile.ZipInfo) -> str | None: + parts: Final = PurePosixPath(member.filename).parts + if not parts or not parts[0].endswith(".dist-info"): + return None + return parts[0] + + +def _wheel_metadata_tags(archive: zipfile.ZipFile, members: tuple[zipfile.ZipInfo, ...]) -> tuple[str, ...]: + if len(members) != 1: + return () + metadata: Final = BytesParser(policy=policy.default).parsebytes(archive.read(members[0])) + tags: Final = cast(list[str], metadata.get_all("Tag", [])) + return tuple(tag.strip() for tag in tags) + + +def _load_native_module(native_path: Path) -> ModuleType | None: module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path) if module_spec is None or module_spec.loader is None: - return False + return None try: native_module: Final = importlib.util.module_from_spec(module_spec) module_spec.loader.exec_module(native_module) except Exception as error: sys.stderr.write(f"native module load failed: {error}\n") - return False - return True + return None + return native_module def main() -> int: @@ -29,8 +52,39 @@ def main() -> int: return 2 wheel: Final = Path(sys.argv[1]) + wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3) + if len(wheel_tags) != 4: + sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n") + return 1 + + wheel_identity: Final = wheel_tags[0].split("-") + if len(wheel_identity) != 2 or wheel_identity[0] != "litellm" or not wheel_identity[1]: + sys.stderr.write(f"unexpected wheel identity: {wheel_tags[0]}\n") + return 1 + + expected_dist_info_directory: Final = f"{wheel_tags[0]}.dist-info" + expected_dist_info_directories: Final = frozenset((expected_dist_info_directory,)) + python_tag: Final = wheel_tags[1] + abi_tag: Final = wheel_tags[2] + platform_tag: Final = wheel_tags[3] + expanded_filename_tags: Final = frozenset( + "-".join(tag) for tag in product(python_tag.split("."), abi_tag.split("."), platform_tag.split(".")) + ) + with zipfile.ZipFile(wheel) as archive: wheel_members: Final = archive.infolist() + dist_info_directories: Final = frozenset( + directory for member in wheel_members if (directory := _dist_info_directory(member)) is not None + ) + required_dist_info_files: Final = ("METADATA", "RECORD", "WHEEL") + dist_info_file_counts: Final = { + filename: sum(member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members) + for filename in required_dist_info_files + } + wheel_metadata_members: Final = tuple( + member for member in wheel_members if member.filename == f"{expected_dist_info_directory}/WHEEL" + ) + wheel_metadata_tags: Final = _wheel_metadata_tags(archive, wheel_metadata_members) native_members: Final = tuple( member for member in wheel_members @@ -52,14 +106,10 @@ def main() -> int: native_path.parent.mkdir(parents=True, exist_ok=True) native_path.write_bytes(archive.read(native_member)) - wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3) - if len(wheel_tags) != 4: - sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n") - return 1 - - python_tag: Final = wheel_tags[1] - abi_tag: Final = wheel_tags[2] - platform_tag: Final = wheel_tags[3] + wheel_metadata_tags_match: Final = ( + len(wheel_metadata_tags) == len(expanded_filename_tags) + and frozenset(wheel_metadata_tags) == expanded_filename_tags + ) commit_sha: Final = os.environ.get("RELEASE_WHEEL_COMMIT_SHA", os.environ.get("GITHUB_SHA", "unknown")) rustc_version: Final = subprocess.run( ("rustc", "--version"), @@ -120,14 +170,26 @@ def main() -> int: text=True, ).stdout extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols - native_module_loads: Final = _loads_native_module(native_path) + native_module: Final = _load_native_module(native_path) + native_module_loads: Final = native_module is not None + panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") native_size_limit: Final = 20_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( + (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), + (f"ABI tag is {EXPECTED_ABI_TAG}", abi_tag == EXPECTED_ABI_TAG), + (f"Platform tag is {EXPECTED_PLATFORM_TAG}", platform_tag == EXPECTED_PLATFORM_TAG), + ("Wheel dist-info directory matches the filename", dist_info_directories == expected_dist_info_directories), + ( + "Required dist-info files are present exactly once", + all(count == 1 for count in dist_info_file_counts.values()), + ), + ("Wheel metadata tags match the filename", wheel_metadata_tags_match), ("Debug sections are absent", debug_sections_absent), ("Static symbol table is absent", static_symbol_table_absent), ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), + ("Production module omits the panic test hook", panic_test_hook_absent), ("Native extension does not exceed 20 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -146,6 +208,26 @@ def main() -> int: sys.stderr.write(f"{native_member.filename} contains a static symbol table\n") if not extension_entry_point_present: sys.stderr.write("native extension does not export PyInit__native\n") + if python_tag != EXPECTED_PYTHON_TAG: + sys.stderr.write(f"unexpected Python tag: expected {EXPECTED_PYTHON_TAG}, found {python_tag}\n") + if abi_tag != EXPECTED_ABI_TAG: + sys.stderr.write(f"unexpected ABI tag: expected {EXPECTED_ABI_TAG}, found {abi_tag}\n") + if platform_tag != EXPECTED_PLATFORM_TAG: + sys.stderr.write(f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}\n") + if dist_info_directories != expected_dist_info_directories: + sys.stderr.write( + f"unexpected dist-info directories: expected {[expected_dist_info_directory]}, " + f"found {sorted(dist_info_directories)}\n" + ) + if any(count != 1 for count in dist_info_file_counts.values()): + sys.stderr.write(f"required dist-info file counts are invalid: {dist_info_file_counts}\n") + elif not wheel_metadata_tags_match: + sys.stderr.write( + f"WHEEL tags do not match filename: expected {sorted(expanded_filename_tags)}, " + f"found {sorted(wheel_metadata_tags)}\n" + ) + if native_module is not None and not panic_test_hook_absent: + sys.stderr.write("production native module exposes _panic_for_test\n") if not native_size_within_limit: sys.stderr.write(f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB\n") if unexpected_members: diff --git a/tests/test_litellm/test_verify_linux_native_wheel.py b/tests/test_litellm/test_verify_linux_native_wheel.py new file mode 100644 index 00000000000..86f5debfe3b --- /dev/null +++ b/tests/test_litellm/test_verify_linux_native_wheel.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import importlib.util +import subprocess +import sys +import zipfile +from collections.abc import Callable +from pathlib import Path +from types import ModuleType +from typing import Final, Protocol, cast + +import pytest + +_REPO_ROOT: Final = Path(__file__).resolve().parents[2] +_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "verify_linux_native_wheel.py" + + +class _VerifierModule(Protocol): + subprocess: ModuleType + _load_native_module: Callable[[Path], ModuleType | None] + main: Callable[[], int] + + +_SPEC: Final = importlib.util.spec_from_file_location("verify_linux_native_wheel", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_LOADED_VERIFIER: Final = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _LOADED_VERIFIER +_SPEC.loader.exec_module(_LOADED_VERIFIER) +verifier: Final = cast(_VerifierModule, _LOADED_VERIFIER) + +_EXPECTED_TAG: Final = "cp310-abi3-linux_x86_64" +_NATIVE_MEMBER: Final = "litellm/rust_bridge/_native.abi3.so" +_DIST_INFO: Final = "litellm-1.100.0.dist-info" + + +def _write_wheel( + tmp_path: Path, + *, + filename_tag: str, + metadata_tags: tuple[str, ...] | None = (_EXPECTED_TAG,), + dist_info: str = _DIST_INFO, + duplicate_wheel: bool = False, +) -> Path: + wheel: Final = tmp_path / f"litellm-1.100.0-{filename_tag}.whl" + with zipfile.ZipFile(wheel, "w", compression=zipfile.ZIP_DEFLATED) as archive: + archive.writestr(_NATIVE_MEMBER, b"synthetic native extension") + archive.writestr( + f"{dist_info}/METADATA", + "Metadata-Version: 2.1\nName: litellm\nVersion: 1.100.0\n", + ) + archive.writestr( + f"{dist_info}/RECORD", + f"{_NATIVE_MEMBER},,\n{dist_info}/WHEEL,,\n", + ) + if metadata_tags is not None: + wheel_metadata: Final = ( + "Wheel-Version: 1.0\nGenerator: regression-test\nRoot-Is-Purelib: false\n" + + "".join(f"Tag: {tag}\n" for tag in metadata_tags) + ) + archive.writestr(f"{dist_info}/WHEEL", wheel_metadata) + if duplicate_wheel: + archive.writestr(f"{dist_info}/WHEEL", wheel_metadata) + return wheel + + +def _fake_subprocess_run(command: tuple[str, ...], **_: object) -> subprocess.CompletedProcess[str]: + if command == ("rustc", "--version"): + stdout = "rustc 1.98.0 (regression-test)\n" + elif "--sections" in command: + stdout = "[ 1] .text PROGBITS\n" + elif "--dyn-syms" in command: + stdout = "PyInit__native\n" + else: + raise AssertionError(f"unexpected subprocess command: {command}") + return subprocess.CompletedProcess(command, 0, stdout=stdout, stderr="") + + +def _run_verifier( + monkeypatch: pytest.MonkeyPatch, + wheel: Path, + *, + exposes_panic: bool = False, +) -> int: + native_module: Final = ModuleType("litellm.rust_bridge._native") + if exposes_panic: + setattr(native_module, "_panic_for_test", lambda: None) + + def _fake_load_native_module(_: Path) -> ModuleType: + return native_module + + monkeypatch.setattr(verifier, "_load_native_module", _fake_load_native_module) + monkeypatch.setattr(verifier.subprocess, "run", _fake_subprocess_run) + monkeypatch.setattr(sys, "argv", [str(_MODULE_PATH), str(wheel)]) + monkeypatch.setenv("GITHUB_STEP_SUMMARY", str(wheel.parent / "summary.md")) + return verifier.main() + + +def test_accepts_expected_release_wheel_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) + + assert _run_verifier(monkeypatch, wheel) == 0 + + +def test_rejects_cp312_version_specific_wheel(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + tag: Final = "cp312-cp312-linux_x86_64" + wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,)) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_non_linux_platform_tag(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + tag: Final = "cp310-abi3-win_amd64" + wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,)) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +@pytest.mark.parametrize( + "metadata_tags", + [None, ("cp312-cp312-linux_x86_64",)], + ids=["missing", "mismatched"], +) +def test_rejects_missing_or_mismatched_wheel_metadata_tag( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + metadata_tags: tuple[str, ...] | None, +) -> None: + wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG, metadata_tags=metadata_tags) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_wheel_metadata_from_wrong_dist_info_directory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + wheel: Final = _write_wheel( + tmp_path, + filename_tag=_EXPECTED_TAG, + dist_info="decoy-1.0.0.dist-info", + ) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_duplicate_wheel_metadata_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + wheel: Final = _write_wheel( + tmp_path, + filename_tag=_EXPECTED_TAG, + metadata_tags=(_EXPECTED_TAG, _EXPECTED_TAG), + ) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + with pytest.warns(UserWarning, match="Duplicate name"): + wheel: Final = _write_wheel( + tmp_path, + filename_tag=_EXPECTED_TAG, + duplicate_wheel=True, + ) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_production_module_exposing_panic_hook(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) + + assert _run_verifier(monkeypatch, wheel, exposes_panic=True) == 1 From ef9a207ed545e740863926647af148e5195285c9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 05:44:35 -0700 Subject: [PATCH 38/81] fix(ci): harden release wheel reporting --- .../workflows/report-rust-release-wheel.yml | 43 +++++++++++++++++-- .github/workflows/test-rust.yml | 26 ++--------- 2 files changed, 43 insertions(+), 26 deletions(-) diff --git a/.github/workflows/report-rust-release-wheel.yml b/.github/workflows/report-rust-release-wheel.yml index 76aefd70ef6..1d93b56f77f 100644 --- a/.github/workflows/report-rust-release-wheel.yml +++ b/.github/workflows/report-rust-release-wheel.yml @@ -9,11 +9,14 @@ on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consum permissions: {} +concurrency: + group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }} + cancel-in-progress: false + jobs: report-release-wheel: name: report release wheel if: >- - github.event.workflow_run.conclusion == 'success' && github.event.workflow_run.event == 'pull_request' && github.event.workflow_run.path == '.github/workflows/test-rust.yml' && github.event.workflow_run.head_repository.full_name == github.repository && @@ -21,7 +24,8 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 permissions: - issues: write + issues: write # PR comments use the issues API + pull-requests: read # Current-head validation rejects stale workflow runs steps: - name: Link release wheel report on PR @@ -32,8 +36,19 @@ jobs: script: | const marker = process.env.COMMENT_MARKER; const workflowRun = context.payload.workflow_run; + const allowedConclusions = new Set([ + "action_required", + "cancelled", + "failure", + "neutral", + "skipped", + "stale", + "startup_failure", + "success", + "timed_out", + ]); if ( - workflowRun.conclusion !== "success" || + !allowedConclusions.has(workflowRun.conclusion) || workflowRun.event !== "pull_request" || workflowRun.path !== ".github/workflows/test-rust.yml" || workflowRun.head_repository?.full_name !== @@ -59,11 +74,15 @@ jobs: const runUrl = `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` + `/actions/runs/${runId}`; + const result = + workflowRun.conclusion === "success" + ? "successfully" + : `with \`${workflowRun.conclusion}\``; const body = [ marker, "## LiteLLM Rust workflow", "", - `Workflow completed successfully for \`${headSha}\``, + `Workflow completed ${result} for \`${headSha}\``, "", `[View workflow run](${runUrl})`, ].join("\n"); @@ -78,6 +97,22 @@ jobs: comment.user?.login === "github-actions[bot]" && comment.body?.startsWith(marker), ); + const currentPullRequest = ( + await github.rest.pulls.get({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: pullRequestNumber, + }) + ).data; + if ( + currentPullRequest.state !== "open" || + currentPullRequest.head.repo?.full_name !== + `${context.repo.owner}/${context.repo.repo}` || + currentPullRequest.head.sha !== headSha + ) { + core.info("source workflow no longer matches the current pull request head"); + return; + } if (existing) { await github.rest.issues.updateComment({ owner: context.repo.owner, diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index b01e7fabe4a..77da358a2fc 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -108,25 +108,9 @@ jobs: rustup toolchain install stable --profile minimal rustup default stable - - name: Cache release build - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 - with: - path: | - ~/.cargo/registry - ~/.cargo/git - litellm-rust/target - key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }} - restore-keys: | - ${{ runner.os }}-cargo-release- - - name: Build release wheel run: uv build --wheel --out-dir dist - - name: Verify stripped native extension - env: - RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl - - name: Build panic contract wheel run: >- uv build --wheel --out-dir panic-dist @@ -135,9 +119,7 @@ jobs: - name: Smoke-test native panic unwinding run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - - name: Upload release wheel - uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 - with: - name: litellm-release-wheel-linux-x86_64 - path: dist/*.whl - if-no-files-found: error + - name: Verify stripped native extension + env: + RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl From 38150dfc2c50526664125c25c6eabf07d06e07ed Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 06:37:39 -0700 Subject: [PATCH 39/81] fix(ci): pin workflow toolchain dependencies --- .circleci/config.yml | 8 ++--- .../actions/setup-uv-with-retries/action.yml | 14 +++----- .github/workflows/test-rust.yml | 12 +++---- rust-toolchain.toml | 4 +++ .../test_circleci_rust_toolchain.py | 33 ++++++++++++++----- 5 files changed, 42 insertions(+), 29 deletions(-) create mode 100644 rust-toolchain.toml diff --git a/.circleci/config.yml b/.circleci/config.yml index 55fa9410845..dfc539fb80e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -112,10 +112,10 @@ commands: node --version npm --version install_rust: - description: "Install pinned rustup (1.28.2) and Rust toolchain (1.97.1) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself." + description: "Install pinned rustup (1.28.2) and Rust toolchain (1.98.0) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself." steps: - run: - name: Install Rust (rustup 1.28.2, toolchain 1.97.1) + name: Install Rust (rustup 1.28.2, toolchain 1.98.0) command: | case "$(uname -m)" in x86_64) @@ -135,7 +135,7 @@ commands: "https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init" echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c - chmod +x /tmp/rustup-init - /tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.97.1 + /tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.98.0 rm -f /tmp/rustup-init echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV" export PATH="$HOME/.cargo/bin:$PATH" @@ -300,7 +300,7 @@ jobs: if ($rustupActual -ne $rustupExpected) { throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual" } - & $rustupInit -y --profile minimal --default-toolchain stable + & $rustupInit -y --profile minimal --default-toolchain 1.98.0 if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } diff --git a/.github/actions/setup-uv-with-retries/action.yml b/.github/actions/setup-uv-with-retries/action.yml index 1627038dc3d..98ff91f0283 100644 --- a/.github/actions/setup-uv-with-retries/action.yml +++ b/.github/actions/setup-uv-with-retries/action.yml @@ -1,11 +1,7 @@ name: "Set up uv with retries" description: >- - Install uv via astral-sh/setup-uv, retrying on transient failures. Even with - an exact pinned version, the action resolves the artifact URL by fetching - https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a - single request with no retry, timeout, or fallback, so one connection-level - network error ("fetch failed") fails the whole job before any test runs. - Retrying the full step covers the manifest fetch and the binary download. + Install uv via astral-sh/setup-uv, retrying the full setup step so manifest + resolution and binary downloads get fresh attempts after transient failures. inputs: version: @@ -18,7 +14,7 @@ runs: - name: Set up uv (attempt 1) id: attempt-1 continue-on-error: true - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} @@ -31,7 +27,7 @@ runs: id: attempt-2 if: steps.attempt-1.outcome == 'failure' continue-on-error: true - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} @@ -42,6 +38,6 @@ runs: - name: Set up uv (attempt 3) if: steps.attempt-2.outcome == 'failure' - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 77da358a2fc..aada0fcf239 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -6,6 +6,7 @@ on: - "litellm-rust/**" - ".cargo/**" - "pyproject.toml" + - "rust-toolchain.toml" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" @@ -19,6 +20,7 @@ on: - "litellm-rust/**" - ".cargo/**" - "pyproject.toml" + - "rust-toolchain.toml" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" @@ -48,9 +50,7 @@ jobs: persist-credentials: false - name: Set up Rust - run: | - rustup toolchain install stable --profile minimal --component clippy,rustfmt - rustup default stable + run: rustup toolchain install - name: Cache Cargo registry and target uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 @@ -59,7 +59,7 @@ jobs: ~/.cargo/registry ~/.cargo/git litellm-rust/target - key: ${{ runner.os }}-cargo-${{ hashFiles('litellm-rust/Cargo.lock') }} + key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }} restore-keys: | ${{ runner.os }}-cargo- @@ -104,9 +104,7 @@ jobs: version: "0.10.9" - name: Set up Rust - run: | - rustup toolchain install stable --profile minimal - rustup default stable + run: rustup toolchain install - name: Build release wheel run: uv build --wheel --out-dir dist diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 00000000000..a1598ccbb34 --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,4 @@ +[toolchain] +channel = "1.98.0" +profile = "minimal" +components = ["rustfmt", "clippy"] diff --git a/tests/test_litellm/test_circleci_rust_toolchain.py b/tests/test_litellm/test_circleci_rust_toolchain.py index c35ced51e16..800ca21b95d 100644 --- a/tests/test_litellm/test_circleci_rust_toolchain.py +++ b/tests/test_litellm/test_circleci_rust_toolchain.py @@ -17,28 +17,27 @@ Two invariants are pinned here: Windows job, so the check accepts either. A new job that syncs without one falls back to the unpinned path, which is exactly the regression a static check catches at PR time and a green CI run does not. - 2. `install_rust` itself pins what it downloads: an explicit rustup version in - the URL, a verified SHA-256, and an exact toolchain version rather than a - channel name. - -The Windows job predates `install_rust` and provisions its toolchain inline, so -invariant 2 is scoped to `install_rust`; invariant 1 covers both. + 2. Both installers pin what they download: an explicit rustup version, a + verified SHA-256, and the exact toolchain in `rust-toolchain.toml`. """ from __future__ import annotations import re from pathlib import Path +from typing import Final import pytest import yaml REPO_ROOT = Path(__file__).resolve().parents[2] CONFIG = REPO_ROOT / ".circleci" / "config.yml" +TOOLCHAIN: Final = REPO_ROOT / "rust-toolchain.toml" BUILDS_WORKSPACE = re.compile(r"\buv\s+(?:sync|build)\b") RUSTUP_ARCHIVE_URL = re.compile(r"https://static\.rust-lang\.org/rustup/archive/\d+\.\d+\.\d+/") -EXACT_TOOLCHAIN = re.compile(r"--default-toolchain\s+\"?\d+\.\d+\.\d+\"?") +EXACT_TOOLCHAIN = re.compile(r"--default-toolchain\s+\"?(\d+\.\d+\.\d+)\"?") +TOOLCHAIN_CHANNEL: Final = re.compile(r'^channel = "(\d+\.\d+\.\d+)"$', re.MULTILINE) def _config() -> dict[str, object]: @@ -57,6 +56,12 @@ def _step_text(step: object) -> str: return "" +def _pinned_toolchain() -> str: + match: Final = TOOLCHAIN_CHANNEL.search(TOOLCHAIN.read_text()) + assert match is not None, "rust-toolchain.toml must pin an exact channel" + return match.group(1) + + def _without_comments(text: str) -> str: return "\n".join(line for line in text.splitlines() if not line.lstrip().startswith("#")) @@ -142,7 +147,17 @@ def test_install_rust_verifies_the_installer_checksum(install_rust_command: str) def test_install_rust_pins_an_exact_toolchain_version(install_rust_command: str) -> None: - assert EXACT_TOOLCHAIN.search(install_rust_command), ( - "install_rust must pin an exact toolchain version (e.g. 1.97.1); a channel name like " + match: Final = EXACT_TOOLCHAIN.search(install_rust_command) + assert match is not None, ( + "install_rust must pin an exact toolchain version (e.g. 1.98.0); a channel name like " "stable/beta/nightly makes the compiler drift with whatever upstream published that day" ) + assert match.group(1) == _pinned_toolchain() + + +def test_windows_installer_matches_the_repo_toolchain() -> None: + windows_steps: Final = _step_lists()["job using_litellm_on_windows"] + windows_command: Final = "\n".join(_step_text(step) for step in windows_steps) + match: Final = EXACT_TOOLCHAIN.search(windows_command) + assert match is not None + assert match.group(1) == _pinned_toolchain() From cbb8a1784db64b8433c07d09af2a114993b717ff Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 06:49:43 -0700 Subject: [PATCH 40/81] chore(ci): extract setup-uv pin --- .github/actions/setup-uv-with-retries/action.yml | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/.github/actions/setup-uv-with-retries/action.yml b/.github/actions/setup-uv-with-retries/action.yml index 98ff91f0283..1627038dc3d 100644 --- a/.github/actions/setup-uv-with-retries/action.yml +++ b/.github/actions/setup-uv-with-retries/action.yml @@ -1,7 +1,11 @@ name: "Set up uv with retries" description: >- - Install uv via astral-sh/setup-uv, retrying the full setup step so manifest - resolution and binary downloads get fresh attempts after transient failures. + Install uv via astral-sh/setup-uv, retrying on transient failures. Even with + an exact pinned version, the action resolves the artifact URL by fetching + https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a + single request with no retry, timeout, or fallback, so one connection-level + network error ("fetch failed") fails the whole job before any test runs. + Retrying the full step covers the manifest fetch and the binary download. inputs: version: @@ -14,7 +18,7 @@ runs: - name: Set up uv (attempt 1) id: attempt-1 continue-on-error: true - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: version: ${{ inputs.version }} @@ -27,7 +31,7 @@ runs: id: attempt-2 if: steps.attempt-1.outcome == 'failure' continue-on-error: true - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: version: ${{ inputs.version }} @@ -38,6 +42,6 @@ runs: - name: Set up uv (attempt 3) if: steps.attempt-2.outcome == 'failure' - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: version: ${{ inputs.version }} From ce0c85ea691922d1e2f06bec93df10b7dc7f43ad Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 12:06:00 -0700 Subject: [PATCH 41/81] refactor(rust): colocate native wheel contract checks --- .github/workflows/test-rust.yml | 12 ++-- .../rust_bridge}/smoke_test_native_wheel.py | 0 .../rust_bridge}/verify_linux_native_wheel.py | 57 ++++++++++++---- .../test_verify_linux_native_wheel.py | 65 +++++++------------ 4 files changed, 74 insertions(+), 60 deletions(-) rename {.github/scripts => litellm/rust_bridge}/smoke_test_native_wheel.py (100%) rename {.github/scripts => litellm/rust_bridge}/verify_linux_native_wheel.py (86%) rename tests/test_litellm/{ => rust_bridge}/test_verify_linux_native_wheel.py (62%) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index aada0fcf239..271c73733a3 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -7,8 +7,8 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" - - ".github/scripts/smoke_test_native_wheel.py" - - ".github/scripts/verify_linux_native_wheel.py" + - "litellm/rust_bridge/smoke_test_native_wheel.py" + - "litellm/rust_bridge/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -21,8 +21,8 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" - - ".github/scripts/smoke_test_native_wheel.py" - - ".github/scripts/verify_linux_native_wheel.py" + - "litellm/rust_bridge/smoke_test_native_wheel.py" + - "litellm/rust_bridge/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" permissions: @@ -115,9 +115,9 @@ jobs: --config-setting "maturin.build-args=--features panic-test,extension-module" - name: Smoke-test native panic unwinding - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl + run: python litellm/rust_bridge/smoke_test_native_wheel.py panic-dist/*.whl - name: Verify stripped native extension env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl + run: python litellm/rust_bridge/verify_linux_native_wheel.py dist/*.whl diff --git a/.github/scripts/smoke_test_native_wheel.py b/litellm/rust_bridge/smoke_test_native_wheel.py similarity index 100% rename from .github/scripts/smoke_test_native_wheel.py rename to litellm/rust_bridge/smoke_test_native_wheel.py diff --git a/.github/scripts/verify_linux_native_wheel.py b/litellm/rust_bridge/verify_linux_native_wheel.py similarity index 86% rename from .github/scripts/verify_linux_native_wheel.py rename to litellm/rust_bridge/verify_linux_native_wheel.py index 61f176ab183..783de87a8b2 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/litellm/rust_bridge/verify_linux_native_wheel.py @@ -6,18 +6,40 @@ import re import subprocess import sys import zipfile +from collections.abc import Callable, Mapping, Sequence from email import policy from email.parser import BytesParser from itertools import product from pathlib import Path, PurePosixPath from types import ModuleType -from typing import Final, cast +from typing import Final, Protocol, cast EXPECTED_PYTHON_TAG: Final = "cp310" EXPECTED_ABI_TAG: Final = "abi3" EXPECTED_PLATFORM_TAG: Final = "linux_x86_64" +class CommandRunner(Protocol): + def __call__( + self, + command: tuple[str, ...], + *, + check: bool, + capture_output: bool, + text: bool, + ) -> subprocess.CompletedProcess[str]: ... + + +def _run_command( + command: tuple[str, ...], + *, + check: bool, + capture_output: bool, + text: bool, +) -> subprocess.CompletedProcess[str]: + return subprocess.run(command, check=check, capture_output=capture_output, text=text) + + def _dist_info_directory(member: zipfile.ZipInfo) -> str | None: parts: Final = PurePosixPath(member.filename).parts if not parts or not parts[0].endswith(".dist-info"): @@ -46,12 +68,19 @@ def _load_native_module(native_path: Path) -> ModuleType | None: return native_module -def main() -> int: - if len(sys.argv) != 2: - sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") +def main( + argv: Sequence[str] | None = None, + environment: Mapping[str, str] | None = None, + load_native_module: Callable[[Path], ModuleType | None] = _load_native_module, + run_command: CommandRunner = _run_command, +) -> int: + arguments: Final = tuple(sys.argv if argv is None else argv) + resolved_environment: Final = os.environ if environment is None else environment + if len(arguments) != 2: + sys.stderr.write(f"usage: {Path(arguments[0]).name} WHEEL\n") return 2 - wheel: Final = Path(sys.argv[1]) + wheel: Final = Path(arguments[1]) wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3) if len(wheel_tags) != 4: sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n") @@ -110,8 +139,10 @@ def main() -> int: len(wheel_metadata_tags) == len(expanded_filename_tags) and frozenset(wheel_metadata_tags) == expanded_filename_tags ) - commit_sha: Final = os.environ.get("RELEASE_WHEEL_COMMIT_SHA", os.environ.get("GITHUB_SHA", "unknown")) - rustc_version: Final = subprocess.run( + commit_sha: Final = resolved_environment.get( + "RELEASE_WHEEL_COMMIT_SHA", resolved_environment.get("GITHUB_SHA", "unknown") + ) + rustc_version: Final = run_command( ("rustc", "--version"), check=True, capture_output=True, @@ -147,14 +178,14 @@ def main() -> int: "", ) ) - summary_path: Final = os.environ.get("GITHUB_STEP_SUMMARY") + summary_path: Final = resolved_environment.get("GITHUB_STEP_SUMMARY") if summary_path is None: sys.stdout.write(size_report) else: Path(summary_path).write_text(size_report) - sections: Final = subprocess.run( - ("readelf", "--sections", "--wide", native_path), + sections: Final = run_command( + ("readelf", "--sections", "--wide", str(native_path)), check=True, capture_output=True, text=True, @@ -163,14 +194,14 @@ def main() -> int: debug_sections_absent: Final = not debug_sections static_symbol_table_absent: Final = ".symtab" not in sections - dynamic_symbols: Final = subprocess.run( - ("readelf", "--dyn-syms", "--wide", native_path), + dynamic_symbols: Final = run_command( + ("readelf", "--dyn-syms", "--wide", str(native_path)), check=True, capture_output=True, text=True, ).stdout extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols - native_module: Final = _load_native_module(native_path) + native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") native_size_limit: Final = 20_000_000 diff --git a/tests/test_litellm/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py similarity index 62% rename from tests/test_litellm/test_verify_linux_native_wheel.py rename to tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index 86f5debfe3b..a4291ce0a65 100644 --- a/tests/test_litellm/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -1,32 +1,16 @@ from __future__ import annotations -import importlib.util import subprocess -import sys import zipfile -from collections.abc import Callable from pathlib import Path from types import ModuleType -from typing import Final, Protocol, cast +from typing import Final import pytest -_REPO_ROOT: Final = Path(__file__).resolve().parents[2] -_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "verify_linux_native_wheel.py" +from litellm.rust_bridge import verify_linux_native_wheel as verifier - -class _VerifierModule(Protocol): - subprocess: ModuleType - _load_native_module: Callable[[Path], ModuleType | None] - main: Callable[[], int] - - -_SPEC: Final = importlib.util.spec_from_file_location("verify_linux_native_wheel", _MODULE_PATH) -assert _SPEC is not None and _SPEC.loader is not None -_LOADED_VERIFIER: Final = importlib.util.module_from_spec(_SPEC) -sys.modules[_SPEC.name] = _LOADED_VERIFIER -_SPEC.loader.exec_module(_LOADED_VERIFIER) -verifier: Final = cast(_VerifierModule, _LOADED_VERIFIER) +_MODULE_PATH: Final = Path(verifier.__file__) _EXPECTED_TAG: Final = "cp310-abi3-linux_x86_64" _NATIVE_MEMBER: Final = "litellm/rust_bridge/_native.abi3.so" @@ -76,7 +60,6 @@ def _fake_subprocess_run(command: tuple[str, ...], **_: object) -> subprocess.Co def _run_verifier( - monkeypatch: pytest.MonkeyPatch, wheel: Path, *, exposes_panic: bool = False, @@ -88,31 +71,33 @@ def _run_verifier( def _fake_load_native_module(_: Path) -> ModuleType: return native_module - monkeypatch.setattr(verifier, "_load_native_module", _fake_load_native_module) - monkeypatch.setattr(verifier.subprocess, "run", _fake_subprocess_run) - monkeypatch.setattr(sys, "argv", [str(_MODULE_PATH), str(wheel)]) - monkeypatch.setenv("GITHUB_STEP_SUMMARY", str(wheel.parent / "summary.md")) - return verifier.main() + environment: Final = {"GITHUB_STEP_SUMMARY": str(wheel.parent / "summary.md")} + return verifier.main( + (str(_MODULE_PATH), str(wheel)), + environment, + _fake_load_native_module, + _fake_subprocess_run, + ) -def test_accepts_expected_release_wheel_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_accepts_expected_release_wheel_tags(tmp_path: Path) -> None: wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) - assert _run_verifier(monkeypatch, wheel) == 0 + assert _run_verifier(wheel) == 0 -def test_rejects_cp312_version_specific_wheel(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_cp312_version_specific_wheel(tmp_path: Path) -> None: tag: Final = "cp312-cp312-linux_x86_64" wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,)) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 -def test_rejects_non_linux_platform_tag(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_non_linux_platform_tag(tmp_path: Path) -> None: tag: Final = "cp310-abi3-win_amd64" wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,)) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 @pytest.mark.parametrize( @@ -122,17 +107,15 @@ def test_rejects_non_linux_platform_tag(tmp_path: Path, monkeypatch: pytest.Monk ) def test_rejects_missing_or_mismatched_wheel_metadata_tag( tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, metadata_tags: tuple[str, ...] | None, ) -> None: wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG, metadata_tags=metadata_tags) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 def test_rejects_wheel_metadata_from_wrong_dist_info_directory( tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, ) -> None: wheel: Final = _write_wheel( tmp_path, @@ -140,20 +123,20 @@ def test_rejects_wheel_metadata_from_wrong_dist_info_directory( dist_info="decoy-1.0.0.dist-info", ) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 -def test_rejects_duplicate_wheel_metadata_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_duplicate_wheel_metadata_tags(tmp_path: Path) -> None: wheel: Final = _write_wheel( tmp_path, filename_tag=_EXPECTED_TAG, metadata_tags=(_EXPECTED_TAG, _EXPECTED_TAG), ) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 -def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path) -> None: with pytest.warns(UserWarning, match="Duplicate name"): wheel: Final = _write_wheel( tmp_path, @@ -161,10 +144,10 @@ def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path, monkeypatch: pyte duplicate_wheel=True, ) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 -def test_rejects_production_module_exposing_panic_hook(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_production_module_exposing_panic_hook(tmp_path: Path) -> None: wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) - assert _run_verifier(monkeypatch, wheel, exposes_panic=True) == 1 + assert _run_verifier(wheel, exposes_panic=True) == 1 From 814204e21f87c65234d5167cff92c92d357a5587 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 16:06:39 -0700 Subject: [PATCH 42/81] fix(rust): satisfy native wheel verifier lint --- .../rust_bridge/verify_linux_native_wheel.py | 93 +++++++++++-------- .../test_verify_linux_native_wheel.py | 44 +++++---- 2 files changed, 80 insertions(+), 57 deletions(-) diff --git a/litellm/rust_bridge/verify_linux_native_wheel.py b/litellm/rust_bridge/verify_linux_native_wheel.py index 783de87a8b2..899e2a211c0 100644 --- a/litellm/rust_bridge/verify_linux_native_wheel.py +++ b/litellm/rust_bridge/verify_linux_native_wheel.py @@ -7,12 +7,10 @@ import subprocess import sys import zipfile from collections.abc import Callable, Mapping, Sequence -from email import policy -from email.parser import BytesParser from itertools import product from pathlib import Path, PurePosixPath -from types import ModuleType -from typing import Final, Protocol, cast +from types import MappingProxyType, ModuleType +from typing import Final, Protocol EXPECTED_PYTHON_TAG: Final = "cp310" EXPECTED_ABI_TAG: Final = "abi3" @@ -50,9 +48,8 @@ def _dist_info_directory(member: zipfile.ZipInfo) -> str | None: def _wheel_metadata_tags(archive: zipfile.ZipFile, members: tuple[zipfile.ZipInfo, ...]) -> tuple[str, ...]: if len(members) != 1: return () - metadata: Final = BytesParser(policy=policy.default).parsebytes(archive.read(members[0])) - tags: Final = cast(list[str], metadata.get_all("Tag", [])) - return tuple(tag.strip() for tag in tags) + lines: Final = archive.read(members[0]).splitlines() + return tuple(line.removeprefix(b"Tag:").strip().decode("ascii") for line in lines if line.startswith(b"Tag:")) def _load_native_module(native_path: Path) -> ModuleType | None: @@ -62,7 +59,7 @@ def _load_native_module(native_path: Path) -> ModuleType | None: try: native_module: Final = importlib.util.module_from_spec(module_spec) module_spec.loader.exec_module(native_module) - except Exception as error: + except Exception as error: # noqa: BLE001 # native module initialization can raise arbitrary exceptions sys.stderr.write(f"native module load failed: {error}\n") return None return native_module @@ -106,10 +103,14 @@ def main( directory for member in wheel_members if (directory := _dist_info_directory(member)) is not None ) required_dist_info_files: Final = ("METADATA", "RECORD", "WHEEL") - dist_info_file_counts: Final = { - filename: sum(member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members) - for filename in required_dist_info_files - } + dist_info_file_counts: Final = MappingProxyType( + { + filename: sum( + member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members + ) + for filename in required_dist_info_files + } + ) wheel_metadata_members: Final = tuple( member for member in wheel_members if member.filename == f"{expected_dist_info_directory}/WHEEL" ) @@ -233,36 +234,46 @@ def main( if summary_path is not None: Path(summary_path).write_text(verified_report) - if debug_sections: - sys.stderr.write(f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}\n") - if not static_symbol_table_absent: - sys.stderr.write(f"{native_member.filename} contains a static symbol table\n") - if not extension_entry_point_present: - sys.stderr.write("native extension does not export PyInit__native\n") - if python_tag != EXPECTED_PYTHON_TAG: - sys.stderr.write(f"unexpected Python tag: expected {EXPECTED_PYTHON_TAG}, found {python_tag}\n") - if abi_tag != EXPECTED_ABI_TAG: - sys.stderr.write(f"unexpected ABI tag: expected {EXPECTED_ABI_TAG}, found {abi_tag}\n") - if platform_tag != EXPECTED_PLATFORM_TAG: - sys.stderr.write(f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}\n") - if dist_info_directories != expected_dist_info_directories: - sys.stderr.write( - f"unexpected dist-info directories: expected {[expected_dist_info_directory]}, " - f"found {sorted(dist_info_directories)}\n" + invalid_dist_info_files: Final = any(count != 1 for count in dist_info_file_counts.values()) + validation_errors: Final = tuple( + message + for failed, message in ( + (bool(debug_sections), f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}"), + (not static_symbol_table_absent, f"{native_member.filename} contains a static symbol table"), + (not extension_entry_point_present, "native extension does not export PyInit__native"), + ( + python_tag != EXPECTED_PYTHON_TAG, + f"unexpected Python tag: expected {EXPECTED_PYTHON_TAG}, found {python_tag}", + ), + (abi_tag != EXPECTED_ABI_TAG, f"unexpected ABI tag: expected {EXPECTED_ABI_TAG}, found {abi_tag}"), + ( + platform_tag != EXPECTED_PLATFORM_TAG, + f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}", + ), + ( + dist_info_directories != expected_dist_info_directories, + f"unexpected dist-info directories: expected {expected_dist_info_directory}, " + f"found {', '.join(sorted(dist_info_directories))}", + ), + (invalid_dist_info_files, f"required dist-info file counts are invalid: {dist_info_file_counts}"), + ( + not invalid_dist_info_files and not wheel_metadata_tags_match, + f"WHEEL tags do not match filename: expected {', '.join(sorted(expanded_filename_tags))}, " + f"found {', '.join(sorted(wheel_metadata_tags))}", + ), + ( + native_module is not None and not panic_test_hook_absent, + "production native module exposes _panic_for_test", + ), + ( + not native_size_within_limit, + f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB", + ), + (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), ) - if any(count != 1 for count in dist_info_file_counts.values()): - sys.stderr.write(f"required dist-info file counts are invalid: {dist_info_file_counts}\n") - elif not wheel_metadata_tags_match: - sys.stderr.write( - f"WHEEL tags do not match filename: expected {sorted(expanded_filename_tags)}, " - f"found {sorted(wheel_metadata_tags)}\n" - ) - if native_module is not None and not panic_test_hook_absent: - sys.stderr.write("production native module exposes _panic_for_test\n") - if not native_size_within_limit: - sys.stderr.write(f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB\n") - if unexpected_members: - sys.stderr.write(f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}\n") + if failed + ) + sys.stderr.write("".join(f"{message}\n" for message in validation_errors)) return 0 if all(passed for _, passed in validations) else 1 diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index a4291ce0a65..8d0082dddc1 100644 --- a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -3,7 +3,7 @@ from __future__ import annotations import subprocess import zipfile from pathlib import Path -from types import ModuleType +from types import MappingProxyType, ModuleType from typing import Final import pytest @@ -47,16 +47,26 @@ def _write_wheel( return wheel -def _fake_subprocess_run(command: tuple[str, ...], **_: object) -> subprocess.CompletedProcess[str]: +def _fake_subprocess_run( + command: tuple[str, ...], + *, + check: bool, + capture_output: bool, + text: bool, +) -> subprocess.CompletedProcess[str]: + assert check and capture_output and text if command == ("rustc", "--version"): - stdout = "rustc 1.98.0 (regression-test)\n" - elif "--sections" in command: - stdout = "[ 1] .text PROGBITS\n" - elif "--dyn-syms" in command: - stdout = "PyInit__native\n" - else: - raise AssertionError(f"unexpected subprocess command: {command}") - return subprocess.CompletedProcess(command, 0, stdout=stdout, stderr="") + return subprocess.CompletedProcess(command, 0, stdout="rustc 1.98.0 (regression-test)\n", stderr="") + if "--sections" in command: + return subprocess.CompletedProcess(command, 0, stdout="[ 1] .text PROGBITS\n", stderr="") + if "--dyn-syms" in command: + return subprocess.CompletedProcess(command, 0, stdout="PyInit__native\n", stderr="") + raise AssertionError(f"unexpected subprocess command: {command}") + + +class _NativeModuleWithPanicHook(ModuleType): + def _panic_for_test(self) -> None: + return None def _run_verifier( @@ -64,14 +74,16 @@ def _run_verifier( *, exposes_panic: bool = False, ) -> int: - native_module: Final = ModuleType("litellm.rust_bridge._native") - if exposes_panic: - setattr(native_module, "_panic_for_test", lambda: None) + native_module: Final = ( + _NativeModuleWithPanicHook("litellm.rust_bridge._native") + if exposes_panic + else ModuleType("litellm.rust_bridge._native") + ) def _fake_load_native_module(_: Path) -> ModuleType: return native_module - environment: Final = {"GITHUB_STEP_SUMMARY": str(wheel.parent / "summary.md")} + environment: Final = MappingProxyType({"GITHUB_STEP_SUMMARY": str(wheel.parent / "summary.md")}) return verifier.main( (str(_MODULE_PATH), str(wheel)), environment, @@ -102,8 +114,8 @@ def test_rejects_non_linux_platform_tag(tmp_path: Path) -> None: @pytest.mark.parametrize( "metadata_tags", - [None, ("cp312-cp312-linux_x86_64",)], - ids=["missing", "mismatched"], + (None, ("cp312-cp312-linux_x86_64",)), + ids=("missing", "mismatched"), ) def test_rejects_missing_or_mismatched_wheel_metadata_tag( tmp_path: Path, From 90eadac40927315831bcc7d7ae6de57f155acbd8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 06:55:19 -0700 Subject: [PATCH 43/81] test(build): keep wheel checks outside package --- .../scripts}/smoke_test_native_wheel.py | 0 .../scripts}/verify_linux_native_wheel.py | 0 .github/workflows/test-rust.yml | 12 +++--- .../test_verify_linux_native_wheel.py | 38 +++++++++++++++++-- 4 files changed, 41 insertions(+), 9 deletions(-) rename {litellm/rust_bridge => .github/scripts}/smoke_test_native_wheel.py (100%) rename {litellm/rust_bridge => .github/scripts}/verify_linux_native_wheel.py (100%) diff --git a/litellm/rust_bridge/smoke_test_native_wheel.py b/.github/scripts/smoke_test_native_wheel.py similarity index 100% rename from litellm/rust_bridge/smoke_test_native_wheel.py rename to .github/scripts/smoke_test_native_wheel.py diff --git a/litellm/rust_bridge/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py similarity index 100% rename from litellm/rust_bridge/verify_linux_native_wheel.py rename to .github/scripts/verify_linux_native_wheel.py diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 271c73733a3..aada0fcf239 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -7,8 +7,8 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" - - "litellm/rust_bridge/smoke_test_native_wheel.py" - - "litellm/rust_bridge/verify_linux_native_wheel.py" + - ".github/scripts/smoke_test_native_wheel.py" + - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -21,8 +21,8 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" - - "litellm/rust_bridge/smoke_test_native_wheel.py" - - "litellm/rust_bridge/verify_linux_native_wheel.py" + - ".github/scripts/smoke_test_native_wheel.py" + - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" permissions: @@ -115,9 +115,9 @@ jobs: --config-setting "maturin.build-args=--features panic-test,extension-module" - name: Smoke-test native panic unwinding - run: python litellm/rust_bridge/smoke_test_native_wheel.py panic-dist/*.whl + run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - name: Verify stripped native extension env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python litellm/rust_bridge/verify_linux_native_wheel.py dist/*.whl + run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index 8d0082dddc1..e449d4392d8 100644 --- a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -1,16 +1,48 @@ from __future__ import annotations +import importlib.util import subprocess +import sys import zipfile +from collections.abc import Callable, Mapping, Sequence from pathlib import Path from types import MappingProxyType, ModuleType -from typing import Final +from typing import Final, Protocol, cast import pytest -from litellm.rust_bridge import verify_linux_native_wheel as verifier -_MODULE_PATH: Final = Path(verifier.__file__) +class _CommandRunner(Protocol): + def __call__( + self, + command: tuple[str, ...], + *, + check: bool, + capture_output: bool, + text: bool, + ) -> subprocess.CompletedProcess[str]: ... + + +class _VerifierModule(Protocol): + main: Callable[ + [ + Sequence[str] | None, + Mapping[str, str] | None, + Callable[[Path], ModuleType | None], + _CommandRunner, + ], + int, + ] + + +_REPO_ROOT: Final = Path(__file__).resolve().parents[3] +_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "verify_linux_native_wheel.py" +_SPEC: Final = importlib.util.spec_from_file_location("verify_linux_native_wheel", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_LOADED_VERIFIER: Final = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _LOADED_VERIFIER +_SPEC.loader.exec_module(_LOADED_VERIFIER) +verifier: Final = cast(_VerifierModule, _LOADED_VERIFIER) _EXPECTED_TAG: Final = "cp310-abi3-linux_x86_64" _NATIVE_MEMBER: Final = "litellm/rust_bridge/_native.abi3.so" From 9de4e84feb128f2248253d9f52ae13049f2c88a5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 15:07:57 -0700 Subject: [PATCH 44/81] refactor(rust): extract domain-neutral Python interop --- litellm-rust/AGENTS.md | 9 +- litellm-rust/CLAUDE.md | 9 +- litellm-rust/Cargo.lock | 109 +++++++++++++++++- litellm-rust/Cargo.toml | 3 + litellm-rust/README.md | 8 +- .../PROVIDER_CODING_STANDARDS.md | 2 +- litellm-rust/crates/ai-gateway/README.md | 7 +- .../core/tests/workspace_crate_allowlist.rs | 16 ++- litellm-rust/crates/python-bridge/AGENTS.md | 2 +- litellm-rust/crates/python-bridge/CLAUDE.md | 5 +- litellm-rust/crates/python-bridge/Cargo.toml | 3 +- .../python-bridge/benches/serialization.rs | 11 +- litellm-rust/crates/python-bridge/src/gil.rs | 32 ----- litellm-rust/crates/python-bridge/src/lib.rs | 16 +-- .../python-bridge/tests/marshal_boundary.rs | 11 +- litellm-rust/crates/python-interop/AGENTS.md | 1 + litellm-rust/crates/python-interop/Cargo.toml | 15 +++ litellm-rust/crates/python-interop/src/gil.rs | 21 ++++ litellm-rust/crates/python-interop/src/lib.rs | 5 + .../src/marshal.rs | 0 .../crates/python-interop/tests/interop.rs | 44 +++++++ 21 files changed, 247 insertions(+), 82 deletions(-) delete mode 100644 litellm-rust/crates/python-bridge/src/gil.rs create mode 100644 litellm-rust/crates/python-interop/AGENTS.md create mode 100644 litellm-rust/crates/python-interop/Cargo.toml create mode 100644 litellm-rust/crates/python-interop/src/gil.rs create mode 100644 litellm-rust/crates/python-interop/src/lib.rs rename litellm-rust/crates/{python-bridge => python-interop}/src/marshal.rs (100%) create mode 100644 litellm-rust/crates/python-interop/tests/interop.rs diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 36a5ad5a8f4..b8b6291283d 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,6 +1,6 @@ # AGENTS.md -litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are MODULES inside the layers. +litellm-rust has four crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. ## Crates @@ -8,9 +8,10 @@ litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes ( |-------|------| | litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | | litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. ## Where a route lives @@ -28,7 +29,7 @@ core/src/messages/ Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched. -Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. +Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index fe6ceedbb86..3dcf1853efc 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -21,12 +21,13 @@ variants of it. The test for a good abstraction is that adding the next provider is a few declarative lines, not a new file of duplicated flow. Only diverge from the base when behavior is genuinely different, and say so explicitly in the PR. -## Crates (exactly three — see AGENTS.md) +## Crates (see AGENTS.md) `litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call. `litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and -`litellm-python-bridge` exposes it to the Python SDK. A crate is a **layer**, not -a route — add modules, not crates. +`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop` +holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate +is a layer or shared foundation, not a route; add modules, not crates. ## Core Boundary @@ -175,7 +176,7 @@ cd litellm-rust cargo fmt --check # the ai-gateway binary + server code is behind the `server` feature cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings -cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings +cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings cargo test --workspace ``` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4388e561026..dd41cf0e84b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -919,6 +919,12 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" +[[package]] +name = "futures-timer" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" + [[package]] name = "futures-util" version = "0.3.33" @@ -972,6 +978,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + [[package]] name = "h2" version = "0.3.27" @@ -1432,14 +1444,24 @@ dependencies = [ "criterion", "litellm-ai-gateway", "litellm-core", + "litellm-python-interop", "pyo3", "pyo3-async-runtimes", - "pythonize", - "serde", "serde_json", "tokio", ] +[[package]] +name = "litellm-python-interop" +version = "0.1.0" +dependencies = [ + "pyo3", + "pythonize", + "rstest", + "serde", + "serde_json", +] + [[package]] name = "litemap" version = "0.8.2" @@ -1627,6 +1649,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -1899,6 +1930,12 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + [[package]] name = "reqwest" version = "0.12.28" @@ -1956,6 +1993,35 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn 2.0.119", + "unicode-ident", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2488,6 +2554,36 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow", +] + [[package]] name = "tower" version = "0.5.3" @@ -2903,6 +2999,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + [[package]] name = "writeable" version = "0.6.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index c17a0605fc7..c447d915abe 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -2,6 +2,7 @@ members = [ "crates/core", "crates/ai-gateway", + "crates/python-interop", "crates/python-bridge", ] resolver = "2" @@ -15,12 +16,14 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } +litellm-python-interop = { path = "crates/python-interop" } axum = "0.7" pyo3 = "0.29.0" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } +rstest = "0.26.1" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" diff --git a/litellm-rust/README.md b/litellm-rust/README.md index bcccf93300b..a0d79c6f0a5 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -26,9 +26,10 @@ coverage and production evidence. |-------|------| | litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. | | litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. ## Layout @@ -38,7 +39,8 @@ crates/ src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client src/providers/anthropic/messages/transformation.rs ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints. - python-bridge/ PyO3 bridge for Python LiteLLM. + python-interop/ Domain-neutral PyO3 conversion and GIL primitives. + python-bridge/ PyO3 API adapter for Python LiteLLM. ``` The folder shape follows the Python provider tree: diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index a1860d8a9c9..4a689cb9579 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -54,6 +54,6 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` cd litellm-rust cargo fmt --check cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings - cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings + cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings cargo test --workspace ``` diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 7a6c620ee84..5cbb47220be 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -6,15 +6,16 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. ## Crates -`litellm-rust` is exactly three crates (a crate is a **layer**, not a route): +`litellm-rust` has four crates. A crate is a layer or shared foundation, not a route: | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs index 656ba033b62..8a8a5ea263a 100644 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -1,7 +1,8 @@ -//! Enforcement: the litellm-rust workspace has exactly three crates. +//! Enforcement: the litellm-rust workspace has exactly four crates. //! -//! `core` (pure translation), `ai-gateway` (routes + all network I/O), and -//! `python-bridge` (the PyO3 cdylib). Adding or removing a crate must be a +//! `core` (the Rust SDK), `ai-gateway` (the HTTP/WebSocket host), +//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the +//! PyO3 cdylib). Adding or removing a crate must be a //! deliberate act: this test fails until the allowlist here is updated, forcing //! whoever changes the crate set to justify the new crate per the rule that a //! crate is a layer needing independent compilation / its own deps / a separate @@ -16,10 +17,15 @@ use std::path::{Path, PathBuf}; /// The one true crate set. Update BOTH this and `litellm-rust/AGENTS.md` when the /// workspace legitimately gains or loses a crate. -const EXPECTED_MEMBERS: &[&str] = &["crates/core", "crates/ai-gateway", "crates/python-bridge"]; +const EXPECTED_MEMBERS: &[&str] = &[ + "crates/core", + "crates/ai-gateway", + "crates/python-interop", + "crates/python-bridge", +]; /// The crate subdirectory names that must exist under `crates/`. -const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-bridge"]; +const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-interop", "python-bridge"]; const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact)."; diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index ad3cddfa5fd..42282ca4da4 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,3 +1,3 @@ -litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over the litellm-core route entrypoints (e.g. `litellm_core::messages::messages`). +litellm-python-bridge is the PyO3 cdylib that exposes LiteLLM Rust APIs to the Python SDK. Keep API registration, domain dependency wiring, request assembly, and Python exception mapping here. Put domain-neutral Python/Serde conversion and GIL primitives in litellm-python-interop. Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint. diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md index 3ce8b8c639a..d25ae5a8130 100644 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -5,8 +5,9 @@ Rules for `litellm-rust/crates/python-bridge`. ## Responsibility `python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. -Keep this crate thin. It adapts Python objects to Rust payloads and returns -Python-compatible dictionaries. +Keep this crate thin. It exposes LiteLLM Rust APIs, assembles domain requests, +maps domain errors to Python exceptions, and delegates generic conversion and +GIL handling to `litellm-python-interop`. ## Bridge Shape diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index b1fdfd7677a..498003de149 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -18,10 +18,9 @@ panic-test = [] [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } +litellm-python-interop.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true -pythonize.workspace = true -serde.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/benches/serialization.rs b/litellm-rust/crates/python-bridge/benches/serialization.rs index 8a90cf667d0..0b9436d0cb7 100644 --- a/litellm-rust/crates/python-bridge/benches/serialization.rs +++ b/litellm-rust/crates/python-bridge/benches/serialization.rs @@ -2,6 +2,7 @@ use std::hint::black_box; use std::time::Duration; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use litellm_python_interop::{from_py, to_py}; use pyo3::prelude::*; use pyo3::types::PyDict; use serde_json::{Value, json}; @@ -25,7 +26,7 @@ fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Va } fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value { - pythonize::depythonize(value).expect("payload should depythonize") + from_py(value).expect("payload should depythonize") } fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { @@ -37,12 +38,10 @@ fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { } fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py { - pythonize::pythonize(py, value) - .expect("response should pythonize") - .unbind() + to_py(py, value).expect("response should pythonize") } -fn serialization(c: &mut Criterion) { +fn bridge_serialization(c: &mut Criterion) { Python::initialize(); Python::attach(|py| { for &(label, payload_bytes) in PAYLOAD_SIZES { @@ -98,6 +97,6 @@ criterion_group! { .sample_size(20) .warm_up_time(Duration::from_secs(1)) .measurement_time(Duration::from_secs(4)); - targets = serialization + targets = bridge_serialization } criterion_main!(benches); diff --git a/litellm-rust/crates/python-bridge/src/gil.rs b/litellm-rust/crates/python-bridge/src/gil.rs deleted file mode 100644 index e887c8ec1e3..00000000000 --- a/litellm-rust/crates/python-bridge/src/gil.rs +++ /dev/null @@ -1,32 +0,0 @@ -//! GIL accounting. -//! -//! A single chokepoint for releasing the GIL around blocking work. Every -//! blocking call in the bridge goes through [`release_gil`] instead of calling -//! `Python::detach` directly, so the release count stays accurate and we -//! have one place to extend later (timing histograms, per-call labels, etc.). - -use std::sync::atomic::{AtomicU64, Ordering}; - -use pyo3::prelude::*; - -/// Number of times the bridge has released the GIL since process start. -static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); - -/// Release the GIL around `f`, recording the release. -/// -/// `f` must not touch any Python state — that is what makes releasing the GIL -/// safe. Returning the value back to Python re-acquires the GIL at the call -/// site, after `f` has finished. -pub fn release_gil(py: Python<'_>, f: F) -> T -where - F: FnOnce() -> T + Send, - T: Send, -{ - GIL_RELEASES.fetch_add(1, Ordering::Relaxed); - py.detach(f) -} - -/// Total GIL releases performed by the bridge so far. -pub fn release_count() -> u64 { - GIL_RELEASES.load(Ordering::Relaxed) -} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 18e0b05cbb5..746e0770f9b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,16 +13,12 @@ use litellm_core::chat_completions::{ use litellm_core::error::CoreError; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; +use litellm_python_interop::{from_py, release_count, release_gil, to_py}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; use serde_json::{Map, Value}; -mod gil; -mod marshal; - -use marshal::{from_py, to_py}; - pyo3::create_exception!( _native, RustBridgeDeclined, @@ -230,7 +226,7 @@ fn ocr( timeout_seconds, )?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { model: &model, document, @@ -318,7 +314,7 @@ fn transcription( }; let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; let timeout = optional_timeout(timeout_seconds); - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( AudioTranscriptionRequest { model: &model, @@ -419,7 +415,7 @@ fn messages( let (body, extra_headers, timeout) = marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest { model: &model, body, @@ -546,7 +542,7 @@ fn chat_completions( timeout_seconds, )?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions( ChatCompletionsRequest { model: &model, @@ -610,7 +606,7 @@ fn achat_completions( #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); - stats.set_item("releases", gil::release_count())?; + stats.set_item("releases", release_count())?; Ok(stats.into_any().unbind()) } diff --git a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs index 6a6ede22e85..d397d20b9fd 100644 --- a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs +++ b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs @@ -1,7 +1,7 @@ use std::fs; use std::path::{Path, PathBuf}; -const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[ +const DISALLOWED_OUTSIDE_INTEROP: &[&str] = &[ "py.import(\"json\")", "pythonize::", "serde_json::to_string", @@ -33,18 +33,15 @@ fn rust_sources(directory: &Path) -> Vec { } #[test] -fn serialization_is_centralized_in_marshal_module() { +fn serialization_uses_the_interop_boundary() { let root = source_root(); for path in rust_sources(&root) { - if path == root.join("marshal.rs") { - continue; - } let source = fs::read_to_string(&path).expect("bridge source should be readable"); - for disallowed in DISALLOWED_OUTSIDE_MARSHAL { + for disallowed in DISALLOWED_OUTSIDE_INTEROP { assert!( !source.contains(disallowed), - "{} bypasses the typed marshal module with `{disallowed}`", + "{} bypasses litellm-python-interop with `{disallowed}`", path.display() ); } diff --git a/litellm-rust/crates/python-interop/AGENTS.md b/litellm-rust/crates/python-interop/AGENTS.md new file mode 100644 index 00000000000..d1d61e5dfa0 --- /dev/null +++ b/litellm-rust/crates/python-interop/AGENTS.md @@ -0,0 +1 @@ +litellm-python-interop is the domain-neutral PyO3 foundation. Keep generic Python/Serde conversion and interpreter primitives here. Do not add LiteLLM domain crates, route types, API registration, or cdylib build features. diff --git a/litellm-rust/crates/python-interop/Cargo.toml b/litellm-rust/crates/python-interop/Cargo.toml new file mode 100644 index 00000000000..9da6af6e2e2 --- /dev/null +++ b/litellm-rust/crates/python-interop/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-python-interop" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +pyo3.workspace = true +pythonize.workspace = true +serde.workspace = true + +[dev-dependencies] +rstest.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/python-interop/src/gil.rs b/litellm-rust/crates/python-interop/src/gil.rs new file mode 100644 index 00000000000..04b966a6002 --- /dev/null +++ b/litellm-rust/crates/python-interop/src/gil.rs @@ -0,0 +1,21 @@ +use std::sync::atomic::{AtomicU64, Ordering}; + +use pyo3::prelude::*; + +static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); + +/// Runs work detached from the interpreter and records the release. +/// +/// `f` must not access Python state while the interpreter is detached. +pub fn release_gil(py: Python<'_>, f: F) -> T +where + F: FnOnce() -> T + Send, + T: Send, +{ + GIL_RELEASES.fetch_add(1, Ordering::Relaxed); + py.detach(f) +} + +pub fn release_count() -> u64 { + GIL_RELEASES.load(Ordering::Relaxed) +} diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs new file mode 100644 index 00000000000..df2bd260fdb --- /dev/null +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -0,0 +1,5 @@ +mod gil; +mod marshal; + +pub use gil::{release_count, release_gil}; +pub use marshal::{from_py, to_py}; diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-interop/src/marshal.rs similarity index 100% rename from litellm-rust/crates/python-bridge/src/marshal.rs rename to litellm-rust/crates/python-interop/src/marshal.rs diff --git a/litellm-rust/crates/python-interop/tests/interop.rs b/litellm-rust/crates/python-interop/tests/interop.rs new file mode 100644 index 00000000000..9c456dcb938 --- /dev/null +++ b/litellm-rust/crates/python-interop/tests/interop.rs @@ -0,0 +1,44 @@ +use pyo3::Python; +use rstest::{fixture, rstest}; +use serde_json::{Value, json}; + +use litellm_python_interop::{from_py, release_count, release_gil, to_py}; + +struct InitializedPython; + +impl InitializedPython { + fn attach(&self, f: F) -> R + where + F: for<'py> FnOnce(Python<'py>) -> R, + { + Python::attach(f) + } +} + +#[fixture] +#[once] +fn initialized_python() -> InitializedPython { + Python::initialize(); + InitializedPython +} + +#[rstest] +fn serde_values_round_trip_through_python(#[from(initialized_python)] python: &InitializedPython) { + python.attach(|py| { + let expected = json!({"model": "test", "items": [1, true, null]}); + let python_value = to_py(py, &expected).expect("value should convert to Python"); + let actual: Value = + from_py(python_value.bind(py)).expect("Python value should convert to serde"); + + assert_eq!(actual, expected); + }); +} + +#[rstest] +fn release_gil_runs_work_and_records_it(#[from(initialized_python)] python: &InitializedPython) { + let before = release_count(); + let result = python.attach(|py| release_gil(py, || 42)); + + assert_eq!(result, 42); + assert_eq!(release_count(), before + 1); +} From 518a2a70f159c1e118d1ea0d387c7da8ba87d642 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 05:47:01 -0700 Subject: [PATCH 45/81] refactor(rust): standardize the core Error type --- .../src/audio_transcription/common_utils.rs | 10 +-- .../src/audio_transcription/handler.rs | 31 ++++---- .../src/audio_transcription/hooks.rs | 58 +++++++------- .../ai-gateway/src/audio_transcription/mod.rs | 4 +- .../crates/ai-gateway/src/io/realtime.rs | 43 +++++------ .../crates/ai-gateway/src/io/realtime_pool.rs | 4 +- .../crates/ai-gateway/src/io/responses_ws.rs | 76 +++++++++---------- .../crates/ai-gateway/src/ocr/common_utils.rs | 73 +++++++++--------- .../crates/ai-gateway/src/ocr/handler.rs | 15 ++-- .../crates/ai-gateway/src/ocr/hooks.rs | 65 ++++++++-------- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 4 +- .../crates/ai-gateway/src/ocr/tests.rs | 8 +- .../crates/ai-gateway/src/python/config.rs | 12 ++- .../ai-gateway/src/routes/messages/mod.rs | 38 +++++----- .../ai-gateway/src/routes/messages/service.rs | 16 ++-- .../ai-gateway/src/routes/realtime/service.rs | 11 ++- .../src/routes/responses/service.rs | 12 +-- .../src/audio_transcription/transformation.rs | 11 ++- .../crates/core/src/call_lifecycle/mod.rs | 34 ++++----- .../core/src/chat_completions/common_utils.rs | 7 +- .../core/src/chat_completions/handler.rs | 30 ++++---- .../crates/core/src/chat_completions/mod.rs | 5 +- .../core/src/chat_completions/prepare.rs | 21 +++-- .../crates/core/src/chat_completions/tests.rs | 40 +++++----- .../src/chat_completions/transformation.rs | 11 ++- litellm-rust/crates/core/src/error.rs | 8 +- litellm-rust/crates/core/src/http_utils.rs | 8 +- litellm-rust/crates/core/src/lib.rs | 2 +- .../crates/core/src/messages/common_utils.rs | 7 +- .../crates/core/src/messages/handler.rs | 25 +++--- litellm-rust/crates/core/src/messages/mod.rs | 7 +- .../crates/core/src/messages/prepare.rs | 12 +-- .../crates/core/src/messages/tests.rs | 10 +-- .../core/src/messages/transformation.rs | 11 ++- .../crates/core/src/ocr/transformation.rs | 11 ++- .../anthropic/chat_completions/tests.rs | 16 ++-- .../chat_completions/transformation.rs | 25 +++--- .../anthropic/messages/transformation.rs | 12 +-- .../azure_ai/messages/transformation.rs | 22 +++--- .../providers/azure_ai/ocr/transformation.rs | 62 +++++++-------- .../providers/bedrock/audio_transcription.rs | 20 ++--- .../core/src/providers/bedrock/aws_base.rs | 38 +++++----- .../bedrock/chat_completions/tests.rs | 14 ++-- .../chat_completions/transformation.rs | 23 +++--- .../providers/mistral/ocr/transformation.rs | 28 +++---- .../openai/realtime/transformation.rs | 10 +-- .../openai/responses/transformation.rs | 6 +- .../providers/vertex_ai/ocr/transformation.rs | 50 ++++++------ .../core/src/realtime/transformation.rs | 6 +- .../core/src/responses/instrumentation.rs | 8 +- .../crates/core/src/responses/websocket.rs | 6 +- litellm-rust/crates/python-bridge/src/lib.rs | 36 ++++----- 52 files changed, 544 insertions(+), 578 deletions(-) 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 index 270d5c2d97a..140bc8aeea8 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs @@ -1,10 +1,8 @@ -use std::collections::BTreeMap; - -use litellm_core::CoreResult; use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; use serde_json::{Map, Value}; +use std::collections::BTreeMap; pub(super) fn audio_transcription_provider_config( provider: &str, @@ -17,7 +15,7 @@ pub(super) fn audio_transcription_provider_config( pub(super) fn string_headers( headers: Option>, -) -> CoreResult> { +) -> Result, Error> { headers .unwrap_or_default() .into_iter() @@ -26,7 +24,7 @@ pub(super) fn string_headers( .as_str() .map(|value| (key.clone(), value.to_string())) .ok_or_else(|| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "audio transcription extra_headers.{key} must be a string" )) }) diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs index 33c13550f58..1bdd4ae72a2 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs @@ -1,11 +1,9 @@ -use std::time::SystemTime; - -use litellm_core::CoreResult; use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; -use litellm_core::error::CoreError; +use litellm_core::error::Error; 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 std::time::SystemTime; use super::common_utils::truncate_error_body; use super::types::ProviderAudioTranscriptionRequest; @@ -13,10 +11,9 @@ 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}")) - })?; +) -> Result { + let body = serde_json::to_vec(&request.body) + .map_err(|error| Error::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); @@ -27,21 +24,20 @@ pub(crate) async fn execute_audio_transcription_provider_call( let response = request_builder .send() .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; let status = response.status(); let text = response .text() .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::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}")) - })?; + let response_json: Value = serde_json::from_str(&text) + .map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?; Ok(request .config .transform_transcription_response(&request.model, response_json)? @@ -51,14 +47,13 @@ pub(crate) async fn execute_audio_transcription_provider_call( pub(crate) async fn sign_request( request: &ProviderAudioTranscriptionRequest, optional_params: &serde_json::Map, -) -> CoreResult { +) -> Result { 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 body = serde_json::to_vec(&request.body) + .map_err(|error| Error::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()); diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 0c9faeda6e7..5e1240de759 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -1,11 +1,9 @@ -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 litellm_core::error::Error; use serde_json::{Map, Value, json}; +use std::future::Future; +use std::pin::Pin; use super::common_utils::{audio_transcription_provider_config, has_header, string_headers}; use super::handler::sign_request; @@ -26,7 +24,7 @@ pub(crate) struct AudioTranscriptionLifecycleHooks { request_metadata: RequestMetadata, } -type AudioFuture<'a, T> = Pin> + Send + 'a>>; +type AudioFuture<'a, T> = Pin> + Send + 'a>>; type AudioLogFuture<'a> = Pin + Send + 'a>>; impl AudioTranscriptionLifecycleHooks { @@ -45,7 +43,7 @@ impl AudioTranscriptionLifecycleHooks { async fn run_pre_call_guardrails( &self, request: PreparedAudioTranscriptionRequest, - ) -> CoreResult { + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(request); } @@ -63,17 +61,17 @@ impl AudioTranscriptionLifecycleHooks { .await .map_err(guardrail_error_to_core_error)?; let Value::Object(mut data) = guardrail_request.data else { - return Err(CoreError::InvalidRequest( + return Err(Error::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()) + Error::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( + return Err(Error::InvalidRequest( "audio transcription optional_params must be an object".to_string(), )); } @@ -89,9 +87,9 @@ impl AudioTranscriptionLifecycleHooks { async fn prepare_provider_request( &self, request: PreparedAudioTranscriptionRequest, - ) -> CoreResult { + ) -> Result { let config = audio_transcription_provider_config(&request.custom_llm_provider) - .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + .ok_or_else(|| Error::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( @@ -135,7 +133,7 @@ impl AudioTranscriptionLifecycleHooks { async fn run_during_call_guardrails( &self, request: ProviderAudioTranscriptionRequest, - ) -> CoreResult { + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(request); } @@ -153,12 +151,12 @@ impl AudioTranscriptionLifecycleHooks { .await .map_err(guardrail_error_to_core_error)?; let Value::Object(mut data) = guardrail_request.data else { - return Err(CoreError::InvalidRequest( + return Err(Error::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()) + Error::InvalidRequest("audio transcription guardrail removed body".to_string()) })?; Ok(ProviderAudioTranscriptionRequest { body, ..request }) } @@ -241,7 +239,7 @@ impl CallLifecycleHooks( &'a self, context: &'a CallLifecycleContext, - error: &'a CoreError, + error: &'a Error, timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -281,22 +279,22 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { } } -fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { - CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +fn guardrail_error_to_core_error(error: GuardrailError) -> Error { + Error::InvalidRequest(format!("{}: {}", error.kind, error.message)) } -fn core_error_kind(error: &CoreError) -> &'static str { +fn core_error_kind(error: &Error) -> &'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::Connect(_) => "ConnectError", - CoreError::Routing(_) => "RoutingError", - CoreError::Unsupported(_) => "UnsupportedRequest", + Error::Auth(_) => "AuthError", + Error::InvalidProvider(_) => "InvalidProvider", + Error::InvalidRequest(_) => "InvalidRequest", + Error::InvalidType { .. } => "InvalidType", + Error::MissingField(_) => "MissingField", + Error::Http { .. } => "HttpError", + Error::InvalidResponse(_) => "InvalidResponse", + Error::Network(_) => "NetworkError", + Error::Connect(_) => "ConnectError", + Error::Routing(_) => "RoutingError", + Error::Unsupported(_) => "UnsupportedRequest", } } diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs index 5d33d912c40..3983846d7b6 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs @@ -1,4 +1,4 @@ -use litellm_core::CoreResult; +use litellm_core::Error; use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; @@ -13,7 +13,7 @@ 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 { +pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { let PreparedAudioTranscriptionCall { request, hooks } = prepare_audio_transcription_call(request); CallLifecycle::default() diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 845e7bf9527..662f7328982 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -15,8 +15,7 @@ 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::error::Error; use litellm_core::realtime::transformation::RealtimeProviderConfig; use litellm_core::realtime::types::RealtimeEvent; use tokio::net::TcpStream; @@ -48,7 +47,7 @@ pub(crate) type UpstreamRx = SplitStream; /// Resolve the OpenAI API key from the explicit param or the environment. /// /// Blank/whitespace values are treated as absent (guard at resolution time). -pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { +pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { api_key .map(str::trim) .filter(|key| !key.is_empty()) @@ -58,7 +57,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { .ok() .filter(|key| !key.trim().is_empty()) }) - .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) } /// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`. @@ -70,24 +69,24 @@ pub(crate) async fn dial_upstream( model: &str, api_key: &str, api_base: Option<&str>, -) -> CoreResult { +) -> Result { let url = OPENAI_REALTIME_CONFIG.complete_url(api_base, model); let mut request = url .as_str() .into_client_request() - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; // GA realtime: only Authorization. The legacy OpenAI-Beta header triggers // beta_api_shape_disabled, so we do not send it. request.headers_mut().insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {api_key}")) - .map_err(|err| CoreError::Auth(err.to_string()))?, + .map_err(|err| Error::Auth(err.to_string()))?, ); let (upstream, _response) = connect_async(request) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; Ok(upstream) } @@ -96,22 +95,22 @@ pub(crate) async fn dial_upstream( /// Used by the pool to pre-read OpenAI's unprompted `session.created`. Returns an /// error on a non-text frame, a closed socket, or undecodable JSON so the pool can /// discard a misbehaving socket rather than warm it. -pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult { +pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> Result { loop { let message = upstream_rx .next() .await - .ok_or_else(|| CoreError::Network("upstream closed before first event".to_string()))? - .map_err(|err| CoreError::Network(err.to_string()))?; + .ok_or_else(|| Error::Network("upstream closed before first event".to_string()))? + .map_err(|err| Error::Network(err.to_string()))?; match message { Message::Text(text) => { return serde_json::from_str(&text) - .map_err(|err| CoreError::InvalidResponse(err.to_string())); + .map_err(|err| Error::InvalidResponse(err.to_string())); } // Ignore protocol frames (ping/pong) while waiting for the first event. Message::Ping(_) | Message::Pong(_) => continue, Message::Close(_) => { - return Err(CoreError::Network( + return Err(Error::Network( "upstream closed before first event".to_string(), )); } @@ -139,7 +138,7 @@ pub(crate) async fn splice( mut observe: impl FnMut(&RealtimeEvent) + Send, mut client_in: In, mut client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -154,7 +153,7 @@ where client_out .send(outbound) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; } } @@ -175,26 +174,26 @@ where // inflate its own spend log. Logging observes upstream events only. for outbound in config.transform_realtime_request(&event, model)?.events { let payload = serde_json::to_string(&outbound) - .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + .map_err(|err| Error::InvalidResponse(err.to_string()))?; upstream_tx .send(Message::Text(payload)) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; } } // upstream -> client upstream_message = upstream_rx.next() => { let Some(message) = upstream_message else { break }; // upstream closed - match message.map_err(|err| CoreError::Network(err.to_string()))? { + match message.map_err(|err| Error::Network(err.to_string()))? { Message::Text(text) => { let event: RealtimeEvent = serde_json::from_str(&text) - .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + .map_err(|err| Error::InvalidResponse(err.to_string()))?; observe(&event); for outbound in config.transform_realtime_response(&event, model)?.events { client_out .send(outbound) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; } } Message::Close(_) => break, @@ -225,7 +224,7 @@ pub async fn realtime( observe: impl FnMut(&RealtimeEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -258,7 +257,7 @@ pub async fn realtime_warm( observe: impl FnMut(&RealtimeEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, 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 4a1a3cd1166..49e9c459a88 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs @@ -28,7 +28,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use futures_util::StreamExt; -use litellm_core::CoreResult; +use litellm_core::Error; use litellm_core::realtime::types::RealtimeEvent; use crate::io::realtime::{ @@ -438,7 +438,7 @@ impl RealtimePool { /// /// `key.api_key` is already resolved (non-blank). The first frame OpenAI sends /// unprompted is `session.created`; we buffer exactly that and read nothing more. -async fn warm_one(key: &UpstreamKey) -> CoreResult { +async fn warm_one(key: &UpstreamKey) -> Result { let upstream: UpstreamWs = dial_upstream(&key.model, &key.api_key, key.api_base.as_deref()).await?; let (tx, mut rx) = upstream.split(); 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 9b51019f4bc..0b01747b1a5 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -4,10 +4,10 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::Error; use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG; use litellm_core::responses::types::ResponsesWsEvent; use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig; -use litellm_core::{CoreError, CoreResult}; use tokio::net::TcpStream; use tokio::sync::Mutex; use tokio_tungstenite::tungstenite::Message; @@ -37,51 +37,49 @@ impl ResponsesWebSocketConnection { url: &str, headers: &HashMap, timeout: Option, - ) -> CoreResult { + ) -> Result { let mut request = url .into_client_request() - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; for (name, value) in headers { let header_name = name .parse::() - .map_err(|error| CoreError::InvalidRequest(error.to_string()))?; + .map_err(|error| Error::InvalidRequest(error.to_string()))?; let header_value = HeaderValue::from_str(value) - .map_err(|error| CoreError::InvalidRequest(error.to_string()))?; + .map_err(|error| Error::InvalidRequest(error.to_string()))?; request.headers_mut().insert(header_name, header_value); } let connect = connect_async(request); let result = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { - CoreError::Network("Responses WebSocket connection timed out".to_string()) + Error::Network("Responses WebSocket connection timed out".to_string()) })?, None => connect.await, }; let (socket, _) = result.map_err(|error| match error { - tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http { + tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), }, - other => CoreError::Network(other.to_string()), + other => Error::Network(other.to_string()), })?; Ok(Self { socket: Arc::new(Mutex::new(Some(socket))), }) } - pub async fn send_text(&self, text: String) -> CoreResult<()> { + pub async fn send_text(&self, text: String) -> Result<(), Error> { let mut socket = self.socket.lock().await; let Some(socket) = socket.as_mut() else { - return Err(CoreError::Network( - "Responses WebSocket is closed".to_string(), - )); + return Err(Error::Network("Responses WebSocket is closed".to_string())); }; socket .send(Message::Text(text)) .await - .map_err(|error| CoreError::Network(error.to_string())) + .map_err(|error| Error::Network(error.to_string())) } - pub async fn recv_text(&self) -> CoreResult> { + pub async fn recv_text(&self) -> Result, Error> { let mut socket_guard = self.socket.lock().await; let Some(socket) = socket_guard.as_mut() else { return Ok(None); @@ -90,27 +88,27 @@ impl ResponsesWebSocketConnection { Some(Ok(Message::Text(text))) => Ok(Some(text)), Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) .map(Some) - .map_err(|error| CoreError::InvalidResponse(error.to_string())), + .map_err(|error| Error::InvalidResponse(error.to_string())), Some(Ok(Message::Close(_))) | None => Ok(None), Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(CoreError::Network(error.to_string())), + Some(Err(error)) => Err(Error::Network(error.to_string())), } } - pub async fn close(&self) -> CoreResult<()> { + pub async fn close(&self) -> Result<(), Error> { let mut socket = self.socket.lock().await; if let Some(socket) = socket.as_mut() { socket .close(None) .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; } *socket = None; Ok(()) } } -pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { +pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { api_key .map(str::trim) .filter(|value| !value.is_empty()) @@ -120,38 +118,38 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { .ok() .filter(|value| !value.trim().is_empty()) }) - .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) } async fn dial_upstream( model: &str, api_key: &str, api_base: Option<&str>, -) -> CoreResult { +) -> Result { let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model); let mut request = url .as_str() .into_client_request() - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; request.headers_mut().insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {api_key}")) - .map_err(|error| CoreError::Auth(error.to_string()))?, + .map_err(|error| Error::Auth(error.to_string()))?, ); let result = tokio::time::timeout( Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS), connect_async(request), ) .await - .map_err(|_| CoreError::Network("Responses WebSocket connection timed out".to_string()))?; + .map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?; result .map(|(socket, _)| socket) .map_err(|error| match error { - tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http { + tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), }, - other => CoreError::Network(other.to_string()), + other => Error::Network(other.to_string()), }) } @@ -166,7 +164,7 @@ impl ResponsesWebSocketStreaming { observe: impl FnMut(&ResponsesWsEvent) + Send, client_in: In, client_out: Out, - ) -> CoreResult<()> + ) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -193,7 +191,7 @@ pub(crate) async fn splice( mut observe: impl FnMut(&ResponsesWsEvent) + Send, mut client_in: In, mut client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -210,18 +208,18 @@ where .events { let payload = serde_json::to_string(&outbound) - .map_err(|error| CoreError::InvalidResponse(error.to_string()))?; + .map_err(|error| Error::InvalidResponse(error.to_string()))?; upstream_tx.send(Message::Text(payload)) .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; } } message = upstream_rx.next() => { let Some(message) = message else { break }; - match message.map_err(|error| CoreError::Network(error.to_string()))? { + match message.map_err(|error| Error::Network(error.to_string()))? { Message::Text(text) => { let event = serde_json::from_str::(&text) - .map_err(|error| CoreError::InvalidResponse(error.to_string()))?; + .map_err(|error| Error::InvalidResponse(error.to_string()))?; observe(&event); for outbound in OPENAI_RESPONSES_WS_CONFIG .transform_ws_response(&event, model)? @@ -229,7 +227,7 @@ where { client_out.send(outbound) .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; } } Message::Close(_) => break, @@ -252,7 +250,7 @@ pub async fn async_responses_websocket( mut observe: impl FnMut(&ResponsesWsEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -267,11 +265,11 @@ where .events { let payload = serde_json::to_string(&outbound) - .map_err(|error| CoreError::InvalidResponse(error.to_string()))?; + .map_err(|error| Error::InvalidResponse(error.to_string()))?; upstream_tx .send(Message::Text(payload)) .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; } } ResponsesWebSocketStreaming::bidirectional_forward( @@ -296,7 +294,7 @@ pub async fn responses_ws( observe: impl FnMut(&ResponsesWsEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -514,7 +512,7 @@ mod tests { ) .await .expect_err("status error"); - assert!(matches!(error, CoreError::Http { status: 401, .. })); + assert!(matches!(error, Error::Http { status: 401, .. })); server.await.expect("server task"); } @@ -543,7 +541,7 @@ mod tests { ) .await .expect_err("status error"); - assert!(matches!(error, CoreError::Http { status: 500, .. })); + assert!(matches!(error, Error::Http { status: 500, .. })); server.await.expect("server task"); } } 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 9bc2818b6e7..e0ce165dc93 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -3,8 +3,7 @@ use std::time::{Duration, Instant}; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use litellm_core::CoreResult; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::ocr::transformation::OcrProviderConfig; use reqwest::Url; use serde_json::{Map, Value}; @@ -56,7 +55,7 @@ fn is_azure_document_intelligence_model(model: &str) -> bool { pub(super) fn string_headers( extra_headers: Option>, -) -> CoreResult> { +) -> Result, Error> { extra_headers .unwrap_or_default() .into_iter() @@ -65,7 +64,7 @@ pub(super) fn string_headers( .as_str() .map(|value| (key.clone(), value.to_string())) .ok_or_else(|| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "OCR extra_headers.{key} must be a string, got {}", litellm_core::error::json_type_name(&value) )) @@ -80,7 +79,7 @@ pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { .any(|(key, _)| key.eq_ignore_ascii_case(name)) } -fn document_url_field(document: &Value) -> CoreResult> { +fn document_url_field(document: &Value) -> Result, Error> { let Some(object) = document.as_object() else { return Ok(None); }; @@ -138,13 +137,13 @@ fn is_blocked_ip(ip: IpAddr) -> bool { } } -fn blocked_url_error(url: &Url) -> CoreError { - CoreError::InvalidRequest(format!( +fn blocked_url_error(url: &Url) -> Error { + Error::InvalidRequest(format!( "OCR document URL rejected by SSRF protection: {url}" )) } -async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> { +async fn validate_safe_fetch_url(url: &Url) -> Result<(), Error> { if !matches!(url.scheme(), "http" | "https") { return Err(blocked_url_error(url)); } @@ -162,7 +161,7 @@ async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> { .ok_or_else(|| blocked_url_error(url))?; let addresses = tokio::net::lookup_host((host, port)) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let mut saw_address = false; for address in addresses { saw_address = true; @@ -176,25 +175,25 @@ async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> { Ok(()) } -fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult { +fn redirect_location(response: &reqwest::Response, url: &Url) -> Result { let location = response .headers() .get(reqwest::header::LOCATION) .and_then(|value| value.to_str().ok()) .ok_or_else(|| { - CoreError::InvalidResponse("OCR document redirect missing Location header".to_string()) + Error::InvalidResponse("OCR document redirect missing Location header".to_string()) })?; url.join(location) - .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}"))) + .map_err(|err| Error::InvalidResponse(format!("invalid OCR document redirect: {err}"))) } -async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> { +async fn safe_get_document_url(url: &str) -> Result<(Url, reqwest::Response), Error> { let client = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) .build() - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let mut current_url = Url::parse(url) - .map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?; + .map_err(|err| Error::InvalidRequest(format!("invalid OCR document URL: {err}")))?; for _ in 0..MAX_SAFE_FETCH_REDIRECTS { validate_safe_fetch_url(¤t_url).await?; @@ -202,28 +201,28 @@ async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response) .get(current_url.clone()) .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !response.status().is_redirection() { return Ok((current_url, response)); } current_url = redirect_location(&response, ¤t_url)?; } - Err(CoreError::InvalidRequest( + Err(Error::InvalidRequest( "Too many redirects while fetching OCR document URL".to_string(), )) } -fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> { +fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Result<(), Error> { if max_bytes == 0 { - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}" ))); } if content_length > max_bytes { let size_mb = content_length as f64 / (1024.0 * 1024.0); let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0); - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}" ))); } @@ -233,7 +232,7 @@ fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Core async fn read_response_with_limit( mut response: reqwest::Response, url: &Url, -) -> CoreResult> { +) -> Result, Error> { let max_bytes = max_document_download_bytes(); if let Some(content_length) = response.content_length() { enforce_download_size(content_length, max_bytes, url)?; @@ -246,7 +245,7 @@ async fn read_response_with_limit( while let Some(chunk) = response .chunk() .await - .map_err(|err| CoreError::Network(err.to_string()))? + .map_err(|err| Error::Network(err.to_string()))? { bytes_downloaded += chunk.len() as u64; enforce_download_size(bytes_downloaded, max_bytes, url)?; @@ -255,7 +254,7 @@ async fn read_response_with_limit( Ok(bytes) } -pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult { +pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result { let Some((field, url)) = document_url_field(&document)? else { return Ok(document); }; @@ -267,7 +266,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes let status = response.status(); if !status.is_success() { let body = response.text().await.unwrap_or_default(); - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&body), }); @@ -290,7 +289,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes let mut transformed = document .as_object() .cloned() - .ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?; + .ok_or_else(|| Error::InvalidRequest("OCR document must be an object".to_string()))?; transformed.insert(field.to_string(), Value::String(data_uri)); Ok(Value::Object(transformed)) } @@ -316,11 +315,11 @@ fn retry_after_secs(response: &reqwest::Response) -> u64 { .unwrap_or(2) } -fn operation_status(response_json: &Value) -> CoreResult<&str> { +fn operation_status(response_json: &Value) -> Result<&str, Error> { let status = response_json .get("status") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("status"))?; + .ok_or(Error::MissingField("status"))?; match status { "succeeded" => Ok("succeeded"), "running" | "notStarted" => Ok("running"), @@ -330,11 +329,11 @@ fn operation_status(response_json: &Value) -> CoreResult<&str> { .and_then(|error| error.get("message")) .and_then(Value::as_str) .unwrap_or("Unknown error"); - Err(CoreError::InvalidResponse(format!( + Err(Error::InvalidResponse(format!( "Azure Document Intelligence analysis failed: {message}" ))) } - other => Err(CoreError::InvalidResponse(format!( + other => Err(Error::InvalidResponse(format!( "Unknown operation status: {other}" ))), } @@ -345,9 +344,9 @@ pub(super) async fn poll_document_intelligence( original_url: &str, headers: &[(String, String)], timeout: Option, -) -> CoreResult { +) -> Result { if !same_origin(operation_url, original_url) { - return Err(CoreError::InvalidResponse( + return Err(Error::InvalidResponse( "Azure Document Intelligence: rejected cross-origin polling URL".to_string(), )); } @@ -358,7 +357,7 @@ pub(super) async fn poll_document_intelligence( )); loop { if start.elapsed() > timeout { - return Err(CoreError::Network(format!( + return Err(Error::Network(format!( "Azure Document Intelligence operation polling timed out after {} seconds", timeout.as_secs() ))); @@ -373,21 +372,21 @@ pub(super) async fn poll_document_intelligence( let response = request_builder .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let retry_after = retry_after_secs(&response); let status = response.status(); let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } let response_json: Value = serde_json::from_str(&text).map_err(|err| { - CoreError::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}")) + Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}")) })?; if operation_status(&response_json)? == "succeeded" { return Ok(response_json); @@ -426,7 +425,7 @@ mod tests { assert!(matches!( error, - CoreError::InvalidRequest(message) + Error::InvalidRequest(message) if message.contains("SSRF protection") )); } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 1de34eb400e..815bc84363a 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -1,5 +1,4 @@ -use litellm_core::CoreResult; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::Value; @@ -7,7 +6,7 @@ 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 { +pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Result { let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -19,7 +18,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co let response = request_builder .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let status = response.status(); if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll @@ -31,7 +30,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co .and_then(|value| value.to_str().ok()) .map(str::to_string) .ok_or_else(|| { - CoreError::InvalidResponse( + Error::InvalidResponse( "Azure Document Intelligence returned 202 but no Operation-Location header found" .to_string(), ) @@ -52,17 +51,17 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } let response_json: Value = serde_json::from_str(&text) - .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; + .map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; Ok(request .config diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 95df566dc53..401e26d3b29 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -1,11 +1,9 @@ -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::error::Error; use litellm_core::ocr::transformation::OcrAuthStrategy; use serde_json::{Map, Value, json}; +use std::future::Future; +use std::pin::Pin; use super::common_utils::{ convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, @@ -27,7 +25,7 @@ pub(crate) struct OcrLifecycleHooks { request_metadata: RequestMetadata, } -type OcrFuture<'a, T> = Pin> + Send + 'a>>; +type OcrFuture<'a, T> = Pin> + Send + 'a>>; type OcrLogFuture<'a> = Pin + Send + 'a>>; impl OcrLifecycleHooks { @@ -46,7 +44,7 @@ impl OcrLifecycleHooks { async fn run_pre_call_guardrails( &self, request: PreparedOcrRequest, - ) -> CoreResult { + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(request); } @@ -74,9 +72,9 @@ impl OcrLifecycleHooks { async fn prepare_provider_request( &self, request: PreparedOcrRequest, - ) -> CoreResult { + ) -> Result { let config = ocr_provider_config(&request.custom_llm_provider, &request.model) - .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + .ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?; let env_lookup = |key: &str| std::env::var(key).ok(); let headers = string_headers(request.extra_headers)?; let auth_strategy = config.auth_strategy(); @@ -120,7 +118,7 @@ impl OcrLifecycleHooks { custom_llm_provider: &str, url: &str, body: Value, - ) -> CoreResult { + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(body); } @@ -217,7 +215,7 @@ impl CallLifecycleHooks for OcrLi fn async_log_failure_event<'a>( &'a self, context: &'a CallLifecycleContext, - error: &'a CoreError, + error: &'a Error, timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -278,19 +276,19 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { fn parse_ocr_pre_call_guardrail_request( request: GuardrailRequest, -) -> CoreResult<(Value, Map)> { +) -> Result<(Value, Map), Error> { let Value::Object(mut data) = request.data else { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "OCR pre_call guardrail must return an object".to_string(), )); }; let document = data.remove("document").ok_or_else(|| { - CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string()) + Error::InvalidRequest("OCR pre_call guardrail removed document".to_string()) })?; let optional_params = match data.remove("optional_params") { Some(Value::Object(params)) => params, Some(_) => { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "OCR pre_call guardrail optional_params must be an object".to_string(), )); } @@ -299,33 +297,32 @@ fn parse_ocr_pre_call_guardrail_request( Ok((document, optional_params)) } -fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult { +fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result { let Value::Object(mut data) = request.data else { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "OCR during_call guardrail must return an object".to_string(), )); }; - data.remove("body").ok_or_else(|| { - CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string()) - }) + data.remove("body") + .ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string())) } -fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { - CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +fn guardrail_error_to_core_error(error: GuardrailError) -> Error { + Error::InvalidRequest(format!("{}: {}", error.kind, error.message)) } -fn core_error_kind(error: &CoreError) -> &'static str { +fn core_error_kind(error: &Error) -> &'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::Connect(_) => "ConnectError", - CoreError::Routing(_) => "RoutingError", - CoreError::Unsupported(_) => "UnsupportedRequest", + Error::Auth(_) => "AuthError", + Error::InvalidProvider(_) => "InvalidProvider", + Error::InvalidRequest(_) => "InvalidRequest", + Error::InvalidType { .. } => "InvalidType", + Error::MissingField(_) => "MissingField", + Error::Http { .. } => "HttpError", + Error::InvalidResponse(_) => "InvalidResponse", + Error::Network(_) => "NetworkError", + Error::Connect(_) => "ConnectError", + Error::Routing(_) => "RoutingError", + Error::Unsupported(_) => "UnsupportedRequest", } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index c4c13e2300c..b59ab626fd3 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,4 +1,4 @@ -use litellm_core::CoreResult; +use litellm_core::Error; use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; @@ -13,7 +13,7 @@ pub use types::OcrRequest; use handler::execute_ocr_provider_call; use prepare::{PreparedOcrCall, prepare_ocr_call}; -pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { +pub async fn ocr(request: OcrRequest<'_>) -> Result { let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); CallLifecycle::default() .run_request(request, &hooks, execute_ocr_provider_call) diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index bb2a6b06501..8c3f0425149 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -395,7 +395,7 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { .await .expect_err("provider error propagates"); - assert!(matches!(err, CoreError::Http { status: 500, .. })); + assert!(matches!(err, Error::Http { status: 500, .. })); server.await.expect("server task completes"); assert_eq!( logger.events(), @@ -439,7 +439,7 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() { .await .expect_err("guardrail blocks request"); - assert!(matches!(err, CoreError::InvalidRequest(_))); + assert!(matches!(err, Error::InvalidRequest(_))); assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]); assert_eq!( logger.events(), @@ -607,7 +607,7 @@ fn string_headers_rejects_non_string_values() { let err = string_headers(Some(headers)).expect_err("non-string header rejected"); assert_eq!( err, - CoreError::InvalidRequest( + Error::InvalidRequest( "OCR extra_headers.x-retry-count must be a string, got number".to_string() ) ); diff --git a/litellm-rust/crates/ai-gateway/src/python/config.rs b/litellm-rust/crates/ai-gateway/src/python/config.rs index c028d3d6b51..d5a4dd69c8d 100644 --- a/litellm-rust/crates/ai-gateway/src/python/config.rs +++ b/litellm-rust/crates/ai-gateway/src/python/config.rs @@ -6,33 +6,31 @@ //! (and recorded in [`crate::gil`]); the realtime hot path never touches Python. //! //! Compiled only under the `python-config` feature. - -use litellm_core::CoreResult; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::router::{Deployment, Router}; use pyo3::prelude::*; use crate::gil; /// Load the router's `model_list` from `config_path` via the Python reader. -pub fn load_router_from_config(config_path: &str) -> CoreResult { +pub fn load_router_from_config(config_path: &str) -> Result { gil::record_acquisition(); Python::attach(|py| { let model_list = py .import("litellm.proxy.read_model_list") .and_then(|module| module.getattr("read_model_list")) .and_then(|reader| reader.call1((config_path,))) - .map_err(|err| CoreError::Routing(format!("read_model_list failed: {err}")))?; + .map_err(|err| Error::Routing(format!("read_model_list failed: {err}")))?; let model_list_json: String = py .import("json") .and_then(|json| json.getattr("dumps")) .and_then(|dumps| dumps.call1((model_list,))) .and_then(|encoded| encoded.extract()) - .map_err(|err| CoreError::Routing(format!("serializing model_list failed: {err}")))?; + .map_err(|err| Error::Routing(format!("serializing model_list failed: {err}")))?; let deployments: Vec = serde_json::from_str(&model_list_json) - .map_err(|err| CoreError::Routing(format!("parsing model_list failed: {err}")))?; + .map_err(|err| Error::Routing(format!("parsing model_list failed: {err}")))?; Ok(Router::new(deployments)) }) 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 7e38d10c6ff..e9f8c477f36 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -9,7 +9,7 @@ use axum::http::StatusCode; use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue}; use axum::response::{IntoResponse, Response}; use axum::routing::post; -use litellm_core::CoreError; +use litellm_core::Error; use serde_json::{Map, Value}; use crate::auth::RequireMasterKey; @@ -46,7 +46,7 @@ fn stream_response(upstream: reqwest::Response) -> Result Result Result>, CoreError> { +fn forwarded_headers(headers: &HeaderMap) -> Result>, Error> { let forwarded = headers .iter() .filter(|(name, _)| { @@ -74,19 +74,19 @@ fn forwarded_headers(headers: &HeaderMap) -> Result>, }) .map(|(name, value)| { let value = value.to_str().map_err(|_| { - CoreError::InvalidRequest(format!("invalid value for header {}", name.as_str())) + Error::InvalidRequest(format!("invalid value for header {}", name.as_str())) })?; Ok((name.to_string(), Value::String(value.to_string()))) }) - .collect::, CoreError>>()?; + .collect::, Error>>()?; Ok((!forwarded.is_empty()).then_some(forwarded)) } #[derive(Debug)] -struct MessagesRouteError(CoreError); +struct MessagesRouteError(Error); -impl From for MessagesRouteError { - fn from(error: CoreError) -> Self { +impl From for MessagesRouteError { + fn from(error: Error) -> Self { Self(error) } } @@ -94,28 +94,28 @@ impl From for MessagesRouteError { impl IntoResponse for MessagesRouteError { fn into_response(self) -> Response { let (status, message) = match self.0 { - CoreError::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message), - CoreError::InvalidProvider(_) | CoreError::Routing(_) => ( + Error::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message), + Error::InvalidProvider(_) | Error::Routing(_) => ( StatusCode::NOT_FOUND, "no messages deployment is configured for this model".to_string(), ), - CoreError::Auth(_) => ( + Error::Auth(_) => ( StatusCode::BAD_GATEWAY, "messages provider authentication failed".to_string(), ), - CoreError::Http { .. } - | CoreError::Network(_) - | CoreError::Connect(_) - | CoreError::InvalidResponse(_) - | CoreError::InvalidType { .. } - | CoreError::MissingField(_) => ( + Error::Http { .. } + | Error::Network(_) + | Error::Connect(_) + | Error::InvalidResponse(_) + | Error::InvalidType { .. } + | Error::MissingField(_) => ( StatusCode::BAD_GATEWAY, "messages provider request failed".to_string(), ), // The gateway has no Python implementation to decline to, so a // request the core cannot serve is reported to the caller. The // reason is a fixed internal string, never provider content. - CoreError::Unsupported(reason) => ( + Error::Unsupported(reason) => ( StatusCode::BAD_REQUEST, format!("messages request is not supported: {reason}"), ), 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 5f4c5fe8de4..4fd29db05d6 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs @@ -1,10 +1,10 @@ use std::sync::Arc; +use litellm_core::Error; use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER; use litellm_core::messages::types::MessagesRequest; use litellm_core::messages::{messages, messages_stream}; use litellm_core::router::Router; -use litellm_core::{CoreError, CoreResult}; use serde_json::{Map, Value}; pub(crate) enum MessagesResponse { @@ -16,16 +16,16 @@ pub async fn run( router: &Arc, body: Value, extra_headers: Option>, -) -> CoreResult { +) -> Result { let model = body .get("model") .and_then(Value::as_str) .map(str::trim) .filter(|model| !model.is_empty()) - .ok_or_else(|| CoreError::InvalidRequest("messages body requires a model".to_string()))?; - let deployment = router.get_available_deployment(model).ok_or_else(|| { - CoreError::Routing(format!("no deployment available for model '{model}'")) - })?; + .ok_or_else(|| Error::InvalidRequest("messages body requires a model".to_string()))?; + let deployment = router + .get_available_deployment(model) + .ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?; let provider_model = deployment.litellm_params.model.as_str(); let upstream_model = provider_model .split_once('/') @@ -37,7 +37,7 @@ pub async fn run( }; let mut body = body; body.as_object_mut() - .ok_or_else(|| CoreError::InvalidRequest("messages body must be an object".to_string()))? + .ok_or_else(|| Error::InvalidRequest("messages body must be an object".to_string()))? .insert( "model".to_string(), Value::String(upstream_model.to_string()), @@ -60,6 +60,6 @@ pub async fn run( serde_json::to_value(response) .map(MessagesResponse::Json) .map_err(|err| { - CoreError::InvalidResponse(format!("failed to serialize messages response: {err}")) + Error::InvalidResponse(format!("failed to serialize messages response: {err}")) }) } 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 4ae8cfe7379..b8ee77c4269 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -11,8 +11,7 @@ use std::time::Duration; 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::error::Error; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::router::Router; @@ -29,15 +28,15 @@ pub async fn run( observe: impl FnMut(&RealtimeEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, >::Error: std::fmt::Display, { - let deployment = router.get_available_deployment(model).ok_or_else(|| { - CoreError::Routing(format!("no deployment available for model '{model}'")) - })?; + let deployment = router + .get_available_deployment(model) + .ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?; let params = &deployment.litellm_params; // Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model. let provider_model = params diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs index 165c95695d3..e8f840c0c8e 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs @@ -2,13 +2,13 @@ use std::sync::Arc; use std::time::Duration; use futures_util::{Sink, Stream}; +use litellm_core::Error; use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext}; use litellm_core::responses::instrumentation::{ ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome, ResponsesWsMetadata, }; use litellm_core::responses::types::ResponsesWsEvent; -use litellm_core::{CoreError, CoreResult}; use crate::integrations::custom_logger::{ CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails, @@ -26,22 +26,22 @@ pub async fn run( metadata: RequestMetadata, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, Out::Error: std::fmt::Display, { - let deployment = router.get_available_deployment(model).ok_or_else(|| { - CoreError::Routing(format!("no deployment available for model '{model}'")) - })?; + let deployment = router + .get_available_deployment(model) + .ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?; let params = &deployment.litellm_params; let provider_model = params .model .strip_prefix("openai/") .unwrap_or(¶ms.model); if params.model.contains('/') && !params.model.starts_with("openai/") { - return Err(CoreError::InvalidProvider( + return Err(Error::InvalidProvider( "Responses WebSocket route supports OpenAI deployments only".to_string(), )); } diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs index eab34c13843..16a28fbcac0 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -1,7 +1,6 @@ +use crate::Error; use serde_json::{Map, Value}; -use crate::CoreResult; - use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData}; #[derive(Clone, Debug, PartialEq, Eq)] @@ -32,13 +31,13 @@ pub trait AudioTranscriptionProviderConfig: Sync { model: &str, audio: Value, optional_params: Map, - ) -> CoreResult; + ) -> Result; fn transform_transcription_response( &self, model: &str, response_json: Value, - ) -> CoreResult; + ) -> Result; fn complete_url( &self, @@ -46,12 +45,12 @@ pub trait AudioTranscriptionProviderConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn auth_strategy( &self, model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; } diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index d9b68a1b726..637c156e192 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -1,7 +1,7 @@ use std::future::Future; use std::time::{Instant, SystemTime, UNIX_EPOCH}; -use crate::{CoreError, CoreResult}; +use crate::Error; pub mod types; @@ -11,14 +11,14 @@ pub use types::{ }; pub trait CallLifecycleHooks: Send + Sync { - type PreCallFuture<'a>: Future> + Send + 'a + type PreCallFuture<'a>: Future> + Send + 'a where Self: 'a, InitialReq: 'a, ProviderReq: 'a, Resp: 'a; - type DuringCallFuture<'a>: Future> + Send + 'a + type DuringCallFuture<'a>: Future> + Send + 'a where Self: 'a, InitialReq: 'a, @@ -56,7 +56,7 @@ pub trait CallLifecycleHooks: Send + Sync { fn async_log_failure_event<'a>( &'a self, context: &'a CallLifecycleContext, - error: &'a CoreError, + error: &'a Error, timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a>; } @@ -86,12 +86,12 @@ impl<'a> CallLifecycle<'a> { request: InitialReq, hooks: &Hooks, provider_call: ProviderCall, - ) -> CoreResult + ) -> Result where InitialReq: CallLifecycleRequest, Hooks: CallLifecycleHooks, ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, + ProviderFuture: Future>, { let context = request.lifecycle_context(); self.run(context, request, hooks, provider_call).await @@ -103,11 +103,11 @@ impl<'a> CallLifecycle<'a> { request: InitialReq, hooks: &Hooks, provider_call: ProviderCall, - ) -> CoreResult + ) -> Result where Hooks: CallLifecycleHooks, ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, + ProviderFuture: Future>, { let call_start = epoch_seconds(); let mut phases = Vec::new(); @@ -166,7 +166,7 @@ impl<'a> CallLifecycle<'a> { &self, context: &CallLifecycleContext, hooks: &Hooks, - error: &CoreError, + error: &Error, call_start: f64, phases: &mut Vec, ) where @@ -251,8 +251,8 @@ mod tests { } impl CallLifecycleHooks for RecordingHooks { - type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; - type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type PreCallFuture<'a> = BoxFuture<'a, Result>; + type DuringCallFuture<'a> = BoxFuture<'a, Result>; type SuccessFuture<'a> = BoxFuture<'a, ()>; type FailureFuture<'a> = BoxFuture<'a, ()>; @@ -294,7 +294,7 @@ mod tests { fn async_log_failure_event<'a>( &'a self, _context: &'a CallLifecycleContext, - _error: &'a CoreError, + _error: &'a Error, _timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -304,8 +304,8 @@ mod tests { } impl CallLifecycleHooks for RecordingHooks { - type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; - type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type PreCallFuture<'a> = BoxFuture<'a, Result>; + type DuringCallFuture<'a> = BoxFuture<'a, Result>; type SuccessFuture<'a> = BoxFuture<'a, ()>; type FailureFuture<'a> = BoxFuture<'a, ()>; @@ -345,7 +345,7 @@ mod tests { fn async_log_failure_event<'a>( &'a self, _context: &'a CallLifecycleContext, - _error: &'a CoreError, + _error: &'a Error, _timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -383,13 +383,13 @@ mod tests { "request".to_string(), &hooks, |_request| async move { - Err::(CoreError::Network("provider down".to_string())) + Err::(Error::Network("provider down".to_string())) }, ) .await .expect_err("call fails"); - assert_eq!(error, CoreError::Network("provider down".to_string())); + assert_eq!(error, Error::Network("provider down".to_string())); assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]); } diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index 36eaf242a5a..ca51471eb7c 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -1,8 +1,7 @@ -use serde_json::{Map, Value}; - -use crate::error::CoreResult; +use crate::Error; use crate::http_utils::string_headers as shared_string_headers; use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; +use serde_json::{Map, Value}; use super::transformation::ChatCompletionsProviderConfig; @@ -23,6 +22,6 @@ pub(super) fn chat_completions_provider_config( pub(super) fn string_headers( extra_headers: Option>, -) -> CoreResult> { +) -> Result, Error> { shared_string_headers(HEADER_CONTEXT, extra_headers) } diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index afc4529fd26..7e2731442cc 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,6 +1,6 @@ use serde_json::Value; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::http_utils::truncate_error_body; use super::client::http_client; @@ -11,9 +11,9 @@ use super::types::{ pub(super) async fn execute_chat_completions_provider_call( request: ProviderChatCompletionsRequest, -) -> CoreResult { +) -> Result { let body = serde_json::to_vec(&request.body).map_err(|err| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "failed to serialize chat completions request: {err}" )) })?; @@ -32,9 +32,9 @@ pub(super) async fn execute_chat_completions_provider_call( // so the host can still serve it. Everything else here, a timeout // above all, may have reached the provider and been answered. if err.is_connect() || err.is_builder() { - CoreError::Connect(err.to_string()) + Error::Connect(err.to_string()) } else { - CoreError::Network(err.to_string()) + Error::Network(err.to_string()) } })?; @@ -42,17 +42,17 @@ pub(super) async fn execute_chat_completions_provider_call( let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } let body: Value = serde_json::from_str(&text).map_err(|err| { - CoreError::InvalidResponse(format!("invalid chat completions response JSON: {err}")) + Error::InvalidResponse(format!("invalid chat completions response JSON: {err}")) })?; request .config @@ -69,10 +69,10 @@ pub(super) async fn execute_chat_completions_provider_call( /// second kind has already been billed, and a host that keeps a reference /// implementation must not retry those, so collapse them to one variant that /// can only mean the provider was already called. -pub(super) fn as_response_error(err: CoreError) -> CoreError { +pub(super) fn as_response_error(err: Error) -> Error { match err { - already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already, - other => CoreError::InvalidResponse(other.to_string()), + already @ (Error::InvalidResponse(_) | Error::Http { .. }) => already, + other => Error::InvalidResponse(other.to_string()), } } @@ -80,7 +80,7 @@ pub(super) fn as_response_error(err: CoreError) -> CoreError { pub(super) async fn signed_headers( request: &ProviderChatCompletionsRequest, body: &[u8], -) -> CoreResult> { +) -> Result, Error> { use std::collections::BTreeMap; use std::time::SystemTime; @@ -101,7 +101,7 @@ pub(super) async fn signed_headers( .iter() .any(|(name, _)| is_sigv4_computed_header(name)) { - return Err(CoreError::Unsupported( + return Err(Error::Unsupported( "request forwards a header AWS SigV4 computes", )); } @@ -137,9 +137,9 @@ pub(super) async fn signed_headers( pub(super) async fn signed_headers( request: &ProviderChatCompletionsRequest, _body: &[u8], -) -> CoreResult> { +) -> Result, Error> { match &request.auth { - ChatCompletionsAuth::AwsSigV4 { .. } => Err(CoreError::Unsupported( + ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported( "AWS SigV4 requires the bedrock-auth feature", )), _ => Ok(request.upstream_headers.clone()), diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index f30ac1a24bf..0d009d36d16 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -6,6 +6,7 @@ //! credentials, and it resolves the provider, translates the conversation, //! calls the provider, and returns a typed OpenAI-shaped response. +use crate::Error; mod client; mod common_utils; pub mod conversation; @@ -17,15 +18,13 @@ pub mod types; use serde_json::{Map, Value}; -use crate::error::CoreResult; - use handler::execute_chat_completions_provider_call; use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config}; use types::{ChatCompletionsRequest, ChatCompletionsResponse}; pub async fn chat_completions( request: ChatCompletionsRequest<'_>, -) -> CoreResult { +) -> Result { execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await } diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 1e1c8d1bafd..142b2f2aaed 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,6 +1,6 @@ use serde_json::Value; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::http_utils::has_header; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; @@ -11,7 +11,7 @@ use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsR pub(super) fn resolve_provider_config<'a>( model: &'a str, custom_llm_provider: Option<&'a str>, -) -> CoreResult<(String, &'static dyn ChatCompletionsProviderConfig)> { +) -> Result<(String, &'static dyn ChatCompletionsProviderConfig), Error> { let provider_info = get_custom_llm_provider(model, custom_llm_provider) .or_else(|| { custom_llm_provider.map(|provider| CustomLlmProvider { @@ -20,35 +20,34 @@ pub(super) fn resolve_provider_config<'a>( }) }) .ok_or_else(|| { - CoreError::InvalidProvider( + Error::InvalidProvider( "unable to resolve custom_llm_provider for chat completions request".to_string(), ) })?; let config = chat_completions_provider_config(provider_info.custom_llm_provider) - .ok_or_else(|| CoreError::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; + .ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; Ok((provider_info.model.to_string(), config)) } -pub(super) fn parse_messages(messages: Value) -> CoreResult> { - serde_json::from_value(messages).map_err(|err| { - CoreError::InvalidRequest(format!("invalid chat completions messages: {err}")) - }) +pub(super) fn parse_messages(messages: Value) -> Result, Error> { + serde_json::from_value(messages) + .map_err(|err| Error::InvalidRequest(format!("invalid chat completions messages: {err}"))) } pub(super) fn prepare_chat_completions_call( request: ChatCompletionsRequest<'_>, -) -> CoreResult { +) -> Result { let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?; let env_lookup = |key: &str| std::env::var(key).ok(); let messages = parse_messages(request.messages)?; if messages.is_empty() { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "chat completions requires at least one message".to_string(), )); } if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) { - return Err(CoreError::Unsupported(reason.0)); + return Err(Error::Unsupported(reason.0)); } let mut headers = string_headers(request.extra_headers)?; diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index e2383723cb0..2858d180e27 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -1,6 +1,6 @@ use serde_json::{Map, Value, json}; -use crate::error::CoreError; +use crate::error::Error; use super::prepare::prepare_chat_completions_call; use super::transformation::ChatCompletionsAuth; @@ -29,7 +29,7 @@ fn request<'a>( /// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers /// carry resolved credentials), so unwrap the failure case by hand. -fn decline(request: ChatCompletionsRequest<'_>) -> CoreError { +fn decline(request: ChatCompletionsRequest<'_>) -> Error { match prepare_chat_completions_call(request) { Err(error) => error, Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), @@ -196,7 +196,7 @@ fn declines_an_unsupported_request_before_resolving_credentials() { call.api_key = None; // No api_key is set and no env is consulted: the gate must run first, so the // error is the decline rather than a missing-credential error. - assert_eq!(decline(call), CoreError::Unsupported("streaming")); + assert_eq!(decline(call), Error::Unsupported("streaming")); } #[test] @@ -208,7 +208,7 @@ fn rejects_an_unknown_provider() { json!([{"role": "user", "content": "hi"}]), json!({}), )), - CoreError::InvalidProvider("openai".to_string()) + Error::InvalidProvider("openai".to_string()) ); } @@ -221,7 +221,7 @@ fn rejects_a_model_with_no_resolvable_provider() { json!([{"role": "user", "content": "hi"}]), json!({}), )), - CoreError::InvalidProvider(_) + Error::InvalidProvider(_) )); } @@ -234,7 +234,7 @@ fn rejects_an_empty_or_malformed_message_list() { json!([]), json!({}), )), - CoreError::InvalidRequest("chat completions requires at least one message".to_string()) + Error::InvalidRequest("chat completions requires at least one message".to_string()) ); assert!(matches!( decline(request( @@ -243,7 +243,7 @@ fn rejects_an_empty_or_malformed_message_list() { json!("not a list"), json!({}), )), - CoreError::InvalidRequest(_) + Error::InvalidRequest(_) )); } @@ -258,7 +258,7 @@ fn rejects_non_string_extra_headers() { call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); assert_eq!( decline(call), - CoreError::InvalidRequest( + Error::InvalidRequest( "chat completions extra_headers.x-trace must be a string, got number".to_string() ) ); @@ -374,7 +374,7 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() { .await .expect_err("{forwarded} should decline instead of being signed"); assert!( - matches!(error, CoreError::Unsupported(_)), + matches!(error, Error::Unsupported(_)), "{forwarded} declined as {error:?}, which the host would not fall back on" ); } @@ -727,7 +727,7 @@ mod round_trip { .expect_err("response cannot be normalized"); handle.await.expect("server task"); assert!( - matches!(err, CoreError::InvalidResponse(_)), + matches!(err, Error::InvalidResponse(_)), "expected a post-send error, got {err:?}" ); } @@ -745,7 +745,7 @@ mod round_trip { .expect_err("response cannot be normalized"); handle.await.expect("server task"); assert!( - matches!(err, CoreError::InvalidResponse(_)), + matches!(err, Error::InvalidResponse(_)), "expected a post-send error, got {err:?}" ); } @@ -763,7 +763,7 @@ mod round_trip { .expect_err("upstream rejects"); handle.await.expect("server task"); assert!( - matches!(err, CoreError::Http { status: 429, .. }), + matches!(err, Error::Http { status: 429, .. }), "expected a 429, got {err:?}" ); } @@ -787,7 +787,7 @@ mod round_trip { .await .expect_err("nothing is listening"); assert!( - matches!(err, CoreError::Connect(_)), + matches!(err, Error::Connect(_)), "expected a pre-send connect failure, got {err:?}" ); } @@ -797,24 +797,24 @@ mod round_trip { use crate::chat_completions::handler::as_response_error; for original in [ - CoreError::MissingField("usage"), - CoreError::Unsupported("non-text response content block"), - CoreError::InvalidRequest("whatever".to_string()), - CoreError::Auth("whatever".to_string()), + Error::MissingField("usage"), + Error::Unsupported("non-text response content block"), + Error::InvalidRequest("whatever".to_string()), + Error::Auth("whatever".to_string()), ] { let label = format!("{original:?}"); assert!( - matches!(as_response_error(original), CoreError::InvalidResponse(_)), + matches!(as_response_error(original), Error::InvalidResponse(_)), "{label} must not stay retryable once the provider has answered" ); } // An upstream status is already unambiguous, so it survives intact. assert!(matches!( - as_response_error(CoreError::Http { + as_response_error(Error::Http { status: 500, body: "boom".to_string() }), - CoreError::Http { status: 500, .. } + Error::Http { status: 500, .. } )); } } diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index a30ce9dc77c..a0868209305 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -1,7 +1,6 @@ +use crate::Error; use serde_json::{Map, Value}; -use crate::error::CoreResult; - use super::types::{ ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData, ProviderChatResponseData, @@ -39,7 +38,7 @@ pub trait ChatCompletionsProviderConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn auth( &self, @@ -47,7 +46,7 @@ pub trait ChatCompletionsProviderConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn default_headers(&self) -> &'static [(&'static str, &'static str)] { &[("content-type", "application/json")] @@ -91,13 +90,13 @@ pub trait ChatCompletionsProviderConfig: Sync { model: &str, messages: Vec, optional_params: Map, - ) -> CoreResult; + ) -> Result; fn transform_response( &self, model: &str, response: ProviderChatResponseData, - ) -> CoreResult; + ) -> Result; } pub fn unsupported_param( diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 739532f8cb5..db3fa2ec704 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -1,9 +1,7 @@ -use thiserror::Error; +use thiserror::Error as ThisError; -pub type CoreResult = Result; - -#[derive(Debug, Error, PartialEq, Eq)] -pub enum CoreError { +#[derive(Debug, ThisError, PartialEq, Eq)] +pub enum Error { #[error("expected {expected}, got {actual}")] InvalidType { expected: &'static str, diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index c541f50275b..10661fadf96 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -3,7 +3,7 @@ use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; /// Bound an upstream error body before it crosses a host boundary, so provider /// bodies stay data-minimized. @@ -18,7 +18,7 @@ pub fn truncate_error_body(body: &str) -> String { pub fn string_headers( context: &'static str, extra_headers: Option>, -) -> CoreResult> { +) -> Result, Error> { extra_headers .unwrap_or_default() .into_iter() @@ -27,7 +27,7 @@ pub fn string_headers( .as_str() .map(|value| (key.clone(), value.to_string())) .ok_or_else(|| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "{context} extra_headers.{key} must be a string, got {}", json_type_name(&value) )) @@ -81,7 +81,7 @@ mod tests { let err = string_headers("chat completions", Some(headers)).expect_err("non-string value"); assert_eq!( err, - CoreError::InvalidRequest( + Error::InvalidRequest( "chat completions extra_headers.x-trace must be a string, got number".to_string() ) ); diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index dce4a425ea0..0e18d24e5d8 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -13,4 +13,4 @@ pub mod responses; pub mod router; pub mod routing_utils; -pub use error::{CoreError, CoreResult}; +pub use error::Error; diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index a14dffbc1fe..8dfdb2e361a 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,9 +1,8 @@ -use serde_json::{Map, Value}; - -use crate::error::CoreResult; +use crate::Error; use crate::http_utils::string_headers as shared_string_headers; use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; +use serde_json::{Map, Value}; use super::transformation::AnthropicMessagesProviderConfig; @@ -23,6 +22,6 @@ pub(super) fn messages_provider_config( pub(super) fn string_headers( extra_headers: Option>, -) -> CoreResult> { +) -> Result, Error> { shared_string_headers(HEADER_CONTEXT, extra_headers) } diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 1c895f66eba..13a65d86131 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,5 +1,5 @@ use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use super::client::http_client; use super::common_utils::truncate_error_body; @@ -7,7 +7,7 @@ use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest}; pub(super) async fn execute_messages_provider_call( request: ProviderMessagesRequest, -) -> CoreResult { +) -> Result { let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -19,32 +19,31 @@ pub(super) async fn execute_messages_provider_call( let response = request_builder .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let status = response.status(); let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } - let response = serde_json::from_str(&text).map_err(|err| { - CoreError::InvalidResponse(format!("invalid messages response JSON: {err}")) - })?; + let response = serde_json::from_str(&text) + .map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?; request.config.transform_response(&request.model, response) } pub(super) async fn execute_messages_provider_stream( request: ProviderMessagesRequest, -) -> CoreResult { +) -> Result { if request.provider != ANTHROPIC_MESSAGES_PROVIDER { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "streaming messages is not supported for this provider".to_string(), )); } @@ -60,14 +59,14 @@ pub(super) async fn execute_messages_provider_stream( let response = request_builder .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let status = response.status(); if !status.is_success() { let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; - return Err(CoreError::Http { + .map_err(|err| Error::Network(err.to_string()))?; + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index acb36d89daf..ee2877e61fc 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -7,6 +7,7 @@ //! is the streaming variant; it hands the raw upstream response back so a host //! can splice the event stream to its own caller. +use crate::Error; mod client; mod common_utils; mod handler; @@ -14,17 +15,15 @@ mod prepare; pub mod transformation; pub mod types; -use crate::error::CoreResult; - use handler::{execute_messages_provider_call, execute_messages_provider_stream}; use prepare::prepare_messages_call; use types::{AnthropicMessagesResponse, MessagesRequest}; -pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { +pub async fn messages(request: MessagesRequest<'_>) -> Result { execute_messages_provider_call(prepare_messages_call(request)?).await } -pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult { +pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { execute_messages_provider_stream(prepare_messages_call(request)?).await } diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 94b5b1eaed7..3b253ac3766 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,4 +1,4 @@ -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; @@ -7,7 +7,7 @@ use super::types::{MessagesRequest, ProviderMessagesRequest}; pub(super) fn prepare_messages_call( request: MessagesRequest<'_>, -) -> CoreResult { +) -> Result { let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) .or_else(|| { request @@ -18,7 +18,7 @@ pub(super) fn prepare_messages_call( }) }) .ok_or_else(|| { - CoreError::InvalidProvider( + Error::InvalidProvider( "unable to resolve custom_llm_provider for messages request".to_string(), ) })?; @@ -26,7 +26,7 @@ pub(super) fn prepare_messages_call( let provider = provider_info.custom_llm_provider; let config = messages_provider_config(provider) - .ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?; + .ok_or_else(|| Error::InvalidProvider(provider.to_string()))?; let env_lookup = |key: &str| std::env::var(key).ok(); let mut headers = string_headers(request.extra_headers)?; @@ -53,11 +53,11 @@ pub(super) fn prepare_messages_call( let url = config.complete_url(request.api_base, &model, &env_lookup)?; let typed_request = serde_json::from_value(request.body).map_err(|err| { - CoreError::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) })?; let transformed = config.transform_request(typed_request)?; let body = serde_json::to_value(transformed).map_err(|err| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "failed to serialize Anthropic messages request: {err}" )) })?; diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index 9fc1763683b..df9f7051011 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -4,7 +4,7 @@ use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use crate::error::CoreError; +use crate::error::Error; use super::common_utils::{ has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, @@ -77,7 +77,7 @@ fn truncate_error_body_caps_long_payloads() { fn string_headers_rejects_non_string_values() { let headers = json!({"x-count": 3}).as_object().unwrap().clone(); let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert!(matches!(err, CoreError::InvalidRequest(_))); + assert!(matches!(err, Error::InvalidRequest(_))); } #[test] @@ -341,7 +341,7 @@ async fn messages_requires_auth_when_no_key_and_no_header() { .await .expect_err("missing auth errors"); - assert!(matches!(err, CoreError::Auth(_))); + assert!(matches!(err, Error::Auth(_))); } #[tokio::test] @@ -420,7 +420,7 @@ async fn messages_maps_provider_error_status_to_http_error() { .await .expect_err("provider error propagates"); - assert!(matches!(err, CoreError::Http { status: 401, .. })); + assert!(matches!(err, Error::Http { status: 401, .. })); } #[tokio::test] @@ -437,5 +437,5 @@ async fn messages_rejects_unsupported_provider() { .await .expect_err("unsupported provider errors"); - assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai")); + assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai")); } diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index b478e20d24b..673a5728aca 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -1,6 +1,5 @@ -use crate::error::CoreResult; - use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse}; +use crate::Error; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MessagesAuthStrategy { @@ -23,13 +22,13 @@ pub trait AnthropicMessagesProviderConfig: Sync { api_base: Option<&str>, model: &str, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn resolve_api_key( &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn auth_strategy(&self) -> MessagesAuthStrategy { MessagesAuthStrategy::Header("x-api-key") @@ -49,7 +48,7 @@ pub trait AnthropicMessagesProviderConfig: Sync { fn transform_request( &self, request: AnthropicMessagesRequest, - ) -> CoreResult { + ) -> Result { Ok(request) } @@ -57,7 +56,7 @@ pub trait AnthropicMessagesProviderConfig: Sync { &self, _model: &str, response: AnthropicMessagesResponse, - ) -> CoreResult { + ) -> Result { Ok(response) } } diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index cb3e735e533..3d3c16c8cb6 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -1,7 +1,6 @@ +use crate::Error; use serde_json::{Map, Value}; -use crate::CoreResult; - use super::types::{OcrRequestData, OcrResponseData}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -43,13 +42,13 @@ pub trait OcrProviderConfig: Sync { model: &str, document: Value, optional_params: Map, - ) -> CoreResult; + ) -> Result; fn transform_ocr_response( &self, model: &str, response_json: Value, - ) -> CoreResult; + ) -> Result; fn complete_url( &self, @@ -57,13 +56,13 @@ pub trait OcrProviderConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn resolve_api_key( &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn auth_strategy(&self) -> OcrAuthStrategy { OcrAuthStrategy::Bearer diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs index 4534ac0182c..b22de6c47de 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::Error; use serde_json::json; fn messages(value: Value) -> Vec { @@ -19,7 +20,7 @@ fn transform(model: &str, msgs: Value, opts: Value) -> Value { .body } -fn transform_response(body: Value) -> CoreResult { +fn transform_response(body: Value) -> Result { ANTHROPIC_CHAT_COMPLETIONS_CONFIG .transform_response("claude-sonnet-4-5", ProviderChatResponseData { body }) } @@ -390,29 +391,26 @@ fn declines_a_response_carrying_a_non_text_block() { "usage": {"input_tokens": 1, "output_tokens": 1} })) .expect_err("non-text block"); - assert_eq!( - err, - CoreError::Unsupported("non-text response content block") - ); + assert_eq!(err, Error::Unsupported("non-text response content block")); } #[test] fn errors_on_a_response_missing_required_fields() { assert_eq!( transform_response(json!("nope")).expect_err("not an object"), - CoreError::InvalidResponse("messages response is not an object".to_string()) + Error::InvalidResponse("messages response is not an object".to_string()) ); assert_eq!( transform_response(json!({"model": "m", "usage": {}})).expect_err("no content"), - CoreError::MissingField("content") + Error::MissingField("content") ); assert_eq!( transform_response(json!({"model": "m", "content": []})).expect_err("no usage"), - CoreError::MissingField("usage") + Error::MissingField("usage") ); assert_eq!( transform_response(json!({"content": [], "usage": {}})).expect_err("no model"), - CoreError::MissingField("model") + Error::MissingField("model") ); } diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index 3658642b539..97cc48aa6f2 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -10,7 +10,7 @@ use crate::chat_completions::types::{ ProviderChatRequestData, ProviderChatResponseData, }; use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::providers::anthropic::messages::transformation::{ complete_anthropic_url, resolve_anthropic_api_key, }; @@ -74,7 +74,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { Ok(complete_anthropic_url(api_base, env_lookup)) } @@ -84,7 +84,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { Ok(ChatCompletionsAuth::Header { name: "x-api-key", value: resolve_anthropic_api_key(api_key, env_lookup)?, @@ -137,7 +137,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { model: &str, messages: Vec, optional_params: Map, - ) -> CoreResult { + ) -> Result { Ok(ProviderChatRequestData { body: anthropic_body(model, &build_conversation(&messages), optional_params), }) @@ -147,15 +147,16 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { &self, _model: &str, response: ProviderChatResponseData, - ) -> CoreResult { - let body = response.body.as_object().ok_or_else(|| { - CoreError::InvalidResponse("messages response is not an object".into()) - })?; + ) -> Result { + let body = response + .body + .as_object() + .ok_or_else(|| Error::InvalidResponse("messages response is not an object".into()))?; let content = body .get("content") .and_then(Value::as_array) - .ok_or(CoreError::MissingField("content"))?; + .ok_or(Error::MissingField("content"))?; // The route declines tool and thinking requests, so a non-text block // means the response carries something this path never asked for. // Decline rather than silently dropping it; the host falls back. @@ -163,7 +164,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { .iter() .any(|block| block.get("type").and_then(Value::as_str) != Some("text")) { - return Err(CoreError::Unsupported("non-text response content block")); + return Err(Error::Unsupported("non-text response content block")); } let text: String = content .iter() @@ -173,7 +174,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { let usage = body .get("usage") .and_then(Value::as_object) - .ok_or(CoreError::MissingField("usage"))?; + .ok_or(Error::MissingField("usage"))?; let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0); Ok(ChatCompletionsResponse { @@ -181,7 +182,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { model: body .get("model") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("model"))? + .ok_or(Error::MissingField("model"))? .to_string(), choices: vec![ChatCompletionsChoice { index: 0, diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index 829f2260d3c..8fcc0f36c7d 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -1,4 +1,4 @@ -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; @@ -17,12 +17,12 @@ pub fn non_empty(value: Option<&str>) -> Option<&str> { pub fn resolve_anthropic_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) .ok_or_else(|| { - CoreError::Auth( + Error::Auth( "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \ environment variable" .to_string(), @@ -52,7 +52,7 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { api_base: Option<&str>, _model: &str, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { Ok(complete_anthropic_url(api_base, env_lookup)) } @@ -60,7 +60,7 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_anthropic_api_key(api_key, env_lookup) } @@ -121,7 +121,7 @@ mod tests { ); assert!(matches!( resolve_anthropic_api_key(None, &|_| None).expect_err("missing key"), - CoreError::Auth(_) + Error::Auth(_) )); } 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 7b958c77ba3..70dad0300f1 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 @@ -1,4 +1,4 @@ -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use crate::messages::types::{ AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock, @@ -28,12 +28,12 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig = pub fn resolve_azure_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty())) .ok_or_else(|| { - CoreError::Auth( + Error::Auth( "Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable" .to_string(), ) @@ -43,12 +43,12 @@ pub fn resolve_azure_api_key( pub fn complete_azure_anthropic_url( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let api_base = non_empty(api_base) .map(str::to_string) .or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty())) .ok_or_else(|| { - CoreError::Auth( + Error::Auth( "Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \ Expected format: https://.services.ai.azure.com/anthropic" .to_string(), @@ -147,7 +147,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { api_base: Option<&str>, _model: &str, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_azure_anthropic_url(api_base, env_lookup) } @@ -155,7 +155,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_azure_api_key(api_key, env_lookup) } @@ -174,7 +174,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { fn transform_request( &self, request: AnthropicMessagesRequest, - ) -> CoreResult { + ) -> Result { let mut request = fold_system_role_messages(request); if let Some(system) = request.system.as_mut() { strip_scope_from_system(system); @@ -190,7 +190,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { &self, model: &str, response: AnthropicMessagesResponse, - ) -> CoreResult { + ) -> Result { self.anthropic.transform_response(model, response) } } @@ -268,7 +268,7 @@ mod tests { "https://env.services.ai.azure.com/anthropic/v1/messages" ); let err = complete_azure_anthropic_url(Some(" "), &|_| None).expect_err("missing base"); - assert!(matches!(err, CoreError::Auth(_))); + assert!(matches!(err, Error::Auth(_))); } #[test] @@ -284,7 +284,7 @@ mod tests { ); assert!(matches!( resolve_azure_api_key(None, &|_| None).expect_err("missing key"), - CoreError::Auth(_) + Error::Auth(_) )); } 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 eabd15677cc..b26a7925e8a 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,6 +1,6 @@ use std::collections::BTreeSet; -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; @@ -32,17 +32,17 @@ fn resolve_value( env_name: &str, env_lookup: &dyn Fn(&str) -> Option, missing_message: &str, -) -> CoreResult { +) -> Result { non_empty(explicit) .map(str::to_string) .or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| CoreError::Auth(missing_message.to_string())) + .ok_or_else(|| Error::Auth(missing_message.to_string())) } pub fn resolve_azure_ai_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { resolve_value( api_key, AZURE_AI_API_KEY_ENV, @@ -54,7 +54,7 @@ pub fn resolve_azure_ai_api_key( pub fn resolve_azure_ai_api_base( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { resolve_value( api_base, AZURE_AI_API_BASE_ENV, @@ -66,7 +66,7 @@ pub fn resolve_azure_ai_api_base( pub fn complete_azure_ai_url( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let base = resolve_azure_ai_api_base(api_base, env_lookup)?; Ok(format!( "{}/providers/mistral/azure/ocr", @@ -77,7 +77,7 @@ pub fn complete_azure_ai_url( pub fn resolve_document_intelligence_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { resolve_value( api_key, AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, @@ -89,7 +89,7 @@ pub fn resolve_document_intelligence_api_key( pub fn resolve_document_intelligence_endpoint( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { resolve_value( api_base, AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, @@ -127,7 +127,7 @@ fn pages_token_is_valid(token: &str) -> bool { } } -fn normalize_pages_param(pages: &Value) -> CoreResult> { +fn normalize_pages_param(pages: &Value) -> Result, Error> { match pages { Value::String(value) => { let normalized = value @@ -138,7 +138,7 @@ fn normalize_pages_param(pages: &Value) -> CoreResult> { if normalized.split(',').all(pages_token_is_valid) { Ok(Some(normalized)) } else { - Err(CoreError::InvalidRequest(format!( + Err(Error::InvalidRequest(format!( "Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'." ))) } @@ -152,7 +152,7 @@ fn normalize_pages_param(pages: &Value) -> CoreResult> { for value in values { let page = value.as_i64().expect("checked is_i64"); if page < 0 { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(), )); } @@ -176,16 +176,16 @@ fn normalize_pages_param(pages: &Value) -> CoreResult> { if normalized.split(',').all(pages_token_is_valid) { return Ok(Some(normalized)); } - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'." ))); } - Err(CoreError::InvalidRequest( + Err(Error::InvalidRequest( "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." .to_string(), )) } - _ => Err(CoreError::InvalidRequest( + _ => Err(Error::InvalidRequest( "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." .to_string(), )), @@ -197,7 +197,7 @@ pub fn complete_document_intelligence_url( model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?; let mut url = format!( "{}/documentintelligence/documentModels/{}:analyze?api-version={}", @@ -216,20 +216,20 @@ pub fn complete_document_intelligence_url( Ok(url) } -fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { - let object = document.as_object().ok_or_else(|| CoreError::InvalidType { +fn document_url_from_mistral_document(document: &Value) -> Result<&str, Error> { + let object = document.as_object().ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(document), })?; let doc_type = object .get("type") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("document.type"))?; + .ok_or(Error::MissingField("document.type"))?; let field_name = match doc_type { "document_url" => "document_url", "image_url" => "image_url", other => { - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "Invalid document type: {other}. Must be 'document_url' or 'image_url'" ))); } @@ -238,7 +238,7 @@ fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { .get(field_name) .and_then(Value::as_str) .filter(|value| !value.is_empty()) - .ok_or(CoreError::MissingField(field_name)) + .ok_or(Error::MissingField(field_name)) } fn extract_base64_from_data_uri(data_uri: &str) -> &str { @@ -290,7 +290,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { model: &str, document: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } @@ -298,7 +298,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } @@ -308,7 +308,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_azure_ai_url(api_base, env_lookup) } @@ -316,7 +316,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_azure_ai_api_key(api_key, env_lookup) } @@ -335,7 +335,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { _model: &str, document: Value, _optional_params: Map, - ) -> CoreResult { + ) -> Result { let document_url = document_url_from_mistral_document(&document)?; let mut data = Map::new(); if document_url.starts_with("data:") { @@ -359,19 +359,19 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { let response = response_json .as_object() - .ok_or_else(|| CoreError::InvalidType { + .ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&response_json), })?; let status = response .get("status") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("status"))?; + .ok_or(Error::MissingField("status"))?; if status != "succeeded" { - return Err(CoreError::InvalidResponse(format!( + return Err(Error::InvalidResponse(format!( "Azure Document Intelligence analysis failed with status: {status}" ))); } @@ -414,7 +414,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_document_intelligence_url(api_base, model, optional_params, env_lookup) } @@ -422,7 +422,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_document_intelligence_api_key(api_key, env_lookup) } diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index 5e885734182..bb4f6afe5f9 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -6,7 +6,7 @@ use crate::audio_transcription::transformation::{ use crate::audio_transcription::types::{ AudioTranscriptionRequestData, AudioTranscriptionResponseData, }; -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region}; use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}; @@ -18,8 +18,8 @@ pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig = pub struct BedrockAudioTranscriptionConfig; -fn audio_fields(audio: Value) -> CoreResult<(String, String)> { - let object = audio.as_object().ok_or_else(|| CoreError::InvalidType { +fn audio_fields(audio: Value) -> Result<(String, String), Error> { + let object = audio.as_object().ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&audio), })?; @@ -27,13 +27,13 @@ fn audio_fields(audio: Value) -> CoreResult<(String, String)> { .get("data") .and_then(Value::as_str) .filter(|value| !value.is_empty()) - .ok_or(CoreError::MissingField("audio.data"))?; + .ok_or(Error::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()) + Error::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string()) })?; Ok((data.to_string(), format.to_string())) } @@ -55,7 +55,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { _model: &str, audio: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { 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") { @@ -87,14 +87,14 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { &self, _model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { 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()) + Error::InvalidResponse("Bedrock response has no output content".to_string()) })?; let mut text = String::new(); for block in content { @@ -111,7 +111,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { 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 @@ -133,7 +133,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { let (_, model_region) = bedrock_model_id_and_region(model); Ok(AudioTranscriptionAuth::AwsSigV4 { region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), 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 b11639aa09b..e5e52bfce95 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -4,7 +4,7 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; use crate::caching::in_memory_cache::InMemoryCache; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ @@ -197,7 +197,7 @@ pub fn classify_auth( pub async fn resolve_credentials( config: AwsAuthConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> CoreResult { +) -> Result { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); match flow { @@ -244,9 +244,10 @@ pub async fn resolve_credentials( 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}")) - }) + provider + .provide_credentials() + .await + .map_err(|error| Error::Auth(format!("AWS profile credentials failed: {error}"))) } AwsAuthFlow::AssumeRole { role, session_name } => { if is_already_running_as_role(&role, &resolved).await? { @@ -260,7 +261,7 @@ pub async fn resolve_credentials( .build() .await; let credentials = provider.provide_credentials().await.map_err(|error| { - CoreError::Auth(format!("AWS default credentials failed: {error}")) + Error::Auth(format!("AWS default credentials failed: {error}")) })?; set_cached_credentials( key, @@ -301,7 +302,7 @@ pub async fn resolve_credentials( provider .provide_credentials() .await - .map_err(|error| CoreError::Auth(format!("AWS role credentials failed: {error}"))) + .map_err(|error| Error::Auth(format!("AWS role credentials failed: {error}"))) } AwsAuthFlow::WebIdentity { token, @@ -325,13 +326,13 @@ pub async fn resolve_credentials( .send() .await .map_err(|error| { - CoreError::Auth(format!("AWS web identity credentials failed: {error}")) + Error::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()) + Error::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}")) + Error::Auth(format!("AWS web identity expiration was invalid: {error}")) })?; Ok(Credentials::new( credentials.access_key_id(), @@ -350,9 +351,10 @@ pub async fn resolve_credentials( 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}")) - })?; + let credentials = provider + .provide_credentials() + .await + .map_err(|error| Error::Auth(format!("AWS default credentials failed: {error}")))?; set_cached_credentials( key, credentials.clone(), @@ -363,7 +365,7 @@ pub async fn resolve_credentials( } } -async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> CoreResult { +async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result { if role_identity(role).is_none() { return Ok(false); } @@ -437,7 +439,7 @@ pub fn sign_bedrock_post( region: &str, credentials: &Credentials, signing_time: SystemTime, -) -> CoreResult> { +) -> Result, Error> { let identity: Identity = credentials.clone().into(); let params = v4::SigningParams::builder() .identity(&identity) @@ -447,14 +449,14 @@ pub fn sign_bedrock_post( .settings(SigningSettings::default()) .build() .map(SigningParams::from) - .map_err(|error| CoreError::Auth(format!("AWS signing parameters failed: {error}")))?; + .map_err(|error| Error::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}")))?; + .map_err(|error| Error::Auth(format!("AWS signable request failed: {error}")))?; let (instructions, _) = sign(request, ¶ms) - .map_err(|error| CoreError::Auth(format!("AWS request signing failed: {error}")))? + .map_err(|error| Error::Auth(format!("AWS request signing failed: {error}")))? .into_parts(); Ok(instructions .headers() diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs index 4b75dcb8e9d..c86f061b9ca 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::Error; use serde_json::json; fn messages(value: Value) -> Vec { @@ -23,7 +24,7 @@ fn transform(msgs: Value, opts: Value) -> Value { .body } -fn transform_response(body: Value) -> CoreResult { +fn transform_response(body: Value) -> Result { BEDROCK_CHAT_COMPLETIONS_CONFIG.transform_response( "anthropic.claude-sonnet-4-5-v1:0", ProviderChatResponseData { body }, @@ -478,25 +479,22 @@ fn declines_a_response_carrying_a_tool_use_block() { "usage": {"inputTokens": 1, "outputTokens": 1} })) .expect_err("tool use block"); - assert_eq!( - err, - CoreError::Unsupported("non-text response content block") - ); + assert_eq!(err, Error::Unsupported("non-text response content block")); } #[test] fn errors_on_a_response_missing_required_fields() { assert_eq!( transform_response(json!("nope")).expect_err("not an object"), - CoreError::InvalidResponse("converse response is not an object".to_string()) + Error::InvalidResponse("converse response is not an object".to_string()) ); assert_eq!( transform_response(json!({"usage": {}})).expect_err("no output"), - CoreError::MissingField("output.message.content") + Error::MissingField("output.message.content") ); assert_eq!( transform_response(json!({"output": {"message": {"content": []}}})).expect_err("no usage"), - CoreError::MissingField("usage") + Error::MissingField("usage") ); } diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index b107950748e..ef5f44b4a14 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -11,7 +11,7 @@ use crate::chat_completions::types::{ ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData, ProviderChatResponseData, }; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region}; use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}; @@ -110,7 +110,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { 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 @@ -137,7 +137,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { // Python reads `api_key` as the Bedrock bearer token and consults the // env only when the caller passed none, so a caller-supplied empty key // falls through to SigV4 without reaching for the environment. An @@ -208,7 +208,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { _model: &str, messages: Vec, optional_params: Map, - ) -> CoreResult { + ) -> Result { Ok(ProviderChatRequestData { body: converse_body(&build_conversation(&messages), &optional_params), }) @@ -218,17 +218,18 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { &self, model: &str, response: ProviderChatResponseData, - ) -> CoreResult { - let body = response.body.as_object().ok_or_else(|| { - CoreError::InvalidResponse("converse response is not an object".into()) - })?; + ) -> Result { + let body = response + .body + .as_object() + .ok_or_else(|| Error::InvalidResponse("converse response is not an object".into()))?; let content = body .get("output") .and_then(|output| output.get("message")) .and_then(|message| message.get("content")) .and_then(Value::as_array) - .ok_or(CoreError::MissingField("output.message.content"))?; + .ok_or(Error::MissingField("output.message.content"))?; // The route declines tool requests, so anything other than a text block // is something this path never asked for. Decline; the host falls back. if content.iter().any(|block| { @@ -236,7 +237,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { .as_object() .is_none_or(|block| block.len() != 1 || !block.contains_key("text")) }) { - return Err(CoreError::Unsupported("non-text response content block")); + return Err(Error::Unsupported("non-text response content block")); } let text: String = content .iter() @@ -246,7 +247,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { let usage = body .get("usage") .and_then(Value::as_object) - .ok_or(CoreError::MissingField("usage"))?; + .ok_or(Error::MissingField("usage"))?; let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0); let computed = usage_from_parts( field("inputTokens"), 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 dc720cc4244..6a8a38204a9 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::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value}; @@ -47,7 +47,7 @@ pub fn complete_url(api_base: Option<&str>) -> String { /// Resolve the Mistral API key from the explicit param or the environment. /// -/// Blank/whitespace values are treated as absent. Returns `CoreError::Auth` +/// Blank/whitespace values are treated as absent. Returns `Error::Auth` /// when no usable key is available. /// /// Note: the env fallback only reads the process environment. Secret-manager @@ -56,13 +56,13 @@ pub fn complete_url(api_base: Option<&str>) -> String { pub fn resolve_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { api_key .map(str::trim) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) } pub struct MistralOcrConfig; @@ -79,9 +79,9 @@ impl OcrProviderConfig for MistralOcrConfig { model: &str, document: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { if !document.is_object() { - return Err(CoreError::InvalidType { + return Err(Error::InvalidType { expected: "object", actual: json_type_name(&document), }); @@ -104,10 +104,10 @@ impl OcrProviderConfig for MistralOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { let response_object = response_json .as_object() - .ok_or_else(|| CoreError::InvalidType { + .ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&response_json), })?; @@ -140,7 +140,7 @@ impl OcrProviderConfig for MistralOcrConfig { _model: &str, _optional_params: &Map, _env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { Ok(complete_url(api_base)) } @@ -148,7 +148,7 @@ impl OcrProviderConfig for MistralOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_api_key(api_key, env_lookup) } } @@ -165,11 +165,11 @@ pub fn transform_ocr_request( model: &str, document: Value, optional_params: Map, -) -> CoreResult { +) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } -pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult { +pub fn transform_ocr_response(model: &str, response_json: Value) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } @@ -250,7 +250,7 @@ mod tests { assert_eq!( err, - CoreError::InvalidType { + Error::InvalidType { expected: "object", actual: "string", } @@ -307,6 +307,6 @@ mod tests { #[test] fn resolve_api_key_errors_when_absent() { let err = resolve_api_key(None, &|_| None).expect_err("missing key should error"); - assert_eq!(err, CoreError::Auth(MISSING_KEY_MESSAGE.to_string())); + assert_eq!(err, Error::Auth(MISSING_KEY_MESSAGE.to_string())); } } 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 b3f6b03b28a..f1985f81b7d 100644 --- a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs @@ -1,4 +1,4 @@ -use crate::CoreResult; +use crate::Error; use crate::realtime::transformation::RealtimeProviderConfig; use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; @@ -72,7 +72,7 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig { &self, event: &RealtimeEvent, _model: &str, - ) -> CoreResult { + ) -> Result { Ok(RealtimeTransformResult::passthrough(event.clone())) } @@ -80,7 +80,7 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig { &self, event: &RealtimeEvent, _model: &str, - ) -> CoreResult { + ) -> Result { Ok(RealtimeTransformResult::passthrough(event.clone())) } } @@ -88,14 +88,14 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig { pub fn transform_realtime_request( event: &RealtimeEvent, model: &str, -) -> CoreResult { +) -> Result { OPENAI_REALTIME_CONFIG.transform_realtime_request(event, model) } pub fn transform_realtime_response( event: &RealtimeEvent, model: &str, -) -> CoreResult { +) -> Result { OPENAI_REALTIME_CONFIG.transform_realtime_response(event, model) } 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 e15197c468c..be86bb90311 100644 --- a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs @@ -1,4 +1,4 @@ -use crate::CoreResult; +use crate::Error; use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model}; @@ -15,7 +15,7 @@ impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig { &self, event: &ResponsesWsEvent, model: &str, - ) -> CoreResult { + ) -> Result { Ok(ResponsesWsTransformResult::passthrough(enforce_model( event, model, ))) @@ -25,7 +25,7 @@ impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig { &self, event: &ResponsesWsEvent, _model: &str, - ) -> CoreResult { + ) -> Result { Ok(ResponsesWsTransformResult::passthrough(event.clone())) } } 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 6300149c237..ee095447028 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,4 +1,4 @@ -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; @@ -43,7 +43,7 @@ pub fn is_deepseek_model(model: &str) -> bool { pub fn resolve_vertex_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { api_key .map(str::trim) .filter(|key| !key.is_empty()) @@ -51,7 +51,7 @@ pub fn resolve_vertex_api_key( .or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) .ok_or_else(|| { - CoreError::Auth( + Error::Auth( "Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers" .to_string(), ) @@ -61,12 +61,12 @@ pub fn resolve_vertex_api_key( fn vertex_project( params: &Map, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { string_param(params, &["vertex_project", "vertex_ai_project"]) .map(str::to_string) .or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty())) .ok_or_else(|| { - CoreError::InvalidRequest( + Error::InvalidRequest( "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" .to_string(), ) @@ -99,7 +99,7 @@ pub fn complete_vertex_mistral_url( model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let project = vertex_project(optional_params, env_lookup)?; let location = vertex_location(optional_params, env_lookup); let base = vertex_mistral_api_base(api_base, &location); @@ -112,7 +112,7 @@ pub fn complete_vertex_deepseek_url( api_base: Option<&str>, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let project = vertex_project(optional_params, env_lookup)?; let location = vertex_location(optional_params, env_lookup); let base = api_base @@ -125,20 +125,20 @@ pub fn complete_vertex_deepseek_url( )) } -fn document_content_item(document: &Value) -> CoreResult { - let object = document.as_object().ok_or_else(|| CoreError::InvalidType { +fn document_content_item(document: &Value) -> Result { + let object = document.as_object().ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(document), })?; let doc_type = object .get("type") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("document.type"))?; + .ok_or(Error::MissingField("document.type"))?; let url_field = match doc_type { "image_url" => "image_url", "document_url" => "document_url", other => { - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "Unsupported document type: {other}. Expected 'image_url' or 'document_url'" ))); } @@ -147,7 +147,7 @@ fn document_content_item(document: &Value) -> CoreResult { .get(url_field) .and_then(Value::as_str) .filter(|value| !value.is_empty()) - .ok_or(CoreError::MissingField(url_field))?; + .ok_or(Error::MissingField(url_field))?; Ok(json!({ "type": "image_url", @@ -163,7 +163,7 @@ fn deepseek_model_name(model: &str) -> String { } } -fn first_choice_content(response: &Value) -> CoreResult { +fn first_choice_content(response: &Value) -> Result { response .get("choices") .and_then(Value::as_array) @@ -176,9 +176,7 @@ fn first_choice_content(response: &Value) -> CoreResult { Value::Object(_) => true, _ => false, }) - .ok_or_else(|| { - CoreError::InvalidResponse("No content in DeepSeek OCR response".to_string()) - }) + .ok_or_else(|| Error::InvalidResponse("No content in DeepSeek OCR response".to_string())) } fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> Value { @@ -219,7 +217,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { model: &str, document: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } @@ -227,7 +225,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } @@ -237,7 +235,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_vertex_mistral_url(api_base, model, optional_params, env_lookup) } @@ -245,7 +243,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_vertex_api_key(api_key, env_lookup) } @@ -264,7 +262,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { model: &str, document: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { let mut data = Map::new(); data.insert( "model".to_string(), @@ -289,10 +287,10 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { let response = response_json .as_object() - .ok_or_else(|| CoreError::InvalidType { + .ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&response_json), })?; @@ -314,7 +312,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { }); } - let object = ocr_data.as_object().ok_or_else(|| CoreError::InvalidType { + let object = ocr_data.as_object().ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&ocr_data), })?; @@ -346,7 +344,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { _model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_vertex_deepseek_url(api_base, optional_params, env_lookup) } @@ -354,7 +352,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_vertex_api_key(api_key, env_lookup) } } diff --git a/litellm-rust/crates/core/src/realtime/transformation.rs b/litellm-rust/crates/core/src/realtime/transformation.rs index 69b88687000..b08084514ef 100644 --- a/litellm-rust/crates/core/src/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/realtime/transformation.rs @@ -1,4 +1,4 @@ -use crate::CoreResult; +use crate::Error; use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; pub trait RealtimeProviderConfig { @@ -11,12 +11,12 @@ pub trait RealtimeProviderConfig { &self, event: &RealtimeEvent, model: &str, - ) -> CoreResult; + ) -> Result; /// Transform a backend → client event before it is forwarded downstream. fn transform_realtime_response( &self, event: &RealtimeEvent, model: &str, - ) -> CoreResult; + ) -> Result; } diff --git a/litellm-rust/crates/core/src/responses/instrumentation.rs b/litellm-rust/crates/core/src/responses/instrumentation.rs index ec04571da14..b1098f4d386 100644 --- a/litellm-rust/crates/core/src/responses/instrumentation.rs +++ b/litellm-rust/crates/core/src/responses/instrumentation.rs @@ -5,9 +5,9 @@ use std::time::{SystemTime, UNIX_EPOCH}; use serde_json::Value; +use crate::Error; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType}; -use crate::{CoreError, CoreResult}; #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct ResponsesWsUsage { @@ -205,7 +205,7 @@ impl ResponsesWsInstrumentation { } } -type LifecycleFuture<'a, T> = Pin> + Send + 'a>>; +type LifecycleFuture<'a, T> = Pin> + Send + 'a>>; impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation { type PreCallFuture<'a> = LifecycleFuture<'a, ()>; @@ -246,7 +246,7 @@ impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation { fn async_log_failure_event<'a>( &'a self, _context: &'a CallLifecycleContext, - _error: &'a CoreError, + _error: &'a Error, _timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -342,7 +342,7 @@ mod tests { ), (), &instrumentation, - |_| async { Ok::<(), CoreError>(()) }, + |_| async { Ok::<(), Error>(()) }, ) .await; diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index 92dc19627a0..5d037e9cf1b 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -1,4 +1,4 @@ -use crate::CoreResult; +use crate::Error; use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult}; @@ -19,13 +19,13 @@ pub trait ResponsesWebSocketProviderConfig: Sync { &self, event: &ResponsesWsEvent, model: &str, - ) -> CoreResult; + ) -> Result; fn transform_ws_response( &self, event: &ResponsesWsEvent, model: &str, - ) -> CoreResult; + ) -> Result; } pub fn complete_websocket_url( diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 746e0770f9b..68aa9436b15 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -10,7 +10,7 @@ use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompleti use litellm_core::chat_completions::{ chat_completions as run_chat_completions, chat_completions_decline_reason, }; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; use litellm_python_interop::{from_py, release_count, release_gil, to_py}; @@ -54,13 +54,13 @@ fn chat_completions_response_to_py( to_py(py, &response) } -fn core_error_to_pyerr(err: CoreError) -> PyErr { +fn core_error_to_pyerr(err: Error) -> PyErr { match err { - CoreError::Auth(message) => PyValueError::new_err(message), - CoreError::InvalidProvider(_) - | CoreError::InvalidRequest(_) - | CoreError::InvalidType { .. } - | CoreError::MissingField(_) => PyValueError::new_err(err.to_string()), + Error::Auth(message) => PyValueError::new_err(message), + Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) => PyValueError::new_err(err.to_string()), other => PyRuntimeError::new_err(other.to_string()), } } @@ -71,22 +71,22 @@ fn core_error_to_pyerr(err: CoreError) -> PyErr { /// Everything raised before the request goes out is safe for the host to retry /// on its own path; anything after it is not, because the provider has already /// done the work and billed for it. -fn chat_completions_error_to_pyerr(err: CoreError) -> PyErr { +fn chat_completions_error_to_pyerr(err: Error) -> PyErr { match err { - CoreError::Unsupported(_) - | CoreError::Auth(_) - | CoreError::InvalidProvider(_) - | CoreError::InvalidRequest(_) - | CoreError::InvalidType { .. } - | CoreError::MissingField(_) - | CoreError::Routing(_) + Error::Unsupported(_) + | Error::Auth(_) + | Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) + | Error::Routing(_) // Nothing reached the provider, so serving it on Python cannot double // bill and is the only way the caller gets an answer at all. - | CoreError::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - CoreError::Http { status, body } => { + | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), + Error::Http { status, body } => { RustUpstreamError::new_err((status, format!("{status}: {body}"))) } - CoreError::Network(message) | CoreError::InvalidResponse(message) => { + Error::Network(message) | Error::InvalidResponse(message) => { RustUpstreamError::new_err((0u16, message)) } } From 9bd870d47a700b183e0ee0d9bf7647aa7c739561 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:36:28 -0700 Subject: [PATCH 46/81] fix(databricks): upgrade legacy thinking to adaptive on adaptive-only Claude models --- .../llms/databricks/chat/transformation.py | 4 ++++ .../test_databricks_chat_transformation.py | 21 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index c587146005f..65622d62af2 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -330,6 +330,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): ) -> dict: is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) mapped_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) + if "claude" in model: + AnthropicConfig.translate_legacy_thinking_for_adaptive_model( + model=model, optional_params=mapped_params, custom_llm_provider="databricks" + ) if "tools" in mapped_params: mapped_params["tools"] = self._map_openai_to_dbrx_tool(model=model, tools=mapped_params["tools"]) if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens: diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 41fb2589655..71661cc532b 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -422,6 +422,27 @@ def test_databricks_config_probes_capabilities_under_databricks_namespace(): assert DatabricksConfig().custom_llm_provider == "databricks" +@pytest.mark.parametrize( + "model, expected_thinking, expected_output_config", + [ + ("databricks-claude-opus-4-8", {"type": "adaptive"}, {"effort": "high"}), + ("databricks-claude-opus-4-6", {"type": "enabled", "budget_tokens": 4096}, None), + ], + ids=["adaptive_only_upgrades_to_adaptive", "legacy_capable_forwards_verbatim"], +) +def test_map_openai_params_upgrades_legacy_thinking_on_adaptive_only_claude( + model, expected_thinking, expected_output_config +): + mapped = DatabricksConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4096}}, + optional_params={}, + model=model, + drop_params=False, + ) + assert mapped["thinking"] == expected_thinking + assert mapped.get("output_config") == expected_output_config + + def _streaming_chunk(usage=None, choices=None): base = { "id": "chatcmpl-test", From 3af19cbf61706f1cf3d3360b8802b5cbc1991827 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:36:29 -0700 Subject: [PATCH 47/81] test(proxy-extras): use the modern optional annotation in the deploy budget test --- tests/proxy_migration_tests/test_prisma_toolchain.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 4c258c4d007..733870f3239 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -22,7 +22,6 @@ import sys import time from collections.abc import Callable from pathlib import Path -from typing import Optional import pytest @@ -302,7 +301,7 @@ def test_db_push_timeout_hint_names_the_per_command_budget( ids=["raised_command_budget_carries_over", "lowered_command_budget_does_not", "override_wins_upward", "override_wins_downward"], ) def test_migrate_deploy_budget_keeps_a_raised_command_budget( - command_timeout: str, deploy_timeout: Optional[str], expected: float, monkeypatch: pytest.MonkeyPatch + command_timeout: str, deploy_timeout: str | None, expected: float, monkeypatch: pytest.MonkeyPatch ) -> None: """Deployments that raised the per-command budget to survive a long deploy keep that budget for deploy.""" monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, command_timeout) From 034ff5855802e1b2b036d4cee6208a3efd8a15fd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:36:32 -0700 Subject: [PATCH 48/81] test(otel): assert Langfuse logger behavior instead of its class --- .../integrations/otel/test_langfuse_logger.py | 30 +++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/integrations/otel/test_langfuse_logger.py b/tests/test_litellm/integrations/otel/test_langfuse_logger.py index af0597517dc..3e35395389e 100644 --- a/tests/test_litellm/integrations/otel/test_langfuse_logger.py +++ b/tests/test_litellm/integrations/otel/test_langfuse_logger.py @@ -12,9 +12,9 @@ pytest.importorskip("opentelemetry") from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter # noqa: E402 +import litellm # noqa: E402 from litellm.caching.dual_cache import DualCache # noqa: E402 -from litellm.integrations.otel.langfuse_logger import LangfuseOpenTelemetryV2 # noqa: E402 -from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger # noqa: E402 +from litellm.integrations.otel.logger import build_otel_v2_logger # noqa: E402 from litellm.integrations.otel.model.config import OpenTelemetryV2Config, is_otel_v2_enabled # noqa: E402 from litellm.integrations.otel.model.spans import LITELLM_PROXY_REQUEST_SPAN_NAME, SpanRole # noqa: E402 from litellm.integrations.otel.plumbing import context as otel_context # noqa: E402 @@ -22,6 +22,7 @@ from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.plumbing.context import set_request_root_span # noqa: E402 from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 # noqa: E402 from litellm.proxy._types import UserAPIKeyAuth # noqa: E402 +from litellm.proxy.utils import ProxyLogging # noqa: E402 from litellm.types.llms.openai import ( # noqa: E402 ResponseCompletedEvent, ResponsesAPIResponse, @@ -294,13 +295,29 @@ def test_unrenderable_output_never_raises_into_the_request(): def test_factory_keeps_the_base_logger_unless_langfuse_content_capture_is_on(capture, mappers): logger, exporter = _logger(capture=capture, mappers=mappers) - assert type(logger) is OpenTelemetryV2 _run_request(logger, CHAT_DATA, "acompletion", ModelResponse()) attrs = _root_attrs(exporter) assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs -def test_langfuse_otel_preset_builds_the_langfuse_logger(monkeypatch): +@pytest.mark.parametrize( + ("capture", "mappers", "relays_streams"), + [ + ("span_only", ("genai", "langfuse"), True), + ("no_content", ("genai", "langfuse"), False), + ("span_only", ("genai",), False), + ], +) +def test_only_langfuse_content_capture_takes_proxy_streams_off_the_fast_path( + monkeypatch, capture, mappers, relays_streams +): + logger, _ = _logger(capture=capture, mappers=mappers) + monkeypatch.setattr(litellm, "callbacks", [logger]) + + assert ProxyLogging._callback_capabilities().has_iterator_override is relays_streams + + +def test_langfuse_otel_preset_builds_a_logger_that_stamps_the_root(monkeypatch): monkeypatch.setenv("LITELLM_OTEL_V2", "true") monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk") monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk") @@ -311,7 +328,10 @@ def test_langfuse_otel_preset_builds_the_langfuse_logger(monkeypatch): loggers: list = [] try: built = _maybe_construct_otel_v2("langfuse_otel", loggers) - assert isinstance(built, LangfuseOpenTelemetryV2) + assert built is not None assert _maybe_construct_otel_v2("langfuse_otel", loggers) is built + root = _start_root(built) + asyncio.run(built.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion")) + assert INPUT_ATTR in dict(root.attributes or {}) finally: is_otel_v2_enabled.cache_clear() From cde9d94c3650dfbc0b704280973d6e4319a7530a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:40:24 -0700 Subject: [PATCH 49/81] feat(agentcore-a2a): derive runtime session id from A2A message.contextId (#39371) Native AgentCore A2A always sent either a fresh generated runtime session id or the single configured runtimeSessionId, so related turns lost context and unrelated callers shared one AgentCore microVM. The runtime session id is now params.message.contextId scoped to the calling key hash, then runtimeSessionId, then generated, and is length-validated (33-256) before the header is signed. Invalid ids surface as JSON-RPC -32602 / HTTP 400 instead of a 500. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock_agentcore/transformation.py | 51 ++++- litellm/a2a_protocol/utils.py | 25 +++ litellm/llms/langflow/a2a.py | 27 +-- .../proxy/agent_endpoints/a2a_endpoints.py | 2 + .../test_bedrock_agentcore_a2a.py | 191 ++++++++++++++++++ .../agent_endpoints/test_a2a_endpoints.py | 56 +++++ 6 files changed, 325 insertions(+), 27 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 32252711997..1e8cc4ff90e 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -10,8 +10,19 @@ from collections.abc import AsyncIterator, Mapping from typing import Any, Final from litellm._logging import verbose_logger +from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, +) +from litellm.a2a_protocol.utils import ( + get_session_id_from_a2a_params, + scope_session_to_principal, +) +from litellm.exceptions import BadRequestError from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig +RUNTIME_SESSION_ID_MIN_LENGTH: Final = 33 +RUNTIME_SESSION_ID_MAX_LENGTH: Final = 256 + # Reserved outbound header names that must never be sourced from per-request # ``agent_extra_headers`` for AgentCore requests. ``agent_extra_headers`` carries # values rewritten from the client-controlled ``x-a2a-{agent}-*`` convention, so @@ -19,8 +30,9 @@ from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreCo # request identity / SigV4 metadata by overwriting headers the proxy sets from # trusted server-side config. # -# The runtime headers (session / user id) are derived server-side from -# ``runtimeSessionId`` / ``runtimeUserId`` in the agent's ``litellm_params``; +# The runtime headers (session / user id) are derived server-side from the A2A +# ``message.contextId`` and ``runtimeSessionId`` / ``runtimeUserId`` in the +# agent's ``litellm_params``; # ``authorization`` is set by the AgentCore signer (JWT or SigV4); ``host`` and # the ``x-amz-*`` family are owned by SigV4 itself. _RESERVED_EXACT_HEADERS: Final = frozenset( @@ -66,6 +78,31 @@ def _filter_reserved_headers( return filtered or None +def _request_scoped_runtime_session_id( + params: Mapping[str, Any], + litellm_params: Mapping[str, Any], +) -> str | None: + context_id: Final = get_session_id_from_a2a_params(params) + if not isinstance(context_id, str) or not context_id: + return None + return scope_session_to_principal(context_id, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)) + + +def _validate_runtime_session_id(session_id: str, model: str) -> str: + if RUNTIME_SESSION_ID_MIN_LENGTH <= len(session_id) <= RUNTIME_SESSION_ID_MAX_LENGTH: + return session_id + raise BadRequestError( + message=( + f"Invalid AgentCore runtime session id {session_id!r}: AWS requires " + f"{RUNTIME_SESSION_ID_MIN_LENGTH}-{RUNTIME_SESSION_ID_MAX_LENGTH} characters. It is built from the A2A " + "message.contextId (prefixed with a 16-hex-char hash of the calling key and '-') when set, " + "otherwise from the agent's configured runtimeSessionId." + ), + model=model, + llm_provider="bedrock", + ) + + class BedrockAgentCoreA2ATransformation: """ Request/response transformation for Bedrock AgentCore A2A agents. @@ -100,7 +137,9 @@ class BedrockAgentCoreA2ATransformation: here to prevent a caller-controlled ``x-a2a-{agent}-*`` header from spoofing the AgentCore runtime user id or other SigV4 metadata. Use ``api_key`` / ``runtimeUserId`` / ``runtimeSessionId`` in litellm_params - (not ``agent_extra_headers``) to override those values. + (not ``agent_extra_headers``) to override those values. The runtime + session id is taken from ``params["message"]["contextId"]`` (scoped to + the calling key) when present, then ``runtimeSessionId``, else generated. Returns: Tuple of (url, signed_headers, signed_body_bytes) @@ -139,7 +178,11 @@ class BedrockAgentCoreA2ATransformation: # Set required AgentCore session headers (normally set by transform_request, # which we skip because it also builds {"prompt": "..."}) headers: Final[dict] = {} - session_id: Final = agentcore_config._get_runtime_session_id(optional_params) + session_id: Final = _validate_runtime_session_id( + _request_scoped_runtime_session_id(params, litellm_params) + or agentcore_config._get_runtime_session_id(optional_params), + model=model, + ) headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = session_id runtime_user_id: Final = agentcore_config._get_runtime_user_id(optional_params) if runtime_user_id: diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index f2e61f66105..7c459daf720 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -2,6 +2,8 @@ Utility functions for A2A protocol. """ +import hashlib +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import litellm @@ -140,6 +142,29 @@ class A2ARequestUtils: return prompt_tokens, completion_tokens, total_tokens +def get_session_id_from_a2a_params(params: Mapping[str, Any]) -> str | None: + message: Final = params.get("message", {}) + if isinstance(message, dict): + return message.get("contextId") + return getattr(message, "contextId", None) + + +def scope_session_to_principal(session_id: str, principal: str | None) -> str: + """ + Bind a client-supplied A2A contextId to the authenticated principal. + + Without this, two distinct keys authorized for the same agent could set the + same contextId and read/append to each other's backend memory. The + principal is hashed (it is already a hashed token) so the raw value is never + sent to the agent backend, while the original contextId is kept as a suffix + for operator-side correlation. + """ + if not principal: + return session_id + principal_prefix: Final = hashlib.sha256(principal.encode("utf-8")).hexdigest()[:16] + return f"{principal_prefix}-{session_id}" + + # Backwards compatibility aliases def extract_text_from_a2a_message(message: Any) -> str: return A2ARequestUtils.extract_text_from_message(message) diff --git a/litellm/llms/langflow/a2a.py b/litellm/llms/langflow/a2a.py index cae750d586e..060dc0a4d05 100644 --- a/litellm/llms/langflow/a2a.py +++ b/litellm/llms/langflow/a2a.py @@ -1,28 +1,9 @@ -import hashlib from typing import Any, Final - -def get_session_id_from_a2a_params(params: dict[str, Any]) -> str | None: - message: Final = params.get("message", {}) - if isinstance(message, dict): - return message.get("contextId") - return getattr(message, "contextId", None) - - -def scope_session_to_principal(session_id: str, principal: str | None) -> str: - """ - Bind a client-supplied A2A contextId to the authenticated principal. - - Without this, two distinct keys authorized for the same LangFlow agent could - set the same contextId and read/append to each other's LangFlow memory. The - principal is hashed (it is already a hashed token) so the raw value is never - sent to the LangFlow backend, while the original contextId is kept as a - suffix for operator-side correlation. - """ - if not principal: - return session_id - principal_prefix: Final = hashlib.sha256(principal.encode("utf-8")).hexdigest()[:16] - return f"{principal_prefix}-{session_id}" +from litellm.a2a_protocol.utils import ( + get_session_id_from_a2a_params, + scope_session_to_principal, +) def merge_a2a_session_into_litellm_params( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 31b05320cd3..28882484db4 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -1019,4 +1019,6 @@ async def invoke_agent_a2a( ) except Exception: pass + if isinstance(e, litellm.BadRequestError): + return _jsonrpc_error(body.get("id"), -32602, e.message, 400) return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500) diff --git a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index 5503a5668bf..a8fe464ec32 100644 --- a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -11,7 +11,9 @@ Verifies that: import json +import httpx import pytest +import respx from unittest.mock import AsyncMock, MagicMock, patch @@ -295,6 +297,195 @@ class TestTransformation: assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") +SESSION_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" +CONTEXT_ID = "conversation-alpha-0001-0000000000000000" +KEY_HASH = "hashed-key-of-caller-one" + + +def _params_with_context(context_id: object) -> dict: + return {"message": {**SAMPLE_PARAMS["message"], "contextId": context_id}} + + +def _scoped(context_id: str, key_hash: str) -> str: + import hashlib + + return f"{hashlib.sha256(key_hash.encode()).hexdigest()[:16]}-{context_id}" + + +def _session_header(params: dict, litellm_params: dict) -> str: + from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( + BedrockAgentCoreA2ATransformation, + ) + + _, headers, _ = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + request_id="req-001", + params=params, + litellm_params=litellm_params, + ) + return headers[SESSION_HEADER] + + +@pytest.fixture +def httpx_transport(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +class TestRequestScopedRuntimeSession: + """message.contextId selects the AgentCore runtime session, scoped to the calling key.""" + + def test_context_id_scoped_to_calling_key(self): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + + litellm_params = {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH} + assert _session_header(_params_with_context(CONTEXT_ID), litellm_params) == _scoped(CONTEXT_ID, KEY_HASH) + + def test_context_id_used_verbatim_without_principal(self): + assert _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS) == CONTEXT_ID + + def test_same_context_id_reuses_session_and_other_context_isolated(self): + first = _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS) + second = _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS) + other = _session_header( + _params_with_context("conversation-beta-00002-0000000000000000"), + SAMPLE_LITELLM_PARAMS, + ) + assert first == second + assert other != first + + def test_same_context_id_from_different_keys_is_isolated(self): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + + params = _params_with_context(CONTEXT_ID) + caller_one = _session_header(params, {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH}) + caller_two = _session_header( + params, {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: "hashed-key-of-caller-two"} + ) + assert caller_one != caller_two + assert caller_one.endswith(f"-{CONTEXT_ID}") + assert caller_two.endswith(f"-{CONTEXT_ID}") + + def test_context_id_takes_precedence_over_configured_session(self): + litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40} + assert _session_header(_params_with_context(CONTEXT_ID), litellm_params) == CONTEXT_ID + + def test_configured_session_is_fallback_without_context_id(self): + litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40} + assert _session_header(SAMPLE_PARAMS, litellm_params) == "a" * 40 + assert _session_header(_params_with_context(""), litellm_params) == "a" * 40 + + def test_no_context_id_and_no_config_generates_new_session_per_request(self): + first = _session_header(SAMPLE_PARAMS, SAMPLE_LITELLM_PARAMS) + second = _session_header(SAMPLE_PARAMS, SAMPLE_LITELLM_PARAMS) + assert first != second + assert 33 <= len(first) <= 256 + + @pytest.mark.parametrize( + "context_id", + [ + "short-context-id", + "x" * 257, + ], + ) + def test_invalid_context_id_rejected_with_clear_error(self, context_id): + import litellm + + with pytest.raises(litellm.BadRequestError, match="Invalid AgentCore runtime session id") as exc_info: + _session_header(_params_with_context(context_id), SAMPLE_LITELLM_PARAMS) + assert exc_info.value.status_code == 400 + assert "33-256" in str(exc_info.value) + + def test_scoped_context_id_shorter_than_33_rejected(self): + import litellm + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + + litellm_params = {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH} + with pytest.raises(litellm.BadRequestError, match=_scoped("c" * 15, KEY_HASH)): + _session_header(_params_with_context("c" * 15), litellm_params) + assert _session_header(_params_with_context("c" * 16), litellm_params) == _scoped("c" * 16, KEY_HASH) + + def test_invalid_configured_session_rejected(self): + import litellm + + litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "too-short"} + with pytest.raises(litellm.BadRequestError, match="Invalid AgentCore runtime session id"): + _session_header(SAMPLE_PARAMS, litellm_params) + + def test_non_string_context_id_falls_back(self): + litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40} + assert _session_header(_params_with_context(12345), litellm_params) == "a" * 40 + + def test_spoofed_session_header_does_not_override_context_id(self): + from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( + BedrockAgentCoreA2ATransformation, + ) + + _, headers, _ = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + request_id="req-001", + params=_params_with_context(CONTEXT_ID), + litellm_params=SAMPLE_LITELLM_PARAMS, + agent_extra_headers={SESSION_HEADER: "s" * 40}, + ) + assert headers[SESSION_HEADER] == CONTEXT_ID + + @pytest.mark.asyncio + async def test_context_id_session_header_on_outbound_non_streaming_post(self, httpx_transport): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + with respx.mock(assert_all_called=True) as router: + route = router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response(200, json={"jsonrpc": "2.0", "id": "req-001", "result": {}}) + ) + await BedrockAgentCoreA2AConfig().handle_non_streaming( + request_id="req-001", + params=_params_with_context(CONTEXT_ID), + litellm_params={**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH}, + ) + + assert route.calls.last.request.headers[SESSION_HEADER] == _scoped(CONTEXT_ID, KEY_HASH) + + @pytest.mark.asyncio + async def test_context_id_session_header_on_outbound_streaming_post(self, httpx_transport): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + with respx.mock(assert_all_called=True) as router: + route = router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response(200, json={"jsonrpc": "2.0", "id": "req-001", "result": {}}) + ) + events = [ + event + async for event in BedrockAgentCoreA2AConfig().handle_streaming( + request_id="req-001", + params=_params_with_context(CONTEXT_ID), + litellm_params={**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH}, + ) + ] + + assert events == [{"jsonrpc": "2.0", "id": "req-001", "result": {}}] + assert route.calls.last.request.headers[SESSION_HEADER] == _scoped(CONTEXT_ID, KEY_HASH) + + class TestNonStreaming: """Test end-to-end non-streaming flow.""" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 2ff38af80b1..43034f889f6 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -918,6 +918,62 @@ async def test_task_method_failure_hook_uses_enriched_request_data(): assert failure_data.get("agent_id") == "test-agent" +@pytest.mark.asyncio +async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400(): + from litellm.proxy._types import UserAPIKeyAuth + + agent = _make_agent_mock() + agent.litellm_params = { + "custom_llm_provider": "bedrock", + "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/demo", + "api_key": "test-jwt-token", + } + mock_request = _make_request_mock( + "message/send", + { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-1", + "contextId": "too-short", + } + }, + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") + + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + stack.enter_context( + patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook test uses; no HTTP call is made because the request is rejected before signing + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert response.status_code == 400 + assert body["id"] == "req-1" + assert body["error"]["code"] == -32602 + assert "Invalid AgentCore runtime session id" in body["error"]["message"] + assert "Internal error" not in body["error"]["message"] + mock_proxy_logging.post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_extended_agent_card_rewrites_url(): from litellm.proxy._types import UserAPIKeyAuth From 987ab769213527d9bd1cfaeb3b038674dbec40f4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:40:59 -0700 Subject: [PATCH 50/81] fix(proxy): share per-model budget counters across replicas through the spend counter cache (#39375) * fix(proxy): share per-model budget counters across replicas through the spend counter cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): keep the shared fake Redis store immutable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/model_max_budget_limiter.py | 35 +++++--- litellm/proxy/proxy_server.py | 2 +- ...test_unit_test_max_model_budget_limiter.py | 85 +++++++++++++++++++ .../proxy/test_redis_auth_cache_flag.py | 24 +++++- 4 files changed, 132 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index c5d10b2749b..efaaab277a9 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,6 +1,6 @@ import json import time -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final @@ -199,18 +199,10 @@ async def build_model_max_budget_usage( ) for budget_model, budget_config in budgets ) - batched: Final = await cache.async_batch_get_cache( - keys=list(spend_keys) # mutable-ok: async_batch_get_cache annotates keys as list, so one must exist here - ) - # async_batch_get_cache returns None if it fails internally, and its result is - # index-aligned with `keys` otherwise. An unusable result reads as a miss, - # which is what a never-written counter already reads as. - current_spends: Final = ( - tuple(batched) if isinstance(batched, list) and len(batched) == len(budgets) else (None,) * len(budgets) - ) + current_spends: Final = await _current_window_spends(cache=cache, spend_keys=spend_keys) return { budget_model: { - "current_spend": round(_as_spend(current_spend), 4), + "current_spend": round(current_spend, 4), "budget_limit": budget_config.max_budget, "time_period": budget_config.budget_duration, } @@ -218,6 +210,22 @@ async def build_model_max_budget_usage( } +async def _current_window_spends(cache: DualCache, spend_keys: Sequence[str]) -> tuple[float, ...]: + """Redis holds the window total across replicas; the in-memory copy is one replica's share.""" + keys: Final = list(spend_keys) # mutable-ok: both batch readers annotate their key argument as list + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + shared: Final = await redis_cache.async_batch_get_cache(key_list=keys) + return tuple(_as_spend(shared.get(key)) for key in keys) + # async_batch_get_cache returns None if it fails internally, and its result is + # index-aligned with `keys` otherwise. An unusable result reads as a miss, + # which is what a never-written counter already reads as. + batched: Final = await cache.async_batch_get_cache(keys=keys) + if not isinstance(batched, list) or len(batched) != len(keys): + return (0.0,) * len(keys) + return tuple(_as_spend(current_spend) for current_spend in batched) + + def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None: try: budget_config: Final = BudgetConfig.model_validate(raw_budget_config) @@ -404,7 +412,10 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return current_spend + _as_spend(await self._cached_spend(legacy_spend_key)) async def _cached_spend(self, spend_key: str) -> float | None: - return await self.dual_cache.async_get_cache(key=spend_key) + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None: + return await self.dual_cache.async_get_cache(key=spend_key) + return await redis_cache.async_get_cache(key=spend_key) async def async_filter_deployments( self, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 85a57e5af2b..d52f05f6166 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2280,7 +2280,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache( ) spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) -model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache) +model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) redis_usage_cache: RedisCache | None = None # redis cache used for tracking spend, tpm/rpm limits polling_via_cache_enabled: Literal["all"] | list[str] | bool = False diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 3785ccdcfba..096efc33aaf 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -1,3 +1,5 @@ +import asyncio +from types import MappingProxyType from unittest.mock import AsyncMock, patch @@ -5,6 +7,7 @@ import pytest import litellm from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache from datetime import datetime, timezone from litellm.litellm_core_utils.duration_parser import duration_in_seconds @@ -1332,3 +1335,85 @@ async def test_the_user_scope_has_no_pre_upgrade_counter_to_carry(): await limiter.is_user_within_model_budget( user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4" ) + + +class _SharedFakeRedis(RedisCache): + """Stand-in for the one Redis every replica's DualCache is attached to. + + Only the methods the limiter and DualCache call are implemented, and + ``super().__init__`` is skipped so no connection is opened. + """ + + def __init__(self): + self._store = MappingProxyType({}) + + async def async_set_cache(self, key, value, **kwargs): + self._store = MappingProxyType({**self._store, key: value}) + + async def async_get_cache(self, key, **kwargs): + return self._store.get(key) + + async def async_batch_get_cache(self, key_list, **kwargs): + return {key: self._store.get(key) for key in key_list} + + async def async_increment_pipeline(self, increment_list, **kwargs): + for op in increment_list: + total = self._store.get(op["key"], 0.0) + op["increment_value"] + self._store = MappingProxyType({**self._store, op["key"]: total}) + return [self._store[op["key"]] for op in increment_list] + + +async def _log_spend(limiter, *, key_hash, model_max_budget, response_cost): + await limiter.async_log_success_event( + _success_kwargs( + model_group="gpt-4", + response_cost=response_cost, + key_hash=key_hash, + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + # The Redis push is scheduled as a task rather than awaited inline. + await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task())) + + +@pytest.mark.asyncio +async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another(): + """ + Each replica increments its own in-memory copy of the per-model counter and + pushes the increment to the shared Redis, so only Redis holds the window's + total. A replica that has served part of the traffic must still enforce and + report the total, not its own share. + + Regression: reads went to the in-memory tier first, so a replica whose local + copy sat under the cap kept admitting requests and /key/info on it reported + that local share, while the shared counter was already over the cap. + """ + shared_redis = _SharedFakeRedis() + replica_a = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + replica_b = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + key_hash = "vk-shared" + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "30d"}} + user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget) + + await _log_spend(replica_b, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.25) + await _log_spend(replica_a, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.5) + await _log_spend(replica_a, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.5) + + with pytest.raises(litellm.BudgetExceededError): + await replica_b.is_key_within_model_budget(user_api_key, "gpt-4") + + usage_on_b = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=key_hash, + model_max_budget=model_max_budget, + cache=replica_b.dual_cache, + ) + assert usage_on_b["gpt-4"]["current_spend"] == 1.25 + + # Control: a replica that never served this key reads the same total. + replica_c = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + with pytest.raises(litellm.BudgetExceededError): + await replica_c.is_key_within_model_budget(user_api_key, "gpt-4") diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index 772b08bc9d0..573bfc40c96 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -5,7 +5,7 @@ Verifies that _init_cache attaches Redis to user_api_key_cache only when the flag is explicitly set to True, and leaves it in-memory-only otherwise. """ -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager import json from unittest.mock import MagicMock, patch @@ -167,3 +167,25 @@ class TestRedisAuthCacheFlag: f"cli_sso_session_cache must always get Redis " f"(enable_redis_auth_cache={flag_value!r})" ) + + def test_flag_absent_still_shares_the_model_budget_counters_over_redis(self): + """ + Per-model budget counters are spend counters: the limiter must be able to + push and read them through Redis without the auth-cache opt-in, or every + worker enforces and reports its own share of a key's spend + """ + fake_redis = _FakeRedisCache() + limiter_cache = ps.model_max_budget_limiter.dual_cache + touched_caches = ( + limiter_cache, + ps.spend_counter_cache, + ps.cli_sso_session_cache, + ps.user_api_key_cache, + ps.litellm_config_cache, + ) + with ExitStack() as detached: + for cache in touched_caches: + detached.enter_context(patch.object(cache, "redis_cache", None)) + ps._attach_redis_usage_cache(fake_redis, enable_redis_auth_cache=False) + assert limiter_cache.redis_cache is fake_redis + assert ps.user_api_key_cache.redis_cache is None From 0346bb265934a09bd5d8eab336facba1cc5bc01b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:58:05 +0000 Subject: [PATCH 51/81] fix(bedrock): upgrade legacy thinking after the invoke response_format stub model swap Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../anthropic_claude3_transformation.py | 5 +++++ ...ations_anthropic_claude3_transformation.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 2a4c38e71ea..07ddf6570f2 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -107,6 +107,11 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Restore original model name model = original_model + # The stub model hides the original model from the parent's legacy thinking upgrade + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=original_model, optional_params=optional_params, custom_llm_provider="bedrock" + ) + # The stub model hides the original model from the parent's forced-tool-use backstop response_format_tool_choice: Final = optional_params.get("tool_choice") if ( diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 41d82e4f960..d136cb6450b 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -671,3 +671,22 @@ def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice assert "output_format" not in result assert "tools" in result assert "tool_choice" not in result + + +def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking(local_model_cost_map): + """Regression: the tool-based ``response_format`` path swaps in a Claude 3 stub + model before the shared Anthropic mapping, which hid the adaptive-only model + from the legacy ``thinking`` upgrade and left ``type=enabled`` on the wire.""" + result = AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": {"type": "json_object"}, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "max_tokens": 8192, + }, + optional_params={}, + model="us.anthropic.claude-fable-5-1", + drop_params=False, + ) + + assert result["thinking"] == {"type": "adaptive"} + assert result["output_config"] == {"effort": "high"} From 5da9b7ef900bb60657cd6c4340b3f54833463be2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:12:05 -0700 Subject: [PATCH 52/81] fix(otel): stamp the Langfuse root observation from the post-guardrail request and response --- litellm/integrations/otel/langfuse_logger.py | 43 ++++++--------- litellm/integrations/otel/logger.py | 5 +- .../integrations/otel/test_langfuse_logger.py | 52 +++++++++++++------ 3 files changed, 57 insertions(+), 43 deletions(-) diff --git a/litellm/integrations/otel/langfuse_logger.py b/litellm/integrations/otel/langfuse_logger.py index 9986eae4d0a..ed47533e700 100644 --- a/litellm/integrations/otel/langfuse_logger.py +++ b/litellm/integrations/otel/langfuse_logger.py @@ -8,41 +8,26 @@ from litellm.integrations.otel.model.request_io import request_input, response_o from litellm.integrations.otel.plumbing.context import request_root_span if TYPE_CHECKING: - from litellm.caching.dual_cache import DualCache from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import CallTypesLiteral, ModelResponseStream - -ROOT_OBSERVATION_IO_CALL_TYPES: Final = frozenset( - {"completion", "acompletion", "responses", "aresponses", "anthropic_messages", "aanthropic_messages"} -) + from litellm.types.utils import ModelResponseStream class LangfuseOpenTelemetryV2(OpenTelemetryV2): """Stamps the request's input and output on the root observation while it is still recording. Langfuse shows a trace's input and output from its root observation. The proxy's root span ends - when the response is sent, before the success callback runs, so the stamps have to come from the - request-task hooks: input at pre-call, output at post-call success or at the end of the stream. + when the response is sent, before the success callback runs, so both stamps come from the + post-call hooks in the request task: the request as it stands after the pre-call chain and the + response as it is returned, for the call types whose response renders as a message. """ - async def async_pre_call_hook( - self, - user_api_key_dict: "UserAPIKeyAuth", - cache: "DualCache", - data: Mapping[str, object], - call_type: "CallTypesLiteral", - ) -> None: - await super().async_pre_call_hook(user_api_key_dict, cache, data, call_type) - if call_type in ROOT_OBSERVATION_IO_CALL_TYPES: - self._stamp_root(LANGFUSE_OBSERVATION_INPUT, lambda: request_input(data)) - async def async_post_call_success_hook( self, data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", response: object, ) -> None: - self._stamp_root(LANGFUSE_OBSERVATION_OUTPUT, lambda: response_output(response)) + self._stamp_root_io(data, lambda: response_output(response)) async def async_post_call_streaming_iterator_hook( self, @@ -54,16 +39,22 @@ class LangfuseOpenTelemetryV2(OpenTelemetryV2): async for chunk in response: relayed.append(chunk) yield chunk - self._stamp_root(LANGFUSE_OBSERVATION_OUTPUT, lambda: stream_output(tuple(relayed), request_data)) + self._stamp_root_io(request_data, lambda: stream_output(tuple(relayed), request_data)) - def _stamp_root(self, key: str, render: Callable[[], str | None]) -> None: + def _stamp_root_io(self, data: Mapping[str, object], render_output: Callable[[], str | None]) -> None: root: Final = request_root_span() if root is None or not root.is_recording(): return try: - value: Final = render() + output: Final = render_output() + if output is None: + return + root.set_attribute(LANGFUSE_OBSERVATION_OUTPUT, output) + rendered_input: Final = request_input(data) except Exception: # noqa: BLE001 # telemetry must never fail the request it describes - verbose_logger.debug("otel v2 langfuse: could not render %s for the root observation", key, exc_info=True) + verbose_logger.debug( + "otel v2 langfuse: could not render the root observation input or output", exc_info=True + ) return - if value is not None: - root.set_attribute(key, value) + if rendered_input is not None: + root.set_attribute(LANGFUSE_OBSERVATION_INPUT, rendered_input) diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 4ab1c738488..a550dca6cc8 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -723,13 +723,14 @@ class OpenTelemetryV2(CustomLogger): self, user_api_key_dict: "UserAPIKeyAuth", cache: "DualCache", - data: Mapping[str, object], + data: dict, call_type: "CallTypesLiteral", - ) -> None: + ) -> dict: self.seed_request_identity( user_api_key_dict, model=model_from_request_data(data), ) + return data def record_error_attributes_on_span( self, diff --git a/tests/test_litellm/integrations/otel/test_langfuse_logger.py b/tests/test_litellm/integrations/otel/test_langfuse_logger.py index 3e35395389e..8f93a9a564f 100644 --- a/tests/test_litellm/integrations/otel/test_langfuse_logger.py +++ b/tests/test_litellm/integrations/otel/test_langfuse_logger.py @@ -31,6 +31,8 @@ from litellm.types.llms.openai import ( # noqa: E402 from litellm.types.utils import ( # noqa: E402 Choices, Delta, + Embedding, + EmbeddingResponse, Message, ModelResponse, ModelResponseStream, @@ -258,26 +260,41 @@ def test_root_observation_io_survives_the_root_ending_before_the_success_callbac assert OUTPUT_ATTR in dict(generation.attributes or {}) +def test_root_input_is_the_request_as_the_pre_call_chain_left_it(): + logger, exporter = _logger() + raw = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]} + masked = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "my ssn is [REDACTED]"}]} + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="noted"))]) + root = _start_root(logger) + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), raw, "acompletion")) + asyncio.run(logger.async_post_call_success_hook(data=masked, user_api_key_dict=UserAPIKeyAuth(), response=response)) + root.end() + + assert json.loads(_root_attrs(exporter)[INPUT_ATTR]) == masked["messages"] + + def test_root_already_ended_is_left_alone(): logger, exporter = _logger() + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) root = _start_root(logger) root.end() - asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion")) - - assert INPUT_ATTR not in _root_attrs(exporter) - - -def test_non_chat_call_types_do_not_stamp_input(): - logger, exporter = _logger() - root = _start_root(logger) - asyncio.run( - logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), {"model": "e", "input": "ping"}, "aembedding") + logger.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response) ) - root.end() - assert INPUT_ATTR not in _root_attrs(exporter) + attrs = _root_attrs(exporter) + assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs + + +def test_responses_without_a_message_body_stamp_neither_input_nor_output(): + logger, exporter = _logger() + embedding = EmbeddingResponse(model="e", data=[Embedding(embedding=[0.1], index=0, object="embedding")]) + + _run_request(logger, {"model": "e", "input": "ping"}, "aembedding", embedding) + + attrs = _root_attrs(exporter) + assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs def test_unrenderable_output_never_raises_into_the_request(): @@ -285,7 +302,8 @@ def test_unrenderable_output_never_raises_into_the_request(): _run_request(logger, CHAT_DATA, "acompletion", object()) - assert OUTPUT_ATTR not in _root_attrs(exporter) + attrs = _root_attrs(exporter) + assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs @pytest.mark.parametrize( @@ -331,7 +349,11 @@ def test_langfuse_otel_preset_builds_a_logger_that_stamps_the_root(monkeypatch): assert built is not None assert _maybe_construct_otel_v2("langfuse_otel", loggers) is built root = _start_root(built) - asyncio.run(built.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion")) - assert INPUT_ATTR in dict(root.attributes or {}) + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) + asyncio.run( + built.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response) + ) + attrs = dict(root.attributes or {}) + assert INPUT_ATTR in attrs and OUTPUT_ATTR in attrs finally: is_otel_v2_enabled.cache_clear() From 6fae4b3c3977edcddfeb9e0080f91718ae5c77d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:32:21 -0700 Subject: [PATCH 53/81] fix(guardrails): keep the presidio output masker from unmasking after an in-memory update --- .../proxy/guardrails/guardrail_hooks/presidio.py | 2 ++ .../guardrails/guardrail_hooks/test_presidio.py | 12 ++++++++++++ .../proxy/guardrails/test_guardrail_registry.py | 13 +++++++++++-- 3 files changed, 25 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index da51a905ae3..70ea21320ee 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1633,6 +1633,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Update the guardrails litellm params in memory """ super().update_in_memory_litellm_params(litellm_params) + if self.apply_to_output: + self.output_parse_pii = False if litellm_params.pii_entities_config: self.pii_entities_config = litellm_params.pii_entities_config if litellm_params.presidio_score_thresholds: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index fcf940afd0d..84f7611c0c0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -3129,6 +3129,18 @@ def test_update_in_memory_applies_analyze_chunk_size(): assert guardrail.presidio_analyze_chunk_size_bytes == 99_000 +def test_update_in_memory_keeps_output_masker_from_unmasking(): + masker = _OPTIONAL_PresidioPIIMasking(mock_testing=True, apply_to_output=True, output_parse_pii=False) + unmasker = _OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True) + params = LitellmParams(guardrail="presidio", mode="pre_call", output_parse_pii=True) + + masker.update_in_memory_litellm_params(params) + unmasker.update_in_memory_litellm_params(params) + + assert (masker.apply_to_output, masker.output_parse_pii) == (True, False) + assert (unmasker.apply_to_output, unmasker.output_parse_pii) == (False, True) + + def test_merge_drops_truncated_same_type_fragment_from_overlap(): """A boundary entity seen truncated by chunk 1 and whole by chunk 2 must merge to the single full span; keeping both overlapping spans corrupts the diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 5cbdef5f92f..24742e1bac2 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -566,7 +566,14 @@ def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_st try: handler.initialize_guardrail(_presidio_db_guardrail({"EMAIL_ADDRESS": "MASK", "IP_ADDRESS": "MASK"})) tracked = _presidio_callbacks_in(litellm.callbacks) - roles_before = [(callback.apply_to_output, callback.event_hook) for callback in tracked] + roles_before = [ + (callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked + ] + assert roles_before == [ + (False, True, [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call]), + (False, True, GuardrailEventHooks.post_call), + (True, False, GuardrailEventHooks.post_call), + ] updated = Guardrail( guardrail_id=PRESIDIO_SIBLINGS_GID, @@ -585,7 +592,9 @@ def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_st handler.update_in_memory_guardrail(guardrail_id=PRESIDIO_SIBLINGS_GID, guardrail=updated) assert [callback.pii_entities_config for callback in tracked] == [{"EMAIL_ADDRESS": "MASK"}] * 3 - assert [(callback.apply_to_output, callback.event_hook) for callback in tracked] == roles_before + assert [ + (callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked + ] == roles_before assert _presidio_callbacks_in(litellm.callbacks) == tracked finally: for cb_list, snapshot in zip(lists, snapshots): From 711430216ea73eb0ad45773a81c61155062e0567 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 2 Sep 2026 13:33:14 -0700 Subject: [PATCH 54/81] fix(ui): preserve full AgentCore runtime ARN in agent edit form (#39382) parseDynamicAgentForForm recovered a credential field's value from a stored model string by splitting both the model_template and the model on "/" and matching by array index. That breaks for any placeholder value that itself contains "/", such as a Bedrock AgentCore runtime ARN resource path (runtime/), silently dropping everything after the first slash when populating the edit form. Saving without touching the field then persisted the truncated ARN. Replace the index-matching split with a non-mutating template parse (split on the placeholder pattern, escape and rejoin the literal segments into a regex) so a placeholder captures everything it needs regardless of embedded slashes. Also add a lightweight ARN-shape validator for the AgentCore runtime ARN field, guarded against a malformed pattern string, so a truncated value is rejected client-side before it reaches the backend. Resolves LIT-6737 --- .../public_endpoints/agent_create_fields.json | 4 +- .../public_endpoints/public_endpoints.py | 2 + .../public_endpoints/test_public_endpoints.py | 40 ++++++ .../agent_info.integration.test.tsx | 97 ++++++++++++++ .../_components/agent_type_utils.test.ts | 73 ++++++++++ .../agents/_components/agent_type_utils.ts | 45 +++++-- .../_components/dynamic_agent_form_fields.tsx | 126 ++++++++++-------- .../src/components/networking.tsx | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 9 files changed, 326 insertions(+), 67 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts diff --git a/litellm/proxy/public_endpoints/agent_create_fields.json b/litellm/proxy/public_endpoints/agent_create_fields.json index 36484cc1065..cc2fc17d759 100644 --- a/litellm/proxy/public_endpoints/agent_create_fields.json +++ b/litellm/proxy/public_endpoints/agent_create_fields.json @@ -107,7 +107,9 @@ "required": true, "field_type": "text", "default_value": null, - "include_in_litellm_params": false + "include_in_litellm_params": false, + "validation_pattern": "^arn:aws[a-zA-Z0-9-]*:bedrock-agentcore:[a-z0-9-]+:[0-9]{12}:runtime/.+$", + "validation_message": "Enter the complete Bedrock AgentCore runtime ARN, including the runtime ID after \"runtime/\" (e.g. arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent-runtime)." } ], "litellm_params_template": { diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index f6ee054ceaa..c7f80a61e0f 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -43,6 +43,8 @@ class AgentCredentialField(BaseModel): options: list[str] | None = None default_value: str | None = None include_in_litellm_params: bool | None = None + validation_pattern: str | None = None + validation_message: str | None = None class AgentCreateInfo(BaseModel): diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 31430da71e8..8006f64ba41 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1,3 +1,4 @@ +import re from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -756,6 +757,45 @@ def test_public_agent_hub_returns_empty_when_no_public_groups(): assert response.json() == [] +# --------------------------------------------------------------------------- +# /public/agents/fields +# --------------------------------------------------------------------------- + + +def test_bedrock_agentcore_runtime_arn_validation_pattern_accepts_full_resource_path(): + """Regression for LIT-6737: the AgentCore agent_runtime_arn field's + validation_pattern must accept a complete runtime ARN whose resource part + is itself multi-segment (``runtime/``), and reject the exact + truncated shape a naive split("/")-by-position parse used to produce (the + ARN cut off right after the ``runtime`` resource type). + """ + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/agents/fields") + assert response.status_code == 200 + agents = response.json() + + bedrock_agentcore = next((a for a in agents if a["agent_type"] == "bedrock_agentcore"), None) + assert bedrock_agentcore is not None, "bedrock_agentcore agent type not found" + assert bedrock_agentcore["model_template"] == "bedrock/agentcore/{agent_runtime_arn}" + + fields_by_key = {f["key"]: f for f in bedrock_agentcore["credential_fields"]} + arn_field = fields_by_key["agent_runtime_arn"] + assert arn_field["required"] is True + assert arn_field["include_in_litellm_params"] is False + + pattern = arn_field.get("validation_pattern") + assert pattern, "agent_runtime_arn must ship a validation_pattern so the UI can reject a truncated ARN" + + full_arn = "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime/hosted_agent_4vm3i-BaTdfOELAs" + truncated_arn = "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime" + + assert re.match(pattern, full_arn), "the validator must accept a complete runtime ARN" + assert not re.match(pattern, truncated_arn), "the validator must reject the truncated ARN" + + # --------------------------------------------------------------------------- # /public/endpoints # --------------------------------------------------------------------------- diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index 79bd2f6a21b..de1eb153f6f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -75,6 +75,40 @@ const langgraphInfo: AgentCreateInfo = { ], }; +const FULL_RUNTIME_ARN = "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime/hosted_agent_4vm3i-BaTdfOELAs"; + +const BEDROCK_AGENTCORE_AGENT = { + agent_id: "agent-3", + agent_name: "bedrock-agent", + agent_card_params: { name: "bedrock-agent", description: "agentcore agent", url: "", version: "1.0.0", skills: [] }, + litellm_params: { + custom_llm_provider: "bedrock", + model: `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + }, +}; + +const bedrockAgentcoreInfo: AgentCreateInfo = { + agent_type: "bedrock_agentcore", + agent_type_display_name: "Bedrock AgentCore", + description: "Bedrock AgentCore runtimes", + logo_url: "/b.png", + use_a2a_form_fields: false, + litellm_params_template: { custom_llm_provider: "bedrock" }, + model_template: "bedrock/agentcore/{agent_runtime_arn}", + credential_fields: [ + { + key: "agent_runtime_arn", + label: "Agent Runtime ARN", + field_type: "text", + required: true, + include_in_litellm_params: false, + validation_pattern: "^arn:aws[a-zA-Z0-9-]*:bedrock-agentcore:[a-z0-9-]+:[0-9]{12}:runtime/.+$", + validation_message: + 'Enter the complete Bedrock AgentCore runtime ARN, including the runtime ID after "runtime/".', + }, + ], +}; + const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); const renderView = () => render(); @@ -247,6 +281,69 @@ describe("AgentInfoView update payload", () => { }); }); + it("preserves the full AgentCore runtime ARN (including the resource id after runtime/) across an unedited save", async () => { + vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([bedrockAgentcoreInfo]); + vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); + const user = setup(); + renderView(); + await openEditor(user); + + expect(await screen.findByLabelText("Agent Runtime ARN")).toHaveValue(FULL_RUNTIME_ARN); + + await save(user); + + expect((patchedPayload().litellm_params as Record).model).toBe( + `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + ); + }); + + it("blocks the save and shows a validation error when the Agent Runtime ARN is truncated", async () => { + vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([bedrockAgentcoreInfo]); + vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); + const user = setup(); + renderView(); + await openEditor(user); + + const arnField = await screen.findByLabelText("Agent Runtime ARN"); + await user.clear(arnField); + fireEvent.change(arnField, { + target: { value: "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime" }, + }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + expect( + await screen.findByText( + 'Enter the complete Bedrock AgentCore runtime ARN, including the runtime ID after "runtime/".', + ), + ).toBeInTheDocument(); + expect(networking.patchAgentCall).not.toHaveBeenCalled(); + }); + + it("renders and saves normally when a field's validation_pattern is not a valid regex", async () => { + const infoWithBadPattern: AgentCreateInfo = { + ...bedrockAgentcoreInfo, + credential_fields: [ + { + ...bedrockAgentcoreInfo.credential_fields[0], + validation_pattern: "(unterminated", + }, + ], + }; + vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([infoWithBadPattern]); + vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); + const user = setup(); + renderView(); + await openEditor(user); + + expect(await screen.findByLabelText("Agent Runtime ARN")).toHaveValue(FULL_RUNTIME_ARN); + + await save(user); + + expect((patchedPayload().litellm_params as Record).model).toBe( + `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + ); + }); + it("reloads the agent and leaves edit mode when the edit is cancelled", async () => { const user = setup(); renderView(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts new file mode 100644 index 00000000000..0c2d2500776 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts @@ -0,0 +1,73 @@ +import { describe, it, expect } from "vitest"; +import { detectAgentType, extractModelTemplateValues, parseDynamicAgentForForm } from "./agent_type_utils"; +import type { AgentCreateInfo } from "@/components/networking"; +import type { Agent } from "@/components/agents/types"; + +const FULL_RUNTIME_ARN = "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime/hosted_agent_4vm3i-BaTdfOELAs"; + +const bedrockAgentcoreInfo: AgentCreateInfo = { + agent_type: "bedrock_agentcore", + agent_type_display_name: "Bedrock AgentCore", + model_template: "bedrock/agentcore/{agent_runtime_arn}", + credential_fields: [ + { + key: "agent_runtime_arn", + label: "Agent Runtime ARN", + required: true, + include_in_litellm_params: false, + }, + ], +}; + +describe("extractModelTemplateValues", () => { + it("recovers a placeholder value that itself contains '/' (an AWS ARN resource path)", () => { + const values = extractModelTemplateValues( + "bedrock/agentcore/{agent_runtime_arn}", + `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + ); + + expect(values.agent_runtime_arn).toBe(FULL_RUNTIME_ARN); + }); + + it("recovers a placeholder value with no '/' (single path segment)", () => { + const values = extractModelTemplateValues("langgraph/{assistant_id}", "langgraph/asst_1"); + + expect(values.assistant_id).toBe("asst_1"); + }); + + it("returns no match when the model does not fit the template", () => { + const values = extractModelTemplateValues("langgraph/{assistant_id}", "azure_ai/agents/asst_1"); + + expect(values).toEqual({}); + }); +}); + +describe("parseDynamicAgentForForm", () => { + it("preserves the full runtime ARN, including the resource id after 'runtime/', when populating the edit form", () => { + const agent = { + agent_id: "agent-1", + agent_name: "bedrock-agent", + agent_card_params: { description: "" }, + litellm_params: { + custom_llm_provider: "bedrock", + model: `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + }, + } as unknown as Agent; + + const values = parseDynamicAgentForForm(agent, bedrockAgentcoreInfo); + + expect(values.agent_runtime_arn).toBe(FULL_RUNTIME_ARN); + }); +}); + +describe("detectAgentType", () => { + it("detects bedrock_agentcore agents from the model prefix", () => { + const agent = { + agent_id: "agent-1", + agent_name: "bedrock-agent", + litellm_params: { model: `bedrock/agentcore/${FULL_RUNTIME_ARN}` }, + } as unknown as Agent; + + expect(detectAgentType(agent)).toBe("bedrock_agentcore"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts index f91590c5732..506f355c519 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts @@ -25,6 +25,29 @@ export const detectAgentType = (agent: Agent): string => { return "a2a"; }; +const escapeRegExp = (segment: string): string => segment.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + +/** + * Reverses a `model_template` (e.g. "bedrock/agentcore/{agent_runtime_arn}") against a stored + * `model` string to recover the placeholder values that produced it. Builds a regex from the + * template's literal segments rather than matching by split("/") position, because a + * placeholder's value can itself contain "/" (an AWS ARN's "runtime/" resource path, + * a Vertex AI reasoning engine's "projects/.../reasoningEngines/..." resource id) and would + * otherwise be cut off at the first one. + */ +export const extractModelTemplateValues = (template: string, model: string): Record => { + // Splitting on a regex with a capturing group interleaves the captured placeholder + // names between the surrounding literal segments, e.g. "a/{x}/b" -> ["a/", "x", "/b"]. + const parts = template.split(/\{([a-zA-Z0-9_]+)\}/g); + const fieldNames = parts.filter((_part, index) => index % 2 === 1); + const pattern = parts.map((part, index) => (index % 2 === 1 ? "(.+)" : escapeRegExp(part))).join(""); + + const match = model.match(new RegExp(`^${pattern}$`)); + if (!match) return {}; + + return Object.fromEntries(fieldNames.map((name, index) => [name, match[index + 1]])); +}; + /** * Parses agent data for dynamic form fields (non-A2A agents). * Extracts values from litellm_params based on the agent type metadata. @@ -35,24 +58,18 @@ export const parseDynamicAgentForForm = (agent: Agent, agentTypeInfo: AgentCreat description: agent.agent_card_params?.description || "", }; + const templateValues = + agentTypeInfo.model_template && agent.litellm_params?.model + ? extractModelTemplateValues(agentTypeInfo.model_template, agent.litellm_params.model) + : {}; + // Extract credential field values from litellm_params for (const field of agentTypeInfo.credential_fields) { if (field.include_in_litellm_params !== false) { values[field.key] = agent.litellm_params?.[field.key] || field.default_value || ""; - } else { - // For fields not in litellm_params (like agent_id), try to extract from model string - if (agentTypeInfo.model_template && agent.litellm_params?.model) { - const model = agent.litellm_params.model; - const templateParts = agentTypeInfo.model_template.split("/"); - const modelParts = model.split("/"); - - // Find the placeholder position and extract the value - templateParts.forEach((part, index) => { - if (part === `{${field.key}}` && modelParts[index]) { - values[field.key] = modelParts[index]; - } - }); - } + } else if (templateValues[field.key] !== undefined) { + // For fields not in litellm_params (like agent_runtime_arn), recover from the model string + values[field.key] = templateValues[field.key]; } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx index 04a8b0df9d9..0f3355fe073 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx @@ -24,58 +24,80 @@ interface DynamicAgentFormFieldsProps { export const unmountedDynamicFieldNames = (mountedPanels: readonly string[]): readonly string[] => mountedPanels.includes(AGENT_FORM_CONFIG.cost.key) ? [] : COST_FIELD_NAMES; -const CredentialField = ({ field }: { field: AgentCredentialFieldMetadata }) => ( - - {({ value, onChange, ref, ...control }) => { - const text = typeof value === "string" ? value : ""; - if (field.field_type === "password") { - return ( - - ); - } - if (field.field_type === "textarea") { - return ( -