feat(llm): Claude prompt caching (system + tools + rolling tail), Bedrock-safe

Enables Anthropic/Bedrock prompt caching for Claude routes via LiteLLM
cache_control_injection_points: system prompt + latest message everywhere,
plus tool_config on Bedrock Converse only (the sole route whose transform
consumes it; elsewhere it leaks as an unknown top-level field that native
Anthropic 400-rejects). Unmapped Bedrock Claude models run uncached instead
of crashing. Adds STRIX_PROMPT_CACHE opt-out (default on).
This commit is contained in:
Ahmed Allam
2026-07-26 23:29:41 +00:00
parent fb6606279f
commit 460ceab075
7 changed files with 198 additions and 154 deletions
+64
View File
@@ -429,3 +429,67 @@ def is_known_openai_bare_model(model_name: str) -> bool:
return False
entry = litellm.model_cost.get(name)
return bool(entry and entry.get("litellm_provider") == "openai")
def is_claude_model(model_name: str) -> bool:
return "claude" in (model_name or "").strip().lower()
def is_bedrock_route(model_name: str) -> bool:
"""Whether ``model_name`` resolves to an AWS Bedrock route.
Matches the ``bedrock/...`` LiteLLM route prefix and bare Bedrock model ids
(``[region.]anthropic.claude-...``).
"""
name = (model_name or "").strip().lower()
return name.startswith("bedrock/") or "anthropic." in name
def _prompt_cache_name_candidates(model_name: str) -> list[str]:
"""Candidate LiteLLM model-map keys for ``model_name``, most→least specific.
LiteLLM keys the same model under several names (``bedrock/global.anthropic.
claude-opus-4-1``, ``anthropic.claude-opus-4-1``, ``claude-opus-4-1``) and not
every provider/region-prefixed variant is present for every model. Strip the
LiteLLM route prefix, then leading dotted segments (region, then provider) so
a prefixed name still resolves to a bare key.
"""
name = (model_name or "").strip().lower()
for prefix in ("litellm/", "bedrock/"):
if name.startswith(prefix):
name = name[len(prefix) :]
break
candidates = [name]
rest = name
while "." in rest:
rest = rest.split(".", 1)[1]
candidates.append(rest)
return candidates
def bedrock_route_supports_prompt_caching(model_name: str) -> bool:
"""Whether LiteLLM can confirm this Bedrock model supports prompt caching.
Bedrock's Converse API rejects unknown request fields outright
(``ValidationException: cache_control_injection_points: Extra inputs are
not permitted``), and LiteLLM only consumes the cache marker for models its
(statically bundled) model map recognises as cache-capable. For a Bedrock
model missing from that map — a just-released model, or any model when the
remote model-map refresh fails and a stale local copy is used — the marker
would pass straight through and fail every call, so callers must withhold
it unless support is confirmed here.
"""
import litellm
checker = getattr(getattr(litellm, "utils", None), "supports_prompt_caching", None)
for cand in _prompt_cache_name_candidates(model_name):
if checker is not None:
# supports_prompt_caching raises for models missing from the map;
# keep checking the remaining name candidates.
with contextlib.suppress(Exception):
if checker(cand):
return True
entry = litellm.model_cost.get(cand)
if entry and entry.get("supports_prompt_caching"):
return True
return False
+4
View File
@@ -40,6 +40,10 @@ class LlmSettings(BaseSettings):
default=False,
alias="STRIX_FORCE_REQUIRED_TOOL_CHOICE",
)
prompt_cache: bool = Field(
default=True,
alias="STRIX_PROMPT_CACHE",
)
timeout: int = Field(default=300, alias="LLM_TIMEOUT")
+78 -135
View File
@@ -10,6 +10,9 @@ from openai.types.shared import Reasoning
from strix.config.models import (
DEFAULT_MODEL_RETRY,
bedrock_route_supports_prompt_caching,
is_bedrock_route,
is_claude_model,
is_known_openai_bare_model,
model_supports_reasoning,
request_timeout_extra_args,
@@ -128,6 +131,7 @@ def make_model_settings(
model_name: str,
force_required_tool_choice: bool = False,
request_timeout: float | None = None,
prompt_cache: bool = True,
) -> ModelSettings:
model_settings = ModelSettings(
parallel_tool_calls=False,
@@ -145,154 +149,93 @@ def make_model_settings(
)
if force_required_tool_choice and _accepts_required_tool_choice(model_name):
model_settings = model_settings.resolve(ModelSettings(tool_choice="required"))
if _is_claude_model(model_name) and not _bedrock_route_without_cache_support(model_name):
# Merge into any existing extra_args rather than relying on resolve()'s
# dict-merge semantics — makes it obvious at the call site that unrelated
# LiteLLM options are preserved (make_model_settings currently builds
# from scratch, so extra_args is None here today, but this keeps the
# invariant local if that changes).
merged_extra_args = {
**(model_settings.extra_args or {}),
**_claude_prompt_cache_extra_args(),
}
cache_extra_args = _prompt_cache_extra_args(model_name) if prompt_cache else None
if cache_extra_args:
# Merge into any existing extra_args (e.g. the request timeout) rather
# than relying on resolve()'s dict-merge semantics, so it is obvious at
# the call site that unrelated LiteLLM options are preserved.
model_settings = model_settings.resolve(
ModelSettings(extra_args=merged_extra_args),
ModelSettings(
extra_args={**(model_settings.extra_args or {}), **cache_extra_args},
),
)
return model_settings
def _is_claude_model(model_name: str) -> bool:
return "claude" in (model_name or "").strip().lower()
def _litellm_name_candidates(model_name: str) -> list[str]:
"""Candidate LiteLLM model-map keys for ``model_name``, most→least specific.
LiteLLM keys the same model under several names (``bedrock/global.anthropic.
claude-opus-4-1``, ``anthropic.claude-opus-4-1``, ``claude-opus-4-1``) and not
every provider/region-prefixed variant is present for every model. Strip the
LiteLLM route prefix, then leading dotted segments (region, then provider) so
a prefixed name still resolves to a bare key.
"""
name = (model_name or "").strip().lower()
for prefix in ("litellm/", "bedrock/"):
if name.startswith(prefix):
name = name[len(prefix) :]
break
candidates = [name]
for cand in list(candidates):
rest = cand
while "." in rest:
rest = rest.split(".", 1)[1]
candidates.append(rest)
return candidates
def _bedrock_route_without_cache_support(model_name: str) -> bool:
"""True for a BEDROCK Claude route that LiteLLM can't confirm supports prompt
caching — the one case where injecting the cache marker HARD-CRASHES the run.
Bedrock's Converse API rejects unknown request fields outright
(``ValidationException: cache_control_injection_points: Extra inputs are not
permitted``). LiteLLM's ``AnthropicCacheControlHook`` strips
``cache_control_injection_points`` from the outgoing call only for models it
recognises as cache-capable via its (statically bundled) model map; for a
model missing from that map the marker passes straight through and Bedrock
500s the first call, failing the whole scan. This bites any Bedrock Claude
model LiteLLM hasn't mapped yet — a just-released model, or ANY model when
LiteLLM can't refresh its remote model map (e.g. behind a TLS-intercepting
corporate proxy) and falls back to a stale local copy.
Scope is deliberately narrow — ONLY Bedrock routes. Anthropic-native,
Vertex, and OpenRouter Claude tolerate/ignore the marker (or LiteLLM maps
them under keys we don't resolve), so gating those on confirmed support
would DISABLE caching for genuinely-capable models — a caching regression,
the opposite of this change's intent. So elsewhere we keep injecting by
model family and only withhold on the provider that actually rejects.
"""
name = (model_name or "").strip().lower()
if not name.startswith("bedrock/") and "anthropic." not in name:
# Not a Bedrock route (bedrock/... or a bare bedrock model id like
# global.anthropic.claude-...); other providers don't hard-reject.
return False
import litellm
checker = getattr(getattr(litellm, "utils", None), "supports_prompt_caching", None)
for cand in _litellm_name_candidates(model_name):
if checker is not None:
try:
if checker(cand):
return False # confirmed cache-capable → safe to inject
except Exception: # noqa: BLE001 — unknown model raises; keep checking
pass
entry = litellm.model_cost.get(cand)
if entry and entry.get("supports_prompt_caching"):
return False
return True # Bedrock route, support unconfirmed → withhold to avoid the 500
def _claude_prompt_cache_extra_args() -> dict[str, Any]:
"""Enable Anthropic/Bedrock prompt caching for Claude models via LiteLLM.
def _prompt_cache_extra_args(model_name: str) -> dict[str, Any] | None:
"""LiteLLM ``extra_args`` that enable Anthropic/Bedrock prompt caching.
A Strix scan is a long, multi-turn agentic loop that re-sends a large,
STABLE prefix every turn — the system prompt plus the tool schemas — AND an
append-only conversation transcript that only grows. Without caching
breakpoints the whole request is re-tokenised and billed at the full input
rate on every turn; on Bedrock Claude that is the single biggest lever on
scan cost (measured here: ``cache-read 0% -> 57%`` on a real scan once these
points are set).
rate on every turn; on Claude that is the single biggest lever on scan cost
(measured: ``cache-read 0% -> ~66%`` on a real scan once these points are set).
LiteLLM already implements this end to end: when
``cache_control_injection_points`` is present in the call kwargs its
``AnthropicCacheControlHook`` fires and emits the provider-appropriate
breakpoint (Anthropic ``cache_control``; Bedrock Converse ``cachePoint``),
honouring Anthropic's 4-breakpoint cap. ``LitellmModel`` forwards
``ModelSettings.extra_args`` straight into ``litellm.acompletion()``, so
passing the injection points there is all that is required.
This mirrors the caching policy of production agent harnesses (e.g.
anomalyco/opencode's ``cache-policy``): cache the tool schemas, the system
prompt, and the latest conversation message, capped at Anthropic's 4
breakpoints. We express it through LiteLLM's ``cache_control_injection_points``
— the ``AnthropicCacheControlHook`` fires on that kwarg and emits the
provider-appropriate breakpoint (Anthropic ``cache_control``; Bedrock
Converse ``cachePoint``). ``LitellmModel`` forwards ``ModelSettings.extra_args``
straight into ``litellm.acompletion()``, so passing the points there is all
that is required; this is the LiteLLM-config-layer approach the Agents SDK
maintainer prescribed over a native ``ModelSettings`` caching flag
(openai/openai-agents-python#3008 / #3009).
This is deliberately kept at the LiteLLM-config layer rather than a general
``ModelSettings`` caching flag: that is the direction the Agents SDK
maintainer prescribed when declining a native ``cache_system_prompt`` field
(openai/openai-agents-python#3008 / #3009) — caching is a LiteLLM/provider
behaviour and a ``ModelSettings`` flag would let strict OpenAI-compatible
paths emit non-standard ``cache_control`` parts. Gating on Claude keeps this
a no-op for every other provider (no injection points -> the hook never
fires), and only Claude-family routes (Anthropic native, Bedrock, Vertex,
OpenRouter -> Claude) honour the marker.
Three breakpoints (3 of the 4 allowed), leaving headroom:
- the system prompt (``role: system``) — the largest repeated span
- the tool schemas (``tool_config``) — sizeable and identical every turn
- the conversation tail (``index: -1``) — a ROLLING breakpoint on the last
message, so the accumulated transcript caches incrementally
The tail breakpoint matters more than it looks. The first two only cache the
FIXED prefix; the transcript is append-only (prior turns are immutable, each
turn just appends the new assistant/tool messages), so on a long scan the
growing body is re-sent at full input price every turn and cache-read decays
as a denominator effect even though the prefix keeps hitting. A breakpoint at
``index: -1`` re-caches the whole immutable prefix-so-far each turn and hits
on the next; the hook resolves the negative index against the live message
list. Measured on a 29-turn Bedrock scan, WITHOUT the tail point the cached
prefix stayed pinned at ~56k tokens while per-turn input grew to ~256k and
cache-read fell from 90% to 22%; adding it lifts modelled cache-read to ~96%
and cuts full-price input ~16x on that scan.
All three points degrade gracefully on older LiteLLM: an unrecognised
location is simply not injected (no error), so a stale pin still gets
whatever caching it supports — the system-prompt point (the widest support)
and the tool_config + message-index points applied by LiteLLM's Bedrock
Converse transform on versions that recognise them (verified on litellm
1.90.1).
Returns ``None`` (a strict no-op — the hook never fires) for every route
that would not benefit or could break:
- non-Claude models, and
- Bedrock Claude routes LiteLLM can't confirm as cache-capable. Bedrock's
Converse API rejects unknown request fields outright
(``ValidationException: cache_control_injection_points: Extra inputs are
not permitted``) and LiteLLM only consumes the marker for models its
model map recognises; an unmapped Bedrock model would pass the marker
straight through and crash the first call. Only Bedrock hard-rejects, so
only Bedrock is guarded — gating Anthropic-native/Vertex/OpenRouter on
confirmed support would needlessly disable caching for capable models.
"""
return {
"cache_control_injection_points": [
{"location": "message", "role": "system"},
{"location": "tool_config"},
{"location": "message", "index": -1},
],
}
if not is_claude_model(model_name):
return None
if is_bedrock_route(model_name) and not bedrock_route_supports_prompt_caching(model_name):
return None
points = _prompt_cache_injection_points(model_name)
return {"cache_control_injection_points": points}
def _prompt_cache_injection_points(model_name: str) -> list[dict[str, Any]]:
"""Cache breakpoints for a Claude route (system + tools + latest message).
At most 3 of Anthropic's 4 allowed breakpoints, leaving headroom:
- system prompt (``role: system``) — the largest repeated span.
- tool schemas (``tool_config``) — Bedrock Converse ONLY. LiteLLM's
``tool_config`` location is implemented solely by the Bedrock Converse
transform (which appends a ``cachePoint`` to the tool list); on any other
route it is not consumed and would leak onto the wire as an unknown
top-level ``cache_control_injection_points`` field. It is also redundant
elsewhere: Anthropic orders tools BEFORE the system prompt, so the system
breakpoint already caches the tool schemas in the shared prefix.
- latest message (``index: -1``) — a ROLLING breakpoint on the last
message. A scan transcript is append-only (prior turns are immutable, each
turn just appends new assistant/tool messages), so without it the growing
body is re-sent at full input price every turn and cache-read decays as a
denominator effect even though the prefix keeps hitting. Re-caching the
whole prefix-so-far each turn keeps cache-read high on long scans. (This
is the Strix analogue of opencode's ``latest-user-message``; ``index: -1``
tracks the true tail because Strix appends tool-role, not user-role,
messages each turn.)
Unrecognised locations degrade gracefully on older LiteLLM — they are simply
not injected (no error).
"""
points: list[dict[str, Any]] = [{"location": "message", "role": "system"}]
if is_bedrock_route(model_name):
points.append({"location": "tool_config"})
points.append({"location": "message", "index": -1})
return points
def child_initial_input(
+1
View File
@@ -236,6 +236,7 @@ async def run_strix_scan(
model_name=resolved_model,
force_required_tool_choice=settings.llm.force_required_tool_choice,
request_timeout=settings.llm.timeout,
prompt_cache=settings.llm.prompt_cache,
)
run_config = RunConfig(
model=resolved_model,
+46 -17
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
from itertools import pairwise
from typing import Any
import litellm
import pytest
from strix.core.inputs import build_root_task, child_initial_input, make_model_settings
@@ -60,23 +61,46 @@ def _cache_points(model_name: str) -> Any:
return extra.get("cache_control_injection_points")
@pytest.mark.parametrize(
"model_name",
[
"bedrock/global.anthropic.claude-opus-4-8",
"anthropic/claude-sonnet-4-5",
"openrouter/anthropic/claude-3.5-sonnet",
],
)
def test_make_model_settings_enables_prompt_cache_for_claude(model_name: str) -> None:
points = _cache_points(model_name)
assert points == [
def test_make_model_settings_enables_prompt_cache_for_bedrock_claude() -> None:
# Bedrock Converse is the only route whose transform consumes the
# ``tool_config`` location, so it gets all three breakpoints.
assert _cache_points("bedrock/global.anthropic.claude-opus-4-8") == [
{"location": "message", "role": "system"},
{"location": "tool_config"},
{"location": "message", "index": -1},
]
@pytest.mark.parametrize(
"model_name",
[
"anthropic/claude-sonnet-4-5",
"openrouter/anthropic/claude-3.5-sonnet",
"vertex_ai/claude-sonnet-4-5",
],
)
def test_make_model_settings_enables_prompt_cache_for_non_bedrock_claude(model_name: str) -> None:
# Non-Bedrock routes get system + tail only: LiteLLM implements the
# ``tool_config`` location solely in the Bedrock Converse transform, so
# sending it elsewhere would leak an unknown top-level field (see the
# dedicated no-leak test). Tools are already cached by the system breakpoint
# there — Anthropic orders tools ahead of the system prompt in the prefix.
assert _cache_points(model_name) == [
{"location": "message", "role": "system"},
{"location": "message", "index": -1},
]
def test_tool_config_point_not_leaked_to_non_bedrock_claude() -> None:
# Regression guard: the ``tool_config`` injection point must NEVER be sent on
# a non-Bedrock route. LiteLLM leaves it on the outgoing body as a top-level
# ``cache_control_injection_points`` field there, which native Anthropic
# would reject as an unknown field.
for model in ("anthropic/claude-sonnet-4-5", "openrouter/anthropic/claude-3.5-sonnet"):
points = _cache_points(model) or []
assert all(p.get("location") != "tool_config" for p in points)
@pytest.mark.parametrize("model_name", ["gpt-5", "vertex_ai/gemini-2.5-pro", "openai/o3"])
def test_make_model_settings_no_prompt_cache_for_non_claude(model_name: str) -> None:
# No injection points for non-Claude models: the LiteLLM cache hook never
@@ -85,6 +109,16 @@ def test_make_model_settings_no_prompt_cache_for_non_claude(model_name: str) ->
assert make_model_settings(None, model_name=model_name).extra_args is None
def test_prompt_cache_can_be_disabled() -> None:
# STRIX_PROMPT_CACHE=false kill switch: no injection points even for Claude.
assert (
make_model_settings(
None, model_name="anthropic/claude-sonnet-4-5", prompt_cache=False
).extra_args
is None
)
def test_no_prompt_cache_for_unmapped_bedrock_claude_model(monkeypatch: Any) -> None:
"""A BEDROCK Claude route LiteLLM has NOT mapped (a new release, or any model
when LiteLLM can't refresh its model map and falls back to a stale local
@@ -93,8 +127,6 @@ def test_no_prompt_cache_for_unmapped_bedrock_claude_model(monkeypatch: Any) ->
Extra inputs are not permitted'); LiteLLM only strips the marker for models
it recognises as cache-capable, so an unmapped model would 500 the first
call and fail the whole run."""
import litellm
unmapped = "bedrock/global.anthropic.claude-brand-new-9"
# Simulate a model LiteLLM doesn't know: no cost-map entry, checker says no.
monkeypatch.setattr(litellm, "model_cost", {}, raising=False)
@@ -111,18 +143,15 @@ def test_prompt_cache_kept_for_non_bedrock_claude_even_if_unmapped(monkeypatch:
LiteLLM maps them under keys we don't resolve, e.g. OpenRouter), so gating
them on confirmed support would DISABLE caching for capable models — a
regression. Only Bedrock hard-rejects, so only Bedrock is guarded."""
import litellm
monkeypatch.setattr(litellm, "model_cost", {}, raising=False)
if getattr(getattr(litellm, "utils", None), "supports_prompt_caching", None):
monkeypatch.setattr(litellm.utils, "supports_prompt_caching", lambda *_a, **_k: False)
# Even with LiteLLM knowing nothing, an Anthropic-native / OpenRouter Claude
# still gets the injection points.
# still gets the injection points (system + tail, no tool_config).
for model in ("anthropic/claude-brand-new-9", "openrouter/anthropic/claude-brand-new"):
assert _cache_points(model) == [
{"location": "message", "role": "system"},
{"location": "tool_config"},
{"location": "message", "index": -1},
]
+1
View File
@@ -38,6 +38,7 @@ async def test_persistent_rate_limit_stops_gracefully(
reasoning_effort="high",
force_required_tool_choice=False,
timeout=300,
prompt_cache=True,
),
runtime=types.SimpleNamespace(max_context_images=3),
)
+4 -2
View File
@@ -17,6 +17,7 @@ import strix.tools.notes.tools as notes_tools
import strix.tools.todo.tools as todo_tools
from strix.core import runner
from strix.core.agents import AgentCoordinator
from strix.runtime import session_manager
def _make_rate_limit_error() -> RateLimitError:
@@ -46,6 +47,7 @@ def _patch_engine_scaffold(
reasoning_effort="high",
force_required_tool_choice=False,
timeout=300,
prompt_cache=True,
),
runtime=types.SimpleNamespace(max_context_images=3),
)
@@ -66,8 +68,8 @@ def _patch_engine_scaffold(
async def _cleanup(*_args: Any, **_kwargs: Any) -> None:
return None
monkeypatch.setattr(runner.session_manager, "create_or_reuse", _create_or_reuse)
monkeypatch.setattr(runner.session_manager, "cleanup", _cleanup)
monkeypatch.setattr(session_manager, "create_or_reuse", _create_or_reuse)
monkeypatch.setattr(session_manager, "cleanup", _cleanup)
monkeypatch.setattr(runner, "build_root_task", lambda _scan_config: "task")
monkeypatch.setattr(runner, "build_scope_context", lambda _scan_config: scope_context)