diff --git a/strix/config/models.py b/strix/config/models.py index 38234d97..b6e6b0f0 100644 --- a/strix/config/models.py +++ b/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 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 de922d07..aef1fe13 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,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, 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 bde83bd9..da914879 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 @@ -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({}) == "" diff --git a/tests/test_runner_rate_limit.py b/tests/test_runner_rate_limit.py index 62242339..482c4730 100644 --- a/tests/test_runner_rate_limit.py +++ b/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: "") 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)