From cd17ced01394f3cd648421250db3a87042db7fc5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:30:01 +0000 Subject: [PATCH] fix(router): keep the prompt caching pin working for response messages and malformed tool calls The router's prompt caching check reads tool calls through their JSON form, so an SDK agent loop that passes the response message back, or a request whose tool_calls is not a list of objects, routes as before instead of raising a validation error. A cache_control under a tool's function that is not an object no longer counts as a 5m mark ahead of a 1h breakpoint, since the Anthropic transform drops it. On routes that never translate prompt_cache_breakpoint, such as OpenRouter and Databricks, a breakpoint on a tool call no longer takes the operator's system injection point --- .../anthropic_cache_control_hook.py | 24 ++++-- .../chat/prompt_cache_breakpoints.py | 13 ++- litellm/router_utils/prompt_caching_cache.py | 10 ++- .../test_prompt_caching_deployment_check.py | 43 ++++++++++ tests/unit/test_main.py | 82 +++++++++++++++++++ 5 files changed, 160 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 7df13541828..fbefce48293 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -313,6 +313,14 @@ def _builds_anthropic_messages(model: str, custom_llm_provider: str | None) -> b return isinstance(chat_config, AnthropicConfig) and not isinstance(chat_config, OpenAIGPTConfig) +def _tool_call_breakpoints_become_marks( + messages: Sequence[object], model: str, custom_llm_provider: str | None +) -> bool: + if not any(_carries_mark(tool_call, "prompt_cache_breakpoint") for tool_call in _tool_calls(messages)): + return True + return _builds_anthropic_messages(model, custom_llm_provider) or _served_by_bedrock(model, custom_llm_provider) + + def _custom_llm_provider(non_default_params: Mapping[str, object], point: Mapping[str, object]) -> str | None: provider: Final = non_default_params.get("custom_llm_provider") or point.get(CUSTOM_LLM_PROVIDER_STAMP) return provider if isinstance(provider, str) else None @@ -475,8 +483,11 @@ class AnthropicCacheControlHook(CustomPromptManagement): answered_tool_ids: Final = _tool_call_ids_sent_as_tool_calls_when_a_server_tool_call_is_marked( processed_messages, model, custom_llm_provider ) + census_dialect: Final = openai_dialect or not _tool_call_breakpoints_become_marks( + processed_messages, model, custom_llm_provider + ) breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints( - processed_messages, answered_tool_ids=answered_tool_ids, openai_dialect=openai_dialect + processed_messages, answered_tool_ids=answered_tool_ids, openai_dialect=census_dialect ) processed_messages = self._apply_message_injections( points=applied_message_points, @@ -484,6 +495,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks, openai_dialect=openai_dialect, answered_tool_ids=answered_tool_ids, + census_dialect=census_dialect, ) if ( openai_dialect @@ -504,7 +516,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): *AnthropicCacheControlHook._points_with_a_slot_left( remaining_points, AnthropicCacheControlHook.count_request_cache_breakpoints( - processed_messages, answered_tool_ids=answered_tool_ids, openai_dialect=openai_dialect + processed_messages, answered_tool_ids=answered_tool_ids, openai_dialect=census_dialect ) + external_breakpoints, openai_dialect, @@ -634,6 +646,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): max_blocks: int, openai_dialect: bool = False, answered_tool_ids: frozenset[str] | None = None, + census_dialect: bool | None = None, ) -> list[AllMessageValues]: """Apply message-level cache control injection points in order. @@ -649,8 +662,9 @@ class AnthropicCacheControlHook(CustomPromptManagement): if answered_tool_ids is None else answered_tool_ids ) + counted_dialect: Final = openai_dialect if census_dialect is None else census_dialect used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints( - messages, answered_tool_ids=client_answers, openai_dialect=openai_dialect + messages, answered_tool_ids=client_answers, openai_dialect=counted_dialect ) limit_reached = False @@ -669,7 +683,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): break if AnthropicCacheControlHook._message_has_cache_control( - messages[target_index], client_answers, openai_dialect + messages[target_index], client_answers, counted_dialect ): # Client already marked this message; don't overwrite it. continue @@ -678,7 +692,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): messages[target_index], control, openai_dialect ) if AnthropicCacheControlHook._message_has_cache_control( - messages[target_index], client_answers, openai_dialect + messages[target_index], client_answers, counted_dialect ): used_blocks += 1 diff --git a/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py b/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py index 0b4081b8348..8baa1fd57a2 100644 --- a/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py +++ b/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py @@ -152,9 +152,13 @@ def _namespace_members(tools: Sequence[object]) -> Iterator[tuple[_Path, _Fields ) +def _function_mark(tool: _Fields) -> _Fields | None: + return _fields((_fields(tool.get("function")) or {}).get("cache_control")) + + def _member_mark(member: _Fields) -> object: own: Final = member.get("cache_control") - return own if own is not None else (_fields(member.get("function")) or {}).get("cache_control") + return own if own is not None else _function_mark(member) def _transform_sends_member(member: _Fields) -> bool: @@ -530,12 +534,13 @@ def _on_wire( ) -> bool: """Claude sends a message-level mark on string content, or on list content only for a tool result, and none on a developer message, which reaches it as a system message without one; a tool's mark only when the route sends - it, and one under ``function`` only off Converse and when the tool carries none of its own; no tool_search call - nobody answered; and no mark on tool history the route sends as text.""" + it, and one under ``function`` only off Converse, as an object, and when the tool carries none of its own; no + tool_search call nobody answered; and no mark on tool history the route sends as text.""" if path[0] == "tools": tool: Final = _fields(tools[int(path[1])]) or {} return tool_mark_reaches_claude(tool, route) and ( - len(path) == 2 or (route != "converse" and tool.get("cache_control") is None) + len(path) == 2 + or (route != "converse" and tool.get("cache_control") is None and _function_mark(tool) is not None) ) message: Final = _fields(messages[int(path[1])]) or {} if tool_history_as_text and _is_tool_history(message): diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 8d83c281bbd..d984edddcec 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -44,7 +44,7 @@ _MARK_KEYS: Final = frozenset({"cache_control", "prompt_cache_breakpoint"}) _ENVELOPE_SKIPPED_KEYS: Final = _MARK_KEYS | {"content"} _PREFIX_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) _TOOLS_ADAPTER: Final = TypeAdapter(tuple[JsonValue, ...]) -_OBJECTS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None) @@ -79,8 +79,12 @@ def _carries_cache_mark(block: Mapping[str, object]) -> bool: def _carries_marked_tool_call(message: Mapping[str, object]) -> bool: - tool_calls: Final = _OBJECTS_ADAPTER.validate_python(message.get("tool_calls") or ()) - return any(_carries_cache_mark(call) for call in tool_calls) + tool_calls: Final = _JSON_ADAPTER.validate_python( + to_jsonable_python(message.get("tool_calls"), serialize_unknown=True, bytes_mode="base64") + ) + return isinstance(tool_calls, list) and any( + isinstance(call, dict) and _carries_cache_mark(call) for call in tool_calls + ) def _unmarked(block: JsonValue) -> JsonValue: diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 018fa808f70..e708b0216e7 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -19,6 +19,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ) from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt, token_counter MODEL_GROUP_ALIAS = "my-claude-group" @@ -860,6 +861,48 @@ async def test_pin_survives_the_marked_message_coming_back_as_string_content(): assert filtered == [deployments[0]] +def _lookup_call() -> ChatCompletionMessageToolCall: + return ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="lookup_clause", arguments="{}") + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "assistant_turn", + [ + Message(content=None, tool_calls=[_lookup_call()]), + {"role": "assistant", "content": None, "tool_calls": [_lookup_call()]}, + {"role": "assistant", "content": None, "tool_calls": "call_1"}, + {"role": "assistant", "content": None, "tool_calls": {"id": "call_1"}}, + {"role": "assistant", "content": None, "tool_calls": [1]}, + {"role": "assistant", "content": None, "tool_calls": ["call_1"]}, + {"role": "assistant", "content": None, "tool_calls": [None]}, + ], + ids=["response_message", "response_tool_calls", "string", "object", "number_list", "string_list", "null_list"], +) +async def test_pin_holds_whatever_shape_the_assistant_tool_calls_come_back_in(assistant_turn: object): + 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": "Find the termination clause."}, + ) + turn_two = cast( + list[AllMessageValues], + [*turn_one, assistant_turn, {"role": "tool", "tool_call_id": "call_1", "content": "Section 9."}], + ) + + 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_lookback_stops_where_the_provider_cache_stops(): """ diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 5d07abd0111..3e54111a8c8 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1148,6 +1148,34 @@ def test_a_breakpoint_drop_params_drops_leaves_key_auto_caching_on(respx_mock: r {"type": "tool_search_tool_regex_20251119", "name": "tool_search", "cache_control": _FIVE_MINUTES}, id="converse-tool-search-regex", ), + pytest.param( + "anthropic/claude-opus-5-5", + r"https://api\.anthropic\.com/v1/messages", + _CLAUDE_REPLY, + {"api_key": "fake"}, + { + "type": "function", + "function": {"name": "spawn", "parameters": {"type": "object"}, "cache_control": "ephemeral"}, + }, + id="anthropic-string-function-mark", + ), + pytest.param( + "bedrock/invoke/us.anthropic.claude-opus-5-5", + _INVOKE_ROUTE, + _CLAUDE_REPLY, + _FAKE_AWS, + { + "type": "namespace", + "name": "agents", + "tools": [ + { + "type": "function", + "function": {"name": "spawn", "parameters": {"type": "object"}, "cache_control": "ephemeral"}, + } + ], + }, + id="invoke-namespace-member-string-function-mark", + ), ], ) def test_a_five_minute_mark_on_a_tool_the_route_sends_unmarked_leaves_the_system_breakpoint_free_to_take_the_hour( @@ -1283,6 +1311,60 @@ def test_a_marked_past_server_tool_call_leaves_the_operator_system_point_in_plac assert "cache_control" in json.dumps(json.loads(route.calls.last.request.content)["messages"][0]) +@pytest.mark.parametrize( + ("model", "api_base", "url"), + [ + ("openrouter/anthropic/claude-opus-5-5", None, "https://openrouter.ai/api/v1/chat/completions"), + ( + "databricks/databricks-claude-opus-4-5", + "https://adb-cache.cloud.databricks.com/serving-endpoints", + "https://adb-cache.cloud.databricks.com/serving-endpoints/chat/completions", + ), + ], + ids=["openrouter", "databricks"], +) +def test_a_breakpoint_on_a_tool_call_the_route_never_translates_leaves_the_operator_system_point_in_place( + respx_mock: respx.MockRouter, + openai_api_response: Mapping[str, object], + model: str, + api_base: str | None, + url: str, +): + route: Final = respx_mock.post(url).respond(json=openai_api_response) + + litellm.completion( + model=model, + messages=[ + {"role": "system", "content": "sys"}, + { + "role": "user", + "content": [ + {"type": "text", "text": str(index), "prompt_cache_breakpoint": _BREAKPOINT} for index in range(3) + ], + }, + { + "role": "assistant", + "content": "Looking it up.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "prompt_cache_breakpoint": _BREAKPOINT, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Sunny"}, + {"role": "user", "content": "Thanks"}, + ], + cache_control_injection_points=[{"location": "message", "role": "system"}], + api_key="fake", + api_base=api_base, + ) + + assert "cache_control" in json.dumps(json.loads(route.calls.last.request.content)["messages"][0]) + + def test_a_tuple_additional_drop_params_places_the_same_breakpoints_as_a_list(respx_mock: respx.MockRouter): route: Final = respx_mock.post(url__regex=_CONVERSE_ROUTE).respond(json=_CONVERSE_REPLY)