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

Co-authored-by: Sean Turner <sean.turner@zerohash.com>
Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>

strix/config/models.py

@@ -429,3 +429,45 @@ 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:
+    name = (model_name or "").strip().lower()
+    return name.startswith("bedrock/") or "anthropic." in name
+
+
+def _prompt_cache_name_candidates(model_name: str) -> list[str]:
+    # LiteLLM's model map keys the same model under several names; strip the
+    # route prefix, then leading dotted segments (region, provider).
+    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:
+    # Bedrock rejects the cache marker for models LiteLLM's map doesn't
+    # recognise as cache-capable, so callers withhold it unless 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:
+            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

strix/config/settings.py

@@ -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")
 
 

strix/core/inputs.py

@@ -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,9 +149,38 @@ 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"))
+
+    cache_extra_args = _prompt_cache_extra_args(model_name) if prompt_cache else None
+    if cache_extra_args:
+        model_settings = model_settings.resolve(
+            ModelSettings(
+                extra_args={**(model_settings.extra_args or {}), **cache_extra_args},
+            ),
+        )
     return model_settings
 
 
+def _prompt_cache_extra_args(model_name: str) -> dict[str, Any] | None:
+    """LiteLLM ``cache_control_injection_points`` for Claude prompt caching.
+
+    System prompt + rolling last-message breakpoint everywhere; ``tool_config``
+    only on Bedrock Converse (the only route whose LiteLLM transform consumes
+    it — elsewhere it leaks onto the wire and native Anthropic 400s). Unmapped
+    Bedrock models get no points at all: Bedrock rejects the passed-through
+    field outright.
+    """
+    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: 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 {"cache_control_injection_points": points}
+
+
 def child_initial_input(
     *,
     name: str,

strix/core/runner.py

@@ -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,

tests/test_inputs.py

@@ -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
@@ -55,6 +56,103 @@ 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")
+
+
+def test_make_model_settings_enables_prompt_cache_for_bedrock_claude() -> None:
+    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:
+    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:
+    # LiteLLM only consumes tool_config on Bedrock; elsewhere it leaks onto the
+    # wire and native Anthropic 400s.
+    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)
+
+
+def test_prompt_cache_can_be_disabled() -> None:
+    assert (
+        make_model_settings(
+            None, model_name="anthropic/claude-sonnet-4-5", prompt_cache=False
+        ).extra_args
+        is None
+    )
+
+
+@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:
+    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 model LiteLLM hasn't mapped must run uncached, not crash.
+    unmapped = "bedrock/global.anthropic.claude-brand-new-9"
+    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)
+
+    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:
+    # Only Bedrock hard-rejects unknown cache fields, so only Bedrock is guarded.
+    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)
+
+    for model in ("anthropic/claude-brand-new-9", "openrouter/anthropic/claude-brand-new"):
+        assert _cache_points(model) == [
+            {"location": "message", "role": "system"},
+            {"location": "message", "index": -1},
+        ]
+
+
+def test_conversation_tail_breakpoint_moves_with_appended_transcript() -> None:
+    # LiteLLM must place the index=-1 cache_control on the last message however
+    # long the transcript grows.
+    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({}) == ""
 

tests/test_runner_rate_limit.py

@@ -14,6 +14,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:
@@ -38,6 +39,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),
     )
@@ -56,8 +58,8 @@ async def test_persistent_rate_limit_stops_gracefully(
     async def _cleanup(*_args: Any, **_kwargs: Any) -> None:
         return None
 
-    monkeypatch.setattr(runner.session_manager, "create_or_reuse", _create_or_reuse)  # type: ignore[attr-defined]
-    monkeypatch.setattr(runner.session_manager, "cleanup", _cleanup)  # type: ignore[attr-defined]
+    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: "")

tests/test_runner_root_prompt.py

@@ -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)