fix(caching): scope automatic breakpoints to supported Claude transports

This commit is contained in:
Tin Chi Lo 2026-09-18 20:44:22 -05:00
parent 1d91fc232d
commit 2b086dc7aa
6 changed files with 237 additions and 57 deletions

View file

@ -25,7 +25,10 @@ from litellm.integrations.prompt_management_base import PromptManagementClient
from litellm.litellm_core_utils.prompt_templates.common_utils import (
with_prompt_cache_breakpoint,
)
from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
from litellm.llms.anthropic.common_utils import (
is_claude_code_one_shot_subagent_request,
supports_anthropic_cache_control,
)
from litellm.types.integrations.anthropic_cache_control_hook import (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@ -574,8 +577,9 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider: str | None,
api_base: object,
prompt_cache_options: object,
request_kwargs: object,
) -> Sequence[Mapping[str, object]] | None:
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control):
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control, request_kwargs):
return None
return AnthropicCacheControlHook._stamped_with_dialect(
points, model, custom_llm_provider, api_base, prompt_cache_options
@ -612,6 +616,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
system: str | list | None,
tools: list | None,
cache_control: object = None,
request_kwargs: object = None,
) -> bool:
"""Whether configured injection points must yield to client-set cache_control.
@ -624,7 +629,9 @@ class AnthropicCacheControlHook(CustomPromptManagement):
"""
if all(point.get("_litellm_judged") for point in points):
return False
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control)
return AnthropicCacheControlHook._request_has_cache_control(
messages, system, tools, cache_control, request_kwargs
)
@staticmethod
def _request_has_cache_control(
@ -632,31 +639,29 @@ class AnthropicCacheControlHook(CustomPromptManagement):
system: str | list | None,
tools: list | None = None,
cache_control: object = None,
request_kwargs: object = None,
) -> bool:
"""Return True if the request already carries any client-supplied cache_control.
When the client (e.g. Claude Code) already marks its own breakpoints we
stand down entirely rather than add more, per the auto-caching contract.
Tools count: they are a breakpoint the client can mark, they count toward
the provider's four-block limit, and caching only the tool definitions is
a common pattern, so injecting alongside them can exceed the cap. Tools
carry the mark either at the top level (Anthropic shape) or nested under
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
"""
if cache_control is not None:
return True
if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
return True
if tools is not None:
return any(
isinstance(tool, dict)
and (
tool.get("cache_control") is not None
or (isinstance(tool.get("function"), dict) and tool["function"].get("cache_control") is not None)
)
for tool in tools
"""Client breakpoints own caching in both the request and its extra_body envelope."""
bodies: Final = (
{"messages": messages, "system": system, "tools": tools, "cache_control": cache_control},
_validated_object_mapping(AnthropicCacheControlHook._request_value(request_kwargs, "extra_body")) or {},
)
return any(
body.get("cache_control") is not None
or AnthropicCacheControlHook.count_request_cache_breakpoints(
_validated_object_list(body.get("messages")) or (), body.get("system")
)
return False
> 0
or any(
AnthropicCacheControlHook._request_value(tool, "cache_control") is not None
or AnthropicCacheControlHook._request_value(
AnthropicCacheControlHook._request_value(tool, "function"), "cache_control"
)
is not None
for tool in (_validated_object_list(body.get("tools")) or ())
)
for body in bodies
)
@staticmethod
def get_default_injection_points(
@ -676,36 +681,19 @@ class AnthropicCacheControlHook(CustomPromptManagement):
even when the global flag is off. Caches the system prompt and the
trailing turn, so the stable prefix (system + tools + history) is
reused while the breakpoint advances with the conversation. Returns []
(stand down) when neither flag is on, the provider does not consume
cache_control breakpoints (only anthropic / bedrock do), the model
lacks prompt-caching support, or the request already carries
client-supplied cache_control.
(stand down) when neither flag is on, the model is not Claude on a
supported explicit-cache transport, the model lacks prompt-caching
support, or the request already carries client-supplied cache_control.
"""
import litellm
if litellm.enable_anthropic_prompt_caching is not True and enable_prompt_caching is not True:
return []
provider = custom_llm_provider
if provider is None:
from litellm.litellm_core_utils.get_llm_provider_logic import (
get_llm_provider,
)
try:
_, provider, _, _ = get_llm_provider(model=model)
except Exception: # noqa: BLE001 # unroutable model must never block the call, just skip auto-caching
return []
if provider not in ("anthropic", "bedrock"):
if not supports_anthropic_cache_control(model, custom_llm_provider):
return []
from litellm.utils import supports_prompt_caching
if not supports_prompt_caching(model=model, custom_llm_provider=provider):
return []
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control):
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control, request_kwargs):
return []
if is_claude_code_one_shot_subagent_request(
@ -737,13 +725,15 @@ class AnthropicCacheControlHook(CustomPromptManagement):
prompt and trailing turn) do not depend on which deployment serves the call. Returns the
input list itself when auto-injection would not apply
"""
import litellm
points: Final = next(
(
candidate
for candidate in (
AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
model=model,
model=litellm.model_alias_map.get(model, model),
custom_llm_provider=None,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
@ -789,6 +779,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
prompt-management gate and the AnthropicCacheControlHook run
unchanged.
"""
import litellm
if non_default_params.get("cache_control_injection_points"):
judged: Final = AnthropicCacheControlHook._judged_configured_points(
non_default_params["cache_control_injection_points"],
@ -799,6 +791,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider,
api_base,
non_default_params.get("prompt_cache_options"),
non_default_params,
)
if judged is None:
non_default_params.pop("cache_control_injection_points")
@ -808,7 +801,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
points: Final = AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
system=None,
model=model,
model=litellm.model_alias_map.get(model, model),
custom_llm_provider=custom_llm_provider,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
@ -925,7 +918,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
)
if configured and AnthropicCacheControlHook._should_stand_down(
configured, typed_messages, system, tools, cache_control
configured, typed_messages, system, tools, cache_control, kwargs
):
return messages, system
injection_points: list[CacheControlInjectionPoint] = configured or []

View file

@ -76,6 +76,21 @@ _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/")
def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None) -> bool:
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.utils import supports_prompt_caching
try:
provider: Final = custom_llm_provider if custom_llm_provider is not None else get_llm_provider(model=model)[1]
except Exception: # noqa: BLE001 # Optional caching must not block an unroutable request
return False
return (
provider in ("anthropic", "bedrock", "vertex_ai", "azure_ai")
and "claude" in model.lower()
and supports_prompt_caching(model=model, custom_llm_provider=provider)
)
def is_claude_code_user_agent(user_agent: str) -> bool:
"""Claude Code sends its API calls through the Anthropic SDK as `claude-cli/<version>` and its own
fetches, such as gateway model discovery, as `claude-code/<version>`"""

View file

@ -1968,7 +1968,7 @@ async def generate_key_fn(
- policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
@ -3291,7 +3291,7 @@ async def update_key_fn(
- policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
- blocked: Optional[bool] - Whether the key is blocked
- aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)

View file

@ -17559,8 +17559,8 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie
"type": "Boolean",
"tab": "prompt_caching",
"description": (
"Auto-adds cache_control to the system prompt and trailing turn for supported Anthropic "
"and Bedrock Claude models. The cache is shared across callers on the same upstream credentials."
"Auto-adds cache_control to the system prompt and trailing turn for supported Claude models on "
"Anthropic, Bedrock, Vertex AI, and Azure AI. The cache is shared across callers on the same upstream credentials."
),
},
"anthropic_prompt_caching_ttl": {

View file

@ -1595,6 +1595,178 @@ class TestEnableAnthropicPromptCaching:
assert supports_prompt_caching(model=model, custom_llm_provider=provider) is True
assert self._points(model=model, provider=provider) == []
@pytest.mark.parametrize("family", ["haiku-4-5", "sonnet-5", "opus-5", "fable-5", "fable-5-1"])
@pytest.mark.parametrize(
"provider, template",
[("anthropic", "{}"), ("vertex_ai", "{}"), ("azure_ai", "{}"), ("bedrock", "us.anthropic.{}-v1:0")],
)
@pytest.mark.parametrize("infer_provider", [False, True])
@pytest.mark.parametrize("supported", [False, True])
def test_claude_transport_defaults(self, monkeypatch, local_model_cost_map, family, provider, template, infer_provider, supported):
from litellm.utils import supports_prompt_caching
model = template.format(f"claude-{family}")
qualified = f"{provider}/{model}"
entry = {"litellm_provider": provider, "mode": "chat", "supports_prompt_caching": supported}
monkeypatch.setitem(litellm.model_cost, model, entry)
monkeypatch.setitem(litellm.model_cost, qualified, entry)
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False)
target = qualified if infer_provider else model
resolved_provider = None if infer_provider else provider
assert supports_prompt_caching(model=target, custom_llm_provider=resolved_provider) is supported
points = AnthropicCacheControlHook.get_default_injection_points(
messages=copy.deepcopy(self.MESSAGES), system=None, model=target,
custom_llm_provider=resolved_provider, enable_prompt_caching=True,
)
assert [point["index"] for point in points] == ([None, -1] if supported else [])
affinity_messages = AnthropicCacheControlHook.messages_with_default_injections(
copy.deepcopy(self.MESSAGES), models=[qualified], enable_prompt_caching=True,
)
assert sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in affinity_messages) == (2 if supported else 0)
@pytest.mark.parametrize(
"provider, model",
[
("bedrock", "us.openai.gpt-6-astra"),
("bedrock", "amazon.nova-pro-v1:0"),
("bedrock", "us.xai.grok-4.6"),
("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque"),
("vertex_ai", "gemini-3.8-flash"),
("azure_ai", "gpt-6-astra"),
("anthropic", "unknown-model"),
],
)
def test_non_claude_caching_capability_does_not_enable_defaults(self, monkeypatch, local_model_cost_map, provider, model):
from litellm.utils import supports_prompt_caching
qualified = f"{provider}/{model}"
entry = {"litellm_provider": provider, "mode": "chat", "supports_prompt_caching": True}
monkeypatch.setitem(litellm.model_cost, model, entry)
monkeypatch.setitem(litellm.model_cost, qualified, entry)
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
assert supports_prompt_caching(model=model, custom_llm_provider=provider)
assert self._points(model=model, provider=provider) == []
assert self._points(model=qualified, provider=None) == []
assert AnthropicCacheControlHook.messages_with_default_injections(self.MESSAGES, [qualified]) == self.MESSAGES
@pytest.mark.parametrize("provider", ["vertex_ai", "azure_ai"])
@pytest.mark.parametrize("client_control", ["none", "message", "system", "tool", "function", "top_level"])
@pytest.mark.parametrize("envelope", ["request", "extra_body"])
@pytest.mark.parametrize("configured", [False, True])
def test_new_transports_preserve_client_controls(self, monkeypatch, local_model_cost_map, provider, client_control, envelope, configured):
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig
model = "claude-sonnet-5"
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
monkeypatch.setitem(litellm.model_cost, f"{provider}/{model}", {
**litellm.model_cost[f"{provider}/{model}"], "supports_prompt_caching": True,
})
control = {"type": "ephemeral"}
messages = [{"role": "user", "content": [{"type": "text", "text": "question", **({"cache_control": control} if client_control == "message" else {})}]}]
system = [{"type": "text", "text": "stable context", **({"cache_control": control} if client_control == "system" else {})}]
tools = [{"name": "lookup", "description": "Lookup", "input_schema": {"type": "object", "properties": {}}, **({"cache_control": control} if client_control == "tool" else {})}]
if client_control == "function":
tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}, "cache_control": control}}]
kwargs = {"metadata": {}, "model_info": {"id": "selected-deployment"}, **({"cache_control": control} if client_control == "top_level" else {})}
if envelope == "extra_body":
kwargs["extra_body"] = {"messages": messages, "system": system, "tools": tools}
if "cache_control" in kwargs:
kwargs["extra_body"]["cache_control"] = kwargs.pop("cache_control")
messages, system, tools = [{"role": "user", "content": "question"}], "stable context", []
if configured:
kwargs["cache_control_injection_points"] = [
{"location": "message", "role": "system", "index": None, "control": control},
{"location": "message", "role": None, "index": -1, "control": control},
]
seeded = copy.deepcopy(kwargs)
original = copy.deepcopy((messages, system, tools))
result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control(
messages, system, kwargs, model, provider, tools=tools,
)
if client_control != "none":
assert (result_messages, result_system, tools) == original
assert kwargs["metadata"] == {}
else:
assert kwargs["metadata"]["litellm_gateway_injected_cache"] == "selected-deployment"
assert sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in result_messages) == 1
assert result_system[0]["cache_control"] == control
if provider == "vertex_ai":
wire = VertexAIAnthropicConfig().transform_request(
model=model, messages=[{"role": "system", "content": result_system}, *result_messages],
optional_params={"max_tokens": 8}, litellm_params={}, headers={},
)
assert wire["system"][0]["cache_control"] == control
assert wire["messages"][-1]["content"][-1]["cache_control"] == control
affinity = AnthropicCacheControlHook.messages_with_default_injections(
[{"role": "system", "content": original[1]}, *original[0]], [f"{provider}/{model}"],
tools=tools, request_kwargs=seeded,
)
if client_control != "none":
assert affinity == [{"role": "system", "content": original[1]}, *original[0]]
AnthropicCacheControlHook.maybe_seed_default_injection_points(
seeded, [{"role": "system", "content": original[1]}, *original[0]], model, provider, tools=tools,
)
assert bool(seeded.get("cache_control_injection_points")) == (client_control == "none")
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("model, target, client_control, expected", [
("vertex_ai/claude-sonnet-5", "bedrock/amazon.nova-pro-v1:0", False, 0),
("azure_ai/gpt-6-astra", "azure_ai/claude-sonnet-5", False, 2),
("azure_ai/claude-sonnet-5", None, False, 2),
("azure_ai/claude-sonnet-5", None, True, 1),
("azure_ai/model_router/claude-replacement", None, False, 2),
])
async def test_public_completion_cache_ownership(self, monkeypatch, local_model_cost_map, asynchronous, model, target, client_control, expected):
import httpx
from litellm.llms.custom_httpx.http_handler import HTTPHandler
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
monkeypatch.setattr(litellm, "model_alias_map", {model: target} if target else {})
for qualified in (model, target):
if qualified:
provider = qualified.split("/")[0]
entry = {"litellm_provider": provider, "mode": "chat", "supports_prompt_caching": True}
monkeypatch.setitem(litellm.model_cost, qualified, entry)
monkeypatch.setitem(litellm.model_cost, qualified.split("/", 1)[-1], entry)
sent = []
def respond(request):
sent.append(json.loads(request.content))
return httpx.Response(200, request=request, json={
"id": "msg-test", "type": "message", "role": "assistant", "model": "claude-sonnet-5",
"content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "stop_sequence": None,
"output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, "stopReason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 1, "inputTokens": 10, "outputTokens": 1, "totalTokens": 11},
})
control = {"type": "ephemeral", "ttl": "1h"}
messages = [{"role": "system", "content": "stable context"}, {"role": "user", "content": "question"}]
metadata = {}
kwargs = {
"model": model, "messages": copy.deepcopy(messages), "max_tokens": 32, "num_retries": 0,
"litellm_metadata": metadata,
"api_base": "https://rig.services.ai.azure.com/anthropic", "api_key": "synthetic-test-key",
"aws_access_key_id": "synthetic", "aws_secret_access_key": "synthetic", "aws_region_name": "us-east-1",
**({"extra_body": {"cache_control": control}} if client_control else {}),
}
if asynchronous:
handler = AsyncHTTPHandler()
await handler.client.aclose()
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
handler.client = client
response = await litellm.acompletion(**kwargs, client=handler)
else:
with httpx.Client(transport=httpx.MockTransport(respond)) as client:
response = litellm.completion(**kwargs, client=HTTPHandler(client=client))
assert response.choices[0].message.content == "ok"
assert len(sent) == 1
assert ("litellm_gateway_injected_cache" in metadata) == (expected == 2)
serialized = json.dumps(sent[0])
assert serialized.count('"cache_control"') + serialized.count('"cachePoint"') == expected
if client_control:
assert sent[0]["cache_control"] == control
affinity = AnthropicCacheControlHook.messages_with_default_injections(messages, [model], request_kwargs=kwargs)
assert AnthropicCacheControlHook.count_request_cache_breakpoints(affinity) == (2 if expected == 2 else 0)
def test_databricks_claude_not_injected_despite_caching_support(self, monkeypatch, local_model_cost_map):
from litellm.utils import supports_prompt_caching

View file

@ -7801,7 +7801,7 @@ export interface paths {
* - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
* - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
* - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only.
* - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
* - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
@ -8282,7 +8282,7 @@ export interface paths {
* - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
* - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
* - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only.
* - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
* - prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
* - blocked: Optional[bool] - Whether the key is blocked
* - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)