From 460ceab075d7727fa9c11fbeca1c4a9693fea1b3 Mon Sep 17 00:00:00 2001 From: Ahmed Allam Date: Sun, 26 Jul 2026 23:29:41 +0000 Subject: [PATCH] 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). --- strix/config/models.py | 64 ++++++++++ strix/config/settings.py | 4 + strix/core/inputs.py | 213 +++++++++++-------------------- strix/core/runner.py | 1 + tests/test_inputs.py | 63 ++++++--- tests/test_runner_rate_limit.py | 1 + tests/test_runner_root_prompt.py | 6 +- 7 files changed, 198 insertions(+), 154 deletions(-) diff --git a/strix/config/models.py b/strix/config/models.py index 38234d97..e38b0009 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -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 diff --git a/strix/config/settings.py b/strix/config/settings.py index 0d910e60..78d52273 100644 --- a/strix/config/settings.py +++ b/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") diff --git a/strix/core/inputs.py b/strix/core/inputs.py index 629883fd..4d6efe79 100644 --- a/strix/core/inputs.py +++ b/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,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( diff --git a/strix/core/runner.py b/strix/core/runner.py index 5776ce42..77a4eff1 100644 --- a/strix/core/runner.py +++ b/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, diff --git a/tests/test_inputs.py b/tests/test_inputs.py index 6169a7a8..dc363abd 100644 --- a/tests/test_inputs.py +++ b/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 @@ -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}, ] diff --git a/tests/test_runner_rate_limit.py b/tests/test_runner_rate_limit.py index 62242339..479c8d4e 100644 --- a/tests/test_runner_rate_limit.py +++ b/tests/test_runner_rate_limit.py @@ -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), ) diff --git a/tests/test_runner_root_prompt.py b/tests/test_runner_root_prompt.py index aaa8ef6b..56d7caa6 100644 --- a/tests/test_runner_root_prompt.py +++ b/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)