feat(llm): enable Bedrock/Anthropic prompt caching for Claude models

strix/core/inputs.py

@@ -141,9 +141,156 @@ 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(),
+        }
+        model_settings = model_settings.resolve(
+            ModelSettings(extra_args=merged_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.
+
+    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).
+
+    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 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).
+    """
+    return {
+        "cache_control_injection_points": [
+            {"location": "message", "role": "system"},
+            {"location": "tool_config"},
+            {"location": "message", "index": -1},
+        ],
+    }
+
+
 def child_initial_input(
     *,
     name: str,

tests/test_inputs.py

@@ -55,6 +55,110 @@ def test_child_initial_input_no_consecutive_same_role(parent_history: list[Any])
     assert all(prev != nxt for prev, nxt in pairwise(roles))
 
 
+def _cache_points(model_name: str) -> Any:
+    extra = make_model_settings(None, model_name=model_name).extra_args or {}
+    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 == [
+        {"location": "message", "role": "system"},
+        {"location": "tool_config"},
+        {"location": "message", "index": -1},
+    ]
+
+
+@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
+    # fires, so this stays a strict no-op (won't emit cache_control to strict
+    # OpenAI-compatible endpoints).
+    assert make_model_settings(None, model_name=model_name).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
+    copy) must run UNCACHED, not crash. Bedrock's Converse API rejects the
+    unknown field outright (ValidationException 'cache_control_injection_points:
+    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)
+    if getattr(getattr(litellm, "utils", None), "supports_prompt_caching", None):
+        monkeypatch.setattr(litellm.utils, "supports_prompt_caching", lambda *_a, **_k: False)
+
+    # Bedrock Claude by name, but unmapped → no injection points, no crash.
+    assert make_model_settings(None, model_name=unmapped).extra_args is None
+
+
+def test_prompt_cache_kept_for_non_bedrock_claude_even_if_unmapped(monkeypatch: Any) -> None:
+    """Non-Bedrock Claude routes must KEEP caching-by-family even when LiteLLM
+    can't confirm support — those providers tolerate/ignore the marker (or
+    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.
+    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},
+        ]
+
+
+def test_conversation_tail_breakpoint_moves_with_appended_transcript() -> None:
+    """The tail breakpoint's premise, end-to-end: LiteLLM's own message-injection
+    logic must place the cache_control on the LAST message for both a short and a
+    long transcript — i.e. it tracks the growing (append-only) tail rather than a
+    fixed position — so the immutable prefix-so-far is cached and re-read next
+    turn.
+
+    Driven through the hook's static ``_apply_message_injections`` primitive
+    (stable across LiteLLM versions) rather than the prompt-manager entrypoint
+    (whose signature drifts).
+    """
+    hook_mod = pytest.importorskip("litellm.integrations.anthropic_cache_control_hook")
+    apply = hook_mod.AnthropicCacheControlHook._apply_message_injections
+    points = _cache_points("bedrock/global.anthropic.claude-opus-4-8")
+    msg_points = [p for p in points if p.get("location") == "message"]
+
+    def last_msg_cache_control(n_turns: int) -> Any:
+        messages: list[dict[str, Any]] = [{"role": "system", "content": "stable prompt"}]
+        for i in range(n_turns):
+            messages.append({"role": "assistant", "content": f"turn {i} action"})
+            messages.append({"role": "user", "content": f"turn {i} tool result"})
+        processed = apply(msg_points, messages, 4)
+        last = processed[-1]
+        content = last.get("content")
+        if isinstance(content, list):
+            return content[-1].get("cache_control")
+        return last.get("cache_control")
+
+    assert last_msg_cache_control(2) == {"type": "ephemeral"}
+    assert last_msg_cache_control(20) == {"type": "ephemeral"}
+
+
 def test_build_root_task_empty_config() -> None:
     assert build_root_task({}) == ""