From 6bbcf1dbbf3d533ee1948b9b3183d8962ace4263 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 28 Jul 2026 12:08:08 +0000
Subject: [PATCH 01/56] fix(mcp): keep the streamable-HTTP routing peek on a
UTF-8 boundary
Fixes https://github.com/BerriAI/litellm/issues/34917
---
.../proxy/_experimental/mcp_server/server.py | 24 +++-
.../mcp_server/test_mcp_server.py | 104 ++++++++++++++++++
2 files changed, 125 insertions(+), 3 deletions(-)
diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 14673cf12c1..af5da275961 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -237,6 +237,24 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool:
return False
+def _utf8_boundary_prefix(data: bytes) -> bytes:
+ """``data`` with any trailing incomplete UTF-8 sequence removed.
+
+ Cutting a body at a fixed byte budget can land in the middle of a multibyte
+ character, and ``json.loads`` on such bytes raises ``UnicodeDecodeError``
+ rather than ``JSONDecodeError``. Trimming to a character boundary keeps the
+ truncated peek decodable so callers only have to handle malformed JSON.
+ """
+ for trailing in range(0, min(3, len(data)) + 1):
+ candidate = data[: len(data) - trailing]
+ try:
+ candidate.decode("utf-8")
+ except UnicodeDecodeError:
+ continue
+ return candidate
+ return data
+
+
def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None:
"""The W3C trace context (``traceparent``/``tracestate``) the MCP client
propagated in the request's ``params._meta`` (SEP-414), or ``None``.
@@ -3411,7 +3429,7 @@ if MCP_AVAILABLE:
try:
data = json.loads(body)
return isinstance(data, dict) and data.get("method") == "initialize"
- except (json.JSONDecodeError, TypeError):
+ except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
return False
async def _read_request_body_for_routing(
@@ -3462,7 +3480,7 @@ if MCP_AVAILABLE:
# directly from the original `receive` via wrapped_receive.
break
- return consumed_messages, b"".join(body_chunks)
+ return consumed_messages, _utf8_boundary_prefix(b"".join(body_chunks))
async def _handle_stale_mcp_session(
scope: Scope,
@@ -4227,7 +4245,7 @@ if MCP_AVAILABLE:
"MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock",
_peeked.get("id"),
)
- except (json.JSONDecodeError, TypeError):
+ except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
# Peek cap truncated the body, so it can't be fully parsed.
# Scan the top-level keys (depth-aware) instead of a flat
# substring search: a response's result payload may nest a
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index 1753b0d92a8..7f79e5aebda 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -1,5 +1,6 @@
import asyncio
import contextvars
+import json
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, MagicMock, patch
@@ -1689,6 +1690,109 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body():
assert total_streamed == len(first_chunk) + sum(len(b) for b in oversized_tail)
+@pytest.mark.asyncio
+async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap():
+ """
+ A tool-call POST whose UTF-8 body is larger than the routing peek cap, with a
+ multibyte character straddling the cap boundary, must still be forwarded
+ intact instead of blowing up with a UnicodeDecodeError 500.
+
+ Regression test for https://github.com/BerriAI/litellm/issues/34917
+ """
+ try:
+ from litellm.proxy._experimental.mcp_server import server as mcp_server
+ from litellm.proxy._experimental.mcp_server.server import (
+ handle_streamable_http_mcp,
+ session_manager_stateful,
+ session_manager_stateless,
+ )
+ except ImportError:
+ pytest.skip("MCP server not available")
+
+ peek_cap = mcp_server._MCP_ROUTING_PEEK_MAX_BYTES
+
+ def _splits_multibyte_at_cap(candidate: bytes) -> bool:
+ try:
+ candidate[:peek_cap].decode("utf-8")
+ except UnicodeDecodeError:
+ return True
+ return False
+
+ def _build_body() -> bytes:
+ for pad in range(4):
+ candidate = json.dumps(
+ {
+ "jsonrpc": "2.0",
+ "id": 1,
+ "method": "tools/call",
+ "params": {
+ "name": "update_full_document" + "x" * pad,
+ "arguments": {"markdown": "щ" * 3000},
+ },
+ },
+ ensure_ascii=False,
+ ).encode("utf-8")
+ if len(candidate) > peek_cap and _splits_multibyte_at_cap(candidate):
+ return candidate
+ raise AssertionError("could not build a body splitting a multibyte char at the peek cap")
+
+ body = _build_body()
+
+ messages = [{"type": "http.request", "body": body, "more_body": False}]
+ receive_calls = {"count": 0}
+
+ async def receive():
+ idx = receive_calls["count"]
+ receive_calls["count"] += 1
+ return messages[idx]
+
+ scope = {
+ "type": "http",
+ "method": "POST",
+ "path": "/mcp/progress_test",
+ "headers": [
+ (b"content-type", b"application/json"),
+ (b"authorization", b"Bearer test-key"),
+ ],
+ }
+ send = AsyncMock()
+
+ streamed_chunks = []
+
+ async def stateless_handle(s, r, se):
+ while True:
+ msg = await r()
+ if msg.get("type") != "http.request":
+ break
+ streamed_chunks.append(msg.get("body", b"") or b"")
+ if not msg.get("more_body", False):
+ break
+
+ async def stateful_handle(s, r, se):
+ raise AssertionError("non-initialize POST should not reach stateful manager")
+
+ with (
+ patch(
+ "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+ new_callable=AsyncMock,
+ return_value=(MagicMock(), None, ["progress_test"], None, None, None),
+ ),
+ patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"),
+ patch(
+ "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+ True,
+ ),
+ patch.object(session_manager_stateless, "handle_request", side_effect=stateless_handle),
+ patch.object(session_manager_stateful, "handle_request", side_effect=stateful_handle),
+ patch.object(session_manager_stateless, "_server_instances", {}),
+ patch.object(session_manager_stateful, "_server_instances", {}),
+ ):
+ await handle_streamable_http_mcp(scope, receive, send)
+
+ assert send.await_count == 0, f"unexpected response emitted by the proxy: {send.await_args_list}"
+ assert b"".join(streamed_chunks) == body
+
+
@pytest.mark.asyncio
async def test_enforce_stateful_session_cap_evicts_oldest_idle_then_rejects():
"""
From 4a6a387ca1e511e35858fee0c92fe3e3415d03ee Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 1 Sep 2026 20:48:22 +0000
Subject: [PATCH 02/56] fix(mcp): follow nextCursor on paginated
tools/prompts/resources list operations
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/constants.py | 1 +
litellm/experimental_mcp_client/client.py | 56 ++++++-----
litellm/experimental_mcp_client/pagination.py | 92 +++++++++++++++++++
litellm/experimental_mcp_client/tools.py | 7 +-
.../mcp_server/rest_endpoints.py | 6 +-
.../test_mcp_client.py | 52 ++++++++++-
.../test_pagination.py | 80 ++++++++++++++++
7 files changed, 258 insertions(+), 36 deletions(-)
create mode 100644 litellm/experimental_mcp_client/pagination.py
create mode 100644 tests/test_litellm/experimental_mcp_client/test_pagination.py
diff --git a/litellm/constants.py b/litellm/constants.py
index 9a50797f517..11f35177636 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -136,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
+MCP_LIST_MAX_PAGES: Final = int(os.getenv("LITELLM_MCP_LIST_MAX_PAGES", "100"))
# Allowlist of commands permitted for MCP stdio transport.
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py
index f0a1bff8fdc..814074d35a4 100644
--- a/litellm/experimental_mcp_client/client.py
+++ b/litellm/experimental_mcp_client/client.py
@@ -48,6 +48,12 @@ from pydantic import AnyUrl
from litellm._logging import verbose_logger
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
+from litellm.experimental_mcp_client.pagination import (
+ list_all_prompts,
+ list_all_resource_templates,
+ list_all_resources,
+ list_all_tools,
+)
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
from litellm.types.llms.custom_http import VerifyTypes
from litellm.types.mcp import (
@@ -603,17 +609,17 @@ class MCPClient:
"""
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
- async def _list_tools_operation(session: ClientSession):
- return await session.list_tools()
+ async def _list_tools_operation(session: ClientSession) -> tuple[MCPTool, ...]:
+ return await list_all_tools(session, self.server_url or "stdio")
try:
- result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
- tool_count: Final = len(result.tools)
- tool_names: Final = [tool.name for tool in result.tools]
+ tools: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
+ tool_count: Final = len(tools)
+ tool_names: Final = [tool.name for tool in tools]
verbose_logger.info(
"MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names
)
- return result.tools
+ return list(tools)
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_tools was cancelled")
raise
@@ -734,17 +740,17 @@ class MCPClient:
"""List available prompts from the server."""
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
- async def _list_prompts_operation(session: ClientSession):
- return await session.list_prompts()
+ async def _list_prompts_operation(session: ClientSession) -> tuple[Prompt, ...]:
+ return await list_all_prompts(session, self.server_url or "stdio")
try:
- result: Final = await self.run_with_session(_list_prompts_operation)
- prompt_count: Final = len(result.prompts)
- prompt_names: Final = [prompt.name for prompt in result.prompts]
+ prompts: Final = await self.run_with_session(_list_prompts_operation)
+ prompt_count: Final = len(prompts)
+ prompt_names: Final = [prompt.name for prompt in prompts]
verbose_logger.info(
- "MCP client listed %s tools from %s: %s", prompt_count, self.server_url or "stdio", prompt_names
+ "MCP client listed %s prompts from %s: %s", prompt_count, self.server_url or "stdio", prompt_names
)
- return result.prompts
+ return list(prompts)
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_prompts was cancelled")
raise
@@ -811,17 +817,17 @@ class MCPClient:
"""List available resources from the server."""
verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio")
- async def _list_resources_operation(session: ClientSession):
- return await session.list_resources()
+ async def _list_resources_operation(session: ClientSession) -> tuple[Resource, ...]:
+ return await list_all_resources(session, self.server_url or "stdio")
try:
- result: Final = await self.run_with_session(_list_resources_operation)
- resource_count: Final = len(result.resources)
- resource_names: Final = [resource.name for resource in result.resources]
+ resources: Final = await self.run_with_session(_list_resources_operation)
+ resource_count: Final = len(resources)
+ resource_names: Final = [resource.name for resource in resources]
verbose_logger.info(
"MCP client listed %s resources from %s: %s", resource_count, self.server_url or "stdio", resource_names
)
- return result.resources
+ return list(resources)
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_resources was cancelled")
raise
@@ -847,20 +853,20 @@ class MCPClient:
"""List available resource templates from the server."""
verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio")
- async def _list_resource_templates_operation(session: ClientSession):
- return await session.list_resource_templates()
+ async def _list_resource_templates_operation(session: ClientSession) -> tuple[ResourceTemplate, ...]:
+ return await list_all_resource_templates(session, self.server_url or "stdio")
try:
- result: Final = await self.run_with_session(_list_resource_templates_operation)
- resource_template_count: Final = len(result.resourceTemplates)
- resource_template_names: Final = [resourceTemplate.name for resourceTemplate in result.resourceTemplates]
+ resource_templates: Final = await self.run_with_session(_list_resource_templates_operation)
+ resource_template_count: Final = len(resource_templates)
+ resource_template_names: Final = [resource_template.name for resource_template in resource_templates]
verbose_logger.info(
"MCP client listed %s resource templates from %s: %s",
resource_template_count,
self.server_url or "stdio",
resource_template_names,
)
- return result.resourceTemplates
+ return list(resource_templates)
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_resource_templates was cancelled")
raise
diff --git a/litellm/experimental_mcp_client/pagination.py b/litellm/experimental_mcp_client/pagination.py
new file mode 100644
index 00000000000..8852715aba5
--- /dev/null
+++ b/litellm/experimental_mcp_client/pagination.py
@@ -0,0 +1,92 @@
+"""
+Follows ``nextCursor`` on the paginated MCP list operations so a multi-page catalog is read in full.
+"""
+
+from collections.abc import Awaitable, Callable, Sequence
+from typing import Final, TypeVar
+
+from mcp import ClientSession, Resource
+from mcp.types import PaginatedRequestParams, PaginatedResult, Prompt, ResourceTemplate
+from mcp.types import Tool as MCPTool
+
+from litellm._logging import verbose_logger
+from litellm.constants import MCP_LIST_MAX_PAGES
+
+TPage = TypeVar("TPage", bound=PaginatedResult)
+TItem = TypeVar("TItem")
+
+
+async def collect_pages(
+ fetch_page: Callable[[PaginatedRequestParams | None], Awaitable[TPage]],
+ items_of: Callable[[TPage], Sequence[TItem]],
+ *,
+ method: str,
+ server: str,
+ cursor: str | None = None,
+ seen_cursors: frozenset[str] = frozenset(),
+) -> tuple[TItem, ...]:
+ page: Final = await fetch_page(None if cursor is None else PaginatedRequestParams(cursor=cursor))
+ items: Final = tuple(items_of(page))
+ next_cursor: Final = page.nextCursor
+ pages_read: Final = len(seen_cursors) + 1
+ if next_cursor is None:
+ return items
+ if next_cursor in seen_cursors:
+ verbose_logger.warning(
+ "MCP %s from %s repeated cursor %r; returning the %s page(s) read so far",
+ method,
+ server,
+ next_cursor,
+ pages_read,
+ )
+ return items
+ if pages_read >= MCP_LIST_MAX_PAGES:
+ verbose_logger.warning(
+ "MCP %s from %s still paginating after %s pages (LITELLM_MCP_LIST_MAX_PAGES); returning what was read",
+ method,
+ server,
+ pages_read,
+ )
+ return items
+ rest: Final = await collect_pages(
+ fetch_page,
+ items_of,
+ method=method,
+ server=server,
+ cursor=next_cursor,
+ seen_cursors=seen_cursors | frozenset((next_cursor,)),
+ )
+ return items + rest
+
+
+async def list_all_tools(session: ClientSession, server: str) -> tuple[MCPTool, ...]:
+ return await collect_pages(
+ lambda params: session.list_tools(params=params), lambda page: page.tools, method="tools/list", server=server
+ )
+
+
+async def list_all_prompts(session: ClientSession, server: str) -> tuple[Prompt, ...]:
+ return await collect_pages(
+ lambda params: session.list_prompts(params=params),
+ lambda page: page.prompts,
+ method="prompts/list",
+ server=server,
+ )
+
+
+async def list_all_resources(session: ClientSession, server: str) -> tuple[Resource, ...]:
+ return await collect_pages(
+ lambda params: session.list_resources(params=params),
+ lambda page: page.resources,
+ method="resources/list",
+ server=server,
+ )
+
+
+async def list_all_resource_templates(session: ClientSession, server: str) -> tuple[ResourceTemplate, ...]:
+ return await collect_pages(
+ lambda params: session.list_resource_templates(params=params),
+ lambda page: page.resourceTemplates,
+ method="resources/templates/list",
+ server=server,
+ )
diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py
index 30d50e2a74b..adaca0888aa 100644
--- a/litellm/experimental_mcp_client/tools.py
+++ b/litellm/experimental_mcp_client/tools.py
@@ -9,6 +9,7 @@ from openai.types.chat import ChatCompletionToolParam
from openai.types.responses.function_tool_param import FunctionToolParam
from openai.types.shared_params.function_definition import FunctionDefinition
+from litellm.experimental_mcp_client.pagination import list_all_tools
from litellm.types.llms.anthropic import AnthropicMessagesTool
from litellm.types.utils import ChatCompletionMessageToolCall
@@ -103,10 +104,10 @@ async def load_mcp_tools(
If format is set to "openai", the tools are converted to OpenAI API compatible tools.
"""
- tools: Final = await session.list_tools()
+ tools: Final = await list_all_tools(session, "upstream")
if format == "openai":
- return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools]
- return tools.tools
+ return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools]
+ return list(tools)
########################################################
diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
index 3efb6429326..3ca6b6c5f90 100644
--- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
@@ -1402,11 +1402,7 @@ if MCP_AVAILABLE:
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
async def _list_tools_operation(client):
- async def _list_tools_session_operation(session):
- return await session.list_tools()
-
- list_tools_response: Final = await client.run_with_session(_list_tools_session_operation)
- list_tools_result: Final[list[MCPTool]] = list_tools_response.tools
+ list_tools_result: Final[list[MCPTool]] = await client.list_tools(raise_on_error=True)
model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result]
return {
"tools": model_dumped_tools,
diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py
index fd7ab3afdab..3f501d3859a 100644
--- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py
+++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py
@@ -20,8 +20,10 @@ from mcp.types import (
JSONRPCError,
JSONRPCMessage,
JSONRPCResponse,
+ ListToolsResult,
ServerCapabilities,
)
+from mcp.types import Tool as MCPTool
# Add the parent directory to the path so we can import litellm
@@ -740,8 +742,14 @@ class _ScriptedUpstream:
error, the shape an upstream application uses to report its own failure.
"""
- def __init__(self, tools_list_error: ErrorData | None = None):
+ def __init__(
+ self,
+ tools_list_error: ErrorData | None = None,
+ tool_pages: tuple[tuple[MCPTool, ...], ...] = (),
+ ):
self._tools_list_error = tools_list_error
+ self._tool_pages = tool_pages
+ self.tools_list_cursors: list[str | None] = []
self._to_client_tx, self._to_client_rx = anyio.create_memory_object_stream(10)
self._from_client_tx, self._from_client_rx = anyio.create_memory_object_stream(10)
self._task_group = None
@@ -778,15 +786,37 @@ class _ScriptedUpstream:
)
elif method == "tools/list" and self._tools_list_error is not None:
await self._send(JSONRPCError(jsonrpc="2.0", id=request.id, error=self._tools_list_error))
+ elif method == "tools/list" and self._tool_pages:
+ cursor = (request.params or {}).get("cursor")
+ self.tools_list_cursors.append(cursor)
+ page_index = int(cursor) if cursor else 0
+ has_more = page_index + 1 < len(self._tool_pages)
+ page = ListToolsResult(
+ tools=list(self._tool_pages[page_index]),
+ nextCursor=str(page_index + 1) if has_more else None,
+ )
+ await self._send(
+ JSONRPCResponse(
+ jsonrpc="2.0",
+ id=request.id,
+ result=page.model_dump(by_alias=True, mode="json", exclude_none=True),
+ )
+ )
class _ScriptedClient(MCPClient):
"""An MCPClient whose transport is a scripted in-memory upstream instead of a real connection,
so the real ``ClientSession`` and its real timeout machinery are what run."""
- def __init__(self, *, timeout: float, tools_list_error: ErrorData | None = None):
+ def __init__(
+ self,
+ *,
+ timeout: float,
+ tools_list_error: ErrorData | None = None,
+ tool_pages: tuple[tuple[MCPTool, ...], ...] = (),
+ ):
super().__init__(server_url="http://upstream.local/mcp", timeout=timeout)
- self._upstream = _ScriptedUpstream(tools_list_error=tools_list_error)
+ self._upstream = _ScriptedUpstream(tools_list_error=tools_list_error, tool_pages=tool_pages)
def _create_transport_context(self):
return self._upstream, None
@@ -821,6 +851,22 @@ async def test_list_tools_fails_on_its_own_timeout_when_the_upstream_never_answe
assert list_fault_http_status(fault) == 504
+@pytest.mark.asyncio
+async def test_list_tools_follows_tools_list_pagination_across_the_whole_catalog():
+ """An upstream that pages tools/list (72 tools, 30 per page) must have every page read within the
+ one session, each request carrying the cursor the previous page returned. Reading only the first
+ page made 42 tools invisible to the proxy and every call to them fail as unknown."""
+ tools = tuple(
+ MCPTool(name=f"tool_{i:02d}", inputSchema={"type": "object", "properties": {}}) for i in range(72)
+ )
+ client = _ScriptedClient(timeout=30, tool_pages=(tools[:30], tools[30:60], tools[60:]))
+
+ listed = await asyncio.wait_for(client.list_tools(raise_on_error=True), timeout=10)
+
+ assert [tool.name for tool in listed] == [tool.name for tool in tools]
+ assert client._upstream.tools_list_cursors == [None, "1", "2"]
+
+
@pytest.mark.asyncio
async def test_upstream_json_rpc_error_408_is_not_reported_as_a_client_timeout():
"""The SDK reports its own elapsed read timeout and relays an upstream JSON-RPC error through
diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py
new file mode 100644
index 00000000000..f588fdd1eee
--- /dev/null
+++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py
@@ -0,0 +1,80 @@
+import logging
+
+import pytest
+from mcp.types import ListToolsResult, PaginatedRequestParams
+from mcp.types import Tool as MCPTool
+
+import litellm.experimental_mcp_client.pagination as pagination_module
+from litellm.experimental_mcp_client.pagination import collect_pages
+
+
+def _tool(index: int) -> MCPTool:
+ return MCPTool(name=f"tool_{index:02d}", inputSchema={"type": "object", "properties": {}})
+
+
+class _PagedTools:
+ """A tools/list upstream serving ``total`` tools ``page_size`` at a time, cursors being offsets."""
+
+ def __init__(self, total: int, page_size: int):
+ self._tools = tuple(_tool(i) for i in range(total))
+ self._page_size = page_size
+ self.cursors_seen: list[str | None] = []
+
+ async def fetch(self, params: PaginatedRequestParams | None) -> ListToolsResult:
+ cursor = params.cursor if params is not None else None
+ self.cursors_seen.append(cursor)
+ start = int(cursor) if cursor else 0
+ end = start + self._page_size
+ return ListToolsResult(
+ tools=list(self._tools[start:end]),
+ nextCursor=str(end) if end < len(self._tools) else None,
+ )
+
+
+@pytest.mark.asyncio
+async def test_collect_pages_follows_next_cursor_until_exhausted():
+ upstream = _PagedTools(total=72, page_size=30)
+
+ tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s")
+
+ assert [t.name for t in tools] == [f"tool_{i:02d}" for i in range(72)]
+ assert upstream.cursors_seen == [None, "30", "60"], "each page must be requested with the cursor the previous one returned"
+
+
+@pytest.mark.asyncio
+async def test_collect_pages_single_page_makes_one_request():
+ upstream = _PagedTools(total=5, page_size=30)
+
+ tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s")
+
+ assert len(tools) == 5
+ assert upstream.cursors_seen == [None]
+
+
+@pytest.mark.asyncio
+async def test_collect_pages_stops_on_a_repeated_cursor_and_keeps_what_it_read(caplog):
+ calls: list[str | None] = []
+
+ async def fetch(params: PaginatedRequestParams | None) -> ListToolsResult:
+ calls.append(params.cursor if params else None)
+ return ListToolsResult(tools=[_tool(len(calls))], nextCursor="same")
+
+ with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+ tools = await collect_pages(fetch, lambda page: page.tools, method="tools/list", server="s")
+
+ assert calls == [None, "same"], "the cursor must be followed once and refused the second time it comes back"
+ assert len(tools) == 2
+ assert any("repeated cursor" in record.getMessage() for record in caplog.records)
+
+
+@pytest.mark.asyncio
+async def test_collect_pages_honors_the_page_cap(monkeypatch, caplog):
+ monkeypatch.setattr(pagination_module, "MCP_LIST_MAX_PAGES", 3)
+ upstream = _PagedTools(total=1000, page_size=10)
+
+ with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+ tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s")
+
+ assert len(upstream.cursors_seen) == 3
+ assert len(tools) == 30
+ assert any("LITELLM_MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records)
From 26b48d58919e5021b9b251339bf5c720a3e3649e Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 1 Sep 2026 21:00:27 +0000
Subject: [PATCH 03/56] refactor(mcp): keep list pagination within
type-discipline budget and ratchet LIT001
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/experimental_mcp_client/client.py | 56 ++++++++++---------
.../mcp_server/rest_endpoints.py | 4 +-
.../test_pagination.py | 4 +-
type-discipline-budget.json | 2 +-
4 files changed, 35 insertions(+), 31 deletions(-)
diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py
index 814074d35a4..46d8823e439 100644
--- a/litellm/experimental_mcp_client/client.py
+++ b/litellm/experimental_mcp_client/client.py
@@ -377,29 +377,31 @@ class MCPClient:
return provided_env
# Minimal allowlist of safe/standard environment variables
- safe_keys: Final = {
- "PATH",
- "HOME",
- "USER",
- "LOGNAME",
- "TMPDIR",
- "TMP",
- "TEMP",
- "SHELL",
- "LANG",
- "LC_ALL",
- # Node/Package manager caches
- "NPM_CONFIG_CACHE",
- "PNPM_HOME",
- "XDG_CACHE_HOME",
- "XDG_CONFIG_HOME",
- "XDG_DATA_HOME",
- # System info
- "SYSTEMROOT",
- "COMSPEC",
- "PATHEXT",
- "WINDIR",
- }
+ safe_keys: Final = frozenset(
+ {
+ "PATH",
+ "HOME",
+ "USER",
+ "LOGNAME",
+ "TMPDIR",
+ "TMP",
+ "TEMP",
+ "SHELL",
+ "LANG",
+ "LC_ALL",
+ # Node/Package manager caches
+ "NPM_CONFIG_CACHE",
+ "PNPM_HOME",
+ "XDG_CACHE_HOME",
+ "XDG_CONFIG_HOME",
+ "XDG_DATA_HOME",
+ # System info
+ "SYSTEMROOT",
+ "COMSPEC",
+ "PATHEXT",
+ "WINDIR",
+ }
+ )
safe_env: Final = {}
for key in safe_keys:
@@ -615,7 +617,7 @@ class MCPClient:
try:
tools: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
tool_count: Final = len(tools)
- tool_names: Final = [tool.name for tool in tools]
+ tool_names: Final = tuple(tool.name for tool in tools)
verbose_logger.info(
"MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names
)
@@ -746,7 +748,7 @@ class MCPClient:
try:
prompts: Final = await self.run_with_session(_list_prompts_operation)
prompt_count: Final = len(prompts)
- prompt_names: Final = [prompt.name for prompt in prompts]
+ prompt_names: Final = tuple(prompt.name for prompt in prompts)
verbose_logger.info(
"MCP client listed %s prompts from %s: %s", prompt_count, self.server_url or "stdio", prompt_names
)
@@ -823,7 +825,7 @@ class MCPClient:
try:
resources: Final = await self.run_with_session(_list_resources_operation)
resource_count: Final = len(resources)
- resource_names: Final = [resource.name for resource in resources]
+ resource_names: Final = tuple(resource.name for resource in resources)
verbose_logger.info(
"MCP client listed %s resources from %s: %s", resource_count, self.server_url or "stdio", resource_names
)
@@ -859,7 +861,7 @@ class MCPClient:
try:
resource_templates: Final = await self.run_with_session(_list_resource_templates_operation)
resource_template_count: Final = len(resource_templates)
- resource_template_names: Final = [resource_template.name for resource_template in resource_templates]
+ resource_template_names: Final = tuple(resource_template.name for resource_template in resource_templates)
verbose_logger.info(
"MCP client listed %s resource templates from %s: %s",
resource_template_count,
diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
index 3ca6b6c5f90..beee7c1db78 100644
--- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
@@ -1,6 +1,6 @@
import asyncio
import importlib
-from collections.abc import Awaitable, Callable, Mapping
+from collections.abc import Awaitable, Callable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal
@@ -1402,7 +1402,7 @@ if MCP_AVAILABLE:
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
async def _list_tools_operation(client):
- list_tools_result: Final[list[MCPTool]] = await client.list_tools(raise_on_error=True)
+ list_tools_result: Final[Sequence[MCPTool]] = await client.list_tools(raise_on_error=True)
model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result]
return {
"tools": model_dumped_tools,
diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py
index f588fdd1eee..a76f410ac3c 100644
--- a/tests/test_litellm/experimental_mcp_client/test_pagination.py
+++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py
@@ -38,7 +38,9 @@ async def test_collect_pages_follows_next_cursor_until_exhausted():
tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s")
assert [t.name for t in tools] == [f"tool_{i:02d}" for i in range(72)]
- assert upstream.cursors_seen == [None, "30", "60"], "each page must be requested with the cursor the previous one returned"
+ assert upstream.cursors_seen == [None, "30", "60"], (
+ "each page must be requested with the cursor the previous one returned"
+ )
@pytest.mark.asyncio
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 3d2e97d55a5..b0c3cc7f9fd 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -1,6 +1,6 @@
{
"LIT001": {
- "limit": 22367
+ "limit": 22366
},
"LIT002": {
"limit": 26777
From d32f8a07c88ed165e55ce944026504a1cd15a327 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 1 Sep 2026 21:15:10 +0000
Subject: [PATCH 04/56] fix(mcp): make list page cap a plain constant and use a
real ListToolsResult in the unit mock
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/constants.py | 2 +-
litellm/experimental_mcp_client/pagination.py | 2 +-
tests/mcp_tests/test_mcp_client_unit.py | 6 ++----
.../test_litellm/experimental_mcp_client/test_pagination.py | 2 +-
4 files changed, 5 insertions(+), 7 deletions(-)
diff --git a/litellm/constants.py b/litellm/constants.py
index 11f35177636..07914934495 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -136,7 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
-MCP_LIST_MAX_PAGES: Final = int(os.getenv("LITELLM_MCP_LIST_MAX_PAGES", "100"))
+MCP_LIST_MAX_PAGES: Final = 100
# Allowlist of commands permitted for MCP stdio transport.
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
diff --git a/litellm/experimental_mcp_client/pagination.py b/litellm/experimental_mcp_client/pagination.py
index 8852715aba5..85f46268f92 100644
--- a/litellm/experimental_mcp_client/pagination.py
+++ b/litellm/experimental_mcp_client/pagination.py
@@ -42,7 +42,7 @@ async def collect_pages(
return items
if pages_read >= MCP_LIST_MAX_PAGES:
verbose_logger.warning(
- "MCP %s from %s still paginating after %s pages (LITELLM_MCP_LIST_MAX_PAGES); returning what was read",
+ "MCP %s from %s still paginating after %s pages (MCP_LIST_MAX_PAGES); returning what was read",
method,
server,
pages_read,
diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py
index aadaadd510e..ef4231fe1d9 100644
--- a/tests/mcp_tests/test_mcp_client_unit.py
+++ b/tests/mcp_tests/test_mcp_client_unit.py
@@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch, ANY
import litellm.experimental_mcp_client.client as mcp_client_module
from litellm.experimental_mcp_client.client import MCPClient
from litellm.types.mcp import MCPAuth, MCPTransport
-from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult
+from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult, ListToolsResult
def test_mcp_client_uses_configurable_default_timeout():
@@ -174,9 +174,7 @@ class TestMCPClientUnitTests:
},
)
]
- mock_result = MagicMock()
- mock_result.tools = mock_tools
- mock_session_instance.list_tools.return_value = mock_result
+ mock_session_instance.list_tools.return_value = ListToolsResult(tools=mock_tools)
client = MCPClient("http://example.com")
result = await client.list_tools()
diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py
index a76f410ac3c..93952c176a8 100644
--- a/tests/test_litellm/experimental_mcp_client/test_pagination.py
+++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py
@@ -79,4 +79,4 @@ async def test_collect_pages_honors_the_page_cap(monkeypatch, caplog):
assert len(upstream.cursors_seen) == 3
assert len(tools) == 30
- assert any("LITELLM_MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records)
+ assert any("MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records)
From f2b6c0da81ca247a3b6e7e52c85c51d313310a78 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 16:41:12 -0700
Subject: [PATCH 05/56] feat(bedrock_mantle): serve /v1/messages for Claude
models on Mantle's native Anthropic Messages API
---
.../messages/handler.py | 3 +-
.../llms/bedrock_mantle/messages/__init__.py | 0
.../bedrock_mantle/messages/transformation.py | 101 +++++
litellm/utils.py | 7 +
..._bedrock_mantle_messages_transformation.py | 346 ++++++++++++++++++
tests/test_litellm/test_utils.py | 22 ++
6 files changed, 478 insertions(+), 1 deletion(-)
create mode 100644 litellm/llms/bedrock_mantle/messages/__init__.py
create mode 100644 litellm/llms/bedrock_mantle/messages/transformation.py
create mode 100644 tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
index 87a4801f987..e1309ea4063 100644
--- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
+++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
@@ -501,6 +501,7 @@ def anthropic_messages_handler(
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
)
+ resolved_api_base: Final = dynamic_api_base if dynamic_api_base is not None else api_base
# Store agentic loop params in logging object for agentic hooks
# This provides original request context needed for follow-up calls
@@ -662,7 +663,7 @@ def anthropic_messages_handler(
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
api_key=api_key,
- api_base=api_base,
+ api_base=resolved_api_base,
stream=stream,
kwargs=kwargs,
)
diff --git a/litellm/llms/bedrock_mantle/messages/__init__.py b/litellm/llms/bedrock_mantle/messages/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py
new file mode 100644
index 00000000000..a4365cfa49b
--- /dev/null
+++ b/litellm/llms/bedrock_mantle/messages/transformation.py
@@ -0,0 +1,101 @@
+from collections.abc import Mapping
+from typing import Final
+
+from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
+ DEFAULT_ANTHROPIC_API_VERSION,
+)
+from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.common_utils import MANTLE_MESSAGES_PATH
+from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig
+from litellm.llms.bedrock_mantle.common_utils import (
+ MANTLE_HOST_RE,
+ BedrockMantleAuthMixin,
+ resolve_mantle_region,
+)
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.router import GenericLiteLLMParams
+
+_BASE_SUFFIXES_TO_STRIP: Final = (
+ MANTLE_MESSAGES_PATH,
+ "/v1/messages",
+ "/messages",
+ "/anthropic/v1",
+ "/openai/v1",
+ "/v1",
+)
+
+
+def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mapping[str, object]) -> str:
+ region: Final = resolve_mantle_region({**litellm_params, "api_base": api_base})
+ configured: Final = (
+ api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") or f"https://bedrock-mantle.{region}.api.aws"
+ ).rstrip("/")
+ stripped: Final = next(
+ (configured[: -len(suffix)] for suffix in _BASE_SUFFIXES_TO_STRIP if configured.endswith(suffix)),
+ configured,
+ )
+ host: Final = f"https://bedrock-mantle.{region}.api.aws" if MANTLE_HOST_RE.match(stripped) else stripped
+ return f"{host}{MANTLE_MESSAGES_PATH}"
+
+
+class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleMessagesConfig):
+ def __init__(self, aws_signer: BaseAWSLLM | None = None) -> None:
+ AmazonMantleMessagesConfig.__init__(self)
+ self._aws_signer = aws_signer or self
+
+ @property
+ def custom_llm_provider(self) -> str | None:
+ return "bedrock_mantle"
+
+ def get_complete_url(
+ self,
+ api_base: str | None,
+ api_key: str | None,
+ model: str,
+ optional_params: dict,
+ litellm_params: dict,
+ stream: bool | None = None,
+ ) -> str:
+ return build_mantle_native_messages_url(api_base=api_base, litellm_params=litellm_params)
+
+ def validate_anthropic_messages_environment(
+ self,
+ headers: dict,
+ model: str,
+ messages: list[dict],
+ optional_params: dict,
+ litellm_params: dict,
+ api_key: str | None = None,
+ api_base: str | None = None,
+ ) -> tuple[dict, str | None]:
+ merged_headers, resolved_api_base = super().validate_anthropic_messages_environment(
+ headers=headers,
+ model=model,
+ messages=messages,
+ optional_params=optional_params,
+ litellm_params=litellm_params,
+ api_key=api_key,
+ api_base=api_base,
+ )
+ if any(name.lower() == "anthropic-version" for name in merged_headers):
+ return merged_headers, resolved_api_base
+ return {**merged_headers, "anthropic-version": DEFAULT_ANTHROPIC_API_VERSION}, resolved_api_base
+
+ def transform_anthropic_messages_request(
+ self,
+ model: str,
+ messages: list[dict],
+ anthropic_messages_optional_request_params: dict,
+ litellm_params: GenericLiteLLMParams,
+ headers: dict,
+ ) -> dict:
+ 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 "anthropic_version" in anthropic_messages_optional_request_params:
+ return request
+ return {key: value for key, value in request.items() if key != "anthropic_version"}
diff --git a/litellm/utils.py b/litellm/utils.py
index b724313641f..3439a21b560 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -8681,6 +8681,13 @@ class ProviderConfigManager:
from litellm.llms.bedrock.common_utils import BedrockModelInfo
return BedrockModelInfo.get_bedrock_provider_config_for_messages_api(model)
+ elif litellm.LlmProviders.BEDROCK_MANTLE == provider:
+ if "claude" in model_lower:
+ from litellm.llms.bedrock_mantle.messages.transformation import (
+ BedrockMantleAnthropicMessagesConfig,
+ )
+
+ return BedrockMantleAnthropicMessagesConfig()
elif litellm.LlmProviders.VERTEX_AI == provider:
if "claude" in model_lower:
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import (
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
new file mode 100644
index 00000000000..2961eee925c
--- /dev/null
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
@@ -0,0 +1,346 @@
+"""
+Unit tests for the bedrock_mantle native Anthropic Messages route.
+
+Mantle serves its Claude models only on `/anthropic/v1/messages` (the OpenAI
+paths reject them), so `bedrock_mantle/anthropic.claude-*` requests on
+/v1/messages must hit that endpoint directly instead of the chat-completions
+bridge. These tests lock the dispatcher gate, the URL derivation from the
+OpenAI-surface base that get_llm_provider pre-fills, the version header, the
+Bearer/SigV4 auth chain, and the wire request through the public entrypoint.
+"""
+
+import json
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+import respx
+
+import litellm
+from litellm.caching.llm_caching_handler import LLMClientCache
+from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock_mantle.messages.transformation import (
+ BedrockMantleAnthropicMessagesConfig,
+ build_mantle_native_messages_url,
+)
+from litellm.types.router import GenericLiteLLMParams
+from litellm.utils import ProviderConfigManager
+
+MESSAGES_PATH = "/anthropic/v1/messages"
+
+
+@pytest.fixture(autouse=True)
+def _httpx_transport_with_fresh_clients(monkeypatch):
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
+
+
+@pytest.fixture(autouse=True)
+def _no_ambient_mantle_env(monkeypatch):
+ monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
+ monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
+ monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
+ monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
+ monkeypatch.delenv("AWS_REGION_NAME", raising=False)
+ monkeypatch.delenv("AWS_REGION", raising=False)
+
+
+def _anthropic_response() -> httpx.Response:
+ return httpx.Response(
+ status_code=200,
+ json={
+ "id": "msg_test",
+ "type": "message",
+ "role": "assistant",
+ "model": "anthropic.claude-sonnet-5",
+ "content": [{"type": "text", "text": "pong"}],
+ "stop_reason": "end_turn",
+ "stop_sequence": None,
+ "usage": {"input_tokens": 3, "output_tokens": 1},
+ },
+ )
+
+
+_SSE_EVENTS = (
+ (
+ "message_start",
+ {
+ "type": "message_start",
+ "message": {
+ "id": "msg_stream",
+ "type": "message",
+ "role": "assistant",
+ "model": "anthropic.claude-sonnet-5",
+ "content": [],
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": {"input_tokens": 3, "output_tokens": 1},
+ },
+ },
+ ),
+ ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
+ ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}}),
+ ("content_block_stop", {"type": "content_block_stop", "index": 0}),
+ ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}),
+ ("message_stop", {"type": "message_stop"}),
+)
+
+
+def _sse_response() -> httpx.Response:
+ body = "".join(f"event: {event}\ndata: {json.dumps(payload)}\n\n" for event, payload in _SSE_EVENTS).encode()
+ return httpx.Response(status_code=200, content=body, headers={"content-type": "text/event-stream"})
+
+
+def _mantle_messages_route(region: str) -> respx.Route:
+ return respx.post(f"https://bedrock-mantle.{region}.api.aws{MESSAGES_PATH}")
+
+
+def _sent_body(route: respx.Route) -> dict:
+ return json.loads(route.calls.last.request.content)
+
+
+class TestDispatch:
+ def test_claude_models_get_the_native_messages_config(self):
+ config = ProviderConfigManager.get_provider_anthropic_messages_config(
+ model="anthropic.claude-sonnet-5", provider=litellm.LlmProviders.BEDROCK_MANTLE
+ )
+ assert isinstance(config, BedrockMantleAnthropicMessagesConfig)
+ assert config.custom_llm_provider == "bedrock_mantle"
+
+ @pytest.mark.parametrize("model", ["openai.gpt-5.6-sol", "openai.gpt-oss-120b-1:0", "google.gemma-4-31b"])
+ def test_non_claude_models_keep_the_bridge(self, model):
+ assert (
+ ProviderConfigManager.get_provider_anthropic_messages_config(
+ model=model, provider=litellm.LlmProviders.BEDROCK_MANTLE
+ )
+ is None
+ )
+
+
+class TestURL:
+ @pytest.mark.parametrize(
+ "api_base",
+ [
+ "https://bedrock-mantle.us-east-1.api.aws/v1",
+ "https://bedrock-mantle.us-east-1.api.aws/openai/v1",
+ "https://bedrock-mantle.us-east-1.api.aws/openai/v1/",
+ "https://bedrock-mantle.us-east-1.api.aws",
+ "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages",
+ ],
+ )
+ def test_prefilled_openai_base_becomes_the_messages_endpoint(self, api_base):
+ url = build_mantle_native_messages_url(api_base, {"aws_region_name": "us-east-1"})
+ assert url == f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}"
+
+ def test_aws_region_name_wins_over_the_prefilled_host_region(self):
+ url = build_mantle_native_messages_url(
+ "https://bedrock-mantle.us-east-1.api.aws/v1", {"aws_region_name": "us-east-2"}
+ )
+ assert url == f"https://bedrock-mantle.us-east-2.api.aws{MESSAGES_PATH}"
+
+ def test_host_region_is_used_when_no_region_param(self):
+ url = build_mantle_native_messages_url("https://bedrock-mantle.eu-west-1.api.aws/v1", {})
+ assert url == f"https://bedrock-mantle.eu-west-1.api.aws{MESSAGES_PATH}"
+
+ def test_custom_host_is_preserved(self):
+ url = build_mantle_native_messages_url("https://vpce-abc.bedrock-mantle.example.com/v1", {})
+ assert url == f"https://vpce-abc.bedrock-mantle.example.com{MESSAGES_PATH}"
+
+ def test_env_base_is_used_without_api_base(self, monkeypatch):
+ monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://mantle-proxy.internal/openai/v1")
+ assert build_mantle_native_messages_url(None, {}) == f"https://mantle-proxy.internal{MESSAGES_PATH}"
+
+ def test_default_host_comes_from_mantle_region_env(self, monkeypatch):
+ monkeypatch.setenv("BEDROCK_MANTLE_REGION", "ap-northeast-1")
+ assert build_mantle_native_messages_url(None, {}) == f"https://bedrock-mantle.ap-northeast-1.api.aws{MESSAGES_PATH}"
+
+ def test_config_get_complete_url_reads_litellm_params(self):
+ config = BedrockMantleAnthropicMessagesConfig()
+ url = config.get_complete_url(
+ api_base="https://bedrock-mantle.us-east-1.api.aws/v1",
+ api_key=None,
+ model="anthropic.claude-sonnet-5",
+ optional_params={},
+ litellm_params={"aws_region_name": "us-west-2"},
+ )
+ assert url == f"https://bedrock-mantle.us-west-2.api.aws{MESSAGES_PATH}"
+
+
+class TestEnvironment:
+ def _validate(self, headers: dict, litellm_params: dict) -> dict:
+ config = BedrockMantleAnthropicMessagesConfig()
+ merged, _ = config.validate_anthropic_messages_environment(
+ headers=headers,
+ model="anthropic.claude-sonnet-5",
+ messages=[],
+ optional_params={},
+ litellm_params=litellm_params,
+ )
+ return merged
+
+ def test_adds_the_anthropic_version_header(self):
+ assert self._validate({}, {})["anthropic-version"] == "2023-06-01"
+
+ def test_keeps_a_caller_supplied_version_header(self):
+ merged = self._validate({"Anthropic-Version": "2024-01-01"}, {})
+ assert merged["Anthropic-Version"] == "2024-01-01"
+ assert "anthropic-version" not in merged
+
+ def test_project_id_becomes_the_workspace_header(self):
+ assert self._validate({}, {"aws_bedrock_project_id": "proj_123"})["anthropic-workspace"] == "proj_123"
+
+
+class TestRequestBody:
+ def test_body_carries_model_and_stream_but_not_the_invoke_version(self):
+ config = BedrockMantleAnthropicMessagesConfig()
+ body = config.transform_anthropic_messages_request(
+ model="anthropic.claude-sonnet-5",
+ messages=[{"role": "user", "content": "ping"}],
+ anthropic_messages_optional_request_params={"max_tokens": 8, "stream": True},
+ litellm_params=GenericLiteLLMParams(),
+ headers={},
+ )
+ assert body["model"] == "anthropic.claude-sonnet-5"
+ assert body["stream"] is True
+ assert body["max_tokens"] == 8
+ assert "anthropic_version" not in body
+
+ def test_body_omits_stream_when_not_streaming(self):
+ config = BedrockMantleAnthropicMessagesConfig()
+ body = config.transform_anthropic_messages_request(
+ model="anthropic.claude-sonnet-5",
+ messages=[{"role": "user", "content": "ping"}],
+ anthropic_messages_optional_request_params={"max_tokens": 8},
+ litellm_params=GenericLiteLLMParams(),
+ headers={},
+ )
+ assert "stream" not in body
+
+
+class TestAuth:
+ def test_bearer_from_api_key_skips_aws_credentials(self):
+ signer = BaseAWSLLM()
+ signer.get_credentials = MagicMock(side_effect=AssertionError("must not resolve AWS credentials"))
+ config = BedrockMantleAnthropicMessagesConfig(aws_signer=signer)
+ headers, signed = config.sign_request(
+ headers={"anthropic-version": "2023-06-01"},
+ optional_params={},
+ request_data={"model": "anthropic.claude-sonnet-5"},
+ api_base=f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}",
+ api_key="arg-bearer",
+ )
+ assert headers["Authorization"] == "Bearer arg-bearer"
+ assert headers["anthropic-version"] == "2023-06-01"
+ assert signed == b'{"model": "anthropic.claude-sonnet-5"}'
+
+ def test_bearer_from_mantle_env_key(self, monkeypatch):
+ monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer")
+ config = BedrockMantleAnthropicMessagesConfig()
+ headers, _ = config.sign_request(
+ headers={},
+ optional_params={},
+ request_data={},
+ api_base=f"https://bedrock-mantle.us-east-1.api.aws{MESSAGES_PATH}",
+ api_key=None,
+ )
+ assert headers["Authorization"] == "Bearer env-bearer"
+
+ def test_sigv4_scope_is_pinned_to_the_url_host_region(self):
+ config = BedrockMantleAnthropicMessagesConfig()
+ headers, signed = config.sign_request(
+ headers={"anthropic-version": "2023-06-01"},
+ optional_params={
+ "aws_access_key_id": "AKIAEXAMPLE",
+ "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
+ "aws_region_name": "us-east-1",
+ },
+ request_data={"model": "anthropic.claude-sonnet-5"},
+ api_base=f"https://bedrock-mantle.us-west-2.api.aws{MESSAGES_PATH}",
+ api_key=None,
+ )
+ assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
+ assert "/us-west-2/bedrock/aws4_request" in headers["Authorization"]
+ assert signed == b'{"model": "anthropic.claude-sonnet-5"}'
+
+
+class TestWireRequest:
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_claude_request_hits_the_native_messages_endpoint(self):
+ route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response())
+
+ response = await litellm.anthropic_messages(
+ model="bedrock_mantle/anthropic.claude-sonnet-5",
+ messages=[{"role": "user", "content": "ping"}],
+ max_tokens=8,
+ api_key="test-bearer",
+ aws_region_name="us-east-1",
+ )
+
+ assert response["content"][0]["text"] == "pong"
+ assert route.call_count == 1
+ sent = route.calls.last.request
+ assert sent.headers["authorization"] == "Bearer test-bearer"
+ assert sent.headers["anthropic-version"] == "2023-06-01"
+ assert "x-api-key" not in sent.headers
+ body = _sent_body(route)
+ assert body["model"] == "anthropic.claude-sonnet-5"
+ assert body["messages"] == [{"role": "user", "content": "ping"}]
+ assert "anthropic_version" not in body
+ assert "stream" not in body
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_region_prefix_selects_the_host_and_is_not_sent_as_model(self):
+ route = _mantle_messages_route("us-east-2").mock(return_value=_anthropic_response())
+
+ await litellm.anthropic_messages(
+ model="bedrock_mantle/us-east-2/anthropic.claude-haiku-4-5",
+ messages=[{"role": "user", "content": "ping"}],
+ max_tokens=8,
+ api_key="test-bearer",
+ )
+
+ assert route.call_count == 1
+ assert _sent_body(route)["model"] == "anthropic.claude-haiku-4-5"
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_streaming_sends_stream_and_passes_the_sse_through(self):
+ route = _mantle_messages_route("us-east-1").mock(return_value=_sse_response())
+
+ response = await litellm.anthropic_messages(
+ model="bedrock_mantle/anthropic.claude-sonnet-5",
+ messages=[{"role": "user", "content": "ping"}],
+ max_tokens=8,
+ stream=True,
+ api_key="test-bearer",
+ aws_region_name="us-east-1",
+ )
+ raw = b"".join([chunk async for chunk in response])
+
+ assert route.call_count == 1
+ assert _sent_body(route)["stream"] is True
+ text = raw.decode()
+ assert "event: message_start" in text
+ assert '"text": "pong"' in text
+ assert "event: message_stop" in text
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_sigv4_request_signs_against_the_messages_url(self):
+ route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response())
+
+ await litellm.anthropic_messages(
+ model="bedrock_mantle/anthropic.claude-sonnet-5",
+ messages=[{"role": "user", "content": "ping"}],
+ max_tokens=8,
+ aws_access_key_id="AKIAEXAMPLE",
+ aws_secret_access_key="c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0",
+ aws_region_name="us-east-1",
+ )
+
+ assert route.call_count == 1
+ authorization = route.calls.last.request.headers["authorization"]
+ assert authorization.startswith("AWS4-HMAC-SHA256")
+ assert "/us-east-1/bedrock/aws4_request" in authorization
diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py
index 3336ad6d33a..a4c06122189 100644
--- a/tests/test_litellm/test_utils.py
+++ b/tests/test_litellm/test_utils.py
@@ -3640,6 +3640,28 @@ class TestGetOptionalParamsTencent:
assert isinstance(config, TencentAnthropicMessagesConfig)
assert config.custom_llm_provider == "tencent"
+ def test_bedrock_mantle_claude_messages_config_routing(self):
+ import litellm
+ from litellm.llms.bedrock_mantle.messages.transformation import (
+ BedrockMantleAnthropicMessagesConfig,
+ )
+
+ config = ProviderConfigManager.get_provider_anthropic_messages_config(
+ model="anthropic.claude-sonnet-5",
+ provider=litellm.LlmProviders.BEDROCK_MANTLE,
+ )
+ assert isinstance(config, BedrockMantleAnthropicMessagesConfig)
+ assert config.custom_llm_provider == "bedrock_mantle"
+
+ def test_bedrock_mantle_openai_models_keep_the_messages_bridge(self):
+ import litellm
+
+ config = ProviderConfigManager.get_provider_anthropic_messages_config(
+ model="openai.gpt-5.6-sol",
+ provider=litellm.LlmProviders.BEDROCK_MANTLE,
+ )
+ assert config is None
+
class TestValidateEnvironmentTencent:
"""Tests that validate_environment resolves TENCENT_API_KEY for the tencent provider."""
From 368a8396400bdf5f986f8379d84ac44c43b808c2 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 18:08:41 -0700
Subject: [PATCH 06/56] fix(bedrock_mantle): send anthropic betas in the header
Mantle reads on /v1/messages
---
litellm/anthropic_beta_headers_config.json | 35 +++++
.../anthropic_claude3_transformation.py | 22 ++--
.../bedrock_mantle/messages/transformation.py | 36 +++--
..._bedrock_mantle_messages_transformation.py | 123 +++++++++++++++++-
4 files changed, 196 insertions(+), 20 deletions(-)
diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json
index eb31cc17a15..1331de4c266 100644
--- a/litellm/anthropic_beta_headers_config.json
+++ b/litellm/anthropic_beta_headers_config.json
@@ -131,6 +131,41 @@
"web-fetch-2025-09-10": null,
"web-search-2025-03-05": null
},
+ "bedrock_mantle": {
+ "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19",
+ "advisor-tool-2026-03-01": null,
+ "bash_20241022": null,
+ "bash_20250124": null,
+ "claude-code-20250219": "claude-code-20250219",
+ "code-execution-2025-08-25": null,
+ "compact-2026-01-12": "compact-2026-01-12",
+ "computer-use-2025-01-24": "computer-use-2025-01-24",
+ "computer-use-2025-11-24": "computer-use-2025-11-24",
+ "context-1m-2025-08-07": "context-1m-2025-08-07",
+ "context-management-2025-06-27": "context-management-2025-06-27",
+ "effort-2025-11-24": "effort-2025-11-24",
+ "fast-mode-2026-02-01": null,
+ "files-api-2025-04-14": null,
+ "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
+ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
+ "mcp-client-2025-04-04": null,
+ "mcp-client-2025-11-20": null,
+ "mcp-servers-2025-12-04": null,
+ "output-128k-2025-02-19": "output-128k-2025-02-19",
+ "per-turn-control-2026-07-01": "per-turn-control-2026-07-01",
+ "prompt-caching-scope-2026-01-05": null,
+ "skills-2025-10-02": null,
+ "structured-output-2024-03-01": null,
+ "structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
+ "text_editor_20241022": null,
+ "text_editor_20250124": null,
+ "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
+ "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
+ "tool-examples-2025-10-29": "tool-examples-2025-10-29",
+ "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
+ "web-fetch-2025-09-10": null,
+ "web-search-2025-03-05": "web-search-2025-03-05"
+ },
"vertex_ai": {
"advisor-tool-2026-03-01": null,
"advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19",
diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
index d2be1ad9156..4b52a3bafe6 100644
--- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
+++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
@@ -1,4 +1,4 @@
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
@@ -445,13 +445,16 @@ class AmazonAnthropicClaudeMessagesConfig(
# Bedrock InvokeModel DOES support ``clear_tool_uses_20250919`` under the
# ``context-management-2025-06-27`` beta. AWS docs:
# https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-tool-use.md
- _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: dict[str, str] = {
- "compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
- "clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
- }
+ _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType(
+ {
+ "compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
+ "clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
+ }
+ )
- @staticmethod
+ @classmethod
def _filter_context_management_for_bedrock_invoke(
+ cls,
anthropic_messages_request: dict,
beta_set: set,
) -> None:
@@ -481,7 +484,7 @@ class AmazonAnthropicClaudeMessagesConfig(
anthropic_messages_request.pop("context_management", None)
return
- supported: Final = AmazonAnthropicClaudeMessagesConfig._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS
+ supported: Final = cls._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS
retained_edits: Final = [e for e in edits if isinstance(e, dict) and e.get("type") in supported]
if not retained_edits:
anthropic_messages_request.pop("context_management", None)
@@ -546,15 +549,16 @@ class AmazonAnthropicClaudeMessagesConfig(
if "tool-search-tool-2025-10-19" in beta_set:
beta_set.add("tool-examples-2025-10-29")
+ beta_provider: Final = self.custom_llm_provider or "bedrock"
filtered_betas: Final = sorted(
filter_and_transform_beta_headers(
beta_headers=list(beta_set),
- provider="bedrock",
+ provider=beta_provider,
)
)
dropped_user_betas: Final = sorted(
- b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider="bedrock")
+ b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider=beta_provider)
)
if dropped_user_betas:
verbose_logger.warning(
diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py
index a4365cfa49b..480c09a0476 100644
--- a/litellm/llms/bedrock_mantle/messages/transformation.py
+++ b/litellm/llms/bedrock_mantle/messages/transformation.py
@@ -1,6 +1,9 @@
from collections.abc import Mapping
+from types import MappingProxyType
from typing import Final
+from pydantic import TypeAdapter
+
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
DEFAULT_ANTHROPIC_API_VERSION,
)
@@ -13,6 +16,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
resolve_mantle_region,
)
from litellm.secret_managers.main import get_secret_str
+from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.router import GenericLiteLLMParams
_BASE_SUFFIXES_TO_STRIP: Final = (
@@ -23,6 +27,9 @@ _BASE_SUFFIXES_TO_STRIP: Final = (
"/openai/v1",
"/v1",
)
+_BODY_FIELDS_MANTLE_READS_FROM_HEADERS: Final = frozenset({"anthropic_version", "anthropic_beta"})
+_ANTHROPIC_BETAS: Final = TypeAdapter(tuple[str, ...])
+_MANTLE_REQUEST: Final = TypeAdapter(dict[str, object])
def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mapping[str, object]) -> str:
@@ -39,6 +46,13 @@ def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mappi
class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleMessagesConfig):
+ _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType(
+ {
+ **AmazonMantleMessagesConfig._BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS,
+ "clear_thinking_20251015": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
+ }
+ )
+
def __init__(self, aws_signer: BaseAWSLLM | None = None) -> None:
AmazonMantleMessagesConfig.__init__(self)
self._aws_signer = aws_signer or self
@@ -89,13 +103,17 @@ class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleM
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> dict:
- 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,
+ request: Final = _MANTLE_REQUEST.validate_python(
+ 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 "anthropic_version" in anthropic_messages_optional_request_params:
- return request
- return {key: value for key, value in request.items() if key != "anthropic_version"}
+ betas: Final = request.get("anthropic_beta")
+ if betas is not None:
+ header_betas: Final = ",".join(_ANTHROPIC_BETAS.validate_python(betas))
+ headers["anthropic-beta"] = header_betas # rebind-ok: the handler signs and sends this same dict
+ return {key: value for key, value in request.items() if key not in _BODY_FIELDS_MANTLE_READS_FROM_HEADERS}
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
index 2961eee925c..3544262996c 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
@@ -79,7 +79,10 @@ _SSE_EVENTS = (
},
),
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
- ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}}),
+ (
+ "content_block_delta",
+ {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}},
+ ),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}),
("message_stop", {"type": "message_stop"}),
@@ -152,7 +155,10 @@ class TestURL:
def test_default_host_comes_from_mantle_region_env(self, monkeypatch):
monkeypatch.setenv("BEDROCK_MANTLE_REGION", "ap-northeast-1")
- assert build_mantle_native_messages_url(None, {}) == f"https://bedrock-mantle.ap-northeast-1.api.aws{MESSAGES_PATH}"
+ assert (
+ build_mantle_native_messages_url(None, {})
+ == f"https://bedrock-mantle.ap-northeast-1.api.aws{MESSAGES_PATH}"
+ )
def test_config_get_complete_url_reads_litellm_params(self):
config = BedrockMantleAnthropicMessagesConfig()
@@ -344,3 +350,116 @@ class TestWireRequest:
authorization = route.calls.last.request.headers["authorization"]
assert authorization.startswith("AWS4-HMAC-SHA256")
assert "/us-east-1/bedrock/aws4_request" in authorization
+
+
+def _sent_betas(route: respx.Route) -> list[str]:
+ return route.calls.last.request.headers["anthropic-beta"].split(",")
+
+
+@pytest.mark.usefixtures("local_beta_headers_config")
+class TestBetaHeadersOnTheWire:
+ async def _send(self, **request_params) -> respx.Route:
+ route = _mantle_messages_route("us-east-1").mock(return_value=_anthropic_response())
+ await litellm.anthropic_messages(
+ model="bedrock_mantle/anthropic.claude-sonnet-5",
+ messages=[{"role": "user", "content": "ping"}],
+ max_tokens=8,
+ api_key="test-bearer",
+ aws_region_name="us-east-1",
+ **request_params,
+ )
+ return route
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_betas_mantle_accepts_reach_it_in_the_header(self):
+ route = await self._send(
+ extra_headers={
+ "anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14,context-management-2025-06-27"
+ }
+ )
+
+ assert _sent_betas(route) == [
+ "claude-code-20250219",
+ "context-management-2025-06-27",
+ "interleaved-thinking-2025-05-14",
+ ]
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_betas_mantle_rejects_are_dropped_before_the_request(self):
+ route = await self._send(
+ extra_headers={"anthropic-beta": "code-execution-2025-08-25,context-1m-2025-08-07,files-api-2025-04-14"}
+ )
+
+ assert _sent_betas(route) == ["context-1m-2025-08-07"]
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_no_beta_header_is_sent_when_every_value_is_rejected(self):
+ route = await self._send(extra_headers={"anthropic-beta": "code-execution-2025-08-25"})
+
+ assert "anthropic-beta" not in route.calls.last.request.headers
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_advanced_tool_use_is_renamed_to_the_beta_mantle_knows(self):
+ route = await self._send(extra_headers={"anthropic-beta": "advanced-tool-use-2025-11-20"})
+
+ assert "tool-search-tool-2025-10-19" in _sent_betas(route)
+ assert "advanced-tool-use-2025-11-20" not in _sent_betas(route)
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_a_feature_beta_joins_the_callers_betas_in_the_header(self):
+ route = await self._send(
+ extra_headers={"anthropic-beta": "context-1m-2025-08-07"},
+ context_management={"edits": [{"type": "clear_tool_uses_20250919"}]},
+ )
+
+ assert _sent_betas(route) == ["context-1m-2025-08-07", "context-management-2025-06-27"]
+ assert _sent_body(route)["context_management"] == {"edits": [{"type": "clear_tool_uses_20250919"}]}
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_betas_and_version_never_travel_in_the_body(self):
+ route = await self._send(
+ extra_headers={"anthropic-beta": "context-1m-2025-08-07"},
+ context_management={"edits": [{"type": "clear_tool_uses_20250919"}]},
+ anthropic_version="bedrock-2023-05-31",
+ )
+
+ body = _sent_body(route)
+ assert "anthropic_beta" not in body
+ assert "anthropic_version" not in body
+ assert route.calls.last.request.headers["anthropic-version"] == "2023-06-01"
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_clear_thinking_edit_is_forwarded_with_thinking_on(self):
+ edits = [{"type": "clear_thinking_20251015", "keep": "all"}, {"type": "clear_tool_uses_20250919"}]
+ route = await self._send(
+ context_management={"edits": edits},
+ thinking={"type": "adaptive"},
+ )
+
+ body = _sent_body(route)
+ assert body["context_management"] == {"edits": edits}
+ assert body["thinking"] == {"type": "adaptive"}
+ assert "context-management-2025-06-27" in _sent_betas(route)
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_tools_reach_mantle_unchanged(self):
+ tools = [
+ {
+ "name": "get_weather",
+ "description": "Look up the weather",
+ "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+ }
+ ]
+ route = await self._send(tools=tools, tool_choice={"type": "auto"})
+
+ body = _sent_body(route)
+ assert body["tools"] == tools
+ assert body["tool_choice"] == {"type": "auto"}
From 94b2fd827b4ea2678fa88ac0788e948f3b899348 Mon Sep 17 00:00:00 2001
From: Tin Chi Lo
Date: Sat, 19 Sep 2026 17:02:56 -0700
Subject: [PATCH 07/56] feat(ui): show prompt caching requests and net savings
---
backend/routes/allowlist.py | 1 +
.../prompt_caching_requests.py | 184 ++++++++++
litellm/proxy/proxy_server.py | 4 +
litellm/proxy/spend_tracking/savings.py | 75 ++--
.../prompt_caching_requests.py | 35 ++
.../test_prompt_caching_requests.py | 321 ++++++++++++++++++
.../proxy/spend_tracking/test_savings.py | 37 ++
.../_components/CacheLeakageCard.tsx | 6 +-
.../CostOptimizationView.activity.test.tsx | 1 +
...tCachingRequestsTable.integration.test.tsx | 248 ++++++++++++++
.../PromptCachingRequestsTable.tsx | 186 ++++++++++
.../_components/PromptCachingTab.test.tsx | 23 +-
.../_components/PromptCachingTab.tsx | 7 +
ui/litellm-dashboard/src/lib/http/schema.d.ts | 95 ++++++
14 files changed, 1195 insertions(+), 28 deletions(-)
create mode 100644 litellm/proxy/management_endpoints/prompt_caching_requests.py
create mode 100644 litellm/types/management_endpoints/prompt_caching_requests.py
create mode 100644 tests/test_litellm/proxy/management_endpoints/test_prompt_caching_requests.py
create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.integration.test.tsx
create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.tsx
diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py
index 00c4e0070e6..c7f389c36a4 100644
--- a/backend/routes/allowlist.py
+++ b/backend/routes/allowlist.py
@@ -51,6 +51,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/cache_settings",
"/coordination_redis/",
"/cost_tracking",
+ "/cost_optimization/",
"/cost/",
"/credentials",
"/credential",
diff --git a/litellm/proxy/management_endpoints/prompt_caching_requests.py b/litellm/proxy/management_endpoints/prompt_caching_requests.py
new file mode 100644
index 00000000000..41255bd49b8
--- /dev/null
+++ b/litellm/proxy/management_endpoints/prompt_caching_requests.py
@@ -0,0 +1,184 @@
+from collections.abc import Callable, Mapping
+from datetime import datetime, timezone
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Annotated, Final
+
+from fastapi import APIRouter, Depends, HTTPException, Query
+from pydantic import BaseModel, Json, TypeAdapter
+
+from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth, user_api_key_has_admin_view
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.spend_tracking.savings import (
+ extract_cache_creation_tokens,
+ extract_cache_read_tokens,
+ marks_gateway_injection,
+ prompt_caching_savings_for_request,
+)
+from litellm.proxy.spend_tracking.spend_tracking_utils import (
+ _query_raw_rows, # pyright: ignore[reportPrivateUsage] # existing typed spend-query adapter; rows validated below
+)
+from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY
+from litellm.types.management_endpoints.prompt_caching_requests import (
+ PromptCachingRequest,
+ PromptCachingRequestCursor,
+ PromptCachingRequestFilter,
+ PromptCachingRequestsResponse,
+)
+
+if TYPE_CHECKING:
+ from litellm.router import Router
+
+router: Final = APIRouter()
+
+
+def _numeric_token_sql(path: str) -> str:
+ value: Final = f"metadata #> '{{usage_object,{path}}}'"
+ return (
+ f"CASE WHEN jsonb_typeof({value}) = 'number' THEN ({value} #>> '{{}}')::numeric "
+ f"WHEN {value} = 'true'::jsonb THEN 1 WHEN {value} = 'false'::jsonb THEN 0 END"
+ )
+
+
+def _cache_tokens_sql(*paths: str) -> str:
+ candidates: Final = ", ".join(f"NULLIF(({_numeric_token_sql(path)}), 0)" for path in paths)
+ return f"TRUNC(COALESCE({candidates}, 0))"
+
+
+_CACHE_READ_SQL: Final = _cache_tokens_sql("cache_read_input_tokens", "prompt_tokens_details,cached_tokens")
+_CACHE_CREATION_SQL: Final = _cache_tokens_sql(
+ "cache_creation_input_tokens",
+ "prompt_tokens_details,cache_write_tokens",
+ "prompt_tokens_details,cache_creation_tokens",
+)
+_GATEWAY_INJECTED_SQL: Final = (
+ f"(jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string' "
+ f"AND (metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = '' "
+ f"OR metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = model_id))"
+)
+_FILTER_SQL: Final = MappingProxyType(
+ {
+ "all": f"({_GATEWAY_INJECTED_SQL} OR {_CACHE_READ_SQL} > 0 OR {_CACHE_CREATION_SQL} > 0)",
+ "injected": _GATEWAY_INJECTED_SQL,
+ "hits": f"{_CACHE_READ_SQL} > 0",
+ }
+)
+
+
+def prompt_caching_requests_sql(filter: PromptCachingRequestFilter) -> str:
+ return f"""
+ SELECT request_id, "startTime" AS start_time, "endTime" AS end_time,
+ model, model_id, custom_llm_provider, spend,
+ CASE WHEN jsonb_typeof(metadata->'usage_object') = 'object'
+ THEN metadata->'usage_object' END AS usage_object,
+ CASE WHEN jsonb_typeof(metadata->'cost_breakdown') = 'object'
+ THEN metadata->'cost_breakdown' END AS cost_breakdown,
+ CASE WHEN jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string'
+ THEN metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' END AS gateway_marker
+ FROM "LiteLLM_SpendLogs"
+ WHERE "startTime" >= ($1::text::timestamptz AT TIME ZONE 'UTC')
+ AND "startTime" <= ($2::text::timestamptz AT TIME ZONE 'UTC')
+ AND COALESCE(LOWER(cache_hit), 'false') != 'true'
+ AND {_FILTER_SQL[filter]}
+ AND ($4::text::timestamptz IS NULL OR
+ ("startTime", request_id) < (($4::text::timestamptz AT TIME ZONE 'UTC'), $5::text))
+ ORDER BY "startTime" DESC, request_id DESC
+ LIMIT $3::integer
+ """
+
+
+class _PromptCachingRow(BaseModel):
+ request_id: str
+ start_time: datetime
+ end_time: datetime
+ model: str
+ model_id: str | None
+ custom_llm_provider: str | None
+ spend: float
+ usage_object: Json[Mapping[str, object]] | Mapping[str, object] | None
+ cost_breakdown: Json[Mapping[str, object]] | Mapping[str, object] | None
+ gateway_marker: str | None
+
+
+_REQUEST_ROWS: Final = TypeAdapter(tuple[_PromptCachingRow, ...])
+
+
+def _request_result(row: _PromptCachingRow, llm_router: "Callable[[], Router | None]") -> PromptCachingRequest:
+ return PromptCachingRequest(
+ request_id=row.request_id,
+ start_time=row.start_time.replace(tzinfo=timezone.utc) if row.start_time.tzinfo is None else row.start_time,
+ model=row.model,
+ gateway_injected=marks_gateway_injection(
+ MappingProxyType({GATEWAY_INJECTED_CACHE_METADATA_KEY: row.gateway_marker}), row.model_id
+ ),
+ cache_read_tokens=extract_cache_read_tokens(row.usage_object),
+ cache_creation_tokens=extract_cache_creation_tokens(row.usage_object),
+ spend=row.spend,
+ net_savings=prompt_caching_savings_for_request(
+ model=row.model,
+ custom_llm_provider=row.custom_llm_provider,
+ usage_object=row.usage_object,
+ model_id=row.model_id,
+ llm_router=llm_router,
+ cost_breakdown=row.cost_breakdown,
+ billed_at=row.end_time,
+ ),
+ )
+
+
+@router.get(
+ "/cost_optimization/prompt_caching/requests",
+ tags=["Cost Optimization"], # mutable-ok: FastAPI's route API requires a list
+ response_model=PromptCachingRequestsResponse,
+)
+async def get_prompt_caching_requests(
+ user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+ start_date: datetime,
+ end_date: datetime,
+ page_size: Annotated[int, Query(ge=1, le=100)] = 50,
+ filter: PromptCachingRequestFilter = "all",
+ cursor_start_time: datetime | None = None,
+ cursor_request_id: Annotated[str | None, Query(min_length=1)] = None,
+) -> PromptCachingRequestsResponse:
+ from litellm.proxy.proxy_server import llm_router, prisma_client
+
+ if not user_api_key_has_admin_view(user_api_key_dict):
+ raise HTTPException(status_code=403, detail="Only proxy admin roles can view prompt caching requests")
+ if (cursor_start_time is None) != (cursor_request_id is None):
+ raise HTTPException(status_code=400, detail="cursor_start_time and cursor_request_id must be provided together")
+ if prisma_client is None:
+ raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
+ start: Final = start_date.replace(tzinfo=timezone.utc) if start_date.tzinfo is None else start_date
+ end: Final = end_date.replace(tzinfo=timezone.utc) if end_date.tzinfo is None else end_date
+ if end < start:
+ raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date")
+ cursor_time: Final = (
+ cursor_start_time.replace(tzinfo=timezone.utc)
+ if cursor_start_time is not None and cursor_start_time.tzinfo is None
+ else cursor_start_time
+ )
+ rows: Final = _REQUEST_ROWS.validate_python(
+ await _query_raw_rows(
+ prisma_client,
+ prompt_caching_requests_sql(filter),
+ start.isoformat(),
+ end.isoformat(),
+ page_size + 1,
+ cursor_time.isoformat() if cursor_time is not None else None,
+ cursor_request_id,
+ )
+ or ()
+ )
+
+ def current_router() -> "Router | None":
+ return llm_router
+
+ requests: Final = tuple(_request_result(row, current_router) for row in rows[:page_size])
+ has_more: Final = len(rows) > page_size
+ return PromptCachingRequestsResponse(
+ requests=requests,
+ page_size=page_size,
+ has_more=has_more,
+ next_cursor=PromptCachingRequestCursor(start_time=requests[-1].start_time, request_id=requests[-1].request_id)
+ if has_more
+ else None,
+ )
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index af25d418a63..f4a56e225cc 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -587,6 +587,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
from litellm.proxy.management_endpoints.organization_endpoints import (
router as organization_router,
)
+from litellm.proxy.management_endpoints.prompt_caching_requests import (
+ router as prompt_caching_requests_router,
+)
from litellm.proxy.management_endpoints.router_settings_endpoints import (
router as router_settings_router,
)
@@ -19183,6 +19186,7 @@ app.include_router(workflow_management_router)
app.include_router(memory_router)
app.include_router(plugin_router)
app.include_router(cost_tracking_settings_router)
+app.include_router(prompt_caching_requests_router)
app.include_router(router_settings_router)
app.include_router(fallback_management_router)
app.include_router(cache_settings_router)
diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py
index b7a2ac62844..fbcf9c78d3e 100644
--- a/litellm/proxy/spend_tracking/savings.py
+++ b/litellm/proxy/spend_tracking/savings.py
@@ -578,6 +578,56 @@ def autorouter_savings_for_logging_payload(
)
+def _request_savings_pricing(
+ model: str | None,
+ custom_llm_provider: str | None,
+ model_id: str | None,
+ llm_router: "Callable[[], Router | None] | None",
+) -> tuple[str | None, ModelInfo | None]:
+ router_instance: Final = llm_router() if llm_router else None
+ identity: Final = _resolve_model(model, custom_llm_provider)
+ pricing: Final = _effective_model_info(router_instance, model_id, model or "") or (
+ _model_info(identity) if identity else None
+ )
+ return identity.provider if identity else custom_llm_provider, pricing
+
+
+def _prompt_caching_savings(
+ pricing: ModelInfo | None,
+ provider: str | None,
+ usage_object: Mapping[str, object] | None,
+ cost_breakdown: Mapping[str, object] | None,
+ billed_at: datetime | str | None,
+) -> float | None:
+ usage: Final = _usage_from_spend_log(usage_object)
+ if pricing is None or usage is None:
+ return None
+ basis: Final = _pricing_basis(cost_breakdown)
+ result: Final = calculate_prompt_caching_savings(
+ model_info=pricing,
+ usage=usage,
+ custom_llm_provider=provider,
+ service_tier=basis.service_tier,
+ data_residency=basis.data_residency,
+ vertex_location=basis.vertex_location,
+ billed_at=_coerce_billed_at(billed_at),
+ )
+ return result if isfinite(result) else None
+
+
+def prompt_caching_savings_for_request(
+ model: str | None,
+ custom_llm_provider: str | None,
+ usage_object: Mapping[str, object] | None,
+ model_id: str | None = None,
+ llm_router: "Callable[[], Router | None] | None" = None,
+ cost_breakdown: Mapping[str, object] | None = None,
+ billed_at: datetime | str | None = None,
+) -> float | None:
+ request_pricing: Final = _request_savings_pricing(model, custom_llm_provider, model_id, llm_router)
+ return _prompt_caching_savings(request_pricing[1], request_pricing[0], usage_object, cost_breakdown, billed_at)
+
+
def compute_savings_spend(
model: str | None,
custom_llm_provider: str | None,
@@ -639,29 +689,12 @@ def compute_savings_spend(
# Deployment rates when the request came through one, public rates otherwise --
# `_effective_model_info` merges a deployment's configured prices over the built-in
# map, so a negotiated price is not silently replaced by the list rate.
- router_instance: Router | None = llm_router() if llm_router else None
- identity: Final = _resolve_model(model, custom_llm_provider)
- pricing: Final = _effective_model_info(router_instance, model_id, model or "") or (
- _model_info(identity) if identity else None
- )
+ request_pricing: Final = _request_savings_pricing(model, custom_llm_provider, model_id, llm_router)
+ provider: Final = request_pricing[0]
+ pricing: Final = request_pricing[1]
input_cost: Final = (_get_cost_per_unit(pricing, "input_cost_per_token") or 0.0) if pricing else 0.0
compression: Final = max(compression_saved_tokens, 0) * input_cost
- usage: Final = _usage_from_spend_log(usage_object)
- basis: Final = _pricing_basis(cost_breakdown)
- billed_at_datetime: Final = _coerce_billed_at(billed_at)
- prompt_caching: Final = (
- calculate_prompt_caching_savings(
- model_info=pricing,
- usage=usage,
- custom_llm_provider=identity.provider if identity else custom_llm_provider,
- service_tier=basis.service_tier,
- data_residency=basis.data_residency,
- vertex_location=basis.vertex_location,
- billed_at=billed_at_datetime,
- )
- if pricing is not None and usage is not None
- else 0.0
- )
+ prompt_caching: Final = _prompt_caching_savings(pricing, provider, usage_object, cost_breakdown, billed_at) or 0.0
gateway_injected_caching: Final = prompt_caching if gateway_injected_cache else 0.0
# The figure the logging path recorded wins, before the usage gate on purpose: a row
diff --git a/litellm/types/management_endpoints/prompt_caching_requests.py b/litellm/types/management_endpoints/prompt_caching_requests.py
new file mode 100644
index 00000000000..e72183a113b
--- /dev/null
+++ b/litellm/types/management_endpoints/prompt_caching_requests.py
@@ -0,0 +1,35 @@
+from datetime import datetime
+from typing import Literal, TypeAlias
+
+from pydantic import BaseModel, ConfigDict
+
+PromptCachingRequestFilter: TypeAlias = Literal["all", "injected", "hits"]
+
+
+class PromptCachingRequest(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ request_id: str
+ start_time: datetime
+ model: str
+ gateway_injected: bool
+ cache_read_tokens: int
+ cache_creation_tokens: int
+ spend: float
+ net_savings: float | None
+
+
+class PromptCachingRequestCursor(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ start_time: datetime
+ request_id: str
+
+
+class PromptCachingRequestsResponse(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ requests: tuple[PromptCachingRequest, ...]
+ page_size: int
+ has_more: bool
+ next_cursor: PromptCachingRequestCursor | None
diff --git a/tests/test_litellm/proxy/management_endpoints/test_prompt_caching_requests.py b/tests/test_litellm/proxy/management_endpoints/test_prompt_caching_requests.py
new file mode 100644
index 00000000000..0995de6c39d
--- /dev/null
+++ b/tests/test_litellm/proxy/management_endpoints/test_prompt_caching_requests.py
@@ -0,0 +1,321 @@
+import json
+from collections.abc import AsyncIterator, Mapping
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from types import SimpleNamespace
+from typing import Final
+
+import httpx
+import psycopg
+import pytest
+import pytest_asyncio
+from fastapi import FastAPI
+from prisma import Prisma
+from pydantic import TypeAdapter
+from pytest_postgresql import factories
+
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.management_endpoints.prompt_caching_requests import router
+from litellm.proxy.spend_tracking.savings import (
+ extract_cache_creation_tokens,
+ extract_cache_read_tokens,
+ marks_gateway_injection,
+)
+from litellm.types.management_endpoints.prompt_caching_requests import (
+ PromptCachingRequestFilter,
+ PromptCachingRequestsResponse,
+)
+
+pytestmark = pytest.mark.usefixtures("local_model_cost_map")
+
+_cache_postgresql_proc: Final = factories.postgresql_proc() # pyright: ignore[reportUnknownMemberType] # third-party fixture factory has incomplete callable types
+_cache_postgresql: Final = factories.postgresql("_cache_postgresql_proc")
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
+_JSON_ROWS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
+_START: Final = "2026-09-01T00:00:00Z"
+_END: Final = "2026-09-02T00:00:00Z"
+_URL: Final = "/cost_optimization/prompt_caching/requests"
+_MODEL: Final = "claude-sonnet-5"
+_MARKER: Final = "litellm_gateway_injected_cache"
+_DDL: Final = """
+ CREATE TABLE "LiteLLM_SpendLogs" (
+ request_id TEXT PRIMARY KEY, "startTime" TIMESTAMP, "endTime" TIMESTAMP,
+ model TEXT, model_id TEXT, custom_llm_provider TEXT, spend DOUBLE PRECISION,
+ metadata JSONB, cache_hit TEXT
+ )
+"""
+
+
+@dataclass(frozen=True)
+class _Case:
+ request_id: str
+ metadata: Mapping[str, object]
+ cache_hit: str | None = None
+ start_time: datetime = datetime(2026, 9, 1, 12, 0, 0, 123456)
+
+ def matches(self, filter: PromptCachingRequestFilter) -> bool:
+ if self.cache_hit is not None and self.cache_hit.lower() == "true":
+ return False
+ if not datetime(2026, 9, 1) <= self.start_time <= datetime(2026, 9, 2):
+ return False
+ usage: Final = self.metadata.get("usage_object")
+ normalized: Final = _JSON_OBJECT.validate_python(usage) if isinstance(usage, Mapping) else None
+ injected: Final = marks_gateway_injection(self.metadata, "dep-a")
+ reads: Final = extract_cache_read_tokens(normalized)
+ writes: Final = extract_cache_creation_tokens(normalized)
+ match filter:
+ case "injected":
+ return injected
+ case "hits":
+ return reads > 0
+ case "all":
+ return injected or reads > 0 or writes > 0
+
+
+_CASES: Final = (
+ _Case("injected-empty", {_MARKER: ""}),
+ _Case("injected-deployment", {_MARKER: "dep-a"}),
+ _Case("wrong-deployment", {_MARKER: "dep-b"}),
+ _Case("legacy-read", {"usage_object": {"cache_read_input_tokens": 100}}),
+ _Case("nested-read", {"usage_object": {"prompt_tokens_details": {"cached_tokens": 100}}}),
+ _Case("write", {"usage_object": {"cache_creation_input_tokens": 100}}),
+ _Case("nested-write", {"usage_object": {"prompt_tokens_details": {"cache_write_tokens": 100}}}),
+ _Case("nested-creation", {"usage_object": {"prompt_tokens_details": {"cache_creation_tokens": 100}}}),
+ _Case(
+ "top-precedence",
+ {"usage_object": {"cache_read_input_tokens": -2, "prompt_tokens_details": {"cached_tokens": 100}}},
+ ),
+ _Case(
+ "zero-fallback",
+ {"usage_object": {"cache_read_input_tokens": 0, "prompt_tokens_details": {"cached_tokens": 100}}},
+ ),
+ _Case(
+ "fractional-precedence",
+ {"usage_object": {"cache_read_input_tokens": 0.5, "prompt_tokens_details": {"cached_tokens": 100}}},
+ ),
+ _Case("malformed-number", {"usage_object": {"cache_read_input_tokens": "100"}}),
+ _Case("malformed-container", {"usage_object": [100]}),
+ _Case("boolean-number", {"usage_object": {"cache_read_input_tokens": True}}),
+ _Case("boolean-marker", {_MARKER: True}),
+ _Case("response-cache", {_MARKER: "", "usage_object": {"cache_read_input_tokens": 100}}, "True"),
+ _Case("outside-before", {_MARKER: ""}, start_time=datetime(2026, 8, 31, 23, 59, 59)),
+ _Case(
+ "outside-after", {"usage_object": {"cache_read_input_tokens": 100}}, start_time=datetime(2026, 9, 2, 0, 0, 1)
+ ),
+)
+
+
+@pytest_asyncio.fixture(loop_scope="function")
+async def _cache_prisma(
+ _cache_postgresql: psycopg.Connection[tuple[object, ...]],
+) -> AsyncIterator[Prisma]:
+ info: Final = _cache_postgresql.info
+ database: Final = Prisma(datasource={
+ "url": f"postgresql://{info.user}@{info.host}:{info.port}/{info.dbname}?connection_limit=1",
+ })
+ await database.connect()
+ try:
+ yield database
+ finally:
+ await database.disconnect()
+
+
+def _seed(connection: psycopg.Connection[tuple[object, ...]], cases: tuple[_Case, ...] = _CASES) -> None:
+ with connection.cursor() as cursor:
+ cursor.execute(_DDL)
+ cursor.executemany(
+ """INSERT INTO "LiteLLM_SpendLogs"
+ VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s)""",
+ tuple(
+ (
+ case.request_id,
+ case.start_time,
+ datetime(2026, 9, 1, 12, 0, 1),
+ _MODEL,
+ "dep-a",
+ "anthropic",
+ 0.01,
+ json.dumps(dict(case.metadata)),
+ case.cache_hit,
+ )
+ for case in cases
+ ),
+ )
+ connection.commit()
+
+
+def _app(role: LitellmUserRoles | None) -> FastAPI:
+ application: Final = FastAPI()
+ application.include_router(router)
+
+ def caller() -> UserAPIKeyAuth:
+ return UserAPIKeyAuth(user_role=role)
+
+ application.dependency_overrides[user_api_key_auth] = caller
+ return application
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("filter", ["all", "injected", "hits"])
+@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
+async def test_request_filters_match_accounting_and_paginate_before_projection(
+ _cache_postgresql: psycopg.Connection[tuple[object, ...]],
+ _cache_prisma: Prisma,
+ monkeypatch: pytest.MonkeyPatch,
+ filter: PromptCachingRequestFilter,
+ role: LitellmUserRoles,
+) -> None:
+ from litellm.proxy import proxy_server
+
+ _seed(_cache_postgresql)
+ monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=_cache_prisma))
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ expected: Final = tuple(sorted((case.request_id for case in _CASES if case.matches(filter)), reverse=True))
+ async with httpx.AsyncClient(transport=httpx.ASGITransport(app=_app(role)), base_url="http://test") as client:
+ first: Final = await client.get(
+ _URL, params={"start_date": _START, "end_date": _END, "filter": filter, "page_size": 2}
+ )
+ assert first.status_code == 200
+ first_page: Final = PromptCachingRequestsResponse.model_validate_json(first.content)
+ assert tuple(row.request_id for row in first_page.requests) == expected[:2]
+ assert first_page.has_more is (len(expected) > 2)
+ assert (first_page.next_cursor is not None) is first_page.has_more
+ if first_page.next_cursor is not None:
+ assert first_page.next_cursor.request_id == expected[1]
+ assert first_page.next_cursor.start_time == first_page.requests[-1].start_time
+ next_response: Final = await client.get(
+ _URL, params={
+ "start_date": _START, "end_date": _END, "filter": filter, "page_size": 2,
+ "cursor_start_time": first_page.next_cursor.start_time.astimezone(
+ timezone(timedelta(hours=-7))
+ ).isoformat(),
+ "cursor_request_id": first_page.next_cursor.request_id,
+ }
+ )
+ assert next_response.status_code == 200
+ next_page: Final = PromptCachingRequestsResponse.model_validate_json(next_response.content)
+ assert tuple(row.request_id for row in next_page.requests) == expected[2:4]
+ assert next_page.has_more is (len(expected) > 4)
+ assert (next_page.next_cursor is not None) is next_page.has_more
+ second: Final = await client.get(
+ _URL, params={"start_date": _START, "end_date": _END, "filter": filter, "page_size": 100}
+ )
+ assert second.status_code == 200
+ complete: Final = PromptCachingRequestsResponse.model_validate_json(second.content)
+ assert tuple(row.request_id for row in complete.requests) == expected
+ assert complete.has_more is False
+ assert complete.next_cursor is None
+ assert all(row.start_time.tzinfo == timezone.utc for row in complete.requests)
+ payload: Final = _JSON_OBJECT.validate_json(second.content)
+ assert set(payload) == {"requests", "page_size", "has_more", "next_cursor"}
+ serialized_rows: Final = _JSON_ROWS.validate_python(payload["requests"])
+ assert set(serialized_rows[0]) == {
+ "request_id",
+ "start_time",
+ "model",
+ "gateway_injected",
+ "cache_read_tokens",
+ "cache_creation_tokens",
+ "spend",
+ "net_savings",
+ }
+ by_id: Final = {row.request_id: row for row in complete.requests}
+ if filter == "all":
+ assert by_id["injected-empty"].gateway_injected is True
+ assert by_id["injected-empty"].net_savings is None
+ assert by_id["legacy-read"].gateway_injected is False
+ assert by_id["legacy-read"].net_savings is not None and by_id["legacy-read"].net_savings > 0
+ assert by_id["write"].net_savings is not None and by_id["write"].net_savings < 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("role", [None, LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY])
+async def test_non_admin_is_denied_before_database_access(
+ role: LitellmUserRoles | None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ from litellm.proxy import proxy_server
+
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ async with httpx.AsyncClient(transport=httpx.ASGITransport(app=_app(role)), base_url="http://test") as client:
+ response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END})
+ assert response.status_code == 403
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("params", [
+ {"filter": "savings"}, {"page_size": 0}, {"page_size": 101}, {"start_date": "invalid"},
+ {"cursor_start_time": "invalid", "cursor_request_id": "request"},
+ {"cursor_start_time": _START, "cursor_request_id": ""},
+])
+async def test_invalid_request_is_rejected(params: Mapping[str, str | int]) -> None:
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test"
+ ) as client:
+ response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, **params})
+ assert response.status_code == 422
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("params", [{"cursor_start_time": _START}, {"cursor_request_id": "request"}])
+async def test_incomplete_cursor_is_rejected(
+ params: Mapping[str, str], monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from litellm.proxy import proxy_server
+
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test"
+ ) as client:
+ response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, **params})
+ assert response.status_code == 400
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("delete_before_cursor", [False, True])
+async def test_cursor_keeps_remaining_requests_once_during_insertions_and_deletions(
+ _cache_postgresql: psycopg.Connection[tuple[object, ...]],
+ _cache_prisma: Prisma,
+ monkeypatch: pytest.MonkeyPatch,
+ delete_before_cursor: bool,
+) -> None:
+ from litellm.proxy import proxy_server
+
+ cases: Final = (*_CASES, _Case(
+ "older-cache-read", {"usage_object": {"cache_read_input_tokens": 100}}, start_time=datetime(2026, 9, 1, 11),
+ ))
+ _seed(_cache_postgresql, cases)
+ monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=_cache_prisma))
+ monkeypatch.setattr(proxy_server, "llm_router", None)
+ expected: Final = (*sorted((case.request_id for case in _CASES if case.matches("all")), reverse=True), "older-cache-read")
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test"
+ ) as client:
+ first: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, "page_size": 2})
+ assert first.status_code == 200
+ first_page: Final = PromptCachingRequestsResponse.model_validate_json(first.content)
+ assert tuple(row.request_id for row in first_page.requests) == expected[:2]
+ assert first_page.next_cursor is not None
+ with _cache_postgresql.cursor() as cursor:
+ cursor.executemany(
+ """INSERT INTO "LiteLLM_SpendLogs"
+ SELECT %s, %s, "endTime", model, model_id, custom_llm_provider, spend, metadata, cache_hit
+ FROM "LiteLLM_SpendLogs" WHERE request_id = %s""",
+ (
+ ("newer-request", datetime(2026, 9, 1, 13), expected[0]),
+ ("zz-higher-id", cases[0].start_time, expected[0]),
+ ),
+ )
+ if delete_before_cursor:
+ cursor.execute('DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (expected[0],))
+ _cache_postgresql.commit()
+ following: Final = await client.get(_URL, params={
+ "start_date": _START, "end_date": _END, "page_size": 100,
+ "cursor_start_time": first_page.next_cursor.start_time.isoformat(),
+ "cursor_request_id": first_page.next_cursor.request_id,
+ })
+ assert following.status_code == 200
+ following_page: Final = PromptCachingRequestsResponse.model_validate_json(following.content)
+ assert tuple(row.request_id for row in following_page.requests) == expected[2:]
+ assert following_page.has_more is False
+ assert following_page.next_cursor is None
diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py
index aae966022e3..004f07da431 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_savings.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py
@@ -11,6 +11,7 @@ from litellm.proxy.spend_tracking.savings import (
compute_autorouter_savings,
compute_savings_spend,
marks_gateway_injection,
+ prompt_caching_savings_for_request,
)
from litellm.router import Router
from litellm.types.utils import Usage
@@ -18,6 +19,42 @@ from litellm.types.utils import Usage
pytestmark = pytest.mark.usefixtures("local_model_cost_map")
+@pytest.mark.parametrize("model,usage", [
+ (None, {"cache_read_input_tokens": 100}),
+ ("claude-sonnet-5", None),
+ ("claude-sonnet-5", {"prompt_tokens": "invalid"}),
+])
+def test_prompt_cache_estimate_distinguishes_unknown_from_zero(model: str | None, usage: dict[str, object] | None) -> None:
+ assert prompt_caching_savings_for_request(model, "anthropic", usage) is None
+ assert compute_savings_spend(model, "anthropic", 0, False, usage_object=usage).prompt_caching == 0
+ assert prompt_caching_savings_for_request("claude-sonnet-5", "anthropic", {"prompt_tokens": 100}) == 0
+
+
+def test_prompt_cache_estimate_uses_the_rollup_pricing_and_retains_write_premiums() -> None:
+ router: Final = Router(model_list=[{
+ "model_name": "negotiated",
+ "litellm_params": {
+ "model": "anthropic/claude-sonnet-5", "input_cost_per_token": 1e-6,
+ "cache_creation_input_token_cost": 1.25e-6, "cache_read_input_token_cost": 1e-7,
+ },
+ "model_info": {"id": "negotiated-cache-prices"},
+ }])
+
+ def current_router() -> Router:
+ return router
+
+ usage: Final = {"cache_read_input_tokens": 1000, "cache_creation_input_tokens": 20000}
+ estimate: Final = prompt_caching_savings_for_request(
+ "claude-sonnet-5", "anthropic", usage, model_id="negotiated-cache-prices", llm_router=current_router,
+ )
+ rollup: Final = compute_savings_spend(
+ "claude-sonnet-5", "anthropic", 0, True, usage_object=usage,
+ model_id="negotiated-cache-prices", llm_router=current_router,
+ )
+ assert estimate == pytest.approx(1000 * (1e-6 - 1e-7) - 20000 * (1.25e-6 - 1e-6))
+ assert estimate == rollup.prompt_caching == rollup.gateway_injected_caching
+
+
@pytest.mark.parametrize("modifier", [{"speed": "fast"}, {"inference_geo": "us"}])
@pytest.mark.parametrize("continuing", [False, True])
def test_baseline_preserves_anthropic_pricing_fields(modifier: dict[str, str], continuing: bool) -> None:
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx
index a0877b04648..f5b71a00061 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx
@@ -3,7 +3,6 @@
import React, { useMemo, useState } from "react";
import { ArrowDown, ArrowUp, ArrowUpDown, Info } from "lucide-react";
-import AdvancedDatePicker from "@/components/shared/advanced_date_picker";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
@@ -81,7 +80,7 @@ const SortableHead = ({
};
const CacheLeakageCard: React.FC = ({ activity }) => {
- const { dateValue, onDateChange, results, loading, isFetchingMore, apiKeyTruncation } = activity;
+ const { results, loading, isFetchingMore, apiKeyTruncation } = activity;
const [dimension, setDimension] = useState("key");
const [sort, setSort] = useState({ column: "potentialSavings", dir: "desc" });
const leakage = useMemo(() => computeCacheLeakage(results, dimension), [results, dimension]);
@@ -111,9 +110,6 @@ const CacheLeakageCard: React.FC = ({ activity }) => {
cached token, after cache-write premiums.
-
setDimension(value === "model" ? "model" : "key")}>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx
index 03250e3e53b..f8336f5ab56 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx
@@ -42,6 +42,7 @@ vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () =>
}));
vi.mock("./PromptCompressionTab", () => ({ __esModule: true, default: () =>
}));
+vi.mock("./PromptCachingRequestsTable", () => ({ default: () =>
}));
import CostOptimizationView from "./CostOptimizationView";
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.integration.test.tsx
new file mode 100644
index 00000000000..833a46ce16f
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.integration.test.tsx
@@ -0,0 +1,248 @@
+import { Profiler } from "react";
+import { act, fireEvent, renderWithProviders, screen, testQueryClient, waitFor, within } from "@/../tests/test-utils";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+import type { components } from "@/lib/http/schema";
+import PromptCachingRequestsTable from "./PromptCachingRequestsTable";
+import type { DateRange } from "./useDailyActivityRange";
+
+type CacheRequest = components["schemas"]["PromptCachingRequest"];
+type RequestsResponse = components["schemas"]["PromptCachingRequestsResponse"];
+const firstCursor = { start_time: "2026-09-01T11:59:59.123456Z", request_id: "first-boundary?&" };
+const secondCursor = { start_time: firstCursor.start_time, request_id: "second-boundary" };
+const fetchMock = vi.fn();
+const dates = { from: new Date(2026, 8, 1, 12), to: new Date(2026, 8, 2, 12) };
+const request = (overrides: Partial = {}): CacheRequest => ({
+ request_id: "request-default",
+ start_time: "2026-09-01T12:00:00Z",
+ model: "cache-test-model",
+ gateway_injected: true,
+ cache_read_tokens: 0,
+ cache_creation_tokens: 1000,
+ spend: 0.0375,
+ net_savings: -0.0075,
+ ...overrides,
+});
+const response = (requests: CacheRequest[], nextCursor: RequestsResponse["next_cursor"] = null) => {
+ const body: RequestsResponse = { requests, has_more: nextCursor !== null, next_cursor: nextCursor, page_size: 50 };
+ return Response.json(body);
+};
+const lastQuery = () => new URL(String(fetchMock.mock.calls.at(-1)?.[0]), "http://localhost").searchParams;
+
+describe("PromptCachingRequestsTable", () => {
+ beforeEach(() => {
+ fetchMock.mockReset();
+ vi.stubGlobal("fetch", fetchMock);
+ });
+
+ afterEach(() => {
+ testQueryClient.clear();
+ vi.unstubAllGlobals();
+ vi.unstubAllEnvs();
+ vi.useRealTimers();
+ });
+
+ it("separates recorded injection from cache hits, retains write premiums and unknown savings, and links each request", async () => {
+ const clientHit = {
+ request_id: "client-hit",
+ gateway_injected: false,
+ cache_read_tokens: 10000,
+ cache_creation_tokens: 0,
+ net_savings: 0.27,
+ };
+ fetchMock.mockResolvedValue(
+ response([
+ request({ request_id: "injected/write?&", net_savings: -0.0075 }),
+ request(clientHit),
+ request({ request_id: "unknown-price", net_savings: null }),
+ request({ request_id: "no-benefit", net_savings: 0 }),
+ ]),
+ );
+ renderWithProviders( );
+
+ const table = await screen.findByRole("table", { name: "Prompt caching requests" });
+ const write = within(table).getByRole("row", { name: /injected\/write/ });
+ expect(within(write).getByText("Recorded")).toBeInTheDocument();
+ expect(within(write).getByText("1,000")).toBeInTheDocument();
+ expect(within(write).getByText("$0.0375")).toBeInTheDocument();
+ expect(within(write).getByText("-$0.0075")).toBeInTheDocument();
+ expect(within(write).getByText(new Date("2026-09-01T12:00:00Z").toLocaleString())).toBeInTheDocument();
+ expect(within(write).getByText("cache-test-model")).toHaveAttribute("title", "cache-test-model");
+ expect(within(write).getByRole("link")).toHaveAttribute("href", "/ui/logs?log_id=injected%2Fwrite%3F%26");
+
+ const hit = within(table).getByRole("row", { name: /client-hit/ });
+ expect(within(hit).getByText("Not recorded")).toBeInTheDocument();
+ expect(within(hit).getByText("10,000")).toBeInTheDocument();
+ expect(within(hit).getByText("$0.2700")).toBeInTheDocument();
+ expect(within(table).getByRole("row", { name: /unknown-price/ })).toHaveTextContent("Unavailable");
+ expect(within(table).getByRole("row", { name: /no-benefit/ })).toHaveTextContent("$0.00");
+ expect(screen.getByText(/after cache-write premiums/)).toBeInTheDocument();
+ expect(lastQuery().get("start_date")).toBe("2026-09-01T00:00:00.000Z");
+ expect(lastQuery().get("end_date")).toBe("2026-09-02T23:59:59.999Z");
+ expect(fetchMock.mock.calls[0][1]?.headers).toEqual(expect.objectContaining({ Authorization: "Bearer token-a" }));
+ });
+
+ it("forwards complete server cursors, goes back to prior cursors, and clears them for each caching filter", async () => {
+ fetchMock.mockImplementation(async (input) => {
+ const query = new URL(String(input), "http://localhost").searchParams;
+ const pages = new Map([
+ [null, 1],
+ [firstCursor.request_id, 2],
+ [secondCursor.request_id, 3],
+ ]);
+ const page = pages.get(query.get("cursor_request_id"));
+ const nextCursor =
+ new Map([
+ [1, firstCursor],
+ [2, secondCursor],
+ ]).get(page ?? 0) ?? null;
+ return response([request({ request_id: `${query.get("filter")}-${page}` })], nextCursor);
+ });
+ renderWithProviders( );
+ await screen.findByRole("link", { name: "all-1" });
+ expect(screen.getByRole("button", { name: "Previous" })).toBeDisabled();
+ expect(lastQuery().has("page")).toBe(false);
+ expect(lastQuery().has("cursor_request_id")).toBe(false);
+
+ fireEvent.click(screen.getByRole("button", { name: "Next" }));
+ await screen.findByRole("link", { name: "all-2" });
+ expect(screen.getByText("Page 2")).toBeInTheDocument();
+ expect(lastQuery().get("cursor_start_time")).toBe(firstCursor.start_time);
+ expect(lastQuery().get("cursor_request_id")).toBe(firstCursor.request_id);
+ fireEvent.click(screen.getByRole("button", { name: "Next" }));
+ await screen.findByRole("link", { name: "all-3" });
+ expect(screen.getByText("Page 3")).toBeInTheDocument();
+ expect(lastQuery().get("cursor_start_time")).toBe(secondCursor.start_time);
+ expect(lastQuery().get("cursor_request_id")).toBe(secondCursor.request_id);
+ expect(screen.getByRole("button", { name: "Next" })).toBeDisabled();
+
+ await testQueryClient.invalidateQueries({ refetchType: "none" });
+ fireEvent.click(screen.getByRole("button", { name: "Previous" }));
+ await screen.findByRole("link", { name: "all-2" });
+ await waitFor(() => expect(lastQuery().get("cursor_request_id")).toBe(firstCursor.request_id));
+ expect(lastQuery().get("cursor_start_time")).toBe(firstCursor.start_time);
+ expect(screen.getByText("Page 2")).toBeInTheDocument();
+ fireEvent.click(screen.getByRole("button", { name: "Previous" }));
+ await screen.findByRole("link", { name: "all-1" });
+ await waitFor(() => expect(lastQuery().has("cursor_request_id")).toBe(false));
+ expect(lastQuery().has("cursor_start_time")).toBe(false);
+ fireEvent.click(screen.getByRole("button", { name: "Next" }));
+ await screen.findByRole("link", { name: "all-2" });
+
+ fireEvent.click(screen.getByRole("tab", { name: "LiteLLM injected" }));
+ await screen.findByRole("link", { name: "injected-1" });
+ expect(screen.queryByRole("link", { name: "all-2" })).not.toBeInTheDocument();
+ expect(lastQuery().get("filter")).toBe("injected");
+ expect(lastQuery().has("cursor_request_id")).toBe(false);
+ expect(lastQuery().has("cursor_start_time")).toBe(false);
+
+ fireEvent.click(screen.getByRole("button", { name: "Next" }));
+ await screen.findByRole("link", { name: "injected-2" });
+ fireEvent.click(screen.getByRole("tab", { name: "Cache hits" }));
+ await screen.findByRole("link", { name: "hits-1" });
+ expect(lastQuery().get("filter")).toBe("hits");
+ expect(lastQuery().get("page_size")).toBe("50");
+ expect(screen.getByText("Page 1")).toBeInTheDocument();
+ });
+
+ it("includes the current UTC day for a range ending today, matching the activity totals", async () => {
+ vi.stubEnv("TZ", "America/Los_Angeles");
+ vi.setSystemTime(new Date("2026-09-20T03:00:00Z"));
+ fetchMock.mockResolvedValue(response([]));
+ const today = { from: new Date(2026, 8, 19), to: new Date() };
+ renderWithProviders( );
+
+ await screen.findByText("No matching prompt caching requests in this range");
+ expect(lastQuery().get("start_date")).toBe("2026-09-19T00:00:00.000Z");
+ expect(lastQuery().get("end_date")).toBe("2026-09-20T23:59:59.999Z");
+ });
+
+ it.each(["date", "authentication"])(
+ "hides every old-scope frame and resets pagination when %s changes",
+ async (change) => {
+ fetchMock.mockResolvedValueOnce(response([request({ request_id: "old-first" })], firstCursor));
+ fetchMock.mockResolvedValueOnce(response([request({ request_id: "old-second" })]));
+ const committedOldRows: boolean[] = [];
+ const snapshot = () => {
+ committedOldRows.push(screen.queryByRole("link", { name: "old-second" }) !== null);
+ };
+ const tree = (accessToken: string, dateValue: DateRange) => (
+
+
+
+ );
+ const { rerender } = renderWithProviders(tree("token-a", dates));
+ await screen.findByRole("link", { name: "old-first" });
+ fireEvent.click(screen.getByRole("button", { name: "Next" }));
+ await screen.findByRole("link", { name: "old-second" });
+
+ const pending = Promise.withResolvers();
+ fetchMock.mockReturnValueOnce(pending.promise);
+ committedOldRows.length = 0;
+ rerender(
+ tree(
+ change === "authentication" ? "token-b" : "token-a",
+ change === "date" ? { ...dates, to: new Date(2026, 8, 3) } : dates,
+ ),
+ );
+
+ expect(screen.getByRole("status")).toHaveTextContent("Loading requests");
+ expect(committedOldRows.length).toBeGreaterThan(0);
+ expect(committedOldRows.every((visible) => !visible)).toBe(true);
+ expect(lastQuery().has("cursor_request_id")).toBe(false);
+ expect(lastQuery().has("cursor_start_time")).toBe(false);
+ if (change === "date") {
+ expect(lastQuery().get("end_date")).toBe("2026-09-03T23:59:59.999Z");
+ } else {
+ expect(fetchMock.mock.calls.at(-1)?.[1]?.headers).toEqual(
+ expect.objectContaining({ Authorization: "Bearer token-b" }),
+ );
+ }
+
+ pending.resolve(response([request({ request_id: "new-first" })]));
+ await screen.findByRole("link", { name: "new-first" });
+ expect(screen.getByText("Page 1")).toBeInTheDocument();
+ expect(committedOldRows.every((visible) => !visible)).toBe(true);
+ },
+ );
+
+ it("ignores a delayed response from the previous caching filter", async () => {
+ const stale = Promise.withResolvers();
+ const current = Promise.withResolvers();
+ fetchMock.mockReturnValueOnce(stale.promise).mockReturnValueOnce(current.promise);
+ renderWithProviders( );
+ fireEvent.click(screen.getByRole("tab", { name: "Cache hits" }));
+ expect(lastQuery().get("filter")).toBe("hits");
+
+ current.resolve(response([request({ request_id: "current-hit" })]));
+ await screen.findByRole("link", { name: "current-hit" });
+ await act(async () => {
+ stale.resolve(response([request({ request_id: "stale-all" })], firstCursor));
+ await stale.promise;
+ });
+
+ expect(screen.getByRole("link", { name: "current-hit" })).toBeInTheDocument();
+ expect(screen.queryByRole("link", { name: "stale-all" })).not.toBeInTheDocument();
+ expect(screen.getByRole("button", { name: "Next" })).toBeDisabled();
+ });
+
+ it("offers retry after a failed read and shows the empty state after it succeeds", async () => {
+ fetchMock.mockRejectedValueOnce(new Error("offline"));
+ fetchMock.mockResolvedValueOnce(response([]));
+ renderWithProviders( );
+
+ expect(await screen.findByRole("alert")).toHaveTextContent("Could not load prompt caching requests");
+ fireEvent.click(screen.getByRole("button", { name: "Retry" }));
+ expect(await screen.findByText("No matching prompt caching requests in this range")).toBeInTheDocument();
+ expect(screen.queryByRole("alert")).not.toBeInTheDocument();
+ expect(screen.getByRole("button", { name: "Next" })).toBeDisabled();
+ expect(fetchMock).toHaveBeenCalledTimes(2);
+ });
+
+ it("does not request data for an incomplete date range", async () => {
+ renderWithProviders( );
+ expect(screen.getByText("Select a date range to view requests")).toBeInTheDocument();
+ expect(screen.queryByRole("status")).not.toBeInTheDocument();
+ await waitFor(() => expect(fetchMock).not.toHaveBeenCalled());
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.tsx
new file mode 100644
index 00000000000..29aa9252e7b
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingRequestsTable.tsx
@@ -0,0 +1,186 @@
+"use client";
+
+import { useQuery, type UseQueryOptions } from "@tanstack/react-query";
+import Link from "next/link";
+import { useState } from "react";
+
+import { apiClient } from "@/components/networking";
+import { Button } from "@/components/ui/button";
+import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
+import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
+import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
+import { LOG_ID_QUERY_PARAM } from "@/components/view_logs/logDetailRouting";
+import type { paths } from "@/lib/http/schema";
+import { formatNumberWithCommas } from "@/utils/dataUtils";
+import { uiHref } from "@/utils/uiHref";
+import { usd } from "./costOptimizationUtils";
+import { benchmarksWindow as activityWindow } from "./useAutoRouterBenchmarks";
+import type { DateRange } from "./useDailyActivityRange";
+
+const REQUESTS_PATH = "/cost_optimization/prompt_caching/requests";
+type RequestsEndpoint = paths[typeof REQUESTS_PATH]["get"];
+type RequestsResponse = RequestsEndpoint["responses"][200]["content"]["application/json"];
+type RequestsQuery = NonNullable;
+type RequestFilter = NonNullable;
+type RequestCursor = RequestsResponse["next_cursor"];
+
+interface PromptCachingRequestsTableProps {
+ accessToken: string;
+ dateValue: DateRange;
+}
+
+export default function PromptCachingRequestsTable({ accessToken, dateValue }: PromptCachingRequestsTableProps) {
+ const [filter, setFilter] = useState("all");
+ const window = activityWindow(dateValue, new Date());
+ const startDate = window.start_date ? `${window.start_date}T00:00:00.000Z` : "";
+ const endDate = window.end_date ? `${window.end_date}T23:59:59.999Z` : "";
+ const scope = JSON.stringify([accessToken, startDate, endDate, filter]);
+ const [pagination, setPagination] = useState<{ scope: string; cursors: readonly RequestCursor[] }>({
+ scope,
+ cursors: [null],
+ });
+ const cursors = pagination.scope === scope ? pagination.cursors : [null];
+ const cursor = cursors.at(-1);
+ const page = cursors.length;
+
+ if (pagination.scope !== scope) {
+ setPagination({ scope, cursors: [null] });
+ }
+
+ const enabled = Boolean(accessToken && startDate && endDate);
+ const query: RequestsQuery = {
+ start_date: startDate,
+ end_date: endDate,
+ filter,
+ page_size: 50,
+ cursor_start_time: cursor?.start_time,
+ cursor_request_id: cursor?.request_id,
+ };
+ const queryOptions: UseQueryOptions = {
+ queryKey: [REQUESTS_PATH, accessToken, query],
+ queryFn: ({ signal }) => apiClient.get(REQUESTS_PATH, { accessToken, query, signal }),
+ enabled,
+ retry: false,
+ };
+ const requests = useQuery(queryOptions);
+ const nextCursor = requests.data?.next_cursor;
+
+ const changeFilter = (value: unknown) => {
+ if (value === "all" || value === "injected" || value === "hits") {
+ setFilter(value);
+ }
+ };
+
+ return (
+
+
+
+
Prompt caching requests
+
+ Requests with recorded LiteLLM injection or provider cache reads or writes. A cache hit alone does not
+ establish LiteLLM injection; older logs may not record it.
+
+
+ Net savings are estimated from logged usage and current configured pricing, after cache-write premiums.
+ Negative values mean caching cost more; unavailable means the request could not be priced.
+
+
+
+
+ All caching
+ LiteLLM injected
+ Cache hits
+
+
+
+
+ {!enabled && Select a date range to view requests
}
+ {enabled && requests.isPending && (
+
+ Loading requests...
+
+ )}
+ {enabled && requests.isError && (
+
+
Could not load prompt caching requests
+
void requests.refetch()} disabled={requests.isFetching}>
+ Retry
+
+
+ )}
+ {enabled && requests.isSuccess && (
+ <>
+ {requests.data.requests.length === 0 ? (
+
+ No matching prompt caching requests in this range
+
+ ) : (
+
+
+
+ Request
+ Model
+ LiteLLM injection
+ Cache reads
+ Cache writes
+ Actual cost
+ Net savings
+
+
+
+ {requests.data.requests.map((request) => (
+
+
+
+ {request.request_id}
+
+
+ {new Date(request.start_time).toLocaleString()}
+
+
+
+
+ {request.model}
+
+
+ {request.gateway_injected ? "Recorded" : "Not recorded"}
+ {formatNumberWithCommas(request.cache_read_tokens)}
+
+ {formatNumberWithCommas(request.cache_creation_tokens)}
+
+ {usd(request.spend)}
+
+ {request.net_savings === null ? "Unavailable" : usd(request.net_savings)}
+
+
+ ))}
+
+
+ )}
+
+ setPagination({ scope, cursors: cursors.slice(0, -1) })}
+ >
+ Previous
+
+ Page {page}
+ nextCursor && setPagination({ scope, cursors: [...cursors, nextCursor] })}
+ >
+ Next
+
+
+ >
+ )}
+
+
+ );
+}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx
index 66db347e70f..35464c5852e 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx
@@ -1,4 +1,4 @@
-import { render, waitFor, screen } from "@testing-library/react";
+import { fireEvent, render, waitFor, screen } from "@testing-library/react";
import { describe, expect, it, vi } from "vitest";
const mockGetGeneralSettingsCall = vi.fn();
@@ -12,6 +12,21 @@ vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () =>
}));
const mockCacheLeakageCard = vi.fn();
+const mockRequestsTable = vi.fn();
+const nextDateRange = { from: new Date(2026, 8, 1), to: new Date(2026, 8, 2) };
+
+vi.mock("./PromptCachingRequestsTable", () => ({
+ default: (props: unknown) => {
+ mockRequestsTable(props);
+ return
;
+ },
+}));
+
+vi.mock("@/components/shared/advanced_date_picker", () => ({
+ default: ({ onValueChange }: { onValueChange: (range: typeof nextDateRange) => void }) => (
+ onValueChange(nextDateRange)}>Change caching dates
+ ),
+}));
vi.mock("./CacheLeakageCard", () => ({
__esModule: true,
@@ -24,7 +39,7 @@ vi.mock("./CacheLeakageCard", () => ({
import PromptCachingTab from "./PromptCachingTab";
describe("PromptCachingTab", () => {
- it("renders the cache leakage table alongside the caching settings", async () => {
+ it("shares the selected dates between requests and cache leakage alongside caching settings", async () => {
mockGetGeneralSettingsCall.mockResolvedValue([]);
const activity = {
@@ -42,6 +57,10 @@ describe("PromptCachingTab", () => {
expect(screen.getByTestId("caching-settings")).toBeInTheDocument();
expect(screen.getByTestId("cache-leakage-card")).toBeInTheDocument();
+ expect(screen.getByTestId("caching-requests")).toBeInTheDocument();
+ expect(mockRequestsTable).toHaveBeenCalledWith({ accessToken: "test-token", dateValue: activity.dateValue });
+ fireEvent.click(screen.getByRole("button", { name: "Change caching dates" }));
+ expect(activity.onDateChange).toHaveBeenCalledWith(nextDateRange);
await waitFor(() => expect(mockCacheLeakageCard).toHaveBeenCalledWith(expect.objectContaining({ activity })));
});
});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx
index 59b38f272e0..4e43317998e 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx
@@ -3,12 +3,14 @@
import React, { useCallback, useEffect, useState } from "react";
import { getGeneralSettingsCall } from "@/components/networking";
+import AdvancedDatePicker from "@/components/shared/advanced_date_picker";
import { toast } from "@/lib/toast";
import {
PromptCachingPanel,
generalSettingsItem,
} from "@/app/(dashboard)/router-settings/_components/general_settings";
import CacheLeakageCard from "./CacheLeakageCard";
+import PromptCachingRequestsTable from "./PromptCachingRequestsTable";
import { DailyActivityRange } from "./useDailyActivityRange";
interface PromptCachingTabProps {
@@ -48,6 +50,11 @@ const PromptCachingTab: React.FC = ({ accessToken, activi
return (
+
+
Date range for requests and cache leakage
+
+
+
);
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 81580c8bfb1..d916509c06f 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -3534,6 +3534,23 @@ export interface paths {
patch?: never;
trace?: never;
};
+ "/cost_optimization/prompt_caching/requests": {
+ parameters: {
+ query?: never;
+ header?: never;
+ path?: never;
+ cookie?: never;
+ };
+ /** Get Prompt Caching Requests */
+ get: operations["get_prompt_caching_requests_cost_optimization_prompt_caching_requests_get"];
+ put?: never;
+ post?: never;
+ delete?: never;
+ options?: never;
+ head?: never;
+ patch?: never;
+ trace?: never;
+ };
"/credentials": {
parameters: {
query?: never;
@@ -35814,6 +35831,48 @@ export interface components {
prompt_id: string;
prompt_info?: components["schemas"]["PromptInfo"] | null;
};
+ /** PromptCachingRequest */
+ PromptCachingRequest: {
+ /** Cache Creation Tokens */
+ cache_creation_tokens: number;
+ /** Cache Read Tokens */
+ cache_read_tokens: number;
+ /** Gateway Injected */
+ gateway_injected: boolean;
+ /** Model */
+ model: string;
+ /** Net Savings */
+ net_savings: number | null;
+ /** Request Id */
+ request_id: string;
+ /** Spend */
+ spend: number;
+ /**
+ * Start Time
+ * Format: date-time
+ */
+ start_time: string;
+ };
+ /** PromptCachingRequestCursor */
+ PromptCachingRequestCursor: {
+ /** Request Id */
+ request_id: string;
+ /**
+ * Start Time
+ * Format: date-time
+ */
+ start_time: string;
+ };
+ /** PromptCachingRequestsResponse */
+ PromptCachingRequestsResponse: {
+ /** Has More */
+ has_more: boolean;
+ next_cursor: components["schemas"]["PromptCachingRequestCursor"] | null;
+ /** Page Size */
+ page_size: number;
+ /** Requests */
+ requests: components["schemas"]["PromptCachingRequest"][];
+ };
/** PromptInfo */
PromptInfo: {
/**
@@ -47238,6 +47297,42 @@ export interface operations {
};
};
};
+ get_prompt_caching_requests_cost_optimization_prompt_caching_requests_get: {
+ parameters: {
+ query: {
+ start_date: string;
+ end_date: string;
+ page_size?: number;
+ filter?: "all" | "injected" | "hits";
+ cursor_start_time?: string | null;
+ cursor_request_id?: string | null;
+ };
+ header?: never;
+ path?: never;
+ cookie?: never;
+ };
+ requestBody?: never;
+ responses: {
+ /** @description Successful Response */
+ 200: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": components["schemas"]["PromptCachingRequestsResponse"];
+ };
+ };
+ /** @description Validation Error */
+ 422: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": components["schemas"]["HTTPValidationError"];
+ };
+ };
+ };
+ };
get_credentials_credentials_get: {
parameters: {
query?: never;
From 875f015e24219110dbad35691d80acd1c1a3c375 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 18:33:56 -0700
Subject: [PATCH 08/56] fix(token_counter): count replayed redacted_thinking
blocks so prompt_caching keeps pinning
A conversation that replays a redacted_thinking block (Anthropic redacted reasoning, or the
/v1/messages bridge's stand-in for a reasoning item that carries no summary) made
_count_content_list raise, is_prompt_caching_valid_prompt swallowed that to False, and the
prompt_caching pre-call check neither recorded nor pinned the serving deployment, so the
conversation bounced across the group and paid a cache write on every deployment. The block
now counts like a thinking block with no text: zero tokens for the encrypted payload.
---
litellm/litellm_core_utils/token_counter.py | 11 ++--
.../litellm_core_utils/test_token_counter.py | 19 +++++++
.../test_prompt_caching_deployment_check.py | 52 +++++++++++++++++++
3 files changed, 79 insertions(+), 3 deletions(-)
diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py
index 6c1b7946394..bf37b1be2e4 100644
--- a/litellm/litellm_core_utils/token_counter.py
+++ b/litellm/litellm_core_utils/token_counter.py
@@ -46,6 +46,8 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionDocumentObject,
ChatCompletionNamedToolChoiceParam,
+ ChatCompletionRedactedThinkingBlock,
+ ChatCompletionThinkingBlock,
ChatCompletionToolParam,
OpenAIMessageContentListBlock,
)
@@ -854,6 +856,8 @@ def _count_content_list(
content_list: str
| Iterable[
OpenAIMessageContentListBlock
+ | ChatCompletionThinkingBlock
+ | ChatCompletionRedactedThinkingBlock
| AnthropicMessagesTextParam
| AnthropicMessagesImageParam
| AnthropicMessagesDocumentParam
@@ -898,9 +902,9 @@ def _count_content_list(
use_default_image_token_count,
default_token_count,
)
- elif c["type"] == "thinking":
+ elif c["type"] in ("thinking", "redacted_thinking"):
# Claude extended thinking content block
- # Count the thinking text and skip signature (opaque signature blob)
+ # Count the thinking text and skip the opaque blobs (signature, redacted data)
thinking_text = str(c.get("thinking", ""))
if thinking_text:
num_tokens += count_function(thinking_text)
@@ -920,7 +924,8 @@ def _count_content_list(
raise ValueError(
f"Invalid content item type: {content_type}. "
f"Expected str or dict with 'type' field "
- f"(text, image_url, image, document, file, tool_use, tool_result, thinking, tool_reference)."
+ f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, "
+ f"tool_reference)."
)
return num_tokens
except Exception as e:
diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py
index ba3a6be609f..f19a8891609 100644
--- a/tests/test_litellm/litellm_core_utils/test_token_counter.py
+++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py
@@ -1257,6 +1257,25 @@ def test_token_counter_with_thinking_content():
), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}"
+
+def test_token_counter_with_redacted_thinking_content():
+ """
+ A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in
+ for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking
+ block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the
+ prompt_caching pre-call check stop pinning the deployment that held the cached prefix.
+ """
+ model = "anthropic/claude-sonnet-4-5-20250929"
+ reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."}
+ redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30}
+ user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]}
+ follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]}
+
+ without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up]
+ with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up]
+
+ assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block)
+
def test_token_counter_with_tool_reference_block():
"""
Regression test: a message containing an Anthropic tool-search
diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
index 333e7b2ff31..267109c9164 100644
--- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
+++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
@@ -197,6 +197,58 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is
AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5"
+@pytest.mark.asyncio
+async def test_replayed_redacted_thinking_block_still_records_and_pins():
+ """
+ A model that returns no reasoning summary (gpt-5.x through the /v1/messages bridge, Anthropic with
+ redacted reasoning) hands the client a `redacted_thinking` block, and the client replays it on every
+ later turn. The token count behind `is_prompt_caching_valid_prompt` raised on that block, the helper
+ swallowed it to False, and the check neither recorded the serving deployment nor pinned it, so the
+ conversation bounced across the group and paid a cache write on each deployment.
+ """
+ cache = DualCache()
+ check = PromptCachingDeploymentCheck(cache=cache)
+ model = "openai/gpt-5.6-sol"
+ deployments = _deployments(model, model, model)
+ messages = cast(
+ List[AllMessageValues],
+ [
+ *_messages(word_count=3000),
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:" + "Z" * 400},
+ {"type": "text", "text": "Draw from the box labeled Mixed."},
+ ],
+ },
+ {"role": "user", "content": "Restate that in one sentence."},
+ ],
+ )
+
+ assert is_prompt_caching_valid_prompt(model=model, messages=messages) is True
+
+ await check.async_log_success_event(
+ kwargs={
+ "standard_logging_object": {
+ "call_type": "anthropic_messages",
+ "model": model,
+ "messages": messages,
+ "model_id": "dep-2",
+ }
+ },
+ response_obj=None,
+ start_time=None,
+ end_time=None,
+ )
+ filtered = await check.async_filter_deployments(
+ model=MODEL_GROUP_ALIAS,
+ healthy_deployments=deployments,
+ messages=messages,
+ )
+
+ assert filtered == [deployments[1]]
+
+
def _auto_caching_messages() -> List[AllMessageValues]:
"""A prompt over the model minimum that carries no client cache_control."""
return cast(
From 3772993032e93d283c9c0b0cf5a80909feae52f3 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 18:42:05 -0700
Subject: [PATCH 09/56] fix(anthropic_messages): only Mantle consumes
get_llm_provider's api_base
The /v1/messages handler passed the api_base get_llm_provider resolved to every
provider's native messages config, which shadowed DEEPSEEK_ANTHROPIC_API_BASE and
TENCENT_ANTHROPIC_API_BASE with the chat default and changed the azure_ai
precedence. Messages configs now opt in through uses_get_llm_provider_api_base(),
true only for Bedrock Mantle, whose region-prefixed model must resolve to a
region host before the prefix is stripped. Also registers
BedrockMantleAnthropicMessagesConfig in the lazy import registry.
---
litellm/__init__.py | 3 ++
litellm/_lazy_imports_registry.py | 5 +++
.../messages/handler.py | 6 ++-
.../anthropic_messages/transformation.py | 3 ++
.../bedrock_mantle/messages/transformation.py | 3 ++
...erimental_pass_through_messages_handler.py | 42 +++++++++++++++++++
6 files changed, 61 insertions(+), 1 deletion(-)
diff --git a/litellm/__init__.py b/litellm/__init__.py
index e17ab613dac..d2bbc107205 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -1684,6 +1684,9 @@ if TYPE_CHECKING:
from .llms.bedrock.messages.mantle_transformation import (
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
)
+ from .llms.bedrock_mantle.messages.transformation import (
+ BedrockMantleAnthropicMessagesConfig as BedrockMantleAnthropicMessagesConfig,
+ )
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
from .llms.together_ai.chat.transformation import (
TogetherAIChatConfig as TogetherAIChatConfig,
diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py
index 9cfcb9e41f7..bca04a17250 100644
--- a/litellm/_lazy_imports_registry.py
+++ b/litellm/_lazy_imports_registry.py
@@ -176,6 +176,7 @@ LLM_CONFIG_NAMES: Final = (
"BedrockClaudePlatformMessagesConfig",
"AmazonAnthropicClaudeMessagesConfig",
"AmazonMantleMessagesConfig",
+ "BedrockMantleAnthropicMessagesConfig",
"TogetherAIConfig",
"TogetherAIChatConfig",
"NLPCloudConfig",
@@ -746,6 +747,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.bedrock.messages.mantle_transformation",
"AmazonMantleMessagesConfig",
),
+ "BedrockMantleAnthropicMessagesConfig": (
+ ".llms.bedrock_mantle.messages.transformation",
+ "BedrockMantleAnthropicMessagesConfig",
+ ),
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
"TogetherAIChatConfig": (
".llms.together_ai.chat.transformation",
diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
index e1309ea4063..d87cb0a64f5 100644
--- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
+++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py
@@ -501,7 +501,6 @@ def anthropic_messages_handler(
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
)
- resolved_api_base: Final = dynamic_api_base if dynamic_api_base is not None else api_base
# Store agentic loop params in logging object for agentic hooks
# This provides original request context needed for follow-up calls
@@ -652,6 +651,11 @@ def anthropic_messages_handler(
"display": "summarized",
}
+ resolved_api_base: Final = (
+ dynamic_api_base
+ if dynamic_api_base is not None and anthropic_messages_provider_config.uses_get_llm_provider_api_base()
+ else api_base
+ )
return base_llm_http_handler.anthropic_messages_handler(
model=model,
messages=strip_provider_specific_fields_from_anthropic_messages(messages),
diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py
index 8e7c22930fa..101a5e6c58c 100644
--- a/litellm/llms/base_llm/anthropic_messages/transformation.py
+++ b/litellm/llms/base_llm/anthropic_messages/transformation.py
@@ -128,6 +128,9 @@ class BaseAnthropicMessagesConfig(ABC):
"""
return True
+ def uses_get_llm_provider_api_base(self) -> bool:
+ return False
+
def get_async_streaming_response_iterator(
self,
model: str,
diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py
index 480c09a0476..480fe82ef4c 100644
--- a/litellm/llms/bedrock_mantle/messages/transformation.py
+++ b/litellm/llms/bedrock_mantle/messages/transformation.py
@@ -61,6 +61,9 @@ class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleM
def custom_llm_provider(self) -> str | None:
return "bedrock_mantle"
+ def uses_get_llm_provider_api_base(self) -> bool:
+ return True
+
def get_complete_url(
self,
api_base: str | None,
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 997a97c6fd3..9fa3ef153be 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
@@ -1438,3 +1438,45 @@ async def test_anthropic_messages_leaves_non_provider_failures_unmapped():
)
assert "Traceback" not in str(excinfo.value)
+
+
+def _recording_client(seen_urls: list[str]) -> AsyncHTTPHandler:
+ def record_and_answer(request: httpx.Request) -> httpx.Response:
+ seen_urls.append(str(request.url))
+ return httpx.Response(
+ 200,
+ json={
+ "id": "msg_test",
+ "type": "message",
+ "role": "assistant",
+ "model": "deepseek-chat",
+ "content": [{"type": "text", "text": "pong"}],
+ "stop_reason": "end_turn",
+ "stop_sequence": None,
+ "usage": {"input_tokens": 3, "output_tokens": 1},
+ },
+ )
+
+ upstream = AsyncHTTPHandler()
+ upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(record_and_answer))
+ return upstream
+
+
+@pytest.mark.asyncio
+async def test_provider_messages_api_base_env_is_not_shadowed_by_the_chat_default(monkeypatch):
+ from litellm.llms.anthropic.experimental_pass_through.messages import handler
+
+ monkeypatch.delenv("DEEPSEEK_API_BASE", raising=False)
+ monkeypatch.setenv("DEEPSEEK_ANTHROPIC_API_BASE", "https://deepseek.internal.example/anthropic")
+ seen_urls: list[str] = []
+
+ await handler.anthropic_messages(
+ max_tokens=16,
+ messages=[{"role": "user", "content": "ping"}],
+ model="deepseek/deepseek-chat",
+ api_key="sk-test",
+ client=_recording_client(seen_urls),
+ )
+
+ assert seen_urls == ["https://deepseek.internal.example/anthropic/v1/messages"]
+
From 517fff5bb7bbbd397ad1942cba5a3a1b35e0640a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 19:52:22 -0700
Subject: [PATCH 10/56] fix(router): keep prompt caching affinity when the
breakpoint moves
The prompt_caching pre-call check keyed a deployment pin on a hash of the
whole cacheable prefix, cache_control markers included. Agent clients
such as Claude Code move the marker to the newest user turn on every
request, so the key changed every turn, the pin never matched, and a
multi-turn session drifted across deployments and lost its provider
cache.
Hash the prefix per content block with the markers stripped, chained so
every block position has a key, and write the pin at the breakpoint
block. Lookup walks back over the last PROMPT_CACHE_LOOKBACK_POSITIONS
positions (a run of tool_use or tool_result blocks counting as one), the
same window the provider probes for a cached prefix, in one batch cache
read. Both sides hash the prefix after base64 truncation so a request
carrying raw image bytes derives the keys the success event stored.
---
litellm/constants.py | 3 +
litellm/router_utils/prompt_caching_cache.py | 250 +++++++++++-----
.../test_prompt_caching_deployment_check.py | 273 +++++++++++++++++-
3 files changed, 450 insertions(+), 76 deletions(-)
diff --git a/litellm/constants.py b/litellm/constants.py
index bbeb4846e27..e4576ad4d5c 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -399,6 +399,9 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = (
if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None
else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
)
+# Anthropic checks at most 20 block positions behind a breakpoint for a cached prefix, a run of tool_use
+# or tool_result blocks counting as one position, so deployment affinity probes the same window
+PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20
DEFAULT_TRIM_RATIO: Final = float(
os.getenv("DEFAULT_TRIM_RATIO", 0.75)
) # default ratio of tokens to trim from the end of a prompt
diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py
index 39708e168f5..0b784e1fa91 100644
--- a/litellm/router_utils/prompt_caching_cache.py
+++ b/litellm/router_utils/prompt_caching_cache.py
@@ -4,12 +4,21 @@ Wrapper around router cache. Meant to store model id when prompt caching support
import hashlib
import json
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import dataclass
+from itertools import accumulate
from typing import TYPE_CHECKING, Any, Final, cast
+from pydantic import JsonValue, TypeAdapter
+from pydantic_core import to_jsonable_python
from typing_extensions import TypedDict
from litellm.caching.caching import DualCache
-from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.constants import PROMPT_CACHE_LOOKBACK_POSITIONS
+from litellm.litellm_core_utils.logging_utils import (
+ truncate_base64_in_messages,
+ truncate_base64_in_messages_async,
+)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
if TYPE_CHECKING:
@@ -28,10 +37,100 @@ class PromptCachingCacheValue(TypedDict):
model_id: str
+PROMPT_CACHE_PIN_TTL_SECONDS: Final = 300
+_TOOL_RUN_BLOCK_TYPES: Final = frozenset({"tool_use", "tool_result"})
+_PREFIX_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...])
+_TOOLS_ADAPTER: Final = TypeAdapter(tuple[JsonValue, ...])
+_PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None)
+
+
+@dataclass(frozen=True, slots=True)
+class PrefixPosition:
+ cache_key: str
+ position: int
+
+
+def _sorted_pairs(pairs: Iterable[tuple[str, JsonValue]]) -> tuple[tuple[str, JsonValue], ...]:
+ return tuple(sorted(pairs, key=lambda pair: pair[0]))
+
+
+def _canonical_bytes(value: object) -> bytes:
+ return json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
+
+
+def _block_unit(
+ envelope: tuple[tuple[str, JsonValue], ...], message_run_type: str | None, block: JsonValue
+) -> tuple[bytes, str | None]:
+ if not isinstance(block, dict):
+ return _canonical_bytes((envelope, block)), message_run_type
+ block_type: Final = block.get("type")
+ block_run_type: Final = block_type if isinstance(block_type, str) and block_type in _TOOL_RUN_BLOCK_TYPES else None
+ stripped: Final = _sorted_pairs(item for item in block.items() if item[0] != "cache_control")
+ return _canonical_bytes((envelope, stripped)), message_run_type or block_run_type
+
+
+def _message_units(message: Mapping[str, JsonValue]) -> tuple[tuple[bytes, str | None], ...]:
+ envelope: Final = _sorted_pairs(item for item in message.items() if item[0] not in ("content", "cache_control"))
+ message_run_type: Final = "tool_result" if message.get("role") == "tool" else None
+ content: Final = message.get("content")
+ if isinstance(content, list) and content:
+ return tuple(_block_unit(envelope, message_run_type, block) for block in content)
+ if isinstance(content, str) and content:
+ return ((_canonical_bytes((envelope, (("text", content), ("type", "text")))), message_run_type),)
+ return ((_canonical_bytes((envelope, None)), message_run_type),)
+
+
+def _chain_digest(digest: bytes, unit: bytes) -> bytes:
+ return hashlib.sha256(digest + unit).digest()
+
+
+def _seed(tools: Sequence[ChatCompletionToolParam] | None) -> bytes:
+ if tools is None:
+ return hashlib.sha256(b"").digest()
+ return hashlib.sha256(
+ _canonical_bytes(_TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True)))
+ ).digest()
+
+
+def _positions_of(
+ prefix: tuple[Mapping[str, JsonValue], ...], tools: Sequence[ChatCompletionToolParam] | None
+) -> tuple[PrefixPosition, ...]:
+ units: Final = tuple(unit for message in prefix for unit in _message_units(message))
+ digests: Final = tuple(accumulate((unit_bytes for unit_bytes, _ in units), _chain_digest, initial=_seed(tools)))[1:]
+ run_types: Final = tuple(run_type for _, run_type in units)
+ positions: Final = accumulate(
+ 0 if run_type is not None and run_type == previous else 1
+ for run_type, previous in zip(run_types, (None, *run_types[:-1]))
+ )
+ return tuple(
+ PrefixPosition(cache_key=f"deployment:{digest.hex()}:prompt_caching", position=position)
+ for digest, position in zip(digests, positions)
+ )
+
+
+def _lookback_keys(positions: tuple[PrefixPosition, ...]) -> tuple[str, ...]:
+ if not positions:
+ return ()
+ oldest_probed_position: Final = positions[-1].position - PROMPT_CACHE_LOOKBACK_POSITIONS
+ return tuple(entry.cache_key for entry in reversed(positions) if entry.position > oldest_probed_position)
+
+
+def _pinned_value(value: JsonValue) -> PromptCachingCacheValue | None:
+ if not isinstance(value, dict):
+ return None
+ model_id: Final = value.get("model_id")
+ return PromptCachingCacheValue(model_id=model_id) if isinstance(model_id, str) else None
+
+
+def _first_pin(values: tuple[JsonValue, ...] | None) -> PromptCachingCacheValue | None:
+ if values is None:
+ return None
+ return next((pin for pin in map(_pinned_value, values) if pin is not None), None)
+
+
class PromptCachingCache:
def __init__(self, cache: DualCache):
self.cache = cache
- self.in_memory_cache = InMemoryCache()
@staticmethod
def serialize_object(obj: Any) -> object:
@@ -140,114 +239,123 @@ class PromptCachingCache:
return cacheable_prefix
@staticmethod
- def get_prompt_caching_cache_key(
+ def prefix_positions(
messages: list[AllMessageValues] | None,
- tools: list[ChatCompletionToolParam] | None,
- ) -> str | None:
- if messages is None and tools is None:
- return None
+ tools: Sequence[ChatCompletionToolParam] | None,
+ ) -> tuple[PrefixPosition, ...]:
+ """
+ One cache key per content block of the cacheable prefix, oldest block first.
- # Extract cacheable prefix from messages (only include up to last cache_control block)
- cacheable_messages = None
- if messages is not None:
- cacheable_messages = PromptCachingCache.extract_cacheable_prefix(messages)
- # If no cacheable prefix found, return None (can't cache)
- if not cacheable_messages:
- return None
+ Each key hashes the prefix content up to and including that block, with cache_control markers
+ left out, so the key of a block is the same whichever turn's breakpoint the prefix ends at.
+ String content hashes like a single text block, which is how the provider treats it and how
+ Claude Code re-sends a previously marked message. `position` counts a run of consecutive
+ tool_use (or tool_result) blocks as one, matching the provider's lookback window.
- # Use serialize_object for consistent and stable serialization
- data_to_hash: Final = {}
- if cacheable_messages is not None:
- serialized_messages: Final = PromptCachingCache.serialize_object(cacheable_messages)
- data_to_hash["messages"] = serialized_messages
- if tools is not None:
- serialized_tools: Final = PromptCachingCache.serialize_object(tools)
- data_to_hash["tools"] = serialized_tools
-
- # Combine serialized data into a single string
- data_to_hash_str: Final = json.dumps(
- data_to_hash,
- sort_keys=True,
- separators=(",", ":"),
+ The prefix is hashed in the shape the success event sees it, with long base64 data URIs
+ already replaced by their size placeholder, so a request carrying the raw image bytes
+ derives the same keys the write side stored.
+ """
+ if not messages:
+ return ()
+ return _positions_of(
+ _PREFIX_ADAPTER.validate_python(
+ to_jsonable_python(
+ truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)),
+ serialize_unknown=True,
+ )
+ ),
+ tools,
)
- # Create a hash of the serialized data for a stable cache key
- hashed_data: Final = hashlib.sha256(data_to_hash_str.encode()).hexdigest()
- return f"deployment:{hashed_data}:prompt_caching"
+ @staticmethod
+ async def async_prefix_positions(
+ messages: list[AllMessageValues] | None,
+ tools: Sequence[ChatCompletionToolParam] | None,
+ ) -> tuple[PrefixPosition, ...]:
+ if not messages:
+ return ()
+ return _positions_of(
+ _PREFIX_ADAPTER.validate_python(
+ to_jsonable_python(
+ await truncate_base64_in_messages_async(PromptCachingCache.extract_cacheable_prefix(messages)),
+ serialize_unknown=True,
+ )
+ ),
+ tools,
+ )
+
+ @staticmethod
+ def get_prompt_caching_cache_key(
+ messages: list[AllMessageValues] | None,
+ tools: Sequence[ChatCompletionToolParam] | None,
+ ) -> str | None:
+ positions: Final = PromptCachingCache.prefix_positions(messages, tools)
+ return positions[-1].cache_key if positions else None
def add_model_id(
self,
model_id: str,
messages: list[AllMessageValues] | None,
- tools: list[ChatCompletionToolParam] | None,
+ tools: Sequence[ChatCompletionToolParam] | None,
) -> None:
- if messages is None and tools is None:
- return
-
cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
- # If no cacheable prefix found, don't cache (can't generate cache key)
if cache_key is None:
return
- self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300)
- return
+ self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=PROMPT_CACHE_PIN_TTL_SECONDS)
async def async_add_model_id(
self,
model_id: str,
messages: list[AllMessageValues] | None,
- tools: list[ChatCompletionToolParam] | None,
+ tools: Sequence[ChatCompletionToolParam] | None,
) -> None:
- if messages is None and tools is None:
- return
-
- cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
- # If no cacheable prefix found, don't cache (can't generate cache key)
- if cache_key is None:
+ positions: Final = await PromptCachingCache.async_prefix_positions(messages, tools)
+ if not positions:
return
await self.cache.async_set_cache(
- cache_key,
+ positions[-1].cache_key,
PromptCachingCacheValue(model_id=model_id),
- ttl=300, # store for 5 minutes
+ ttl=PROMPT_CACHE_PIN_TTL_SECONDS,
)
- return
async def async_get_model_id(
self,
messages: list[AllMessageValues] | None,
- tools: list[ChatCompletionToolParam] | None,
+ tools: Sequence[ChatCompletionToolParam] | None,
) -> PromptCachingCacheValue | None:
"""
- Get model ID from cache using the cacheable prefix.
-
- The cache key is based on the cacheable prefix (everything up to and including
- the last cache_control block), so requests with the same cacheable prefix but
- different user messages will have the same cache key.
+ Find the deployment that last served this prefix, walking back from the breakpoint the
+ same way the provider cache does, so a breakpoint that moved forward since the last
+ turn still lands on the deployment whose cache holds the earlier prefix.
"""
- if messages is None and tools is None:
+ cache_keys: Final = _lookback_keys(await PromptCachingCache.async_prefix_positions(messages, tools))
+ if not cache_keys:
return None
- # Generate cache key using cacheable prefix
- cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
- if cache_key is None:
- return None
-
- # Perform cache lookup
- cache_result: Final = await self.cache.async_get_cache(key=cache_key)
- return cache_result
+ return _first_pin(
+ _PINS_ADAPTER.validate_python(
+ await self.cache.async_batch_get_cache(
+ keys=list(cache_keys), # mutable-ok: DualCache.async_batch_get_cache only takes a list
+ )
+ )
+ )
def get_model_id(
self,
messages: list[AllMessageValues] | None,
- tools: list[ChatCompletionToolParam] | None,
+ tools: Sequence[ChatCompletionToolParam] | None,
) -> PromptCachingCacheValue | None:
- if messages is None and tools is None:
+ cache_keys: Final = _lookback_keys(PromptCachingCache.prefix_positions(messages, tools))
+ if not cache_keys:
return None
- cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools)
- # If no cacheable prefix found, return None (can't cache)
- if cache_key is None:
- return None
-
- return self.cache.get_cache(cache_key)
+ return _first_pin(
+ _PINS_ADAPTER.validate_python(
+ self.cache.batch_get_cache(
+ keys=list(cache_keys), # mutable-ok: DualCache.batch_get_cache only takes a list
+ )
+ )
+ )
diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
index 333e7b2ff31..d0a9223dfa7 100644
--- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
+++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
@@ -1,5 +1,6 @@
import asyncio
import copy
+import functools
from typing import List, cast
import pytest
@@ -7,7 +8,7 @@ import pytest
import litellm
from litellm.caching.dual_cache import DualCache
-from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
+from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import (
@@ -30,7 +31,6 @@ def _local_model_cost_map_autouse(local_model_cost_map):
yield
-
def _deployments(*models: str) -> List[dict]:
return [
{
@@ -84,7 +84,9 @@ def test_write_gate_is_what_prevents_a_pin_below_the_model_minimum():
"""
messages = _messages(word_count=1400)
- token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True)
+ token_count = token_counter(
+ messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True
+ )
assert 1024 < token_count < 4096
assert is_prompt_caching_valid_prompt(model="anthropic/claude-opus-4-5", messages=messages) is False
@@ -110,7 +112,9 @@ async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minim
deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6")
messages = _messages(word_count=1400)
- token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True)
+ token_count = token_counter(
+ messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True
+ )
assert DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT < token_count < OPUS_4_6_MIN_TOKENS
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None)
@@ -136,7 +140,9 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum():
deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6")
messages = _messages(word_count=5000)
- token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True)
+ token_count = token_counter(
+ messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True
+ )
assert token_count > OPUS_4_6_MIN_TOKENS
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None)
@@ -539,3 +545,260 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop():
"model_id": "dep-1"
}
assert_loop_stayed_free(took, lags)
+
+
+LONG_PROMPT = "word " * 3000
+ONE_PIXEL_PNG = (
+ "data:image/png;base64,"
+ "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
+)
+
+
+def _turn(*messages: dict) -> List[AllMessageValues]:
+ return cast(List[AllMessageValues], list(messages))
+
+
+def _text(text: str) -> dict:
+ return {"type": "text", "text": text}
+
+
+def _marked(text: str) -> dict:
+ return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}
+
+
+@pytest.mark.asyncio
+async def test_pin_survives_the_breakpoint_moving_to_the_next_turn():
+ """
+ The regression. Claude Code marks only the newest user message each turn, so the last breakpoint
+ moves forward every turn. The key hashed the prefix up to that moving breakpoint, markers
+ included, so no turn after the first ever found the pin the previous turn wrote, and a
+ multi-deployment group re-rolled the deployment mid-session, paying a cache write on a
+ deployment whose provider cache held nothing of the conversation.
+ """
+ cache = DualCache()
+ check = PromptCachingDeploymentCheck(cache=cache)
+ deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
+ turn_one = _turn({"role": "user", "content": [_marked(LONG_PROMPT)]})
+ turn_two = _turn(
+ {"role": "user", "content": [_text(LONG_PROMPT)]},
+ {"role": "assistant", "content": "ok"},
+ {"role": "user", "content": [_marked("next")]},
+ )
+
+ await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=turn_one, tools=None)
+
+ filtered = await check.async_filter_deployments(
+ model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two
+ )
+
+ assert filtered == [deployments[1]]
+
+
+@pytest.mark.asyncio
+async def test_pin_survives_the_marked_message_coming_back_as_string_content():
+ """
+ Claude Code sends the message that carries a breakpoint as a one-block content list and re-sends
+ it next turn as plain string content once the marker has moved on. The provider caches both
+ shapes identically, so the key has to as well, or the walk-back never lands on the turn-one write.
+ """
+ cache = DualCache()
+ check = PromptCachingDeploymentCheck(cache=cache)
+ deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
+ turn_one = _turn(
+ {"role": "system", "content": [_marked(LONG_PROMPT)]},
+ {"role": "user", "content": [_marked("hello")]},
+ )
+ turn_two = _turn(
+ {"role": "system", "content": LONG_PROMPT},
+ {"role": "user", "content": "hello"},
+ {"role": "assistant", "content": "hi"},
+ {"role": "user", "content": [_marked("again")]},
+ )
+
+ await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-1", messages=turn_one, tools=None)
+
+ filtered = await check.async_filter_deployments(
+ model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two
+ )
+
+ assert filtered == [deployments[0]]
+
+
+@pytest.mark.asyncio
+async def test_lookback_stops_where_the_provider_cache_stops():
+ """
+ Anthropic finds a cached prefix at most PROMPT_CACHE_LOOKBACK_POSITIONS block positions behind a
+ breakpoint, the breakpoint block included. Probing further would pin to a deployment whose cache
+ the provider will not consult, and probing less would drop pins the provider still honors.
+ """
+ prompt_cache = PromptCachingCache(cache=DualCache())
+ await prompt_cache.async_add_model_id(
+ model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("block 0")]}), tools=None
+ )
+
+ def turn_with_blocks_after(count: int) -> List[AllMessageValues]:
+ later = [_text(f"block {index}") for index in range(1, count)] + [_marked(f"block {count}")]
+ return _turn({"role": "user", "content": [_text("block 0"), *later]})
+
+ inside_window = turn_with_blocks_after(PROMPT_CACHE_LOOKBACK_POSITIONS - 1)
+ past_window = turn_with_blocks_after(PROMPT_CACHE_LOOKBACK_POSITIONS)
+
+ assert await prompt_cache.async_get_model_id(messages=inside_window, tools=None) == {"model_id": "dep-1"}
+ assert prompt_cache.get_model_id(messages=inside_window, tools=None) == {"model_id": "dep-1"}
+ assert await prompt_cache.async_get_model_id(messages=past_window, tools=None) is None
+ assert prompt_cache.get_model_id(messages=past_window, tools=None) is None
+
+
+@pytest.mark.asyncio
+async def test_a_run_of_tool_blocks_counts_as_one_lookback_position():
+ """
+ The provider counts consecutive tool_use blocks as one lookback position, and consecutive
+ tool_result blocks as one, in both the Anthropic and the OpenAI message shapes. An agent turn that
+ fans out into many tool calls would otherwise push the previous breakpoint out of the window
+ after a single turn, which is exactly when the conversation is longest and the cache matters most.
+ """
+ prompt_cache = PromptCachingCache(cache=DualCache())
+ await prompt_cache.async_add_model_id(
+ model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("task")]}), tools=None
+ )
+ fan_out = PROMPT_CACHE_LOOKBACK_POSITIONS + 5
+
+ def anthropic_shaped(tool_use_type: str, tool_result_type: str) -> List[AllMessageValues]:
+ return _turn(
+ {"role": "user", "content": [_text("task")]},
+ {
+ "role": "assistant",
+ "content": [
+ {"type": tool_use_type, "id": f"call-{index}", "name": "read", "input": {"index": index}}
+ for index in range(fan_out)
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ *(
+ {"type": tool_result_type, "tool_use_id": f"call-{index}", "content": "ok"}
+ for index in range(fan_out)
+ ),
+ _marked("continue"),
+ ],
+ },
+ )
+
+ openai_shaped = _turn(
+ {"role": "user", "content": [_text("task")]},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {"id": f"call-{index}", "type": "function", "function": {"name": "read", "arguments": "{}"}}
+ for index in range(fan_out)
+ ],
+ },
+ *({"role": "tool", "tool_call_id": f"call-{index}", "content": "ok"} for index in range(fan_out)),
+ {"role": "user", "content": [_marked("continue")]},
+ )
+
+ assert await prompt_cache.async_get_model_id(messages=anthropic_shaped("tool_use", "tool_result"), tools=None) == {
+ "model_id": "dep-1"
+ }
+ assert await prompt_cache.async_get_model_id(messages=openai_shaped, tools=None) == {"model_id": "dep-1"}
+ assert await prompt_cache.async_get_model_id(messages=anthropic_shaped("text", "text"), tools=None) is None
+
+
+@pytest.mark.asyncio
+async def test_an_edited_earlier_block_does_not_inherit_the_pin():
+ """Walking back must still bind every block's content, or an edited conversation pins to a stale cache."""
+ prompt_cache = PromptCachingCache(cache=DualCache())
+ await prompt_cache.async_add_model_id(
+ model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("original")]}), tools=None
+ )
+ edited = _turn(
+ {"role": "user", "content": [_text("edited")]},
+ {"role": "assistant", "content": "ok"},
+ {"role": "user", "content": [_marked("next")]},
+ )
+
+ assert await prompt_cache.async_get_model_id(messages=edited, tools=None) is None
+
+
+class _BrokenBatchReadCache(DualCache):
+ async def async_batch_get_cache(self, keys, parent_otel_span=None, local_only=False, **kwargs):
+ return None
+
+
+@pytest.mark.asyncio
+async def test_a_failed_batch_read_pins_nothing():
+ """DualCache answers None rather than a list when the batch read raises, and routing must fall through."""
+ prompt_cache = PromptCachingCache(cache=_BrokenBatchReadCache())
+
+ assert (
+ await prompt_cache.async_get_model_id(messages=_turn({"role": "user", "content": [_marked("x")]}), tools=None)
+ is None
+ )
+
+
+@pytest.mark.asyncio
+async def test_pin_matches_when_the_success_event_truncated_an_image_payload(monkeypatch, local_model_cost_map):
+ """
+ The success event only ever sees the standard logging payload, whose long base64 data URIs are
+ replaced by size placeholders, while routing sees the raw request. Hashing the raw bytes on the
+ read side would key every image-carrying session past its own pin.
+ """
+ capture = _SentMessagesCapture()
+ monkeypatch.setattr(litellm, "callbacks", [capture])
+ image = {"type": "image_url", "image_url": {"url": ONE_PIXEL_PNG}}
+ turn_one = _turn({"role": "user", "content": [image, _marked(LONG_PROMPT)]})
+
+ await litellm.acompletion(
+ model=AUTO_CACHING_MODEL, messages=copy.deepcopy(turn_one), mock_response="ok", api_key="sk-fake"
+ )
+ logged = await _eventually(lambda: capture.messages)
+ assert logged is not None
+ assert logged != turn_one
+
+ cache = DualCache()
+ await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=logged, tools=None)
+ turn_two = _turn(
+ {"role": "user", "content": [image, _text(LONG_PROMPT)]},
+ {"role": "assistant", "content": "ok"},
+ {"role": "user", "content": [_marked("next")]},
+ )
+ deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)
+
+ filtered = await PromptCachingDeploymentCheck(cache=cache).async_filter_deployments(
+ model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two
+ )
+
+ assert filtered == [deployments[1]]
+
+
+@pytest.mark.asyncio
+async def test_claude_code_style_session_stays_on_one_deployment_across_turns(local_model_cost_map):
+ """
+ End to end over the router with a client that marks only the newest user message each turn, the
+ way Claude Code does. Every turn has to land on the deployment that served the first one.
+ """
+ router = litellm.Router(
+ model_list=[
+ {
+ "model_name": MODEL_GROUP_ALIAS,
+ "litellm_params": {"model": AUTO_CACHING_MODEL, "api_key": "sk-fake"},
+ "model_info": {"id": model_id},
+ }
+ for model_id in ("dep-1", "dep-2", "dep-3")
+ ],
+ optional_pre_call_checks=["prompt_caching"],
+ )
+ user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 6))]
+ history: List[AllMessageValues] = []
+ served: List[str] = []
+ for text in user_turns:
+ request = cast(List[AllMessageValues], [*history, {"role": "user", "content": [_marked(text)]}])
+ response = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=request, mock_response="ok")
+ served.append(response._hidden_params["model_id"])
+ pin_key = PromptCachingCache.get_prompt_caching_cache_key(request, None)
+ assert await _eventually(functools.partial(router.cache.get_cache, key=pin_key)) is not None
+ history = [*history, {"role": "user", "content": [_text(text)]}, {"role": "assistant", "content": "ok"}]
+
+ assert served == [served[0]] * len(user_turns)
From c13dcb0abfe3de7b6722e18d7acf0f59eaa39fc8 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 20:11:14 -0700
Subject: [PATCH 11/56] fix(proxy): forward a client's anthropic-beta and
anthropic-version headers to bedrock_mantle
---
litellm/proxy/litellm_pre_call_utils.py | 7 +++++-
..._bedrock_mantle_messages_transformation.py | 19 +++++++++++++++
.../proxy/test_litellm_pre_call_utils.py | 24 ++++++++++++++++++-
3 files changed, 48 insertions(+), 2 deletions(-)
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index 9a973755894..e415a78f412 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -3418,7 +3418,12 @@ async def add_guardrails_from_policy_engine(
_ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join(
- (LlmProviders.ANTHROPIC.value, LlmProviders.BEDROCK.value, LlmProviders.VERTEX_AI.value)
+ (
+ LlmProviders.ANTHROPIC.value,
+ LlmProviders.BEDROCK.value,
+ LlmProviders.BEDROCK_MANTLE.value,
+ LlmProviders.VERTEX_AI.value,
+ )
)
_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = LlmProviders.ANTHROPIC.value
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
index 3544262996c..6bacf8f3d94 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py
@@ -385,6 +385,25 @@ class TestBetaHeadersOnTheWire:
"interleaved-thinking-2025-05-14",
]
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_betas_a_proxy_client_sends_reach_mantle_filtered(self):
+ from litellm.proxy.litellm_pre_call_utils import add_provider_specific_headers_to_request
+
+ proxy_request_data: dict = {}
+ add_provider_specific_headers_to_request(
+ data=proxy_request_data,
+ headers={
+ "anthropic-beta": "claude-code-20250219,fast-mode-2026-02-01,interleaved-thinking-2025-05-14",
+ "anthropic-version": "2023-06-01",
+ "user-agent": "claude-cli/2.1.239",
+ },
+ )
+
+ route = await self._send(**proxy_request_data)
+
+ assert _sent_betas(route) == ["claude-code-20250219", "interleaved-thinking-2025-05-14"]
+
@pytest.mark.asyncio
@respx.mock
async def test_betas_mantle_rejects_are_dropped_before_the_request(self):
diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py
index 88d38d74f49..9257a2dd23d 100644
--- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py
+++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py
@@ -7249,7 +7249,7 @@ CROSS_ACCOUNT_AUTHORIZATION = "Bearer deliberately-configured-pass-through-token
SIGV4_PREFIX = "AWS4-HMAC-SHA256"
AUTHORIZATION_HEADER_CASINGS = ["authorization", "Authorization", "AUTHORIZATION"]
-LEAK_TARGET_PROVIDERS = ["bedrock", "bedrock_converse", "vertex_ai"]
+LEAK_TARGET_PROVIDERS = ["bedrock", "bedrock_converse", "bedrock_mantle", "vertex_ai"]
BEDROCK_ENDPOINT = (
"https://bedrock-runtime.us-west-2.amazonaws.com/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke"
@@ -7342,6 +7342,28 @@ def test_oauth_credential_entry_is_scoped_to_anthropic_alone():
assert [entry["custom_llm_provider"] for entry in credential_entries] == ["anthropic"]
+@pytest.mark.parametrize("custom_llm_provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"])
+def test_client_anthropic_api_headers_reach_every_anthropic_messages_provider(custom_llm_provider):
+ client_headers = {
+ "anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14",
+ "anthropic-version": "2023-06-01",
+ "user-agent": "claude-cli/2.1.239",
+ }
+
+ forwarded = _headers_forwarded_to(client_headers, custom_llm_provider)
+
+ assert forwarded == {
+ "anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14",
+ "anthropic-version": "2023-06-01",
+ }
+
+
+def test_client_anthropic_api_headers_stay_off_openai_compatible_providers():
+ forwarded = _headers_forwarded_to({"anthropic-beta": "claude-code-20250219"}, "openai")
+
+ assert forwarded == {}
+
+
def test_no_provider_specific_header_when_client_sends_nothing_anthropic():
data: dict = {}
add_provider_specific_headers_to_request(
From 0f0c0fe499fc12856273f6094e622a8f9dc72311 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 20:13:59 -0700
Subject: [PATCH 12/56] fix: drop a blank anthropic-beta header before it
reaches the provider
---
litellm/anthropic_beta_headers_manager.py | 2 +-
.../test_anthropic_beta_headers_filtering.py | 18 ++++++++++++++++++
2 files changed, 19 insertions(+), 1 deletion(-)
diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py
index abce47c191e..7e7099a53b0 100644
--- a/litellm/anthropic_beta_headers_manager.py
+++ b/litellm/anthropic_beta_headers_manager.py
@@ -334,7 +334,7 @@ def update_headers_with_filtered_beta(
Updated headers dict
"""
existing_beta: Final = headers.get("anthropic-beta")
- if not existing_beta:
+ if existing_beta is None:
return headers
# Parse existing beta headers
diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py
index 3c967283abf..d404edb1281 100644
--- a/tests/test_litellm/test_anthropic_beta_headers_filtering.py
+++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py
@@ -18,6 +18,7 @@ import pytest
import litellm
from litellm.anthropic_beta_headers_manager import (
filter_and_transform_beta_headers,
+ update_headers_with_filtered_beta,
update_request_with_filtered_beta,
)
@@ -511,3 +512,20 @@ class TestAnthropicBetaHeadersFiltering:
assert (
"unknown-header-123" not in filtered
), f"Unknown header should not be in result for {provider}"
+
+ @pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"])
+ def test_blank_anthropic_beta_header_is_removed(self, provider):
+ headers = {"anthropic-beta": "", "anthropic-version": "2023-06-01"}
+
+ assert update_headers_with_filtered_beta(headers, provider) == {"anthropic-version": "2023-06-01"}
+
+ @pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"])
+ def test_whitespace_only_anthropic_beta_header_is_removed(self, provider):
+ headers = {"anthropic-beta": " , ", "anthropic-version": "2023-06-01"}
+
+ assert update_headers_with_filtered_beta(headers, provider) == {"anthropic-version": "2023-06-01"}
+
+ def test_absent_anthropic_beta_header_is_left_alone(self):
+ headers = {"anthropic-version": "2023-06-01"}
+
+ assert update_headers_with_filtered_beta(headers, "bedrock_mantle") == {"anthropic-version": "2023-06-01"}
From f24208f9ca8c0c5842e92eba09d6bc9b35b8a66f Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 20:17:55 -0700
Subject: [PATCH 13/56] fix(bedrock_mantle): price region-prefixed Claude
responses from the bare Bedrock row
---
litellm/utils.py | 17 ++++++++++++---
tests/test_litellm/test_cost_calculator.py | 25 ++++++++++++++++++++++
2 files changed, 39 insertions(+), 3 deletions(-)
diff --git a/litellm/utils.py b/litellm/utils.py
index 3439a21b560..f3b9fcfd1ed 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -5624,6 +5624,12 @@ def _get_model_info_from_generalization(
return None
+def _strip_mantle_region_prefix(model: str) -> str:
+ from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix
+
+ return split_mantle_region_prefix(model)[1]
+
+
def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> PotentialModelNamesAndCustomLLMProvider:
if custom_llm_provider is None:
# Get custom_llm_provider
@@ -5656,17 +5662,22 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P
split_model = strip_bedrock_routing_prefix(split_model)
+ region_free_split_model: Final = (
+ _strip_mantle_region_prefix(split_model) if custom_llm_provider == "bedrock_mantle" else split_model
+ )
provider_model_info: Final = (
- ProviderConfigManager.get_provider_model_info(model=split_model, provider=LlmProviders(custom_llm_provider))
+ ProviderConfigManager.get_provider_model_info(
+ model=region_free_split_model, provider=LlmProviders(custom_llm_provider)
+ )
if custom_llm_provider in LlmProvidersSet
else None
)
provider_cost_key: Final = (
- provider_model_info.get_model_cost_key(split_model) if provider_model_info is not None else None
+ provider_model_info.get_model_cost_key(region_free_split_model) if provider_model_info is not None else None
)
return PotentialModelNamesAndCustomLLMProvider(
- split_model=split_model,
+ split_model=region_free_split_model,
combined_model_name=combined_model_name,
stripped_model_name=stripped_model_name,
combined_stripped_model_name=combined_stripped_model_name,
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index aef17f3d5d0..fe52b9993f2 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -3522,6 +3522,31 @@ def test_cost_per_token_region_name_applies_to_provider_prefixed_model(_local_mo
)
+def test_completion_cost_mantle_native_messages_prices_claude_from_the_bedrock_row(_local_model_cost_map):
+ """Mantle's native Messages API answers with Anthropic's canonical model name and the proxy
+ resolves a Mantle region for every call, so the first cost candidate is
+ bedrock_mantle//claude-sonnet-5. That name has no row of its own and must fall through to
+ the deployment's bare Bedrock row instead of stopping on an unpriced capability rule at $0."""
+
+ response = litellm.ModelResponse(
+ id="msg_x",
+ choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
+ model="claude-sonnet-5",
+ usage={"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110},
+ )
+ row = litellm.model_cost["anthropic.claude-sonnet-5"]
+ expected = 100 * row["input_cost_per_token"] + 10 * row["output_cost_per_token"]
+ assert expected > 0
+
+ for region_name in ("us-east-1", None):
+ assert litellm.completion_cost(
+ completion_response=response,
+ model="bedrock_mantle/anthropic.claude-sonnet-5",
+ custom_llm_provider="bedrock_mantle",
+ region_name=region_name,
+ ) == pytest.approx(expected)
+
+
def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map):
"""An explicit base_model keeps pricing on that model's own key even when the request carries a
region with different regional rates, so the private provider model never widens region pricing."""
From 3ffe6272c96c08f54f972ef43a2541d73222f2ba Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 20:22:28 -0700
Subject: [PATCH 14/56] fix(router): hash the prompt caching affinity prefix
off the event loop
Offload the per-block hashing through offload_token_count on both the pre-call
read and the success-event write, hash raw bytes as base64 instead of raising,
drop the unused serialize_object helper, and bind the chained digest, the
message envelope, and the bytes path in the regression tests
---
litellm/constants.py | 2 -
litellm/router_utils/prompt_caching_cache.py | 38 +++------------
.../test_router_prompt_caching.py | 48 -------------------
.../test_prompt_caching_deployment_check.py | 40 ++++++++++++++--
4 files changed, 43 insertions(+), 85 deletions(-)
diff --git a/litellm/constants.py b/litellm/constants.py
index e4576ad4d5c..215f25bccd1 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -399,8 +399,6 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = (
if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None
else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
)
-# Anthropic checks at most 20 block positions behind a breakpoint for a cached prefix, a run of tool_use
-# or tool_result blocks counting as one position, so deployment affinity probes the same window
PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20
DEFAULT_TRIM_RATIO: Final = float(
os.getenv("DEFAULT_TRIM_RATIO", 0.75)
diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py
index 0b784e1fa91..78fc5e3fe6d 100644
--- a/litellm/router_utils/prompt_caching_cache.py
+++ b/litellm/router_utils/prompt_caching_cache.py
@@ -15,10 +15,8 @@ from typing_extensions import TypedDict
from litellm.caching.caching import DualCache
from litellm.constants import PROMPT_CACHE_LOOKBACK_POSITIONS
-from litellm.litellm_core_utils.logging_utils import (
- truncate_base64_in_messages,
- truncate_base64_in_messages_async,
-)
+from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages
+from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
if TYPE_CHECKING:
@@ -88,7 +86,9 @@ def _seed(tools: Sequence[ChatCompletionToolParam] | None) -> bytes:
if tools is None:
return hashlib.sha256(b"").digest()
return hashlib.sha256(
- _canonical_bytes(_TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True)))
+ _canonical_bytes(
+ _TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True, bytes_mode="base64"))
+ )
).digest()
@@ -132,23 +132,6 @@ class PromptCachingCache:
def __init__(self, cache: DualCache):
self.cache = cache
- @staticmethod
- def serialize_object(obj: Any) -> object:
- """Helper function to serialize Pydantic objects, dictionaries, or fallback to string."""
- if hasattr(obj, "dict"):
- # If the object is a Pydantic model, use its `dict()` method
- return obj.dict()
- elif isinstance(obj, dict):
- # If the object is a dictionary, serialize it with sorted keys
- return json.dumps(obj, sort_keys=True, separators=(",", ":")) # Standardize serialization
-
- elif isinstance(obj, list):
- # Serialize lists by ensuring each element is handled properly
- return [PromptCachingCache.serialize_object(item) for item in obj]
- elif isinstance(obj, (int, float, bool)):
- return obj # Keep primitive types as-is
- return str(obj)
-
@staticmethod
def extract_cacheable_prefix(
messages: list[AllMessageValues],
@@ -263,6 +246,7 @@ class PromptCachingCache:
to_jsonable_python(
truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)),
serialize_unknown=True,
+ bytes_mode="base64",
)
),
tools,
@@ -275,15 +259,7 @@ class PromptCachingCache:
) -> tuple[PrefixPosition, ...]:
if not messages:
return ()
- return _positions_of(
- _PREFIX_ADAPTER.validate_python(
- to_jsonable_python(
- await truncate_base64_in_messages_async(PromptCachingCache.extract_cacheable_prefix(messages)),
- serialize_unknown=True,
- )
- ),
- tools,
- )
+ return await offload_token_count(PromptCachingCache.prefix_positions)(messages, tools)
@staticmethod
def get_prompt_caching_cache_key(
diff --git a/tests/router_unit_tests/test_router_prompt_caching.py b/tests/router_unit_tests/test_router_prompt_caching.py
index 5c36c30e818..879264ca502 100644
--- a/tests/router_unit_tests/test_router_prompt_caching.py
+++ b/tests/router_unit_tests/test_router_prompt_caching.py
@@ -11,57 +11,9 @@ from unittest.mock import patch, MagicMock, AsyncMock
from create_mock_standard_logging_payload import create_standard_logging_payload
from litellm.types.utils import StandardLoggingPayload
import unittest
-from pydantic import BaseModel
from litellm.router_utils.prompt_caching_cache import PromptCachingCache
-class ExampleModel(BaseModel):
- field1: str
- field2: int
-
-
-def test_serialize_pydantic_object():
- model = ExampleModel(field1="value", field2=42)
- serialized = PromptCachingCache.serialize_object(model)
- assert serialized == {"field1": "value", "field2": 42}
-
-
-def test_serialize_dict():
- obj = {"b": 2, "a": 1}
- serialized = PromptCachingCache.serialize_object(obj)
- assert serialized == '{"a":1,"b":2}' # JSON string with sorted keys
-
-
-def test_serialize_nested_dict():
- obj = {"z": {"b": 2, "a": 1}, "x": [1, 2, {"c": 3}]}
- serialized = PromptCachingCache.serialize_object(obj)
- expected = '{"x":[1,2,{"c":3}],"z":{"a":1,"b":2}}' # JSON string with sorted keys
- assert serialized == expected
-
-
-def test_serialize_list():
- obj = ["item1", {"a": 1, "b": 2}, 42]
- serialized = PromptCachingCache.serialize_object(obj)
- expected = ["item1", '{"a":1,"b":2}', 42]
- assert serialized == expected
-
-
-def test_serialize_fallback():
- obj = 12345 # Simple non-serializable object
- serialized = PromptCachingCache.serialize_object(obj)
- assert serialized == 12345
-
-
-def test_serialize_non_serializable():
- class CustomClass:
- def __str__(self):
- return "custom_object"
-
- obj = CustomClass()
- serialized = PromptCachingCache.serialize_object(obj)
- assert serialized == "custom_object" # Fallback to string conversion
-
-
@pytest.mark.asyncio
async def test_router_prompt_caching_same_cacheable_prefix_routes_to_same_deployment():
"""
diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
index d0a9223dfa7..ad92f442a6e 100644
--- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
+++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
@@ -708,7 +708,10 @@ async def test_a_run_of_tool_blocks_counts_as_one_lookback_position():
@pytest.mark.asyncio
async def test_an_edited_earlier_block_does_not_inherit_the_pin():
- """Walking back must still bind every block's content, or an edited conversation pins to a stale cache."""
+ """
+ Every key must bind the whole prefix before its block, not the block alone, or a conversation
+ that repeats a pinned block after an edit walks back onto a cache the provider no longer holds.
+ """
prompt_cache = PromptCachingCache(cache=DualCache())
await prompt_cache.async_add_model_id(
model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("original")]}), tools=None
@@ -716,12 +719,41 @@ async def test_an_edited_earlier_block_does_not_inherit_the_pin():
edited = _turn(
{"role": "user", "content": [_text("edited")]},
{"role": "assistant", "content": "ok"},
- {"role": "user", "content": [_marked("next")]},
+ {"role": "user", "content": [_marked("original")]},
)
assert await prompt_cache.async_get_model_id(messages=edited, tools=None) is None
+@pytest.mark.asyncio
+async def test_swapped_roles_do_not_inherit_the_pin():
+ """The message envelope is part of what the provider caches, so the same blocks under other roles key apart."""
+ prompt_cache = PromptCachingCache(cache=DualCache())
+ pinned = _turn(
+ {"role": "user", "content": [_text("question")]},
+ {"role": "assistant", "content": [_marked("answer")]},
+ )
+ swapped = _turn(
+ {"role": "assistant", "content": [_text("question")]},
+ {"role": "user", "content": [_marked("answer")]},
+ )
+ await prompt_cache.async_add_model_id(model_id="dep-1", messages=pinned, tools=None)
+
+ assert await prompt_cache.async_get_model_id(messages=pinned, tools=None) == {"model_id": "dep-1"}
+ assert await prompt_cache.async_get_model_id(messages=swapped, tools=None) is None
+
+
+@pytest.mark.asyncio
+async def test_raw_bytes_in_a_block_hash_instead_of_failing_the_request():
+ """A block carrying raw bytes must key like any other block rather than raising out of the router filter."""
+ prompt_cache = PromptCachingCache(cache=DualCache())
+ binary_block = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": b"\xff\xfe"}}
+ turn = _turn({"role": "user", "content": [binary_block, _marked("describe")]})
+ await prompt_cache.async_add_model_id(model_id="dep-1", messages=turn, tools=None)
+
+ assert await prompt_cache.async_get_model_id(messages=turn, tools=None) == {"model_id": "dep-1"}
+
+
class _BrokenBatchReadCache(DualCache):
async def async_batch_get_cache(self, keys, parent_otel_span=None, local_only=False, **kwargs):
return None
@@ -786,11 +818,11 @@ async def test_claude_code_style_session_stays_on_one_deployment_across_turns(lo
"litellm_params": {"model": AUTO_CACHING_MODEL, "api_key": "sk-fake"},
"model_info": {"id": model_id},
}
- for model_id in ("dep-1", "dep-2", "dep-3")
+ for model_id in (f"dep-{number}" for number in range(1, 7))
],
optional_pre_call_checks=["prompt_caching"],
)
- user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 6))]
+ user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 9))]
history: List[AllMessageValues] = []
served: List[str] = []
for text in user_turns:
From 0c68c58eb1d63f0d857bba7ebdd8c4c5dbea992a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 20:39:53 -0700
Subject: [PATCH 15/56] test(proxy): expect bedrock_mantle in the anthropic
header provider list
---
tests/proxy_unit_tests/test_proxy_utils.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py
index 160753e3442..c62aab11930 100644
--- a/tests/proxy_unit_tests/test_proxy_utils.py
+++ b/tests/proxy_unit_tests/test_proxy_utils.py
@@ -2004,7 +2004,7 @@ def test_provider_specific_header():
)
# Verify multi-provider support: anthropic headers work across multiple providers
assert data["provider_specific_header"] == {
- "custom_llm_provider": "anthropic,bedrock,vertex_ai",
+ "custom_llm_provider": "anthropic,bedrock,bedrock_mantle,vertex_ai",
"extra_headers": {
"anthropic-beta": "prompt-caching-2024-07-31",
},
@@ -2076,7 +2076,7 @@ def test_provider_specific_header_multi_provider():
assert "provider_specific_header" in data
assert (
data["provider_specific_header"]["custom_llm_provider"]
- == "anthropic,bedrock,vertex_ai"
+ == "anthropic,bedrock,bedrock_mantle,vertex_ai"
)
assert data["provider_specific_header"]["extra_headers"] == {
"anthropic-beta": "context-1m-2025-08-07",
From 7fc114c24f393d9329c5b6cd919cee459a96fba7 Mon Sep 17 00:00:00 2001
From: yucheng
Date: Sun, 20 Sep 2026 08:24:47 +0000
Subject: [PATCH 16/56] feat(policy_engine): add default fallback policy
attachments
A policy attachment with default: true applies only when no non-default
attachment matches the request, so an opt-in guardrail policy replaces the
fallback one instead of running alongside it. Supported in config.yaml,
/policies/attachments, the Admin UI Attachments tab and the resolver
(matched_via is prefixed with default:).
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../migration.sql | 1 +
.../litellm_proxy_extras/schema.prisma | 1 +
litellm/proxy/_lazy_openapi_snapshot.json | 18 +++
.../policy_engine/attachment_registry.py | 24 ++-
.../proxy/policy_engine/policy_endpoints.py | 1 +
litellm/proxy/schema.prisma | 1 +
.../types/proxy/policy_engine/policy_types.py | 4 +
.../proxy/policy_engine/resolver_types.py | 8 +
schema.prisma | 1 +
.../policy_engine/test_attachment_registry.py | 152 ++++++++++++------
.../_components/AttachmentTable.test.tsx | 13 ++
.../_components/AttachmentTableColumns.tsx | 14 ++
.../_components/add_attachment_form.test.tsx | 15 ++
.../_components/add_attachment_form.tsx | 18 +++
.../_components/build_attachment_data.test.ts | 10 ++
.../_components/build_attachment_data.ts | 2 +
.../src/components/policies/types.ts | 2 +
ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 ++
18 files changed, 241 insertions(+), 56 deletions(-)
create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql
new file mode 100644
index 00000000000..a6c45448d03
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260920041500_add_policy_attachment_is_default/migration.sql
@@ -0,0 +1 @@
+ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "is_default" BOOLEAN NOT NULL DEFAULT false;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
index d2032cec0d0..2d7e557a9d1 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
@@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
+ is_default Boolean @default(false) // Applied only when no non-default attachment matches
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json
index 06e157498aa..6f9a2d8c96d 100644
--- a/litellm/proxy/_lazy_openapi_snapshot.json
+++ b/litellm/proxy/_lazy_openapi_snapshot.json
@@ -34982,6 +34982,12 @@
"PolicyAttachmentCreateRequest": {
"description": "Request body for creating a policy attachment.",
"properties": {
+ "default": {
+ "default": false,
+ "description": "Apply this attachment only when no non-default attachment matches the request.",
+ "title": "Default",
+ "type": "boolean"
+ },
"keys": {
"anyOf": [
{
@@ -35113,6 +35119,12 @@
"description": "Who created the attachment.",
"title": "Created By"
},
+ "default": {
+ "default": false,
+ "description": "Apply this attachment only when no non-default attachment matches the request.",
+ "title": "Default",
+ "type": "boolean"
+ },
"definition_location": {
"default": "db",
"description": "Where this attachment is defined: 'db' (database) or 'config' (config.yaml).",
@@ -37141,6 +37153,12 @@
"PolicyAttachmentCreateRequest": {
"description": "Request body for creating a policy attachment.",
"properties": {
+ "default": {
+ "default": false,
+ "description": "Apply this attachment only when no non-default attachment matches the request.",
+ "title": "Default",
+ "type": "boolean"
+ },
"keys": {
"anyOf": [
{
diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py
index 3735c335bd4..04009151487 100644
--- a/litellm/proxy/policy_engine/attachment_registry.py
+++ b/litellm/proxy/policy_engine/attachment_registry.py
@@ -119,6 +119,7 @@ class AttachmentRegistry:
models=attachment_data.get("models"),
tags=attachment_data.get("tags"),
priority=attachment_data.get("priority"),
+ default=attachment_data.get("default", False),
)
def get_attached_policies(self, context: PolicyMatchContext) -> list[str]:
@@ -142,12 +143,14 @@ class AttachmentRegistry:
"""
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
+ in_scope: Final = tuple(
+ attachment
+ for attachment in self._attachments
+ if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
+ )
+ non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default)
matching_attachments: Final = sorted(
- (
- attachment
- for attachment in self._attachments
- if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
- ),
+ non_default or tuple(attachment for attachment in in_scope if attachment.default),
key=_attachment_sort_key,
)
broadest_attachment_by_policy: Final = MappingProxyType(
@@ -169,6 +172,11 @@ class AttachmentRegistry:
@staticmethod
def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str:
"""Describe why an attachment matched the context."""
+ reason: Final = AttachmentRegistry._describe_scope_match(attachment, context)
+ return f"default:{reason}" if attachment.default else reason
+
+ @staticmethod
+ def _describe_scope_match(attachment: PolicyAttachment, context: PolicyMatchContext) -> str:
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
if attachment.is_global():
@@ -324,6 +332,7 @@ class AttachmentRegistry:
"models": attachment_request.models or [],
"tags": attachment_request.tags or [],
"priority": attachment_request.priority,
+ "is_default": attachment_request.default,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
"created_by": created_by,
@@ -340,6 +349,7 @@ class AttachmentRegistry:
models=attachment_request.models,
tags=attachment_request.tags,
priority=attachment_request.priority,
+ default=attachment_request.default,
)
self.add_attachment(attachment)
@@ -352,6 +362,7 @@ class AttachmentRegistry:
models=created_attachment.models or [],
tags=created_attachment.tags or [],
priority=created_attachment.priority,
+ default=created_attachment.is_default,
created_at=created_attachment.created_at,
updated_at=created_attachment.updated_at,
created_by=created_attachment.created_by,
@@ -429,6 +440,7 @@ class AttachmentRegistry:
models=attachment.models or [],
tags=attachment.tags or [],
priority=attachment.priority,
+ default=attachment.is_default,
created_at=attachment.created_at,
updated_at=attachment.updated_at,
created_by=attachment.created_by,
@@ -468,6 +480,7 @@ class AttachmentRegistry:
models=a.models or [],
tags=a.tags or [],
priority=a.priority,
+ default=a.is_default,
created_at=a.created_at,
updated_at=a.updated_at,
created_by=a.created_by,
@@ -502,6 +515,7 @@ class AttachmentRegistry:
models=(attachment_response.models if attachment_response.models else None),
tags=attachment_response.tags if attachment_response.tags else None,
priority=attachment_response.priority,
+ default=attachment_response.default,
)
for attachment_response in attachments
]
diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py
index 1e30238c8b4..f4b38bea14e 100644
--- a/litellm/proxy/policy_engine/policy_endpoints.py
+++ b/litellm/proxy/policy_engine/policy_endpoints.py
@@ -61,6 +61,7 @@ def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment)
models=attachment.models or [],
tags=attachment.tags or [],
priority=attachment.priority,
+ default=attachment.default,
definition_location="config",
)
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index d2032cec0d0..2d7e557a9d1 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
+ is_default Boolean @default(false) // Applied only when no non-default attachment matches
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
diff --git a/litellm/types/proxy/policy_engine/policy_types.py b/litellm/types/proxy/policy_engine/policy_types.py
index 66e5fbb4b49..73eeffa3585 100644
--- a/litellm/types/proxy/policy_engine/policy_types.py
+++ b/litellm/types/proxy/policy_engine/policy_types.py
@@ -294,6 +294,10 @@ class PolicyAttachment(BaseModel):
le=2147483647,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
+ default: bool = Field(
+ default=False,
+ description="Apply this attachment only when no non-default attachment matches the request.",
+ )
model_config = ConfigDict(extra="forbid")
diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py
index e6f501ed4b5..ebdedb98b12 100644
--- a/litellm/types/proxy/policy_engine/resolver_types.py
+++ b/litellm/types/proxy/policy_engine/resolver_types.py
@@ -311,6 +311,10 @@ class PolicyAttachmentCreateRequest(BaseModel):
le=2147483647,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
+ default: bool = Field(
+ default=False,
+ description="Apply this attachment only when no non-default attachment matches the request.",
+ )
class PolicyAttachmentDBResponse(BaseModel):
@@ -327,6 +331,10 @@ class PolicyAttachmentDBResponse(BaseModel):
default=None,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
+ default: bool = Field(
+ default=False,
+ description="Apply this attachment only when no non-default attachment matches the request.",
+ )
created_at: datetime | None = Field(default=None, description="When the attachment was created.")
updated_at: datetime | None = Field(default=None, description="When the attachment was last updated.")
created_by: str | None = Field(default=None, description="Who created the attachment.")
diff --git a/schema.prisma b/schema.prisma
index d2032cec0d0..2d7e557a9d1 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -1419,6 +1419,7 @@ model LiteLLM_PolicyAttachmentTable {
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
+ is_default Boolean @default(false) // Applied only when no non-default attachment matches
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py
index 089bec59583..b419f3db060 100644
--- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py
+++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py
@@ -30,9 +30,7 @@ class TestGetAttachedPolicies:
)
# Should match any context
- context = PolicyMatchContext(
- team_alias="any-team", key_alias="any-key", model="any-model"
- )
+ context = PolicyMatchContext(team_alias="any-team", key_alias="any-key", model="any-model")
attached = registry.get_attached_policies(context)
assert "global-baseline" in attached
@@ -46,15 +44,11 @@ class TestGetAttachedPolicies:
)
# Match
- context = PolicyMatchContext(
- team_alias="healthcare-team", key_alias="key", model="gpt-4"
- )
+ context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4")
assert "healthcare-policy" in registry.get_attached_policies(context)
# No match - different team
- context_other = PolicyMatchContext(
- team_alias="finance-team", key_alias="key", model="gpt-4"
- )
+ context_other = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4")
assert "healthcare-policy" not in registry.get_attached_policies(context_other)
def test_key_wildcard_pattern_attachment(self):
@@ -67,15 +61,11 @@ class TestGetAttachedPolicies:
)
# Match - key starts with dev-key-
- context = PolicyMatchContext(
- team_alias="team", key_alias="dev-key-123", model="gpt-4"
- )
+ context = PolicyMatchContext(team_alias="team", key_alias="dev-key-123", model="gpt-4")
assert "dev-policy" in registry.get_attached_policies(context)
# No match - different prefix
- context_prod = PolicyMatchContext(
- team_alias="team", key_alias="prod-key-123", model="gpt-4"
- )
+ context_prod = PolicyMatchContext(team_alias="team", key_alias="prod-key-123", model="gpt-4")
assert "dev-policy" not in registry.get_attached_policies(context_prod)
def test_model_specific_attachment(self):
@@ -92,9 +82,7 @@ class TestGetAttachedPolicies:
assert "gpt4-policy" in registry.get_attached_policies(context)
# No match
- context_other = PolicyMatchContext(
- team_alias="team", key_alias="key", model="gpt-3.5"
- )
+ context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-3.5")
assert "gpt4-policy" not in registry.get_attached_policies(context_other)
def test_model_wildcard_pattern(self):
@@ -107,15 +95,11 @@ class TestGetAttachedPolicies:
)
# Match
- context = PolicyMatchContext(
- team_alias="team", key_alias="key", model="bedrock/claude-3"
- )
+ context = PolicyMatchContext(team_alias="team", key_alias="key", model="bedrock/claude-3")
assert "bedrock-policy" in registry.get_attached_policies(context)
# No match
- context_other = PolicyMatchContext(
- team_alias="team", key_alias="key", model="openai/gpt-4"
- )
+ context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="openai/gpt-4")
assert "bedrock-policy" not in registry.get_attached_policies(context_other)
def test_multiple_attachments_match_same_context(self):
@@ -129,9 +113,7 @@ class TestGetAttachedPolicies:
]
)
- context = PolicyMatchContext(
- team_alias="healthcare-team", key_alias="key", model="gpt-4"
- )
+ context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4")
attached = registry.get_attached_policies(context)
# All three should match
@@ -277,9 +259,7 @@ class TestGetAttachedPolicies:
]
)
- context = PolicyMatchContext(
- team_alias="healthcare-team", key_alias="key", model="gpt-4"
- )
+ context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4")
attached = registry.get_attached_policies(context)
# Should only appear once
@@ -288,9 +268,7 @@ class TestGetAttachedPolicies:
def test_many_distinct_policies_resolve_in_linear_time(self):
policy_count = 20_000
registry = AttachmentRegistry()
- registry.load_attachments(
- [{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)]
- )
+ registry.load_attachments([{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)])
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
started = time.perf_counter()
@@ -318,9 +296,7 @@ class TestGetAttachedPolicies:
]
)
- context = PolicyMatchContext(
- team_alias="finance-team", key_alias="key", model="gpt-4"
- )
+ context = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4")
attached = registry.get_attached_policies(context)
assert attached == []
@@ -338,23 +314,15 @@ class TestGetAttachedPolicies:
)
# Match - both team and model match
- context = PolicyMatchContext(
- team_alias="healthcare-team", key_alias="key", model="gpt-4"
- )
+ context = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-4")
assert "strict-policy" in registry.get_attached_policies(context)
# No match - team matches but model doesn't
- context_wrong_model = PolicyMatchContext(
- team_alias="healthcare-team", key_alias="key", model="gpt-3.5"
- )
- assert "strict-policy" not in registry.get_attached_policies(
- context_wrong_model
- )
+ context_wrong_model = PolicyMatchContext(team_alias="healthcare-team", key_alias="key", model="gpt-3.5")
+ assert "strict-policy" not in registry.get_attached_policies(context_wrong_model)
# No match - model matches but team doesn't
- context_wrong_team = PolicyMatchContext(
- team_alias="finance-team", key_alias="key", model="gpt-4"
- )
+ context_wrong_team = PolicyMatchContext(team_alias="finance-team", key_alias="key", model="gpt-4")
assert "strict-policy" not in registry.get_attached_policies(context_wrong_team)
@@ -527,6 +495,79 @@ class TestMatchAttribution:
assert "catch-all" in attached
+class TestDefaultAttachments:
+ """`default: true` attachments apply only when no non-default attachment matches."""
+
+ @staticmethod
+ def _registry() -> AttachmentRegistry:
+ registry = AttachmentRegistry()
+ registry.load_attachments(
+ [
+ {"policy": "guardrail-y", "scope": "*", "default": True},
+ {"policy": "guardrail-x", "tags": ["opt-in"]},
+ ]
+ )
+ return registry
+
+ def test_opted_in_request_gets_only_the_opt_in_policy(self):
+ context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"])
+
+ assert self._registry().get_attached_policies(context) == ["guardrail-x"]
+
+ def test_request_without_opt_in_falls_back_to_default_policy(self):
+ context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2")
+
+ assert self._registry().get_attached_policies(context) == ["guardrail-y"]
+
+ def test_default_attachment_still_honors_its_own_scope(self):
+ registry = AttachmentRegistry()
+ registry.load_attachments([{"policy": "team-default", "teams": ["team-a"], "default": True}])
+
+ assert registry.get_attached_policies(PolicyMatchContext(team_alias="team-a", key_alias="k", model="m")) == [
+ "team-default"
+ ]
+ assert registry.get_attached_policies(PolicyMatchContext(team_alias="team-b", key_alias="k", model="m")) == []
+
+ def test_all_matching_defaults_apply_when_nothing_else_matches(self):
+ registry = AttachmentRegistry()
+ registry.load_attachments(
+ [
+ {"policy": "default-a", "scope": "*", "default": True},
+ {"policy": "default-b", "teams": ["team-a"], "default": True},
+ {"policy": "opt-in", "tags": ["opt-in"]},
+ ]
+ )
+ context = PolicyMatchContext(team_alias="team-a", key_alias="k", model="m")
+
+ assert registry.get_attached_policies(context) == ["default-a", "default-b"]
+
+ def test_non_default_attachments_remain_additive(self):
+ registry = AttachmentRegistry()
+ registry.load_attachments(
+ [
+ {"policy": "baseline", "scope": "*"},
+ {"policy": "opt-in", "tags": ["opt-in"]},
+ {"policy": "fallback", "scope": "*", "default": True},
+ ]
+ )
+ context = PolicyMatchContext(team_alias="t", key_alias="k", model="m", tags=["opt-in"])
+
+ assert registry.get_attached_policies(context) == ["baseline", "opt-in"]
+
+ def test_default_match_reason_is_labelled(self):
+ context = PolicyMatchContext(team_alias="t", key_alias="k", model="m")
+
+ results = self._registry().get_attached_policies_with_reasons(context)
+
+ assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}]
+
+ def test_default_defaults_to_false_when_omitted(self):
+ registry = AttachmentRegistry()
+ registry.load_attachments([{"policy": "p"}])
+
+ assert registry.get_all_attachments()[0].default is False
+
+
class TestAttachmentRegistrySingleton:
"""Test global singleton behavior."""
@@ -557,6 +598,7 @@ def _make_db_attachment_row(
scope: str | None = None,
teams: list[str] | None = None,
priority: int | None = None,
+ is_default: bool = False,
) -> MagicMock:
row = MagicMock()
row.attachment_id = attachment_id
@@ -567,6 +609,7 @@ def _make_db_attachment_row(
row.models = []
row.tags = []
row.priority = priority
+ row.is_default = is_default
row.created_at = datetime.now(timezone.utc)
row.updated_at = datetime.now(timezone.utc)
row.created_by = None
@@ -576,9 +619,7 @@ def _make_db_attachment_row(
def _prisma_with_attachment_rows(rows: list[MagicMock]) -> MagicMock:
prisma = MagicMock()
- prisma.configure_mock(
- **{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)}
- )
+ prisma.configure_mock(**{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)})
return prisma
@@ -629,6 +670,15 @@ class TestConfigAttachmentsPreservedAcrossDbSync:
assert registry.get_all_attachments()[0].priority == 7
+ @pytest.mark.asyncio
+ async def test_sync_round_trips_db_attachment_default_flag(self):
+ registry = AttachmentRegistry()
+ db_row = _make_db_attachment_row(is_default=True)
+
+ await registry.sync_attachments_from_db(_prisma_with_attachment_rows([db_row]))
+
+ assert registry.get_all_attachments()[0].default is True
+
@pytest.mark.asyncio
async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self):
registry = AttachmentRegistry()
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx
index 43ad6a7cc9e..be83f73bb2e 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx
@@ -65,6 +65,19 @@ describe("AttachmentTable", () => {
);
});
+ it("should show a Default badge only for default attachments", () => {
+ const attachments = [
+ makeAttachment({ attachment_id: "att-def00001", policy_name: "fallback", default: true }),
+ makeAttachment({ attachment_id: "att-def00002", policy_name: "regular" }),
+ ];
+ renderWithProviders( );
+ const rows = screen.getAllByRole("row").slice(1);
+ const fallbackRow = rows.find((row) => within(row).queryByText("fallback"));
+ const regularRow = rows.find((row) => within(row).queryByText("regular"));
+ expect(within(fallbackRow!).getByText("Default")).toBeInTheDocument();
+ expect(within(regularRow!).queryByText("Default")).not.toBeInTheDocument();
+ });
+
it("should show skeleton rows when isLoading is true", () => {
renderWithProviders( );
expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0);
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx
index 9a190401d08..3265b9db834 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx
@@ -181,6 +181,20 @@ export const getAttachmentTableColumns = ({
{row.original.priority}
),
},
+ {
+ id: "default",
+ accessorFn: (row) => (row.default ? 1 : 0),
+ meta: { title: "Default" },
+ header: ({ column }) => ,
+ size: 100,
+ enableSorting: true,
+ cell: ({ row }) =>
+ row.original.default ? (
+
+ ) : (
+ -
+ ),
+ },
{
id: "created_at",
accessorFn: (row) => row.created_at ?? "",
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx
index dfc023d428e..14af4a2b8f3 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.test.tsx
@@ -237,6 +237,21 @@ describe("AddAttachmentForm", () => {
expect(createAttachment).toHaveBeenCalledWith("test-token", { policy_name: "policy-alpha", scope: "*" });
});
+ it("sends default: true when the Default switch is turned on", async () => {
+ const user = userEvent.setup();
+ const createAttachment = vi.fn().mockResolvedValue({});
+ renderWithProviders( );
+ await selectPolicy(user, "policy-alpha");
+ await user.click(screen.getByRole("switch", { name: /default/i }));
+ await submit(user);
+ await waitFor(() => expect(createAttachment).toHaveBeenCalledTimes(1));
+ expect(createAttachment).toHaveBeenCalledWith("test-token", {
+ policy_name: "policy-alpha",
+ scope: "*",
+ default: true,
+ });
+ });
+
it.each([
["2147483648", /at most 2147483647/i],
["-2147483649", /at least -2147483648/i],
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx
index 02463a89139..74c4978392f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx
@@ -11,6 +11,7 @@ import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
import { Separator } from "@/components/ui/separator";
+import { Switch } from "@/components/ui/switch";
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { useZodForm } from "@/lib/forms/useZodForm";
@@ -38,6 +39,7 @@ interface AttachmentFormValues {
models: string[];
tags: string[];
priority: number | null;
+ default: boolean;
}
const EMPTY_VALUES: AttachmentFormValues = {
@@ -47,6 +49,7 @@ const EMPTY_VALUES: AttachmentFormValues = {
models: [],
tags: [],
priority: null,
+ default: false,
};
const INT32_MIN = -2147483648;
@@ -64,6 +67,7 @@ const attachmentShape = {
.min(INT32_MIN, `Priority must be at least ${INT32_MIN}`)
.max(INT32_MAX, `Priority must be at most ${INT32_MAX}`)
.nullable(),
+ default: z.boolean(),
};
const buildAttachmentSchema = (scopeType: ScopeType, teamsLoaded: boolean, availableTeams: string[]) =>
@@ -453,6 +457,20 @@ const AddAttachmentForm: React.FC = ({
/>
)}
+
+
+ {({ value, onChange, ref, ...field }) => (
+
+ )}
+
{impactResult && }
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.test.ts
index 930e755f242..80617a40f13 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.test.ts
@@ -80,6 +80,16 @@ describe("buildAttachmentData", () => {
});
});
+ describe("default", () => {
+ it.each(["global", "specific"] as const)("should send default: true for a %s scope", (scopeType) => {
+ expect(buildAttachmentData({ policy_name: "p", default: true }, scopeType).default).toBe(true);
+ });
+
+ it.each([undefined, false])("should omit default when it is %s", (value) => {
+ expect(buildAttachmentData({ policy_name: "p", default: value }, "specific")).not.toHaveProperty("default");
+ });
+ });
+
describe("priority", () => {
it.each(["global", "specific"] as const)("should include priority for a %s scope", (scopeType) => {
expect(buildAttachmentData({ policy_name: "p", priority: 0 }, scopeType).priority).toBe(0);
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.ts
index 8b21142df74..8b50cd7bdc1 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/build_attachment_data.ts
@@ -7,6 +7,7 @@ export interface AttachmentFormInput {
models?: string[];
tags?: string[];
priority?: number | null;
+ default?: boolean;
}
export function buildAttachmentData(
@@ -25,5 +26,6 @@ export function buildAttachmentData(
if (formValues.tags && formValues.tags.length > 0) data.tags = formValues.tags;
}
if (typeof formValues.priority === "number") data.priority = formValues.priority;
+ if (formValues.default === true) data.default = true;
return data;
}
diff --git a/ui/litellm-dashboard/src/components/policies/types.ts b/ui/litellm-dashboard/src/components/policies/types.ts
index 9f3ef02ba5d..430864f93df 100644
--- a/ui/litellm-dashboard/src/components/policies/types.ts
+++ b/ui/litellm-dashboard/src/components/policies/types.ts
@@ -45,6 +45,7 @@ export interface PolicyAttachment {
models: string[];
tags: string[];
priority?: number | null;
+ default?: boolean;
created_at?: string;
updated_at?: string;
created_by?: string;
@@ -80,6 +81,7 @@ export interface PolicyAttachmentCreateRequest {
models?: string[];
tags?: string[];
priority?: number;
+ default?: boolean;
}
export interface PolicyListResponse {
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index e33764c3d1a..ccb5fda2149 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -35164,6 +35164,12 @@ export interface components {
* @description Request body for creating a policy attachment.
*/
PolicyAttachmentCreateRequest: {
+ /**
+ * Default
+ * @description Apply this attachment only when no non-default attachment matches the request.
+ * @default false
+ */
+ default: boolean;
/**
* Keys
* @description Key aliases or patterns this attachment applies to.
@@ -35220,6 +35226,12 @@ export interface components {
* @description Who created the attachment.
*/
created_by?: string | null;
+ /**
+ * Default
+ * @description Apply this attachment only when no non-default attachment matches the request.
+ * @default false
+ */
+ default: boolean;
/**
* Definition Location
* @description Where this attachment is defined: 'db' (database) or 'config' (config.yaml).
From ef55eb6bdf0ccb69533d9541266c5108e0176609 Mon Sep 17 00:00:00 2001
From: yucheng
Date: Sun, 20 Sep 2026 09:07:08 +0000
Subject: [PATCH 17/56] fix(policy_engine): ignore inapplicable non-default
attachments when selecting defaults
A non-default attachment whose policy is missing or whose condition does not match the request
no longer suppresses default attachments. The impact preview marks default counts as an upper bound
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/proxy/litellm_pre_call_utils.py | 4 ++-
.../policy_engine/attachment_registry.py | 18 ++++++++--
litellm/proxy/policy_engine/policy_matcher.py | 15 ++++++++
.../policy_engine/policy_resolve_endpoints.py | 4 ++-
.../proxy/policy_engine/response_retrieval.py | 4 ++-
.../policy_engine/test_attachment_registry.py | 35 ++++++++++++++++++-
.../_components/add_attachment_form.tsx | 2 +-
.../_components/impact_preview_alert.test.tsx | 11 ++++++
.../_components/impact_preview_alert.tsx | 11 ++++--
9 files changed, 94 insertions(+), 10 deletions(-)
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index 9a973755894..8aa861e574f 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -3216,7 +3216,9 @@ def _match_and_track_policies(
attachment_registry: Final = (
attachment_registry_override if attachment_registry_override is not None else get_attachment_registry()
)
- matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(context)
+ matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(
+ context, PolicyMatcher.policy_applies(context, policies_override)
+ )
matching_policy_names: Final = [m["policy_name"] for m in matches_with_reasons]
policy_reasons: Final = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons}
diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py
index 04009151487..d81471b3c1a 100644
--- a/litellm/proxy/policy_engine/attachment_registry.py
+++ b/litellm/proxy/policy_engine/attachment_registry.py
@@ -5,6 +5,7 @@ Attachments define WHERE policies apply, separate from the policy definitions.
This allows the same policy to be attached to multiple scopes.
"""
+from collections.abc import Callable
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypedDict
@@ -122,24 +123,34 @@ class AttachmentRegistry:
default=attachment_data.get("default", False),
)
- def get_attached_policies(self, context: PolicyMatchContext) -> list[str]:
+ def get_attached_policies(
+ self,
+ context: PolicyMatchContext,
+ policy_applies: Callable[[str], bool] | None = None,
+ ) -> list[str]:
"""
Get list of policy names attached to the given context.
Args:
context: The request context to match against
+ policy_applies: Optional predicate; attachments whose policy does not apply are ignored
Returns:
List of policy names that are attached to matching scopes
"""
- return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)]
+ return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context, policy_applies)]
- def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[PolicyAttachmentMatch]:
+ def get_attached_policies_with_reasons(
+ self,
+ context: PolicyMatchContext,
+ policy_applies: Callable[[str], bool] | None = None,
+ ) -> list[PolicyAttachmentMatch]:
"""
Get list of policy names and match reasons for the given context.
Returns a list of dicts with 'policy_name' and 'matched_via' keys.
The 'matched_via' describes which dimension caused the match.
+ Attachments whose policy fails `policy_applies` are dropped before defaults are considered.
"""
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
@@ -147,6 +158,7 @@ class AttachmentRegistry:
attachment
for attachment in self._attachments
if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
+ and (policy_applies is None or policy_applies(attachment.policy))
)
non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default)
matching_attachments: Final = sorted(
diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py
index 001e4115374..2f7fcd23b75 100644
--- a/litellm/proxy/policy_engine/policy_matcher.py
+++ b/litellm/proxy/policy_engine/policy_matcher.py
@@ -7,6 +7,7 @@ apply to a given request based on team alias, key alias, and model.
Policies are matched via policy_attachments which define WHERE each policy applies.
"""
+from collections.abc import Callable
from typing import Final
from litellm._logging import verbose_proxy_logger
@@ -130,6 +131,20 @@ class PolicyMatcher:
"""
return PolicyMatcher.get_matching_policies(context=context)
+ @staticmethod
+ def policy_applies(
+ context: PolicyMatchContext,
+ policies: dict[str, Policy] | None = None,
+ ) -> Callable[[str], bool]:
+ """Predicate telling whether a policy exists and its condition matches the context."""
+ return lambda policy_name: bool(
+ PolicyMatcher.get_policies_with_matching_conditions(
+ policy_names=[policy_name], # mutable-ok: the matcher takes a list
+ context=context,
+ policies=policies,
+ )
+ )
+
@staticmethod
def get_policies_with_matching_conditions(
policy_names: list[str],
diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py
index a8a9856b833..898e42635c5 100644
--- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py
+++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py
@@ -265,7 +265,9 @@ async def resolve_policies_for_context(
)
# Get matching policies with reasons
- match_results: Final = get_attachment_registry().get_attached_policies_with_reasons(context=context)
+ match_results: Final = get_attachment_registry().get_attached_policies_with_reasons(
+ context=context, policy_applies=PolicyMatcher.policy_applies(context)
+ )
if not match_results:
return PolicyResolveResponse(
diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py
index d284c44397e..0f373b08056 100644
--- a/litellm/proxy/policy_engine/response_retrieval.py
+++ b/litellm/proxy/policy_engine/response_retrieval.py
@@ -84,7 +84,9 @@ def _retrieval_context(
def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]:
- matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context)
+ matches: Final = get_attachment_registry().get_attached_policies_with_reasons(
+ context, PolicyMatcher.policy_applies(context)
+ )
if not matches:
return (), MappingProxyType({})
applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions(
diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py
index b419f3db060..faa8d67fe3a 100644
--- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py
+++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py
@@ -14,7 +14,8 @@ from litellm.proxy.policy_engine.attachment_registry import (
AttachmentRegistry,
get_attachment_registry,
)
-from litellm.types.proxy.policy_engine import PolicyMatchContext
+from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
+from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails, PolicyMatchContext
class TestGetAttachedPolicies:
@@ -561,6 +562,38 @@ class TestDefaultAttachments:
assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}]
+ def test_inapplicable_opt_in_policy_does_not_suppress_default(self):
+ context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"])
+ policies = {
+ "guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])),
+ "guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="claude.*")),
+ }
+
+ results = self._registry().get_attached_policies_with_reasons(
+ context, PolicyMatcher.policy_applies(context, policies)
+ )
+
+ assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}]
+
+ def test_attachment_to_missing_policy_does_not_suppress_default(self):
+ context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"])
+ policies = {"guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"]))}
+
+ assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [
+ "guardrail-y"
+ ]
+
+ def test_applicable_opt_in_policy_still_wins_with_predicate(self):
+ context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"])
+ policies = {
+ "guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])),
+ "guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="gpt.*")),
+ }
+
+ assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [
+ "guardrail-x"
+ ]
+
def test_default_defaults_to_false_when_omitted(self):
registry = AttachmentRegistry()
registry.load_attachments([{"policy": "p"}])
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx
index 74c4978392f..5cd240a0838 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_attachment_form.tsx
@@ -473,7 +473,7 @@ const AddAttachmentForm: React.FC = ({
- {impactResult && }
+ {impactResult && }
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/impact_preview_alert.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/impact_preview_alert.test.tsx
index 6c405d7492a..6074508b52b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/impact_preview_alert.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/impact_preview_alert.test.tsx
@@ -69,6 +69,17 @@ describe("ImpactPreviewAlert", () => {
expect(screen.getByText(/1 key\b/i)).toBeInTheDocument();
});
+ it("should present the counts as an upper bound for a default attachment", () => {
+ renderWithProviders( );
+ expect(screen.getByText(/would affect up to/i)).toBeInTheDocument();
+ expect(screen.getByText(/no non-default attachment matches/i)).toBeInTheDocument();
+ });
+
+ it("should not qualify the counts for a non-default attachment", () => {
+ renderWithProviders( );
+ expect(screen.queryByText(/up to/i)).not.toBeInTheDocument();
+ });
+
it("should not show a key section when there are no sample keys", () => {
const noKeys = { affected_keys_count: 0, affected_teams_count: 2, sample_keys: [], sample_teams: ["t1", "t2"] };
renderWithProviders( );
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/impact_preview_alert.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/impact_preview_alert.tsx
index a9264eb9d0d..df8a6987d6f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/impact_preview_alert.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/impact_preview_alert.tsx
@@ -12,6 +12,7 @@ interface ImpactResult {
interface ImpactPreviewAlertProps {
impactResult: ImpactResult;
+ isDefault?: boolean;
}
interface SampleListProps {
@@ -32,8 +33,9 @@ const SampleList: React.FC = ({ label, samples, totalCount }) =
);
-const ImpactPreviewAlert: React.FC = ({ impactResult }) => {
+const ImpactPreviewAlert: React.FC = ({ impactResult, isDefault = false }) => {
const isGlobal = impactResult.affected_keys_count === -1;
+ const qualifier = isDefault ? "up to " : "";
return (
@@ -47,7 +49,7 @@ const ImpactPreviewAlert: React.FC = ({ impactResult })
) : (
- This attachment would affect{" "}
+ This attachment would affect {qualifier}
{impactResult.affected_keys_count} key{impactResult.affected_keys_count !== 1 ? "s" : ""}
{" "}
@@ -57,6 +59,11 @@ const ImpactPreviewAlert: React.FC = ({ impactResult })
.
+ {isDefault && (
+
+ Default attachments only apply to requests no non-default attachment matches, so fewer may be affected.
+
+ )}
{impactResult.sample_keys.length > 0 && (
Date: Sun, 20 Sep 2026 09:09:27 +0000
Subject: [PATCH 18/56] refactor(policy_engine): accept any sequence of policy
names in condition matching
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/proxy/policy_engine/policy_matcher.py | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py
index 2f7fcd23b75..2b54b5dbe41 100644
--- a/litellm/proxy/policy_engine/policy_matcher.py
+++ b/litellm/proxy/policy_engine/policy_matcher.py
@@ -7,7 +7,7 @@ apply to a given request based on team alias, key alias, and model.
Policies are matched via policy_attachments which define WHERE each policy applies.
"""
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
from typing import Final
from litellm._logging import verbose_proxy_logger
@@ -139,7 +139,7 @@ class PolicyMatcher:
"""Predicate telling whether a policy exists and its condition matches the context."""
return lambda policy_name: bool(
PolicyMatcher.get_policies_with_matching_conditions(
- policy_names=[policy_name], # mutable-ok: the matcher takes a list
+ policy_names=(policy_name,),
context=context,
policies=policies,
)
@@ -147,7 +147,7 @@ class PolicyMatcher:
@staticmethod
def get_policies_with_matching_conditions(
- policy_names: list[str],
+ policy_names: Sequence[str],
context: PolicyMatchContext,
policies: dict[str, Policy] | None = None,
) -> list[str]:
From 95cf7066d16186e94a8f27828ac38db55ef45cf7 Mon Sep 17 00:00:00 2001
From: Yujong Lee
Date: Sun, 20 Sep 2026 20:38:13 -0700
Subject: [PATCH 19/56] refactor(cache): use static dispatch and typed backend
codecs
---
litellm-rust/Cargo.lock | 1 +
litellm-rust/crates/cache-memory/src/cache.rs | 23 +--
.../crates/cache-memory/tests/cache.rs | 47 ++++-
litellm-rust/crates/cache-redis/Cargo.toml | 2 +-
litellm-rust/crates/cache-redis/src/cache.rs | 171 ++++++++----------
.../crates/cache-redis/tests/cache.rs | 152 +++++++++++++++-
litellm-rust/crates/cache/Cargo.toml | 1 +
litellm-rust/crates/cache/src/base_cache.rs | 53 +++---
litellm-rust/crates/cache/src/caching.rs | 16 +-
litellm-rust/crates/cache/src/codec.rs | 42 +++++
litellm-rust/crates/cache/src/lib.rs | 6 +-
litellm-rust/crates/cache/tests/caching.rs | 70 ++++++-
litellm-rust/crates/cache/tests/codec.rs | 53 ++++++
13 files changed, 480 insertions(+), 157 deletions(-)
create mode 100644 litellm-rust/crates/cache/src/codec.rs
create mode 100644 litellm-rust/crates/cache/tests/codec.rs
diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock
index 8c35a0be0b4..ccfaee44f50 100644
--- a/litellm-rust/Cargo.lock
+++ b/litellm-rust/Cargo.lock
@@ -2462,6 +2462,7 @@ dependencies = [
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
+ "tokio",
]
[[package]]
diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs
index 1908ff44a81..974cdbe9760 100644
--- a/litellm-rust/crates/cache-memory/src/cache.rs
+++ b/litellm-rust/crates/cache-memory/src/cache.rs
@@ -4,8 +4,7 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::{
- BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
- Error,
+ BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheKwargs, Error,
};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
@@ -214,8 +213,8 @@ impl InMemoryCache {
}
}
-impl BaseCache for InMemoryCache {
- type Value = CacheEntry;
+impl BaseCache for InMemoryCache {
+ type Value = V;
fn default_ttl(&self) -> Duration {
self.default_ttl
@@ -238,17 +237,15 @@ impl BaseCache for InMemoryCache {
self.flush_cache()
}
- fn disconnect(&self) -> CacheFuture<'_, ()> {
- Box::pin(async { Ok(()) })
+ async fn disconnect(&self) -> Result<(), Error> {
+ Ok(())
}
- fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
- Box::pin(async {
- Ok(CacheConnectionResult {
- status: CacheConnectionStatus::Success,
- message: "In-memory cache connection test successful".into(),
- error: None,
- })
+ async fn test_connection(&self) -> Result {
+ Ok(CacheConnectionResult {
+ status: CacheConnectionStatus::Success,
+ message: "In-memory cache connection test successful".into(),
+ error: None,
})
}
}
diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs
index aaf82641db7..ffac9d8ae64 100644
--- a/litellm-rust/crates/cache-memory/tests/cache.rs
+++ b/litellm-rust/crates/cache-memory/tests/cache.rs
@@ -2,7 +2,10 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
-use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error};
+use litellm_cache::{
+ BaseCache, CacheBackend, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, get_cache,
+ set_cache,
+};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
@@ -156,3 +159,45 @@ async fn connection_test_matches_python_result_contract() {
})
);
}
+
+#[tokio::test]
+async fn generic_consumers_share_typed_values_and_honor_expiration() {
+ let clock = clock();
+ let cache: CacheBackend> = Arc::new(cache(clock.clone(), 4));
+ let reader = Arc::clone(&cache);
+ let kwargs = CacheKwargs {
+ ttl: Some(Duration::from_secs(5)),
+ ..Default::default()
+ };
+ set_cache(cache.as_ref(), "sync", "first".into(), kwargs.clone()).unwrap();
+ assert_eq!(
+ get_cache(reader.as_ref(), "sync", &kwargs).unwrap(),
+ Some("first".into())
+ );
+ cache
+ .batch_cache_write("async", "second".into(), kwargs.clone())
+ .await
+ .unwrap();
+ cache
+ .async_set_cache_pipeline(vec![("batch".into(), "third".into())], kwargs.clone())
+ .await
+ .unwrap();
+ drop(cache);
+ for (key, value) in [("sync", "first"), ("async", "second"), ("batch", "third")] {
+ assert_eq!(
+ reader.async_get_cache(key, &kwargs).await.unwrap(),
+ Some(value.into())
+ );
+ }
+ reader.async_delete_cache("async").await.unwrap();
+ assert_eq!(
+ reader.async_get_cache("async", &kwargs).await.unwrap(),
+ None
+ );
+ clock.store(106, Ordering::SeqCst);
+ assert_eq!(get_cache(reader.as_ref(), "sync", &kwargs).unwrap(), None);
+ assert_eq!(
+ reader.async_get_cache("batch", &kwargs).await.unwrap(),
+ None
+ );
+}
diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml
index 933b0feaae4..a60813b6260 100644
--- a/litellm-rust/crates/cache-redis/Cargo.toml
+++ b/litellm-rust/crates/cache-redis/Cargo.toml
@@ -8,8 +8,8 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
redis = "1.7.0"
-serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"
+serde_json.workspace = true
diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs
index 69dee6c6363..8d92fdc8f75 100644
--- a/litellm-rust/crates/cache-redis/src/cache.rs
+++ b/litellm-rust/crates/cache-redis/src/cache.rs
@@ -2,35 +2,37 @@ use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;
use litellm_cache::{
- BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
- Error,
+ BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error,
};
use redis::Commands;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const KEY_PREFIX: &str = "litellm-cache:";
-pub struct RedisCache {
+pub struct RedisCache {
connection: Arc>,
default_ttl: Duration,
+ codec: S,
}
-impl RedisCache {
- pub fn new(url: &str, default_ttl: Option) -> Result {
+impl RedisCache {
+ pub fn new(url: &str, default_ttl: Option, codec: S) -> Result {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
- Ok(Self::with_connection(connection, default_ttl))
+ Ok(Self::with_connection(connection, default_ttl, codec))
}
}
-impl RedisCache
+impl RedisCache
where
+ S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
- fn with_connection(connection: C, default_ttl: Option) -> Self {
+ pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self {
Self {
connection: Arc::new(Mutex::new(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
+ codec,
}
}
@@ -47,48 +49,39 @@ where
PATTERN
}
- fn encode(value: &CacheEntry) -> Result, Error> {
- serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
- }
-
- fn decode(value: Vec) -> Result {
- serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
- }
-
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
- fn run_blocking(connection: Arc>, operation: F) -> CacheFuture<'static, T>
+ async fn run_blocking(connection: Arc>, operation: F) -> Result
where
T: Send + 'static,
F: FnOnce(&mut C) -> Result + Send + 'static,
{
- Box::pin(async move {
- tokio::task::spawn_blocking(move || {
- let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
- operation(&mut connection)
- })
- .await
- .map_err(|_| Error::Unavailable)?
+ tokio::task::spawn_blocking(move || {
+ let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
+ operation(&mut connection)
})
+ .await
+ .map_err(|_| Error::Unavailable)?
}
}
-impl BaseCache for RedisCache
+impl BaseCache for RedisCache
where
+ S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
- type Value = CacheEntry;
+ type Value = S::Value;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
- let payload = Self::encode(&value)?;
+ let payload = self.codec.encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
self.connection()?
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
@@ -96,11 +89,11 @@ where
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> {
- self.connection()?
+ let bytes = self
+ .connection()?
.get::<_, Option>>(Self::namespaced_key(key))
- .map_err(|_| Error::Unavailable)?
- .map(Self::decode)
- .transpose()
+ .map_err(|_| Error::Unavailable)?;
+ bytes.map(|bytes| self.codec.decode(&bytes)).transpose()
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
@@ -125,86 +118,87 @@ where
.map_err(|_| Error::Unavailable)
}
- fn async_set_cache<'a>(
- &'a self,
- key: &'a str,
+ async fn async_set_cache(
+ &self,
+ key: &str,
value: Self::Value,
kwargs: CacheKwargs,
- ) -> CacheFuture<'a, ()> {
- let payload = Self::encode(&value);
+ ) -> Result<(), Error> {
+ let payload = self.codec.encode(&value)?;
let key = Self::namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
- .set_ex::<_, _, ()>(key, payload?, ttl)
+ .set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
})
+ .await
}
- fn async_get_cache<'a>(
- &'a self,
- key: &'a str,
- _: &'a CacheKwargs,
- ) -> CacheFuture<'a, Option> {
+ async fn async_get_cache(
+ &self,
+ key: &str,
+ _: &CacheKwargs,
+ ) -> Result, Error> {
let key = Self::namespaced_key(key);
- Box::pin(async move {
- Self::run_blocking(Arc::clone(&self.connection), move |connection| {
- connection
- .get::<_, Option>>(key)
- .map_err(|_| Error::Unavailable)
- })
- .await?
- .map(Self::decode)
- .transpose()
+ Self::run_blocking(Arc::clone(&self.connection), move |connection| {
+ connection
+ .get::<_, Option>>(key)
+ .map_err(|_| Error::Unavailable)
})
+ .await?
+ .map(|bytes| self.codec.decode(&bytes))
+ .transpose()
}
- fn async_set_cache_pipeline<'a>(
- &'a self,
+ async fn async_set_cache_pipeline(
+ &self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
- ) -> CacheFuture<'a, ()> {
+ ) -> Result<(), Error> {
let entries = cache_list
.into_iter()
.map(|(key, value)| {
- Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
+ self.codec
+ .encode(&value)
+ .map(|payload| (Self::namespaced_key(&key), payload))
})
- .collect::, _>>();
+ .collect::, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
- for (key, payload) in entries? {
+ for (key, payload) in entries {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)?;
}
Ok(())
})
+ .await
}
- fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
+ async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = Self::namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
+ .await
}
- fn disconnect(&self) -> CacheFuture<'_, ()> {
- Box::pin(async { Ok(()) })
+ async fn disconnect(&self) -> Result<(), Error> {
+ Ok(())
}
- fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
- Box::pin(async move {
- Self::run_blocking(Arc::clone(&self.connection), |connection| {
- redis::cmd("PING")
- .query::(connection)
- .map_err(|_| Error::Unavailable)
- })
- .await?;
- Ok(CacheConnectionResult {
- status: CacheConnectionStatus::Success,
- message: "Redis cache connection test successful".into(),
- error: None,
- })
+ async fn test_connection(&self) -> Result {
+ Self::run_blocking(Arc::clone(&self.connection), |connection| {
+ redis::cmd("PING")
+ .query::(connection)
+ .map_err(|_| Error::Unavailable)
+ })
+ .await?;
+ Ok(CacheConnectionResult {
+ status: CacheConnectionStatus::Success,
+ message: "Redis cache connection test successful".into(),
+ error: None,
})
}
}
@@ -212,7 +206,7 @@ where
#[cfg(test)]
mod tests {
use super::RedisCache;
- use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
+ use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKwargs, JsonCodec};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use std::time::Duration;
@@ -224,33 +218,18 @@ mod tests {
}
}
- #[test]
- fn cache_entries_round_trip_through_json() {
- let entry = entry();
- let encoded = RedisCache::::encode(&entry).unwrap();
- assert_eq!(
- RedisCache::::decode(encoded).unwrap(),
- entry
- );
- }
-
- #[test]
- fn invalid_json_is_rejected() {
- assert!(RedisCache::::decode(b"not json".to_vec()).is_err());
- }
-
#[test]
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
assert_eq!(
- RedisCache::::ttl_seconds(Duration::ZERO),
+ RedisCache::>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
- RedisCache::::ttl_seconds(Duration::from_millis(1500)),
+ RedisCache::>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
- RedisCache::::ttl_seconds(Duration::from_secs(15)),
+ RedisCache::>::ttl_seconds(Duration::from_secs(15)),
15
);
}
@@ -258,7 +237,7 @@ mod tests {
#[test]
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
let value = entry();
- let payload = RedisCache::::encode(&value).unwrap();
+ let payload = JsonCodec::::new().encode(&value).unwrap();
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
@@ -271,7 +250,7 @@ mod tests {
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
- let cache = RedisCache::with_connection(connection, None);
+ let cache = RedisCache::with_connection(connection, None, JsonCodec::::new());
cache
.set_cache("key", value.clone(), CacheKwargs::default())
@@ -296,7 +275,7 @@ mod tests {
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
- let cache = RedisCache::with_connection(connection, None);
+ let cache = RedisCache::with_connection(connection, None, JsonCodec::::new());
cache.flush_cache().unwrap();
}
@@ -305,7 +284,7 @@ mod tests {
async fn test_connection_runs_ping_off_executor() {
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
.assert_all_commands_consumed();
- let cache = RedisCache::with_connection(connection, None);
+ let cache = RedisCache::with_connection(connection, None, JsonCodec::::new());
assert_eq!(
cache.test_connection().await.unwrap().status,
diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs
index 76f73145da8..fe15fcd975b 100644
--- a/litellm-rust/crates/cache-redis/tests/cache.rs
+++ b/litellm-rust/crates/cache-redis/tests/cache.rs
@@ -1,6 +1,156 @@
+use std::time::Duration;
+
+use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, Error, JsonCodec, get_cache, set_cache};
use litellm_cache_redis::RedisCache;
+use redis_test::{MockCmd, MockRedisConnection};
+
+struct TaggedByteCodec(u8);
+
+impl CacheCodec for TaggedByteCodec {
+ type Value = u8;
+
+ fn encode(&self, value: &u8) -> Result, Error> {
+ if *value > 127 {
+ return Err(Error::InvalidEntry);
+ }
+ Ok(vec![self.0, *value])
+ }
+
+ fn decode(&self, bytes: &[u8]) -> Result {
+ match bytes {
+ [tag, value] if *tag == self.0 => Ok(*value),
+ _ => Err(Error::InvalidEntry),
+ }
+ }
+}
#[test]
fn constructor_rejects_invalid_urls() {
- assert!(RedisCache::new("not a redis url", None).is_err());
+ assert!(RedisCache::new("not a redis url", None, JsonCodec::::new()).is_err());
+}
+
+#[test]
+fn generic_helpers_use_the_injected_codec_and_ttl() {
+ let connection = MockRedisConnection::new([
+ MockCmd::new(
+ redis::cmd("SETEX")
+ .arg("litellm-cache:counter")
+ .arg(2)
+ .arg([42u8, 7].as_slice()),
+ Ok("OK"),
+ ),
+ MockCmd::new(
+ redis::cmd("GET").arg("litellm-cache:counter"),
+ Ok(vec![42u8, 7]),
+ ),
+ ])
+ .assert_all_commands_consumed();
+ let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
+ let kwargs = CacheKwargs {
+ ttl: Some(Duration::from_millis(1500)),
+ ..Default::default()
+ };
+ set_cache(&cache, "counter", 7, kwargs.clone()).unwrap();
+ assert_eq!(get_cache(&cache, "counter", &kwargs).unwrap(), Some(7));
+}
+
+#[tokio::test]
+async fn async_operations_preserve_codec_ttl_and_missing_values() {
+ let connection = MockRedisConnection::new([
+ MockCmd::new(
+ redis::cmd("SETEX")
+ .arg("litellm-cache:counter")
+ .arg(9)
+ .arg([42u8, 7].as_slice()),
+ Ok("OK"),
+ ),
+ MockCmd::new(
+ redis::cmd("GET").arg("litellm-cache:counter"),
+ Ok(vec![42u8, 7]),
+ ),
+ MockCmd::new(
+ redis::cmd("SETEX")
+ .arg("litellm-cache:batch")
+ .arg(2)
+ .arg([42u8, 8].as_slice()),
+ Ok("OK"),
+ ),
+ MockCmd::new(redis::cmd("DEL").arg("litellm-cache:counter"), Ok(1u32)),
+ MockCmd::new(
+ redis::cmd("GET").arg("litellm-cache:counter"),
+ Ok(redis::Value::Nil),
+ ),
+ ])
+ .assert_all_commands_consumed();
+ let cache = RedisCache::with_connection(
+ connection,
+ Some(Duration::from_secs(9)),
+ TaggedByteCodec(42),
+ );
+ let kwargs = CacheKwargs::default();
+ cache
+ .batch_cache_write("counter", 7, kwargs.clone())
+ .await
+ .unwrap();
+ assert_eq!(
+ cache.async_get_cache("counter", &kwargs).await.unwrap(),
+ Some(7)
+ );
+ cache
+ .async_set_cache_pipeline(
+ vec![("batch".into(), 8)],
+ CacheKwargs {
+ ttl: Some(Duration::from_millis(1500)),
+ ..Default::default()
+ },
+ )
+ .await
+ .unwrap();
+ cache.async_delete_cache("counter").await.unwrap();
+ assert_eq!(
+ cache.async_get_cache("counter", &kwargs).await.unwrap(),
+ None
+ );
+}
+
+#[tokio::test]
+async fn codec_errors_propagate_without_writing_partial_batches() {
+ let connection = MockRedisConnection::new([
+ MockCmd::new(
+ redis::cmd("GET").arg("litellm-cache:invalid"),
+ Ok(vec![99u8, 7]),
+ ),
+ MockCmd::new(
+ redis::cmd("GET").arg("litellm-cache:invalid"),
+ Ok(vec![99u8, 7]),
+ ),
+ ])
+ .assert_all_commands_consumed();
+ let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
+ let kwargs = CacheKwargs::default();
+ assert_eq!(
+ cache.set_cache("invalid", 255, kwargs.clone()),
+ Err(Error::InvalidEntry)
+ );
+ assert_eq!(
+ cache.async_set_cache("invalid", 255, kwargs.clone()).await,
+ Err(Error::InvalidEntry)
+ );
+ assert_eq!(
+ cache
+ .async_set_cache_pipeline(
+ vec![("valid".into(), 7), ("invalid".into(), 255)],
+ kwargs.clone(),
+ )
+ .await,
+ Err(Error::InvalidEntry)
+ );
+ assert_eq!(
+ cache.get_cache("invalid", &kwargs),
+ Err(Error::InvalidEntry)
+ );
+ assert_eq!(
+ cache.async_get_cache("invalid", &kwargs).await,
+ Err(Error::InvalidEntry)
+ );
}
diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml
index a14c4294aa0..350db4b1adb 100644
--- a/litellm-rust/crates/cache/Cargo.toml
+++ b/litellm-rust/crates/cache/Cargo.toml
@@ -13,3 +13,4 @@ thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
+tokio.workspace = true
diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs
index 2ba8ff92ebd..1891e417cb0 100644
--- a/litellm-rust/crates/cache/src/base_cache.rs
+++ b/litellm-rust/crates/cache/src/base_cache.rs
@@ -1,5 +1,4 @@
use std::future::Future;
-use std::pin::Pin;
use std::time::Duration;
use serde::{Deserialize, Serialize};
@@ -7,8 +6,6 @@ use serde_json::{Map, Value};
use crate::Error;
-pub type CacheFuture<'a, T> = Pin> + Send + 'a>>;
-
#[derive(Clone, Debug, Default, PartialEq)]
pub struct CacheKwargs {
pub ttl: Option,
@@ -45,54 +42,54 @@ pub trait BaseCache: Send + Sync {
fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error>;
- fn async_set_cache<'a>(
- &'a self,
- key: &'a str,
+ fn async_set_cache(
+ &self,
+ key: &str,
value: Self::Value,
kwargs: CacheKwargs,
- ) -> CacheFuture<'a, ()> {
- Box::pin(async move { self.set_cache(key, value, kwargs) })
+ ) -> impl Future> + Send {
+ async move { self.set_cache(key, value, kwargs) }
}
- fn async_get_cache<'a>(
- &'a self,
- key: &'a str,
- kwargs: &'a CacheKwargs,
- ) -> CacheFuture<'a, Option> {
- Box::pin(async move { self.get_cache(key, kwargs) })
+ fn async_get_cache(
+ &self,
+ key: &str,
+ kwargs: &CacheKwargs,
+ ) -> impl Future, Error>> + Send {
+ async move { self.get_cache(key, kwargs) }
}
- fn async_set_cache_pipeline<'a>(
- &'a self,
+ fn async_set_cache_pipeline(
+ &self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
- ) -> CacheFuture<'a, ()> {
- Box::pin(async move {
+ ) -> impl Future> + Send {
+ async move {
for (key, value) in cache_list {
- self.set_cache(&key, value, kwargs.clone())?;
+ self.async_set_cache(&key, value, kwargs.clone()).await?;
}
Ok(())
- })
+ }
}
- fn batch_cache_write<'a>(
- &'a self,
- key: &'a str,
+ fn batch_cache_write(
+ &self,
+ key: &str,
value: Self::Value,
kwargs: CacheKwargs,
- ) -> CacheFuture<'a, ()> {
+ ) -> impl Future> + Send {
self.async_set_cache(key, value, kwargs)
}
fn delete_cache(&self, key: &str) -> Result<(), Error>;
- fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
- Box::pin(async move { self.delete_cache(key) })
+ fn async_delete_cache(&self, key: &str) -> impl Future> + Send {
+ async move { self.delete_cache(key) }
}
fn flush_cache(&self) -> Result<(), Error>;
- fn disconnect(&self) -> CacheFuture<'_, ()>;
+ fn disconnect(&self) -> impl Future> + Send;
- fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult>;
+ fn test_connection(&self) -> impl Future> + Send;
}
diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs
index 1aab6ee8e91..21ebcce29bb 100644
--- a/litellm-rust/crates/cache/src/caching.rs
+++ b/litellm-rust/crates/cache/src/caching.rs
@@ -146,21 +146,21 @@ impl CacheEntry {
}
}
-pub fn get_cache(
- cache: &dyn BaseCache,
+pub fn get_cache(
+ cache: &B,
key: &str,
kwargs: &CacheKwargs,
-) -> Result, Error> {
+) -> Result , Error> {
cache.get_cache(key, kwargs)
}
-pub fn set_cache(
- cache: &dyn BaseCache,
+pub fn set_cache(
+ cache: &B,
key: &str,
- entry: CacheEntry,
+ value: B::Value,
kwargs: CacheKwargs,
) -> Result<(), Error> {
- cache.set_cache(key, entry, kwargs)
+ cache.set_cache(key, value, kwargs)
}
-pub type CacheBackend = Arc>;
+pub type CacheBackend = Arc;
diff --git a/litellm-rust/crates/cache/src/codec.rs b/litellm-rust/crates/cache/src/codec.rs
new file mode 100644
index 00000000000..09bee6032f6
--- /dev/null
+++ b/litellm-rust/crates/cache/src/codec.rs
@@ -0,0 +1,42 @@
+use std::marker::PhantomData;
+
+use serde::{Serialize, de::DeserializeOwned};
+
+use crate::Error;
+
+pub trait CacheCodec: Send + Sync {
+ type Value: Clone + Send + Sync + 'static;
+
+ fn encode(&self, value: &Self::Value) -> Result, Error>;
+
+ fn decode(&self, bytes: &[u8]) -> Result;
+}
+
+pub struct JsonCodec(PhantomData V>);
+
+impl Default for JsonCodec {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl JsonCodec {
+ pub const fn new() -> Self {
+ Self(PhantomData)
+ }
+}
+
+impl CacheCodec for JsonCodec
+where
+ V: Clone + Send + Sync + Serialize + DeserializeOwned + 'static,
+{
+ type Value = V;
+
+ fn encode(&self, value: &Self::Value) -> Result, Error> {
+ serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
+ }
+
+ fn decode(&self, bytes: &[u8]) -> Result {
+ serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry)
+ }
+}
diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs
index d0fe3de15cd..a1d9d1402bb 100644
--- a/litellm-rust/crates/cache/src/lib.rs
+++ b/litellm-rust/crates/cache/src/lib.rs
@@ -1,12 +1,12 @@
mod base_cache;
mod caching;
+mod codec;
mod error;
-pub use base_cache::{
- BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs,
-};
+pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs};
pub use caching::{
Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput,
CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache,
};
+pub use codec::{CacheCodec, JsonCodec};
pub use error::Error;
diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs
index 1192fc9a2b0..5c250c6b3c9 100644
--- a/litellm-rust/crates/cache/tests/caching.rs
+++ b/litellm-rust/crates/cache/tests/caching.rs
@@ -1,12 +1,13 @@
use litellm_cache::{
- BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext,
- CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key,
+ BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField,
+ CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key,
};
use sha2::{Digest, Sha256};
-use std::time::Duration;
+use std::{sync::Mutex, time::Duration};
struct TestCache {
default_ttl: Duration,
+ writes: Mutex>,
}
impl BaseCache for TestCache {
@@ -17,6 +18,22 @@ impl BaseCache for TestCache {
}
fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> {
+ Err(Error::Unavailable)
+ }
+
+ async fn async_set_cache(
+ &self,
+ key: &str,
+ value: Self::Value,
+ kwargs: CacheKwargs,
+ ) -> Result<(), Error> {
+ if key == "unavailable" {
+ return Err(Error::Unavailable);
+ }
+ self.writes
+ .lock()
+ .unwrap()
+ .push((key.into(), value, kwargs));
Ok(())
}
@@ -32,11 +49,11 @@ impl BaseCache for TestCache {
Ok(())
}
- fn disconnect(&self) -> CacheFuture<'_, ()> {
- Box::pin(async { Ok(()) })
+ async fn disconnect(&self) -> Result<(), Error> {
+ Ok(())
}
- fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
+ async fn test_connection(&self) -> Result {
unreachable!()
}
}
@@ -45,6 +62,7 @@ impl BaseCache for TestCache {
fn ttl_uses_default_and_allows_per_call_override() {
let cache = TestCache {
default_ttl: Duration::from_secs(60),
+ writes: Mutex::default(),
};
assert_eq!(
cache.get_ttl(&CacheKwargs::default()),
@@ -59,6 +77,46 @@ fn ttl_uses_default_and_allows_per_call_override() {
);
}
+#[tokio::test]
+async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
+ let cache = TestCache {
+ default_ttl: Duration::from_secs(60),
+ writes: Mutex::default(),
+ };
+ let entry = CacheEntry {
+ timestamp: 123.0,
+ response: serde_json::json!("cached"),
+ };
+ let kwargs = CacheKwargs {
+ ttl: Some(Duration::from_secs(5)),
+ ..Default::default()
+ };
+ cache
+ .batch_cache_write("single", entry.clone(), kwargs.clone())
+ .await
+ .unwrap();
+ assert_eq!(
+ cache
+ .async_set_cache_pipeline(
+ vec![
+ ("first".into(), entry.clone()),
+ ("unavailable".into(), entry.clone()),
+ ("skipped".into(), entry.clone()),
+ ],
+ kwargs.clone(),
+ )
+ .await,
+ Err(Error::Unavailable)
+ );
+ assert_eq!(
+ *cache.writes.lock().unwrap(),
+ vec![
+ ("single".into(), entry.clone(), kwargs.clone()),
+ ("first".into(), entry, kwargs),
+ ]
+ );
+}
+
#[test]
fn keys_match_python_order_groups_files_presets_and_namespaces() {
let mut input = CacheKeyInput {
diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs
new file mode 100644
index 00000000000..dad5398a879
--- /dev/null
+++ b/litellm-rust/crates/cache/tests/codec.rs
@@ -0,0 +1,53 @@
+use std::collections::BTreeMap;
+
+use litellm_cache::{CacheCodec, CacheEntry, Error, JsonCodec};
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+struct RoutingState {
+ deployment: String,
+ cooldown_seconds: u64,
+}
+
+#[test]
+fn json_codec_round_trips_typed_domain_values() {
+ let codec = JsonCodec::::new();
+ let value = RoutingState {
+ deployment: "deployment-a".into(),
+ cooldown_seconds: 30,
+ };
+ let bytes = codec.encode(&value).unwrap();
+ assert_eq!(codec.decode(&bytes).unwrap(), value);
+ assert_eq!(
+ serde_json::from_slice::(&bytes).unwrap(),
+ json!({"deployment": "deployment-a", "cooldown_seconds": 30})
+ );
+}
+
+#[test]
+fn response_entries_preserve_the_existing_json_representation() {
+ let codec = JsonCodec::::new();
+ let entry = CacheEntry {
+ timestamp: 123.0,
+ response: json!({"choices": [{"text": "cached"}]}),
+ };
+ let bytes = codec.encode(&entry).unwrap();
+ assert_eq!(bytes, serde_json::to_vec(&entry).unwrap());
+ assert_eq!(codec.decode(&bytes).unwrap(), entry);
+}
+
+#[test]
+fn json_codec_rejects_malformed_and_wrongly_typed_entries() {
+ let codec = JsonCodec::::new();
+ for bytes in [b"not json".as_slice(), br#"{"deployment":12}"#.as_slice()] {
+ assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry);
+ }
+}
+
+#[test]
+fn json_codec_propagates_encoding_errors() {
+ let codec = JsonCodec::>::new();
+ let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]);
+ assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry);
+}
From 081c93908f1e2750c37beb0aa4e87660aa983d4e Mon Sep 17 00:00:00 2001
From: Yujong Lee
Date: Sun, 20 Sep 2026 20:55:45 -0700
Subject: [PATCH 20/56] feat(cache): add native response cache and Python
binding foundations
---
litellm-rust/Cargo.lock | 100 ++++-
litellm-rust/Cargo.toml | 2 +
litellm-rust/crates/cache-redis/src/cache.rs | 83 +++--
.../crates/cache-redis/tests/cache.rs | 78 ++--
litellm-rust/crates/cache-response/Cargo.toml | 19 +
.../crates/cache-response/src/codec.rs | 100 +++++
litellm-rust/crates/cache-response/src/lib.rs | 7 +
.../crates/cache-response/src/native.rs | 97 +++++
.../crates/cache-response/src/response.rs | 124 +++++++
.../crates/cache-response/tests/response.rs | 290 +++++++++++++++
litellm-rust/crates/cache/src/caching.rs | 1 +
litellm-rust/crates/cache/src/error.rs | 2 +
litellm-rust/crates/python-bridge/Cargo.toml | 3 +
.../crates/python-bridge/src/cache/facade.rs | 210 +++++++++++
.../crates/python-bridge/src/cache/mod.rs | 350 ++++++++++++++++++
litellm-rust/crates/python-bridge/src/lib.rs | 6 +
litellm/rust_bridge/_native.pyi | 49 ++-
tests/test_litellm_rust/test_cache.py | 231 ++++++++++++
18 files changed, 1702 insertions(+), 50 deletions(-)
create mode 100644 litellm-rust/crates/cache-response/Cargo.toml
create mode 100644 litellm-rust/crates/cache-response/src/codec.rs
create mode 100644 litellm-rust/crates/cache-response/src/lib.rs
create mode 100644 litellm-rust/crates/cache-response/src/native.rs
create mode 100644 litellm-rust/crates/cache-response/src/response.rs
create mode 100644 litellm-rust/crates/cache-response/tests/response.rs
create mode 100644 litellm-rust/crates/python-bridge/src/cache/facade.rs
create mode 100644 litellm-rust/crates/python-bridge/src/cache/mod.rs
create mode 100644 tests/test_litellm_rust/test_cache.py
diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock
index ccfaee44f50..aa62e4f3770 100644
--- a/litellm-rust/Cargo.lock
+++ b/litellm-rust/Cargo.lock
@@ -2486,6 +2486,21 @@ dependencies = [
"tokio",
]
+[[package]]
+name = "litellm-cache-response"
+version = "0.1.0"
+dependencies = [
+ "litellm-cache",
+ "litellm-cache-memory",
+ "litellm-cache-redis",
+ "py_literal",
+ "redis",
+ "redis-test",
+ "serde",
+ "serde_json",
+ "tokio",
+]
+
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@@ -2649,6 +2664,8 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-gcp",
+ "litellm-cache",
+ "litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@@ -2660,6 +2677,7 @@ dependencies = [
"pyo3",
"pyo3-async-runtimes",
"rstest",
+ "serde",
"serde_json",
"tokio",
"tokio-tungstenite",
@@ -2948,6 +2966,16 @@ dependencies = [
"minimal-lexical",
]
+[[package]]
+name = "num-bigint"
+version = "0.4.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367"
+dependencies = [
+ "num-integer",
+ "num-traits",
+]
+
[[package]]
name = "num-bigint"
version = "0.5.1"
@@ -2958,6 +2986,15 @@ dependencies = [
"num-traits",
]
+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
[[package]]
name = "num-conv"
version = "0.2.2"
@@ -3131,6 +3168,48 @@ version = "2.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
+[[package]]
+name = "pest"
+version = "2.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d45aeb61b4bf818e12d4205f2466f8c4748f85f4fce0146d1c03d69d753f0ad"
+dependencies = [
+ "memchr",
+ "ucd-trie",
+]
+
+[[package]]
+name = "pest_derive"
+version = "2.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "89cc5a242e25ed4e7704d0be240f2cfbe20a8c27e7e252d94835be93d92dc39f"
+dependencies = [
+ "pest",
+ "pest_generator",
+]
+
+[[package]]
+name = "pest_generator"
+version = "2.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7abf21475cc3820fe4b2ca2dc2142902f67a02189f3b5b3a229f4febc01a43e5"
+dependencies = [
+ "pest",
+ "pest_meta",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "pest_meta"
+version = "2.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adba4db388f687393c18c51348d44a41d870ca9df71a2c98172ea3035dc6936e"
+dependencies = [
+ "pest",
+]
+
[[package]]
name = "pin-project"
version = "1.1.13"
@@ -3305,6 +3384,19 @@ dependencies = [
"prost",
]
+[[package]]
+name = "py_literal"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1"
+dependencies = [
+ "num-bigint 0.4.8",
+ "num-complex",
+ "num-traits",
+ "pest",
+ "pest_derive",
+]
+
[[package]]
name = "pyo3"
version = "0.29.2"
@@ -3604,7 +3696,7 @@ dependencies = [
"arcstr",
"combine",
"itoa",
- "num-bigint",
+ "num-bigint 0.5.1",
"percent-encoding",
"ryu",
"sha1_smol",
@@ -4955,6 +5047,12 @@ dependencies = [
"syn 2.0.119",
]
+[[package]]
+name = "ucd-trie"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
+
[[package]]
name = "unarray"
version = "0.1.4"
diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml
index 2f6f5feb4ad..570d0dd3568 100644
--- a/litellm-rust/Cargo.toml
+++ b/litellm-rust/Cargo.toml
@@ -28,6 +28,8 @@ litellm-types = { path = "crates/types" }
litellm-core-utils = { path = "crates/core-utils" }
litellm-cache = { path = "crates/cache" }
litellm-cache-memory = { path = "crates/cache-memory" }
+litellm-cache-redis = { path = "crates/cache-redis" }
+litellm-cache-response = { path = "crates/cache-response" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }
diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs
index 8d92fdc8f75..0faca6cdaaf 100644
--- a/litellm-rust/crates/cache-redis/src/cache.rs
+++ b/litellm-rust/crates/cache-redis/src/cache.rs
@@ -7,12 +7,12 @@ use litellm_cache::{
use redis::Commands;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
-const KEY_PREFIX: &str = "litellm-cache:";
pub struct RedisCache {
connection: Arc>,
default_ttl: Duration,
codec: S,
+ namespace: Option,
}
impl RedisCache {
@@ -33,6 +33,7 @@ where
connection: Arc::new(Mutex::new(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
+ namespace: None,
}
}
@@ -40,13 +41,44 @@ where
self.connection.lock().map_err(|_| Error::Unavailable)
}
- fn namespaced_key(key: &str) -> String {
- format!("{KEY_PREFIX}{key}")
+ pub fn with_namespace(self, namespace: Option) -> Self {
+ Self {
+ namespace: namespace.filter(|value| !value.is_empty()),
+ ..self
+ }
}
- fn namespaced_pattern() -> &'static str {
- const PATTERN: &str = "litellm-cache:*";
- PATTERN
+ fn namespaced_key(&self, key: &str) -> String {
+ match &self.namespace {
+ Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
+ format!("{namespace}:{key}")
+ }
+ _ => key.into(),
+ }
+ }
+
+ fn namespaced_pattern(&self) -> Result {
+ let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
+ let escaped: String = namespace
+ .chars()
+ .flat_map(|ch| {
+ if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
+ vec!['\\', ch]
+ } else {
+ vec![ch]
+ }
+ })
+ .collect();
+ Ok(format!("{escaped}:*"))
+ }
+
+ fn decode_response(&self, value: redis::Value) -> Result, Error> {
+ match value {
+ redis::Value::Nil => Ok(None),
+ redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
+ redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
+ _ => Err(Error::InvalidEntry),
+ }
}
fn ttl_seconds(ttl: Duration) -> u64 {
@@ -84,28 +116,29 @@ where
let payload = self.codec.encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
self.connection()?
- .set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
+ .set_ex::<_, _, ()>(self.namespaced_key(key), payload, ttl)
.map_err(|_| Error::Unavailable)
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result , Error> {
- let bytes = self
+ let value = self
.connection()?
- .get::<_, Option>>(Self::namespaced_key(key))
+ .get::<_, redis::Value>(self.namespaced_key(key))
.map_err(|_| Error::Unavailable)?;
- bytes.map(|bytes| self.codec.decode(&bytes)).transpose()
+ self.decode_response(value)
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.connection()?
- .del::<_, ()>(Self::namespaced_key(key))
+ .del::<_, ()>(self.namespaced_key(key))
.map_err(|_| Error::Unavailable)
}
fn flush_cache(&self) -> Result<(), Error> {
+ let pattern = self.namespaced_pattern()?;
let mut connection = self.connection()?;
let keys = connection
- .scan_match(Self::namespaced_pattern())
+ .scan_match(pattern)
.map_err(|_| Error::Unavailable)?
.collect::>>()
.map_err(|_| Error::Unavailable)?;
@@ -125,7 +158,7 @@ where
kwargs: CacheKwargs,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
- let key = Self::namespaced_key(key);
+ let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
@@ -140,15 +173,14 @@ where
key: &str,
_: &CacheKwargs,
) -> Result, Error> {
- let key = Self::namespaced_key(key);
- Self::run_blocking(Arc::clone(&self.connection), move |connection| {
+ let key = self.namespaced_key(key);
+ let value = Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
- .get::<_, Option>>(key)
+ .get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
})
- .await?
- .map(|bytes| self.codec.decode(&bytes))
- .transpose()
+ .await?;
+ self.decode_response(value)
}
async fn async_set_cache_pipeline(
@@ -161,7 +193,7 @@ where
.map(|(key, value)| {
self.codec
.encode(&value)
- .map(|payload| (Self::namespaced_key(&key), payload))
+ .map(|payload| (self.namespaced_key(&key), payload))
})
.collect::, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
@@ -177,7 +209,7 @@ where
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
- let key = Self::namespaced_key(key);
+ let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
@@ -250,7 +282,8 @@ mod tests {
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
- let cache = RedisCache::with_connection(connection, None, JsonCodec::::new());
+ let cache = RedisCache::with_connection(connection, None, JsonCodec::::new())
+ .with_namespace(Some("litellm-cache".into()));
cache
.set_cache("key", value.clone(), CacheKwargs::default())
@@ -275,7 +308,8 @@ mod tests {
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
- let cache = RedisCache::with_connection(connection, None, JsonCodec::::new());
+ let cache = RedisCache::with_connection(connection, None, JsonCodec::::new())
+ .with_namespace(Some("litellm-cache".into()));
cache.flush_cache().unwrap();
}
@@ -284,7 +318,8 @@ mod tests {
async fn test_connection_runs_ping_off_executor() {
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
.assert_all_commands_consumed();
- let cache = RedisCache::with_connection(connection, None, JsonCodec::::new());
+ let cache = RedisCache::with_connection(connection, None, JsonCodec::::new())
+ .with_namespace(Some("litellm-cache".into()));
assert_eq!(
cache.test_connection().await.unwrap().status,
diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs
index fe15fcd975b..d5bba19a8bd 100644
--- a/litellm-rust/crates/cache-redis/tests/cache.rs
+++ b/litellm-rust/crates/cache-redis/tests/cache.rs
@@ -34,15 +34,12 @@ fn generic_helpers_use_the_injected_codec_and_ttl() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
- .arg("litellm-cache:counter")
+ .arg("counter")
.arg(2)
.arg([42u8, 7].as_slice()),
Ok("OK"),
),
- MockCmd::new(
- redis::cmd("GET").arg("litellm-cache:counter"),
- Ok(vec![42u8, 7]),
- ),
+ MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
@@ -59,27 +56,21 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
- .arg("litellm-cache:counter")
+ .arg("counter")
.arg(9)
.arg([42u8, 7].as_slice()),
Ok("OK"),
),
- MockCmd::new(
- redis::cmd("GET").arg("litellm-cache:counter"),
- Ok(vec![42u8, 7]),
- ),
+ MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])),
MockCmd::new(
redis::cmd("SETEX")
- .arg("litellm-cache:batch")
+ .arg("batch")
.arg(2)
.arg([42u8, 8].as_slice()),
Ok("OK"),
),
- MockCmd::new(redis::cmd("DEL").arg("litellm-cache:counter"), Ok(1u32)),
- MockCmd::new(
- redis::cmd("GET").arg("litellm-cache:counter"),
- Ok(redis::Value::Nil),
- ),
+ MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)),
+ MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(
@@ -116,14 +107,8 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() {
#[tokio::test]
async fn codec_errors_propagate_without_writing_partial_batches() {
let connection = MockRedisConnection::new([
- MockCmd::new(
- redis::cmd("GET").arg("litellm-cache:invalid"),
- Ok(vec![99u8, 7]),
- ),
- MockCmd::new(
- redis::cmd("GET").arg("litellm-cache:invalid"),
- Ok(vec![99u8, 7]),
- ),
+ MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
+ MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
@@ -154,3 +139,48 @@ async fn codec_errors_propagate_without_writing_partial_batches() {
Err(Error::InvalidEntry)
);
}
+
+#[test]
+fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() {
+ let connection = MockRedisConnection::new([
+ MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
+ MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)),
+ ])
+ .assert_all_commands_consumed();
+ let cache = RedisCache::with_connection(connection, None, JsonCodec::::new())
+ .with_namespace(Some("team".into()));
+ assert_eq!(
+ cache.get_cache("key", &CacheKwargs::default()).unwrap(),
+ None
+ );
+ assert_eq!(
+ cache
+ .get_cache("team:key", &CacheKwargs::default())
+ .unwrap(),
+ None
+ );
+}
+
+#[test]
+fn flush_requires_a_namespace_and_escapes_glob_metacharacters() {
+ let unscoped = RedisCache::with_connection(
+ MockRedisConnection::new([]).assert_all_commands_consumed(),
+ None,
+ JsonCodec::::new(),
+ );
+ assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush));
+ let connection = MockRedisConnection::new([
+ MockCmd::new(
+ redis::cmd("SCAN")
+ .cursor_arg(0)
+ .arg("MATCH")
+ .arg("team\\*:*"),
+ Ok(redis_test::redis_value!(["0", ["team*:key"]])),
+ ),
+ MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)),
+ ])
+ .assert_all_commands_consumed();
+ let scoped = RedisCache::with_connection(connection, None, JsonCodec::::new())
+ .with_namespace(Some("team*".into()));
+ scoped.flush_cache().unwrap();
+}
diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml
new file mode 100644
index 00000000000..a0c4a1f74ef
--- /dev/null
+++ b/litellm-rust/crates/cache-response/Cargo.toml
@@ -0,0 +1,19 @@
+[package]
+name = "litellm-cache-response"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+litellm-cache.workspace = true
+litellm-cache-memory.workspace = true
+litellm-cache-redis.workspace = true
+py_literal = "0.4.0"
+redis = "1.7.0"
+serde.workspace = true
+serde_json.workspace = true
+
+[dev-dependencies]
+redis-test = "1.0.4"
+tokio.workspace = true
diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs
new file mode 100644
index 00000000000..137cf61267a
--- /dev/null
+++ b/litellm-rust/crates/cache-response/src/codec.rs
@@ -0,0 +1,100 @@
+use litellm_cache::{CacheCodec, CacheEntry, Error};
+use serde_json::Value;
+
+pub struct ResponseCacheCodec;
+
+impl CacheCodec for ResponseCacheCodec {
+ type Value = CacheEntry;
+
+ fn encode(&self, value: &CacheEntry) -> Result, Error> {
+ if !value.timestamp.is_finite() {
+ return Err(Error::InvalidEntry);
+ }
+ serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
+ }
+
+ fn decode(&self, bytes: &[u8]) -> Result {
+ let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?;
+ let entry: CacheEntry =
+ serde_json::from_value(decode_value(text)?).map_err(|_| Error::InvalidEntry)?;
+ if !entry.timestamp.is_finite() {
+ return Err(Error::InvalidEntry);
+ }
+ Ok(entry)
+ }
+}
+
+pub(crate) fn decode_value(text: &str) -> Result {
+ if let Ok(value) = serde_json::from_str(text) {
+ return Ok(value);
+ }
+ check_literal_depth(text)?;
+ let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?;
+ literal_value(literal, 0)
+}
+
+fn literal_value(value: py_literal::Value, depth: usize) -> Result {
+ use py_literal::Value as Literal;
+ if depth > 128 {
+ return Err(Error::InvalidEntry);
+ }
+ match value {
+ Literal::String(text) => Ok(Value::String(text)),
+ Literal::Boolean(value) => Ok(Value::Bool(value)),
+ Literal::None => Ok(Value::Null),
+ Literal::Integer(value) => {
+ serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry)
+ }
+ Literal::Float(value) => serde_json::Number::from_f64(value)
+ .map(Value::Number)
+ .ok_or(Error::InvalidEntry),
+ Literal::List(values) | Literal::Tuple(values) => values
+ .into_iter()
+ .map(|value| literal_value(value, depth + 1))
+ .collect::, _>>()
+ .map(Value::Array),
+ Literal::Dict(entries) => entries
+ .into_iter()
+ .map(|(key, value)| {
+ let Literal::String(key) = key else {
+ return Err(Error::InvalidEntry);
+ };
+ Ok((key, literal_value(value, depth + 1)?))
+ })
+ .collect::, _>>()
+ .map(Value::Object),
+ _ => Err(Error::InvalidEntry),
+ }
+}
+
+fn check_literal_depth(text: &str) -> Result<(), Error> {
+ let mut quote = None;
+ let mut escaped = false;
+ let mut depth = 0usize;
+ for ch in text.chars() {
+ if escaped {
+ escaped = false;
+ continue;
+ }
+ if let Some(delimiter) = quote {
+ if ch == '\\' {
+ escaped = true;
+ } else if ch == delimiter {
+ quote = None;
+ }
+ continue;
+ }
+ match ch {
+ '\'' | '"' => quote = Some(ch),
+ '[' | '{' | '(' => {
+ depth += 1;
+ if depth > 128 {
+ return Err(Error::InvalidEntry);
+ }
+ }
+ ']' | '}' | ')' => depth = depth.saturating_sub(1),
+ _ => {}
+ }
+ }
+ Ok(())
+}
diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs
new file mode 100644
index 00000000000..454a82e76de
--- /dev/null
+++ b/litellm-rust/crates/cache-response/src/lib.rs
@@ -0,0 +1,7 @@
+mod codec;
+mod native;
+mod response;
+
+pub use codec::ResponseCacheCodec;
+pub use native::NativeResponseCache;
+pub use response::{ResponseCache, ResponseCacheRequest};
diff --git a/litellm-rust/crates/cache-response/src/native.rs b/litellm-rust/crates/cache-response/src/native.rs
new file mode 100644
index 00000000000..c25c61cee82
--- /dev/null
+++ b/litellm-rust/crates/cache-response/src/native.rs
@@ -0,0 +1,97 @@
+use std::{sync::Arc, time::Duration};
+
+use litellm_cache::{CacheEntry, Error};
+use litellm_cache_memory::InMemoryCache;
+use litellm_cache_redis::RedisCache;
+use serde_json::Value;
+
+use crate::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest};
+
+pub enum NativeResponseCache
+where
+ C: redis::ConnectionLike + Send + 'static,
+{
+ Memory(Arc>>),
+ Redis(Arc>>),
+}
+
+impl Clone for NativeResponseCache {
+ fn clone(&self) -> Self {
+ match self {
+ Self::Memory(cache) => Self::Memory(Arc::clone(cache)),
+ Self::Redis(cache) => Self::Redis(Arc::clone(cache)),
+ }
+ }
+}
+
+impl NativeResponseCache {
+ pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
+ Self::Memory(Arc::new(ResponseCache::new(Arc::new(
+ InMemoryCache::response_cache(capacity, ttl, max_entry_bytes),
+ ))))
+ }
+
+ pub fn redis(
+ url: &str,
+ ttl: Option,
+ namespace: Option,
+ ) -> Result {
+ let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace);
+ Ok(Self::Redis(Arc::new(ResponseCache::new(Arc::new(backend)))))
+ }
+}
+
+impl NativeResponseCache {
+ pub fn kind(&self) -> &'static str {
+ match self {
+ Self::Memory(_) => "memory",
+ Self::Redis(_) => "redis",
+ }
+ }
+
+ pub fn lookup(
+ &self,
+ request: &ResponseCacheRequest,
+ now: Duration,
+ ) -> Result, Error> {
+ match self {
+ Self::Memory(cache) => cache.lookup(request, now),
+ Self::Redis(cache) => cache.lookup(request, now),
+ }
+ }
+
+ pub fn store(
+ &self,
+ request: &ResponseCacheRequest,
+ response: Value,
+ now: Duration,
+ ) -> Result<(), Error> {
+ match self {
+ Self::Memory(cache) => cache.store(request, response, now),
+ Self::Redis(cache) => cache.store(request, response, now),
+ }
+ }
+
+ pub async fn async_lookup(
+ &self,
+ request: &ResponseCacheRequest,
+ now: Duration,
+ ) -> Result , Error> {
+ match self {
+ Self::Memory(cache) => cache.async_lookup(request, now).await,
+ Self::Redis(cache) => cache.async_lookup(request, now).await,
+ }
+ }
+
+ pub async fn async_store(
+ &self,
+ request: &ResponseCacheRequest,
+ response: Value,
+ now: Duration,
+ ) -> Result<(), Error> {
+ match self {
+ Self::Memory(cache) => cache.async_store(request, response, now).await,
+ Self::Redis(cache) => cache.async_store(request, response, now).await,
+ }
+ }
+}
diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs
new file mode 100644
index 00000000000..3e9807b6d1b
--- /dev/null
+++ b/litellm-rust/crates/cache-response/src/response.rs
@@ -0,0 +1,124 @@
+use std::{sync::Arc, time::Duration};
+
+use litellm_cache::{
+ BaseCache, CacheControls, CacheEntry, CacheKeyInput, CacheKwargs, Error, cache_key,
+};
+use serde_json::Value;
+
+#[derive(Clone)]
+pub struct ResponseCacheRequest {
+ pub key: CacheKeyInput,
+ pub controls: CacheControls,
+ pub kwargs: CacheKwargs,
+ pub max_age: Option,
+}
+
+impl ResponseCacheRequest {
+ pub fn new(key: CacheKeyInput) -> Self {
+ Self {
+ key,
+ controls: CacheControls {
+ configured: true,
+ supported_call_type: true,
+ native_backend: true,
+ default_on: true,
+ ..Default::default()
+ },
+ kwargs: CacheKwargs::default(),
+ max_age: None,
+ }
+ }
+}
+
+pub struct ResponseCache> {
+ backend: Arc,
+}
+
+impl> ResponseCache {
+ pub fn new(backend: Arc) -> Self {
+ Self { backend }
+ }
+
+ pub fn lookup(
+ &self,
+ request: &ResponseCacheRequest,
+ now: Duration,
+ ) -> Result, Error> {
+ if !request.controls.reads() {
+ return Ok(None);
+ }
+ let entry = self
+ .backend
+ .get_cache(&cache_key(&request.key), &request.kwargs)?;
+ Self::fresh_response(entry, now, request.max_age)
+ }
+
+ pub async fn async_lookup(
+ &self,
+ request: &ResponseCacheRequest,
+ now: Duration,
+ ) -> Result , Error> {
+ if !request.controls.reads() {
+ return Ok(None);
+ }
+ let entry = self
+ .backend
+ .async_get_cache(&cache_key(&request.key), &request.kwargs)
+ .await?;
+ Self::fresh_response(entry, now, request.max_age)
+ }
+
+ pub fn store(
+ &self,
+ request: &ResponseCacheRequest,
+ response: Value,
+ now: Duration,
+ ) -> Result<(), Error> {
+ if !request.controls.writes() {
+ return Ok(());
+ }
+ self.backend.set_cache(
+ &cache_key(&request.key),
+ CacheEntry {
+ timestamp: now.as_secs_f64(),
+ response,
+ },
+ request.kwargs.clone(),
+ )
+ }
+
+ pub async fn async_store(
+ &self,
+ request: &ResponseCacheRequest,
+ response: Value,
+ now: Duration,
+ ) -> Result<(), Error> {
+ if !request.controls.writes() {
+ return Ok(());
+ }
+ self.backend
+ .async_set_cache(
+ &cache_key(&request.key),
+ CacheEntry {
+ timestamp: now.as_secs_f64(),
+ response,
+ },
+ request.kwargs.clone(),
+ )
+ .await
+ }
+
+ fn fresh_response(
+ entry: Option,
+ now: Duration,
+ max_age: Option,
+ ) -> Result, Error> {
+ entry
+ .filter(|entry| entry.fresh(now, max_age))
+ .map(|entry| match entry.response {
+ Value::String(text) => crate::codec::decode_value(&text),
+ value => Ok(value),
+ })
+ .transpose()
+ }
+}
diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs
new file mode 100644
index 00000000000..a4beda9fa4c
--- /dev/null
+++ b/litellm-rust/crates/cache-response/tests/response.rs
@@ -0,0 +1,290 @@
+use std::{
+ sync::{
+ Arc,
+ atomic::{AtomicU64, Ordering},
+ },
+ time::Duration,
+};
+
+use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKeyField, CacheKeyInput, Error};
+use litellm_cache_memory::InMemoryCache;
+use litellm_cache_redis::RedisCache;
+use litellm_cache_response::{
+ NativeResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
+};
+use redis_test::{MockCmd, MockRedisConnection};
+use serde_json::json;
+
+fn request() -> ResponseCacheRequest {
+ ResponseCacheRequest::new(CacheKeyInput {
+ preset: Some("tenant:key".into()),
+ ..Default::default()
+ })
+}
+
+#[tokio::test]
+async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
+ let clock = Arc::new(AtomicU64::new(100));
+ let backend = Arc::new(InMemoryCache::with_clock(
+ Some(8),
+ Some(Duration::from_secs(600)),
+ {
+ let clock = clock.clone();
+ move || Duration::from_secs(clock.load(Ordering::SeqCst))
+ },
+ ));
+ let cache = ResponseCache::new(backend.clone());
+ let mut request = request();
+ request.kwargs.ttl = Some(Duration::from_secs(10));
+ request.max_age = Some(Duration::from_secs(5));
+ cache
+ .store(
+ &request,
+ json!({"choices": [1], "usage": {"total_tokens": 7}}),
+ Duration::from_secs(100),
+ )
+ .unwrap();
+ assert_eq!(
+ backend.expires_at("tenant:key").unwrap(),
+ Some(Duration::from_secs(110))
+ );
+ assert!(
+ cache
+ .async_lookup(&request, Duration::from_secs(105))
+ .await
+ .unwrap()
+ .is_some()
+ );
+ assert_eq!(
+ cache.lookup(&request, Duration::from_secs(106)).unwrap(),
+ None
+ );
+ request.max_age = None;
+ assert_eq!(
+ cache
+ .lookup(&request, Duration::from_secs(106))
+ .unwrap()
+ .unwrap()["usage"]["total_tokens"],
+ 7
+ );
+ clock.store(111, Ordering::SeqCst);
+ assert_eq!(
+ cache
+ .async_lookup(&request, Duration::from_secs(111))
+ .await
+ .unwrap(),
+ None
+ );
+ cache
+ .async_store(&request, json!({"choices": [2]}), Duration::from_secs(111))
+ .await
+ .unwrap();
+ assert_eq!(
+ cache.lookup(&request, Duration::from_secs(111)).unwrap(),
+ Some(json!({"choices": [2]}))
+ );
+}
+
+#[tokio::test]
+async fn directives_skip_io_and_keep_reads_and_writes_independent() {
+ let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024);
+ let mut request = request();
+ let now = Duration::from_secs(100);
+ request.controls.no_store = true;
+ cache
+ .async_store(&request, json!({"v": 1}), now)
+ .await
+ .unwrap();
+ assert_eq!(cache.lookup(&request, now).unwrap(), None);
+ request.controls.no_store = false;
+ request.controls.no_cache = true;
+ cache.store(&request, json!({"v": 2}), now).unwrap();
+ assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None);
+ request.controls.no_cache = false;
+ assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
+ request.controls.default_on = false;
+ cache.store(&request, json!({"v": 3}), now).unwrap();
+ assert_eq!(cache.lookup(&request, now).unwrap(), None);
+ request.controls.use_cache = true;
+ assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
+ request.controls.supported_call_type = false;
+ assert_eq!(cache.lookup(&request, now).unwrap(), None);
+}
+
+#[tokio::test]
+async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_json() {
+ let connection = MockRedisConnection::new([
+ MockCmd::new(
+ redis::cmd("GET").arg("tenant:key"),
+ Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()),
+ ),
+ MockCmd::new(
+ redis::cmd("GET").arg("tenant:key"),
+ Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()),
+ ),
+ MockCmd::new(
+ redis::cmd("SETEX")
+ .arg("tenant:key")
+ .arg(600)
+ .arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()),
+ Ok("OK"),
+ ),
+ ])
+ .assert_all_commands_consumed();
+ let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec)
+ .with_namespace(Some("tenant".into()));
+ let cache = NativeResponseCache::Redis(Arc::new(ResponseCache::new(Arc::new(backend))));
+ let request = request();
+ let expected = json!({"ok": true, "text": "cached"});
+ assert_eq!(
+ cache.lookup(&request, Duration::from_secs(101)).unwrap(),
+ Some(expected.clone())
+ );
+ assert_eq!(
+ cache
+ .async_lookup(&request, Duration::from_secs(101))
+ .await
+ .unwrap(),
+ Some(expected.clone())
+ );
+ cache
+ .async_store(&request, expected, Duration::from_secs(100))
+ .await
+ .unwrap();
+}
+
+#[tokio::test]
+async fn captured_enum_keeps_the_selected_backend_for_background_writes() {
+ let original = NativeResponseCache::memory(8, Duration::from_secs(600), 1024);
+ let captured = original.clone();
+ let replacement = NativeResponseCache::memory(8, Duration::from_secs(600), 1024);
+ let request = request();
+ let writer = tokio::spawn({
+ let request = request.clone();
+ async move {
+ captured
+ .async_store(
+ &request,
+ json!({"selected": "original"}),
+ Duration::from_secs(100),
+ )
+ .await
+ }
+ });
+ writer.await.unwrap().unwrap();
+ assert_eq!(
+ original.lookup(&request, Duration::from_secs(100)).unwrap(),
+ Some(json!({"selected":"original"}))
+ );
+ assert_eq!(
+ replacement
+ .lookup(&request, Duration::from_secs(100))
+ .unwrap(),
+ None
+ );
+}
+
+#[test]
+fn generated_keys_preserve_namespace_and_explicit_keys() {
+ let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024);
+ let key = CacheKeyInput {
+ fields: vec![CacheKeyField {
+ name: "model".into(),
+ value: Some("a".into()),
+ api_parameter: true,
+ internal_parameter: false,
+ }],
+ namespace: Some("tenant".into()),
+ ..Default::default()
+ };
+ let generated = ResponseCacheRequest::new(key.clone());
+ let explicit = ResponseCacheRequest::new(CacheKeyInput {
+ preset: Some(litellm_cache::cache_key(&key)),
+ ..Default::default()
+ });
+ cache
+ .store(&generated, json!({"value": 7}), Duration::from_secs(100))
+ .unwrap();
+ assert_eq!(
+ cache.lookup(&explicit, Duration::from_secs(100)).unwrap(),
+ Some(json!({"value":7}))
+ );
+}
+
+#[test]
+fn response_codec_accepts_python_literals_without_executing_code() {
+ let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#;
+ let entry = ResponseCacheCodec.decode(bytes).unwrap();
+ assert_eq!(
+ entry.response,
+ json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]})
+ );
+ for bytes in [
+ b"__import__('os').system('false')".as_slice(),
+ b"{'timestamp': 'invalid', 'response': {}}",
+ b"{'timestamp': 1e9999, 'response': {}}",
+ ] {
+ assert_eq!(
+ ResponseCacheCodec.decode(bytes).unwrap_err(),
+ Error::InvalidEntry
+ );
+ }
+ let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000));
+ assert_eq!(
+ ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(),
+ Error::InvalidEntry
+ );
+ assert_eq!(
+ ResponseCacheCodec
+ .encode(&CacheEntry {
+ timestamp: f64::NAN,
+ response: json!({})
+ })
+ .unwrap_err(),
+ Error::InvalidEntry
+ );
+}
+
+#[tokio::test]
+async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redis() {
+ let connection = MockRedisConnection::new([MockCmd::new(
+ redis::cmd("GET").arg("tenant:key"),
+ Ok(b"invalid".to_vec()),
+ )])
+ .assert_all_commands_consumed();
+ let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec);
+ let cache = ResponseCache::new(Arc::new(backend));
+ let mut request = request();
+ request.controls.no_cache = true;
+ assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
+ request.controls.no_cache = false;
+ assert_eq!(
+ cache
+ .async_lookup(&request, Duration::ZERO)
+ .await
+ .unwrap_err(),
+ Error::InvalidEntry
+ );
+}
+
+#[test]
+fn malformed_memory_entries_are_rejected_by_the_response_consumer() {
+ let backend = Arc::new(InMemoryCache::default());
+ BaseCache::set_cache(
+ backend.as_ref(),
+ "tenant:key",
+ CacheEntry {
+ timestamp: 100.0,
+ response: json!("not a serialized response"),
+ },
+ Default::default(),
+ )
+ .unwrap();
+ let cache = ResponseCache::new(backend);
+ assert_eq!(
+ cache
+ .lookup(&request(), Duration::from_secs(100))
+ .unwrap_err(),
+ Error::InvalidEntry
+ );
+}
diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs
index 21ebcce29bb..1d1df966ba2 100644
--- a/litellm-rust/crates/cache/src/caching.rs
+++ b/litellm-rust/crates/cache/src/caching.rs
@@ -27,6 +27,7 @@ pub struct CacheKeyField {
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+#[serde(default)]
pub struct CacheKeyInput {
pub fields: Vec,
pub preset: Option,
diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs
index d447c80f62d..ff3ff6572d4 100644
--- a/litellm-rust/crates/cache/src/error.rs
+++ b/litellm-rust/crates/cache/src/error.rs
@@ -4,4 +4,6 @@ pub enum Error {
Unavailable,
#[error("invalid cache entry")]
InvalidEntry,
+ #[error("flushing Redis requires an explicit namespace")]
+ UnscopedFlush,
}
diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml
index a76b069935f..308dfd2dd7a 100644
--- a/litellm-rust/crates/python-bridge/Cargo.toml
+++ b/litellm-rust/crates/python-bridge/Cargo.toml
@@ -20,6 +20,9 @@ tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
bytes.workspace = true
+litellm-cache.workspace = true
+litellm-cache-response.workspace = true
+serde.workspace = true
litellm-auth.workspace = true
litellm-callbacks-legacy-python.workspace = true
litellm-core.workspace = true
diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs
new file mode 100644
index 00000000000..32360e72469
--- /dev/null
+++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs
@@ -0,0 +1,210 @@
+use litellm_cache_response::NativeResponseCache;
+use litellm_host_python::from_py;
+use pyo3::{
+ PyTraverseError, PyVisit,
+ exceptions::PyTypeError,
+ prelude::*,
+ types::{PyDict, PyTuple, PyType},
+};
+use serde_json::Value;
+
+use super::NativeCacheHandle;
+
+struct ClassGuard {
+ class: Py,
+ attributes: Vec<(String, Py)>,
+}
+
+struct ObjectGuard {
+ reference: Py,
+ classes: Vec,
+ config_names: &'static [&'static str],
+ config: Vec,
+}
+
+pub(super) struct FacadeGuard {
+ outer: ObjectGuard,
+ backend: ObjectGuard,
+}
+
+impl ObjectGuard {
+ fn capture(
+ py: Python<'_>,
+ object: &Bound<'_, PyAny>,
+ config_names: &'static [&'static str],
+ ) -> PyResult {
+ let classes = object
+ .get_type()
+ .getattr("__mro__")?
+ .cast_into::()?
+ .iter()
+ .map(|class| {
+ let class = class.cast_into::()?;
+ let attributes = class
+ .getattr("__dict__")?
+ .call_method0("items")?
+ .try_iter()?
+ .map(|item| item?.extract::<(String, Py)>())
+ .collect::>>()?;
+ Ok(ClassGuard {
+ class: class.unbind(),
+ attributes,
+ })
+ })
+ .collect::>>()?;
+ let guard = Self {
+ reference: py
+ .import("weakref")?
+ .getattr("ref")?
+ .call1((object,))?
+ .unbind(),
+ classes,
+ config_names,
+ config: Self::config(object, config_names)?,
+ };
+ if !guard.matches(py, object)? {
+ return Err(PyTypeError::new_err(
+ "native facade registration requires unmodified built-in methods",
+ ));
+ }
+ Ok(guard)
+ }
+
+ fn config(object: &Bound<'_, PyAny>, names: &[&str]) -> PyResult> {
+ names
+ .iter()
+ .map(|name| match object.getattr(*name) {
+ Ok(value) => from_py(&value),
+ Err(error)
+ if error.is_instance_of::(object.py()) =>
+ {
+ Ok(Value::Null)
+ }
+ Err(error) => Err(error),
+ })
+ .collect()
+ }
+
+ fn matches(&self, py: Python<'_>, object: &Bound<'_, PyAny>) -> PyResult {
+ if !self.reference.bind(py).call0()?.is(object) {
+ return Ok(false);
+ }
+ let mro = object
+ .get_type()
+ .getattr("__mro__")?
+ .cast_into::()?;
+ if mro.len() != self.classes.len() {
+ return Ok(false);
+ }
+ let instance = object.getattr("__dict__")?.cast_into::()?;
+ for (class, expected) in mro.iter().zip(&self.classes) {
+ if !class.is(expected.class.bind(py)) {
+ return Ok(false);
+ }
+ let attributes = class.getattr("__dict__")?;
+ if attributes.len()? != expected.attributes.len() {
+ return Ok(false);
+ }
+ for (name, value) in &expected.attributes {
+ if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) {
+ return Ok(false);
+ }
+ }
+ }
+ Ok(Self::config(object, self.config_names)? == self.config)
+ }
+
+ fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+ visit.call(&self.reference)?;
+ for class in &self.classes {
+ visit.call(&class.class)?;
+ for (_, value) in &class.attributes {
+ visit.call(value)?;
+ }
+ }
+ Ok(())
+ }
+}
+
+impl FacadeGuard {
+ pub(super) fn capture(py: Python<'_>, facade: &Bound<'_, PyAny>, kind: &str) -> PyResult {
+ let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?;
+ if !facade.get_type().is(&cache_type) {
+ return Err(PyTypeError::new_err(
+ "only exact built-in Cache facades can be registered",
+ ));
+ }
+ let (module, name, cache_kind) = match kind {
+ "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
+ "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"),
+ _ => unreachable!(),
+ };
+ let backend = facade.getattr("cache")?;
+ if facade.getattr("type")?.extract::()? != cache_kind
+ || !backend.get_type().is(&py.import(module)?.getattr(name)?)
+ {
+ return Err(PyTypeError::new_err(
+ "facade and native backend types must match",
+ ));
+ }
+ Ok(Self {
+ outer: ObjectGuard::capture(
+ py,
+ facade,
+ &[
+ "type",
+ "mode",
+ "ttl",
+ "namespace",
+ "supported_call_types",
+ "redis_flush_size",
+ ],
+ )?,
+ backend: ObjectGuard::capture(
+ py,
+ &backend,
+ &[
+ "namespace",
+ "default_ttl",
+ "max_size_in_memory",
+ "max_size_per_item",
+ ],
+ )?,
+ })
+ }
+
+ fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult {
+ Ok(self.outer.matches(py, facade)?
+ && self.backend.matches(py, &facade.getattr("cache")?)?)
+ }
+
+ pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+ self.outer.traverse(&visit)?;
+ self.backend.traverse(&visit)
+ }
+}
+
+pub(super) fn resolve(
+ py: Python<'_>,
+ facade: &Bound<'_, PyAny>,
+) -> PyResult> {
+ let Ok(dict) = facade
+ .getattr("__dict__")
+ .and_then(|dict| dict.cast_into::().map_err(Into::into))
+ else {
+ return Ok(None);
+ };
+ let Some(handle) = dict.get_item("_native_cache_handle")? else {
+ return Ok(None);
+ };
+ let Ok(handle) = handle.extract::>() else {
+ return Ok(None);
+ };
+ let Some(guard) = &handle.guard else {
+ return Ok(None);
+ };
+ if !guard.matches(py, facade).unwrap_or(false) {
+ return Ok(None);
+ }
+ handle.service().map(Some)
+}
diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs
new file mode 100644
index 00000000000..f00cceeb86e
--- /dev/null
+++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs
@@ -0,0 +1,350 @@
+mod facade;
+
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use litellm_cache::{CacheControls, CacheKeyInput, Error};
+use litellm_cache_response::{NativeResponseCache, ResponseCacheRequest};
+use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py};
+use pyo3::{
+ PyTraverseError, PyVisit,
+ exceptions::{PyRuntimeError, PyTypeError, PyValueError},
+ prelude::*,
+ types::PyDict,
+};
+use serde::Deserialize;
+use serde_json::Value;
+
+use facade::FacadeGuard;
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+struct RequestInput {
+ key: CacheKeyInput,
+ controls: Option,
+ ttl_seconds: Option,
+ max_age_seconds: Option,
+}
+
+fn request(value: &Bound<'_, PyAny>) -> PyResult {
+ let input: RequestInput = from_py(value)?;
+ let mut request = ResponseCacheRequest::new(input.key);
+ if let Some(controls) = input.controls {
+ request.controls = controls;
+ }
+ request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?;
+ request.max_age = input.max_age_seconds.map(duration).transpose()?;
+ Ok(request)
+}
+
+fn duration(seconds: f64) -> PyResult {
+ Duration::try_from_secs_f64(seconds)
+ .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative"))
+}
+
+fn now() -> Duration {
+ SystemTime::now()
+ .duration_since(UNIX_EPOCH)
+ .unwrap_or_default()
+}
+
+fn cache_error(error: Error) -> PyErr {
+ match error {
+ Error::InvalidEntry => PyValueError::new_err(error.to_string()),
+ _ => PyRuntimeError::new_err(error.to_string()),
+ }
+}
+
+#[pyclass(frozen)]
+pub(crate) struct NativeCacheHandle {
+ service: NativeResponseCache,
+ guard: Option,
+ pid: u32,
+}
+
+impl NativeCacheHandle {
+ fn service(&self) -> PyResult {
+ if self.pid != std::process::id() {
+ return Err(PyRuntimeError::new_err(
+ "native cache handles must be recreated after fork",
+ ));
+ }
+ Ok(self.service.clone())
+ }
+}
+
+#[pymethods]
+impl NativeCacheHandle {
+ #[staticmethod]
+ #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))]
+ fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult {
+ Ok(Self {
+ service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes),
+ guard: None,
+ pid: std::process::id(),
+ })
+ }
+
+ #[staticmethod]
+ #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))]
+ fn redis(
+ py: Python<'_>,
+ url: String,
+ ttl_seconds: Option,
+ namespace: Option,
+ ) -> PyResult {
+ let ttl = ttl_seconds.map(duration).transpose()?;
+ let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace))
+ .map_err(cache_error)?;
+ Ok(Self {
+ service,
+ guard: None,
+ pid: std::process::id(),
+ })
+ }
+
+ #[getter]
+ fn backend(&self) -> &'static str {
+ self.service.kind()
+ }
+
+ fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
+ let service = self.service()?;
+ let guard = FacadeGuard::capture(py, facade, self.backend())?;
+ let handle = Py::new(
+ py,
+ Self {
+ service,
+ guard: Some(guard),
+ pid: self.pid,
+ },
+ )?;
+ facade.setattr("_native_cache_handle", handle)
+ }
+
+ fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+ if let Some(guard) = &self.guard {
+ guard.traverse(visit)?;
+ }
+ Ok(())
+ }
+}
+
+enum CacheBinding {
+ Disabled,
+ Native(NativeResponseCache),
+ PythonCallback(Py),
+}
+
+#[pyclass(frozen, name = "CacheBinding")]
+pub(crate) struct ResolvedCache {
+ binding: CacheBinding,
+ pid: u32,
+}
+
+impl ResolvedCache {
+ fn check_process(&self) -> PyResult<()> {
+ if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() {
+ return Err(PyRuntimeError::new_err(
+ "native cache bindings must be resolved again after fork",
+ ));
+ }
+ Ok(())
+ }
+
+ pub(crate) fn lookup_step(
+ &self,
+ py: Python<'_>,
+ input: &Bound<'_, PyAny>,
+ kwargs: Option<&Bound<'_, PyDict>>,
+ ) -> PyResult {
+ self.check_process()?;
+ let awaitable = match &self.binding {
+ CacheBinding::Disabled => ready_none(py)?,
+ CacheBinding::Native(service) => {
+ let request = request(input)?;
+ let service = service.clone();
+ run_async(
+ py,
+ async move { service.async_lookup(&request, now()).await },
+ cache_error,
+ )?
+ }
+ CacheBinding::PythonCallback(object) => object.bind(py).call_method(
+ "async_get_cache",
+ (),
+ Some(callback_kwargs(kwargs)?),
+ )?,
+ };
+ Ok(ExecutionStep::Await(awaitable.unbind()))
+ }
+}
+
+#[pymethods]
+impl ResolvedCache {
+ #[getter]
+ fn kind(&self) -> &'static str {
+ match self.binding {
+ CacheBinding::Disabled => "disabled",
+ CacheBinding::Native(_) => "native",
+ CacheBinding::PythonCallback(_) => "python_callback",
+ }
+ }
+
+ #[pyo3(signature = (request, *, callback_kwargs=None))]
+ fn lookup(
+ &self,
+ py: Python<'_>,
+ request: &Bound<'_, PyAny>,
+ callback_kwargs: Option<&Bound<'_, PyDict>>,
+ ) -> PyResult> {
+ self.check_process()?;
+ match &self.binding {
+ CacheBinding::Disabled => Ok(py.None()),
+ CacheBinding::Native(service) => {
+ let request = self::request(request)?;
+ let service = service.clone();
+ let response = release_gil(py, move || service.lookup(&request, now()))
+ .map_err(cache_error)?;
+ to_py(py, &response)
+ }
+ CacheBinding::PythonCallback(object) => object
+ .bind(py)
+ .call_method(
+ "get_cache",
+ (),
+ Some(self::callback_kwargs(callback_kwargs)?),
+ )
+ .map(Bound::unbind),
+ }
+ }
+
+ #[pyo3(signature = (request, response, *, callback_kwargs=None))]
+ fn store(
+ &self,
+ py: Python<'_>,
+ request: &Bound<'_, PyAny>,
+ response: &Bound<'_, PyAny>,
+ callback_kwargs: Option<&Bound<'_, PyDict>>,
+ ) -> PyResult<()> {
+ self.check_process()?;
+ match &self.binding {
+ CacheBinding::Disabled => Ok(()),
+ CacheBinding::Native(service) => {
+ let request = self::request(request)?;
+ let response: Value = from_py(response)?;
+ let service = service.clone();
+ release_gil(py, move || service.store(&request, response, now()))
+ .map_err(cache_error)
+ }
+ CacheBinding::PythonCallback(object) => object
+ .bind(py)
+ .call_method(
+ "add_cache",
+ (response,),
+ Some(self::callback_kwargs(callback_kwargs)?),
+ )
+ .map(|_| ()),
+ }
+ }
+
+ #[pyo3(signature = (request, *, callback_kwargs=None))]
+ fn async_lookup<'py>(
+ &self,
+ py: Python<'py>,
+ request: &Bound<'py, PyAny>,
+ callback_kwargs: Option<&Bound<'py, PyDict>>,
+ ) -> PyResult> {
+ let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)?
+ else {
+ unreachable!()
+ };
+ Ok(awaitable.into_bound(py))
+ }
+
+ #[pyo3(signature = (request, response, *, callback_kwargs=None))]
+ fn async_store<'py>(
+ &self,
+ py: Python<'py>,
+ request: &Bound<'py, PyAny>,
+ response: &Bound<'py, PyAny>,
+ callback_kwargs: Option<&Bound<'py, PyDict>>,
+ ) -> PyResult> {
+ self.check_process()?;
+ match &self.binding {
+ CacheBinding::Disabled => ready_none(py),
+ CacheBinding::Native(service) => {
+ let request = self::request(request)?;
+ let response: Value = from_py(response)?;
+ let service = service.clone();
+ run_async(
+ py,
+ async move { service.async_store(&request, response, now()).await },
+ cache_error,
+ )
+ }
+ CacheBinding::PythonCallback(object) => object.bind(py).call_method(
+ "async_add_cache",
+ (response,),
+ Some(self::callback_kwargs(callback_kwargs)?),
+ ),
+ }
+ }
+
+ fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+ if let CacheBinding::PythonCallback(object) = &self.binding {
+ visit.call(object)?;
+ }
+ Ok(())
+ }
+}
+
+fn callback_kwargs<'a, 'py>(
+ kwargs: Option<&'a Bound<'py, PyDict>>,
+) -> PyResult<&'a Bound<'py, PyDict>> {
+ kwargs.ok_or_else(|| {
+ PyTypeError::new_err("Python cache callbacks require their original callback_kwargs")
+ })
+}
+
+fn ready_none(py: Python<'_>) -> PyResult> {
+ let future = py
+ .import("asyncio")?
+ .call_method0("get_running_loop")?
+ .call_method0("create_future")?;
+ future.call_method1("set_result", (py.None(),))?;
+ Ok(future)
+}
+
+#[pyclass(frozen)]
+pub(crate) struct CacheResolver {
+ namespace: Py,
+}
+
+#[pymethods]
+impl CacheResolver {
+ #[new]
+ fn new(namespace: Py) -> Self {
+ Self { namespace }
+ }
+
+ pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult {
+ let object = self.namespace.bind(py).getattr("cache")?;
+ let binding = if object.is_none() {
+ CacheBinding::Disabled
+ } else if let Ok(handle) = object.extract::>() {
+ CacheBinding::Native(handle.service()?)
+ } else if let Some(service) = facade::resolve(py, &object)? {
+ CacheBinding::Native(service)
+ } else {
+ CacheBinding::PythonCallback(object.unbind())
+ };
+ Ok(ResolvedCache {
+ binding,
+ pid: std::process::id(),
+ })
+ }
+
+ fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+ visit.call(&self.namespace)
+ }
+}
diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs
index 46f98736aa1..621c111a35b 100644
--- a/litellm-rust/crates/python-bridge/src/lib.rs
+++ b/litellm-rust/crates/python-bridge/src/lib.rs
@@ -1,3 +1,4 @@
+mod cache;
mod credentials;
mod diagnostics;
mod errors;
@@ -9,6 +10,8 @@ mod token_counter;
#[pymodule(gil_used = true)]
mod _native {
+ #[pymodule_export]
+ use crate::cache::{CacheResolver, NativeCacheHandle, ResolvedCache};
#[cfg(feature = "panic-test")]
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
@@ -65,6 +68,9 @@ mod tests {
"achat_completions",
"ResponsesWebSocketConnection",
"TokenCounter",
+ "CacheResolver",
+ "NativeCacheHandle",
+ "CacheBinding",
"gil_stats",
"process_state_started",
"reserve_process_for_forking",
diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi
index 05a6df6d5af..4fd2f0829a3 100644
--- a/litellm/rust_bridge/_native.pyi
+++ b/litellm/rust_bridge/_native.pyi
@@ -1,5 +1,5 @@
from asyncio import Future
-from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
+from collections.abc import AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence
from typing import Never, final
from litellm.llms.base_llm.ocr.transformation import OCRResponse
@@ -93,6 +93,50 @@ class ResponsesWebSocketConnection:
def recv_text(self) -> Future[str | None]: ...
def close(self) -> Future[None]: ...
+@final
+class NativeCacheHandle:
+ def __new__(cls, _uninstantiable: Never, /) -> Never: ...
+ @staticmethod
+ def memory(
+ *, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576
+ ) -> NativeCacheHandle: ...
+ @staticmethod
+ def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> NativeCacheHandle: ...
+ @property
+ def backend(self) -> str: ...
+ def bind_facade(self, facade: object) -> None: ...
+
+@final
+class CacheResolver:
+ def __new__(cls, namespace: object) -> CacheResolver: ...
+ def resolve(self) -> CacheBinding: ...
+
+@final
+class CacheBinding:
+ def __new__(cls, _uninstantiable: Never, /) -> Never: ...
+ @property
+ def kind(self) -> str: ...
+ def lookup(
+ self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None
+ ) -> object: ...
+ def store(
+ self,
+ request: Mapping[str, object] | None,
+ response: object,
+ *,
+ callback_kwargs: dict[str, object] | None = None,
+ ) -> None: ...
+ def async_lookup(
+ self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None
+ ) -> Awaitable[object]: ...
+ def async_store(
+ self,
+ request: Mapping[str, object] | None,
+ response: object,
+ *,
+ callback_kwargs: dict[str, object] | None = None,
+ ) -> Awaitable[object]: ...
+
@final
class TokenCounter:
def __new__(cls, tokenizer_json: str) -> TokenCounter: ...
@@ -109,7 +153,10 @@ def process_state_started() -> bool: ...
def reserve_process_for_forking() -> None: ...
__all__ = [
+ "CacheBinding",
+ "CacheResolver",
"ForkedAfterNativeRuntimeStarted",
+ "NativeCacheHandle",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",
"RustBridgeDeclined",
diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py
new file mode 100644
index 00000000000..d1cea860cb0
--- /dev/null
+++ b/tests/test_litellm_rust/test_cache.py
@@ -0,0 +1,231 @@
+import asyncio
+import contextvars
+import gc
+import json
+import threading
+import time
+import weakref
+from collections.abc import Generator
+from types import SimpleNamespace
+from typing import Final, Protocol, cast
+
+import fakeredis
+import pytest
+import redis
+
+import litellm
+from litellm.caching.caching import Cache
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.rust_bridge import _native
+from litellm.types.caching import LiteLLMCacheType
+from tests.test_litellm_rust.support.isolation import rebound
+
+pytestmark: Final = pytest.mark.requires_rust_extension
+
+
+class CacheLookup(Protocol):
+ def get_cache(self, **kwargs: object) -> object: ...
+
+
+def request(key: str = "key") -> dict[str, object]:
+ return {"key": {"preset": key}}
+
+
+@pytest.fixture
+def redis_url() -> Generator[str]:
+ server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis")
+ worker: Final = threading.Thread(target=server.serve_forever, daemon=True)
+ worker.start()
+ try:
+ yield f"redis://127.0.0.1:{server.server_address[1]}"
+ finally:
+ server.shutdown()
+ server.server_close()
+ worker.join(timeout=5)
+
+
+def test_existing_constructor_and_global_are_unchanged() -> None:
+ facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
+ assert type(facade.cache) is InMemoryCache
+ assert "_native_cache_handle" not in vars(facade)
+ with rebound(litellm, "cache", facade):
+ resolver: Final = _native.CacheResolver(litellm)
+ assert resolver.resolve().kind == "python_callback"
+ resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"})
+ assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7}
+
+
+async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None:
+ namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.memory())
+ resolver: Final = _native.CacheResolver(namespace)
+ selected: Final = resolver.resolve()
+ assert selected.kind == "native"
+ selected.store(request(), {"answer": 1})
+ assert await selected.async_lookup(request()) == {"answer": 1}
+ with rebound(namespace, "cache", _native.NativeCacheHandle.memory()):
+ replacement: Final = resolver.resolve()
+ await selected.async_store(request(), {"answer": 2})
+ assert replacement.lookup(request()) is None
+ assert selected.lookup(request()) == {"answer": 2}
+ with rebound(namespace, "cache", None):
+ disabled: Final = resolver.resolve()
+ assert disabled.kind == "disabled"
+ assert disabled.lookup(None) is None
+ await disabled.async_store(None, object())
+ assert await disabled.async_lookup(None) is None
+ assert selected.lookup(request()) == {"answer": 2}
+
+
+async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None:
+ context: Final = contextvars.ContextVar("cache_context", default="caller")
+ caller: Final = asyncio.current_task()
+ sentinel: Final = object()
+ failure: Final = RuntimeError("callback failed")
+
+ class CustomCache:
+ async def async_get_cache(self, *, marker: object) -> object:
+ assert marker is sentinel
+ assert asyncio.current_task() is caller
+ context.set("callback")
+ return marker
+
+ async def async_add_cache(self, response: object, *, marker: object) -> None:
+ assert response is sentinel
+ assert marker is sentinel
+ raise failure
+
+ namespace: Final = SimpleNamespace(cache=CustomCache())
+ binding: Final = _native.CacheResolver(namespace).resolve()
+ assert binding.kind == "python_callback"
+ assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel
+ assert context.get() == "callback"
+ with pytest.raises(RuntimeError) as caught:
+ await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel})
+ assert caught.value is failure
+
+
+async def test_callback_cancellation_stays_in_the_callers_task() -> None:
+ entered: Final = asyncio.Event()
+ finished: Final = asyncio.Event()
+
+ class CustomCache:
+ async def async_get_cache(self) -> None:
+ entered.set()
+ try:
+ await asyncio.Future()
+ finally:
+ finished.set()
+
+ binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve()
+
+ async def lookup() -> object:
+ return await binding.async_lookup(None, callback_kwargs={})
+
+ task: Final = asyncio.create_task(lookup())
+ await entered.wait()
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ assert finished.is_set()
+
+
+def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None:
+ facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
+ handle: Final = _native.NativeCacheHandle.memory()
+ handle.bind_facade(facade)
+ resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade))
+ native: Final = resolver.resolve()
+ assert native.kind == "native"
+ native.store(request(), {"source": "native"})
+ assert native.lookup(request()) == {"source": "native"}
+ assert cast(CacheLookup, facade).get_cache(cache_key="key") is None
+ sentinel: Final = object()
+
+ def outer_override(**_kwargs: object) -> object:
+ return sentinel
+
+ def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]:
+ return {"source": "override"}
+
+ with rebound(facade, "get_cache", outer_override):
+ fallback: Final = resolver.resolve()
+ assert fallback.kind == "python_callback"
+ assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel
+ assert resolver.resolve().kind == "python_callback"
+ delattr(facade, "get_cache")
+ assert resolver.resolve().kind == "native"
+ with rebound(facade.cache, "get_cache", backend_override):
+ backend_fallback: Final = resolver.resolve()
+ assert backend_fallback.kind == "python_callback"
+ assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"}
+
+
+def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None:
+ class CustomCache(Cache):
+ pass
+
+ handle: Final = _native.NativeCacheHandle.memory()
+ with pytest.raises(TypeError):
+ handle.bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL))
+ facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
+ handle.bind_facade(facade)
+ resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade))
+ with rebound(facade, "cache", InMemoryCache()):
+ assert resolver.resolve().kind == "python_callback"
+ with rebound(facade, "ttl", 12):
+ assert resolver.resolve().kind == "python_callback"
+
+ def custom_key(**_kwargs: object) -> str:
+ return "custom"
+
+ with rebound(facade, "get_cache_key", custom_key):
+ assert resolver.resolve().kind == "python_callback"
+ assert resolver.resolve().kind == "python_callback"
+ delattr(facade, "get_cache_key")
+ assert resolver.resolve().kind == "native"
+
+
+def test_resolver_and_callback_cycles_can_be_collected() -> None:
+ class CustomCache:
+ pass
+
+ def cyclic_reference() -> weakref.ReferenceType[CustomCache]:
+ callback: Final = CustomCache()
+ namespace: Final = SimpleNamespace(cache=callback)
+ binding: Final = _native.CacheResolver(namespace).resolve()
+ setattr(callback, "binding", binding)
+ return weakref.ref(callback)
+
+ reference: Final = cyclic_reference()
+ gc.collect()
+ assert reference() is None
+
+
+async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None:
+ client: Final = redis.Redis.from_url(redis_url)
+ namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url, namespace="team"))
+ binding: Final = _native.CacheResolver(namespace).resolve()
+ response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
+ envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)}
+ client.set("team:sync", str(envelope))
+ client.set("team:async", json.dumps({"timestamp": time.time(), "response": response}))
+ assert binding.lookup(request("sync")) == response
+ assert await binding.async_lookup(request("team:async")) == response
+ await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response)
+ stored: Final = client.get("team:native")
+ assert isinstance(stored, bytes)
+ assert json.loads(stored)["response"] == response
+ assert 0 < client.ttl("team:native") <= 12
+ assert client.get("litellm-cache:team:native") is None
+ assert client.get("team:team:async") is None
+ client.close()
+
+
+def test_invalid_duration_and_request_shape_fail_before_storage() -> None:
+ binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve()
+ for seconds in (-1.0, float("nan"), float("inf")):
+ with pytest.raises(ValueError):
+ binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1})
+ assert binding.lookup(request()) is None
+ with pytest.raises(ValueError):
+ _native.NativeCacheHandle.memory(ttl_seconds=-1)
From 0c3a0a208948e97d0da05ec6c7e28672205c2f00 Mon Sep 17 00:00:00 2001
From: Yujong Lee
Date: Sun, 20 Sep 2026 21:11:48 -0700
Subject: [PATCH 21/56] refactor(cache): separate response policy and host
selection
---
litellm-rust/Cargo.lock | 4 +-
litellm-rust/crates/cache-memory/Cargo.toml | 2 +-
litellm-rust/crates/cache-memory/src/cache.rs | 47 +-----
.../crates/cache-memory/tests/cache.rs | 80 ++++------
litellm-rust/crates/cache-redis/src/cache.rs | 34 +++--
litellm-rust/crates/cache-response/Cargo.toml | 7 +-
litellm-rust/crates/cache-response/README.md | 55 +++++++
.../crates/cache-response/src/caching.rs | 143 ++++++++++++++++++
.../crates/cache-response/src/codec.rs | 4 +-
litellm-rust/crates/cache-response/src/lib.rs | 7 +-
.../crates/cache-response/src/response.rs | 6 +-
.../crates/cache-response/tests/caching.rs | 83 ++++++++++
.../crates/cache-response/tests/response.rs | 40 +++--
litellm-rust/crates/cache/Cargo.toml | 1 -
litellm-rust/crates/cache/src/caching.rs | 143 ------------------
litellm-rust/crates/cache/src/lib.rs | 5 +-
litellm-rust/crates/cache/tests/caching.rs | 94 +-----------
litellm-rust/crates/cache/tests/codec.rs | 14 +-
litellm-rust/crates/python-bridge/Cargo.toml | 2 +
.../crates/python-bridge/src/cache/facade.rs | 3 +-
.../crates/python-bridge/src/cache/mod.rs | 6 +-
.../src => python-bridge/src/cache}/native.rs | 33 ++--
tests/test_litellm_rust/test_cache.py | 20 ++-
23 files changed, 426 insertions(+), 407 deletions(-)
create mode 100644 litellm-rust/crates/cache-response/README.md
create mode 100644 litellm-rust/crates/cache-response/src/caching.rs
create mode 100644 litellm-rust/crates/cache-response/tests/caching.rs
rename litellm-rust/crates/{cache-response/src => python-bridge/src/cache}/native.rs (75%)
diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock
index aa62e4f3770..8b299f5455b 100644
--- a/litellm-rust/Cargo.lock
+++ b/litellm-rust/Cargo.lock
@@ -2460,7 +2460,6 @@ dependencies = [
"rstest",
"serde",
"serde_json",
- "sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
]
@@ -2498,6 +2497,7 @@ dependencies = [
"redis-test",
"serde",
"serde_json",
+ "sha2 0.10.9",
"tokio",
]
@@ -2665,6 +2665,8 @@ dependencies = [
"litellm-auth",
"litellm-auth-gcp",
"litellm-cache",
+ "litellm-cache-memory",
+ "litellm-cache-redis",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core",
diff --git a/litellm-rust/crates/cache-memory/Cargo.toml b/litellm-rust/crates/cache-memory/Cargo.toml
index d4487573a9a..86ab01564c8 100644
--- a/litellm-rust/crates/cache-memory/Cargo.toml
+++ b/litellm-rust/crates/cache-memory/Cargo.toml
@@ -7,8 +7,8 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
-serde_json.workspace = true
[dev-dependencies]
+serde_json.workspace = true
rstest.workspace = true
tokio.workspace = true
diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs
index 974cdbe9760..43186faf3f7 100644
--- a/litellm-rust/crates/cache-memory/src/cache.rs
+++ b/litellm-rust/crates/cache-memory/src/cache.rs
@@ -3,15 +3,12 @@ use std::collections::{BinaryHeap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
-use litellm_cache::{
- BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheKwargs, Error,
-};
+use litellm_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
type ValueMeasure = Arc Result + Send + Sync>;
-type ValueValidator = Arc Result<(), Error> + Send + Sync>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CacheWrite {
@@ -32,7 +29,6 @@ pub struct InMemoryCache {
default_ttl: Duration,
max_entry_bytes: Option,
measure_value: Option>,
- validate_value: Option>,
now: Arc Duration + Send + Sync>,
}
@@ -76,7 +72,6 @@ impl InMemoryCache {
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
max_entry_bytes,
measure_value,
- validate_value: None,
now: Arc::new(now),
}
}
@@ -90,9 +85,6 @@ impl InMemoryCache {
if self.max_size_in_memory == 0 {
return Ok(CacheWrite::Disabled);
}
- if let Some(validate) = &self.validate_value {
- validate(&value)?;
- }
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&value)? > limit
{
@@ -176,43 +168,6 @@ impl InMemoryCache {
}
}
-impl InMemoryCache {
- pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
- Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || {
- SystemTime::now()
- .duration_since(UNIX_EPOCH)
- .unwrap_or_default()
- })
- }
-
- pub fn response_cache_with_clock(
- capacity: usize,
- ttl: Duration,
- max_entry_bytes: usize,
- now: impl Fn() -> Duration + Send + Sync + 'static,
- ) -> Self {
- let mut cache = Self::with_clock_and_size_measurement(
- Some(capacity),
- Some(ttl),
- Some(max_entry_bytes),
- Some(Arc::new(|entry: &CacheEntry| {
- serde_json::to_vec(entry)
- .map(|bytes| bytes.len())
- .map_err(|_| Error::InvalidEntry)
- })),
- now,
- );
- cache.validate_value = Some(Arc::new(|entry: &CacheEntry| {
- entry
- .timestamp
- .is_finite()
- .then_some(())
- .ok_or(Error::InvalidEntry)
- }));
- cache
- }
-}
-
impl BaseCache for InMemoryCache {
type Value = V;
diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs
index ffac9d8ae64..370145bebef 100644
--- a/litellm-rust/crates/cache-memory/tests/cache.rs
+++ b/litellm-rust/crates/cache-memory/tests/cache.rs
@@ -3,8 +3,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use litellm_cache::{
- BaseCache, CacheBackend, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, get_cache,
- set_cache,
+ BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, Error, get_cache, set_cache,
};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
@@ -87,66 +86,49 @@ fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc
}
#[test]
-fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
- let disabled = InMemoryCache::::response_cache(0, Duration::from_secs(60), 80);
+fn disabled_size_limited_and_validated_writes_are_observable() {
+ let cache = |capacity| {
+ InMemoryCache::with_clock_and_size_measurement(
+ Some(capacity),
+ Some(Duration::from_secs(60)),
+ Some(4),
+ Some(Arc::new(|value: &String| {
+ if value.is_empty() {
+ return Err(Error::InvalidEntry);
+ }
+ Ok(value.len())
+ })),
+ || Duration::from_secs(100),
+ )
+ };
+ let disabled = cache(0);
assert_eq!(
- disabled
- .set_cache(
- "a",
- CacheEntry {
- timestamp: 1.0,
- response: serde_json::json!("x")
- },
- None
- )
- .unwrap(),
+ disabled.set_cache("a", "x".into(), None).unwrap(),
CacheWrite::Disabled
);
- let cache = InMemoryCache::::response_cache(2, Duration::from_secs(60), 80);
+ let cache = cache(2);
assert_eq!(
- cache
- .set_cache(
- "large",
- CacheEntry {
- timestamp: 1.0,
- response: serde_json::json!("x".repeat(100))
- },
- None
- )
- .unwrap(),
+ cache.set_cache("large", "oversized".into(), None).unwrap(),
CacheWrite::TooLarge
);
- cache
- .set_cache(
- "small",
- CacheEntry {
- timestamp: 1.0,
- response: serde_json::json!("ok"),
- },
- None,
- )
- .unwrap();
- assert!(cache.get_cache("small").unwrap().is_some());
+ assert_eq!(cache.get_cache("large").unwrap(), None);
assert_eq!(
- cache
- .set_cache(
- "invalid",
- CacheEntry {
- timestamp: f64::NAN,
- response: serde_json::json!("bad"),
- },
- None,
- )
- .unwrap_err(),
- Error::InvalidEntry
+ cache.set_cache("small", "ok".into(), None).unwrap(),
+ CacheWrite::Stored
);
+ assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into()));
+ assert_eq!(
+ cache.set_cache("invalid", String::new(), None),
+ Err(Error::InvalidEntry)
+ );
+ assert_eq!(cache.get_cache("invalid").unwrap(), None);
cache.delete_cache("small").unwrap();
- cache.flush_cache().unwrap();
+ assert_eq!(cache.get_cache("small").unwrap(), None);
}
#[tokio::test]
async fn connection_test_matches_python_result_contract() {
- let cache = InMemoryCache::::default();
+ let cache = InMemoryCache::::default();
let result = BaseCache::test_connection(&cache).await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "In-memory cache connection test successful");
diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs
index 0faca6cdaaf..d4ca0cf0522 100644
--- a/litellm-rust/crates/cache-redis/src/cache.rs
+++ b/litellm-rust/crates/cache-redis/src/cache.rs
@@ -238,30 +238,27 @@ where
#[cfg(test)]
mod tests {
use super::RedisCache;
- use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKwargs, JsonCodec};
+ use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use std::time::Duration;
- fn entry() -> CacheEntry {
- CacheEntry {
- timestamp: 123.0,
- response: json!({"choices": [{"text": "cached"}]}),
- }
+ fn entry() -> serde_json::Value {
+ json!({"deployment": "model-a", "cooldown_seconds": 30})
}
#[test]
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
assert_eq!(
- RedisCache::>::ttl_seconds(Duration::ZERO),
+ RedisCache::>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
- RedisCache::>::ttl_seconds(Duration::from_millis(1500)),
+ RedisCache::>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
- RedisCache::>::ttl_seconds(Duration::from_secs(15)),
+ RedisCache::>::ttl_seconds(Duration::from_secs(15)),
15
);
}
@@ -269,7 +266,9 @@ mod tests {
#[test]
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
let value = entry();
- let payload = JsonCodec::