fix(router): read tool call marks without walking their values

The router's prompt caching check now reads each tool call's own cache_control and prompt_cache_breakpoint
fields instead of serializing the whole tool call, so a tool call that carries a deeply nested value routes as
before instead of answering 500. The cache control hook no longer asks the provider resolver whether a GitHub
Copilot or ChatGPT model is served by Bedrock, so a breakpoint on a tool call never starts their device login.
Tests cover tool calls in response objects and tuples, a function form namespace member mark, the Copilot
login, and the two 30m e2e cases now set drop_params on their own deployment
This commit is contained in:
mateo-berri 2026-10-16 21:26:57 +00:00
parent cd17ced013
commit 3f1e289711
6 changed files with 131 additions and 18 deletions

View file

@ -283,9 +283,11 @@ def _marks_a_server_tool_call(messages: Sequence[object]) -> bool:
def _served_by_bedrock(model: str, custom_llm_provider: str | None) -> bool:
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider
from litellm.types.utils import LlmProviders
if declared_authenticating_provider(model, custom_llm_provider) is not None:
return False
try:
_, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
except (BadRequestError, ValueError):

View file

@ -9,7 +9,7 @@ from dataclasses import dataclass
from itertools import accumulate, chain
from typing import TYPE_CHECKING, Any, Final, cast
from pydantic import JsonValue, TypeAdapter
from pydantic import BaseModel, JsonValue, TypeAdapter
from pydantic_core import to_jsonable_python
from typing_extensions import TypedDict
@ -44,7 +44,6 @@ _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, ...])
_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None)
@ -78,12 +77,18 @@ def _carries_cache_mark(block: Mapping[str, object]) -> bool:
return _mark_pin_ttl_seconds(block) is not None
def _tool_call_fields(call: object) -> Mapping[str, object]:
if isinstance(call, BaseModel):
return dict(call)
if isinstance(call, Mapping):
return call
return {}
def _carries_marked_tool_call(message: Mapping[str, object]) -> bool:
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
tool_calls: Final = message.get("tool_calls")
return isinstance(tool_calls, (list, tuple)) and any(
_carries_cache_mark(_tool_call_fields(call)) for call in tool_calls
)

View file

@ -722,7 +722,7 @@ class TestPromptCacheBreakpointOnClaude:
def test_a_30m_breakpoint_still_caches(
self, client: PassthroughClient, resources: ResourceManager, scoped_key: str, route: ClaudeRoute
) -> None:
model: Final = _register_claude(client, resources, route)
model: Final = _register_claude(client, resources, route, drop_params=True)
body: Final = ChatBody(
model=model,
max_tokens=64,
@ -760,7 +760,7 @@ class TestPromptCacheBreakpointOnClaude:
def test_a_block_carrying_both_keys_still_caches(
self, client: PassthroughClient, resources: ResourceManager, scoped_key: str, route: ClaudeRoute
) -> None:
model: Final = _register_claude(client, resources, route)
model: Final = _register_claude(client, resources, route, drop_params=True)
both_keys: Final = TextContentPart(
text=_cacheable_prefix(),
cache_control=CacheControl(),

View file

@ -12,6 +12,7 @@ from typing import Final, List, Optional, Tuple, cast
from unittest.mock import MagicMock, patch
import pytest
import respx
from openai.types.chat import ChatCompletionMessage
from pydantic import BaseModel, ConfigDict
@ -1523,6 +1524,49 @@ def test_cache_control_hook_fills_the_mark_budget_every_claude_route_sends_after
assert _marks_the_route_sends(model, custom_llm_provider, processed) == 4
def test_a_breakpoint_on_a_copilot_tool_call_never_starts_the_github_device_login(
monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
):
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://github.test/login/device/code")
messages: Final = cast(
list[AllMessageValues],
[
{"role": "system", "content": "You are a weather assistant."},
{"role": "user", "content": "What is the weather in Paris?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
"prompt_cache_breakpoint": {"mode": "explicit"},
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": "Sunny"},
],
)
with respx.mock(assert_all_called=False) as github:
device_login: Final = github.post("https://github.test/login/device/code").respond(status_code=503)
AnthropicCacheControlHook().get_chat_completion_prompt(
model="claude-sonnet-4.6",
messages=messages,
non_default_params={
"custom_llm_provider": "github_copilot",
"cache_control_injection_points": [{"location": "message", "role": "system"}],
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
assert not device_login.called
@pytest.mark.parametrize("tool_call_type", ["function", "custom", None])
@pytest.mark.parametrize(
"mark",

View file

@ -1,6 +1,7 @@
import asyncio
from functools import reduce
from itertools import chain
from typing import Final
from typing import Final, cast
import pytest
@ -12,6 +13,7 @@ from litellm.router_utils.prompt_caching_cache import (
PromptCachingCache,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message
class _ManualClock:
@ -321,6 +323,61 @@ def test_a_one_hour_breakpoint_on_a_tool_call_keeps_the_pin_after_the_mark_moves
"model_id": "deployment-1h"
}
def _marked_lookup_call() -> ChatCompletionMessageToolCall:
return ChatCompletionMessageToolCall(
id="call_1",
type="function",
function=Function(name="lookup_clause", arguments="{}"),
prompt_cache_breakpoint={"mode": "explicit", "ttl": "1h"},
)
@pytest.mark.parametrize(
"marked_turn",
[
Message(content=None, tool_calls=[_marked_lookup_call()]),
{"role": "assistant", "content": None, "tool_calls": [_marked_lookup_call()]},
{
"role": "assistant",
"content": None,
"tool_calls": ({**_tool_call("call_1"), "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}},),
},
],
ids=["response_message", "response_tool_calls", "tuple_of_dicts"],
)
def test_a_one_hour_breakpoint_keeps_the_pin_whatever_shape_the_marked_tool_calls_come_in(marked_turn: object) -> None:
clock: Final = _ManualClock()
cache: Final = _pin_cache(clock)
messages: Final = cast(list[AllMessageValues], [*_tool_call_turn(marked_call_id="call_1")[:-1], marked_turn])
cache.add_model_id(model_id="deployment-1h", messages=messages, tools=None, honors_one_hour_ttl=True)
clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1)
assert cache.get_model_id(messages=messages, tools=None) == {"model_id": "deployment-1h"}
def test_a_deeply_nested_value_on_a_tool_call_after_the_mark_keeps_the_pin() -> None:
cache: Final = _pin_cache(_ManualClock())
deep_call: Final = {**_tool_call("call_1"), "x": reduce(lambda inner, _: {"x": inner}, range(300), "leaf")}
messages: Final = cast(
list[AllMessageValues],
[
{
"role": "system",
"content": "Here is the full text of a complex legal agreement" * 200,
"cache_control": {"type": "ephemeral"},
},
{"role": "user", "content": "Find the termination clause."},
{"role": "assistant", "content": None, "tool_calls": [deep_call]},
{"role": "tool", "tool_call_id": "call_1", "content": "Section 9."},
],
)
cache.add_model_id(model_id="deployment-1", messages=messages, tools=None)
assert cache.get_model_id(messages=messages, tools=None) == {"model_id": "deployment-1"}
async def _first_non_none(read):
while True:
value = await read()

View file

@ -1204,16 +1204,21 @@ def test_a_five_minute_mark_on_a_tool_the_route_sends_unmarked_leaves_the_system
assert _sent_cache_ttls(json.loads(route.calls.last.request.content)) == ("1h", "1h")
@pytest.mark.parametrize(
"member",
[
{"type": "function", "name": "spawn", "parameters": {"type": "object"}, "cache_control": _FIVE_MINUTES},
{
"type": "function",
"function": {"name": "spawn", "parameters": {"type": "object"}, "cache_control": _FIVE_MINUTES},
},
],
ids=["flat", "function_form"],
)
def test_a_five_minute_mark_on_a_namespace_member_leaves_no_hour_after_it_under_drop_params(
respx_mock: respx.MockRouter,
respx_mock: respx.MockRouter, member: Mapping[str, object]
):
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").respond(json=_CLAUDE_REPLY)
member: Final = {
"type": "function",
"name": "spawn",
"parameters": {"type": "object"},
"cache_control": _FIVE_MINUTES,
}
litellm.completion(
model="anthropic/claude-opus-5-5",