diff --git a/strix/core/inputs.py b/strix/core/inputs.py index e53ae321a..e7b091bb3 100644 --- a/strix/core/inputs.py +++ b/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, diff --git a/tests/test_inputs.py b/tests/test_inputs.py index 486c7e549..9f7d64890 100644 --- a/tests/test_inputs.py +++ b/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({}) == ""