From d159b6e0c3563208faf87b23d5fd56d46290ef3a Mon Sep 17 00:00:00 2001 From: Kyzcreig <9063726+Kyzcreig@users.noreply.github.com> Date: Fri, 25 Sep 2026 22:34:33 -0700 Subject: [PATCH] fix(cache): cache-parity forks get a derived cache scope on xAI #109964 made same-model cache-parity forks (background review, /btw) inherit the parent's resolved cache scope. That is the right call for content-addressed caches (Anthropic, DeepSeek, Gemini) and for OpenAI's routing-only prompt_cache_key. xAI is different: x-grok-conv-id / prompt_cache_key select ONE server-side conversation slot, so a fork of similar size that diverges early evicts the parent's slot and the parent's next call reads cold. Measured on grok-4.7 (xai-oauth Responses), ~160k context, fork of similar size diverging right after the system prompt, parent's next call after the fork: shared scope (current): 1,152 / 162,239 read (2/2 runs) derived scope (this PR): 162,176 / 162,239 read (2/2 runs) no fork (control): 177,536 / 177,639 read (2/2 runs) build_cache_parity_fork tags the fork (_prompt_cache_fork_tag). The resolver derives "::" for tagged agents on slot-keyed routes only (xai, xai-oauth, api.x.ai, OpenRouter x-ai/grok-*). The inherited scope is kept everywhere else. The tag is applied outside the memo, so a mid-run provider fallback re-evaluates. OpenRouter's Grok x-grok-conv-id now honours a fork scope over the ambient affinity/conversation scope (chat_completions threads cache_scope_id to the profile). session_id and transcript identity are unchanged. (cherry picked from commit 40fbb39d978648c76940cfb7835f77a42d3719b0) --- agent/background_review.py | 6 + agent/prompt_cache_scope.py | 51 ++++- agent/transports/chat_completions.py | 2 +- .../model-providers/openrouter/__init__.py | 5 + .../test_background_review_cache_parity.py | 29 +++ tests/agent/test_review_fork_cache_scope.py | 190 ++++++++++++++++++ 6 files changed, 279 insertions(+), 4 deletions(-) create mode 100644 tests/agent/test_review_fork_cache_scope.py diff --git a/agent/background_review.py b/agent/background_review.py index f346767634..8f6aa4f6c4 100644 --- a/agent/background_review.py +++ b/agent/background_review.py @@ -998,6 +998,12 @@ def build_cache_parity_fork( review_agent._end_session_on_close = False review_agent._session_db = None review_agent.session_id = agent.session_id + # Slot-keyed caches (xAI) must not see the fork under the parent's key: the fork's divergent + # stream evicts the parent's conversation slot. The resolver derives ``::`` for + # slot-keyed providers only; content-addressed ones keep the shared scope below. + review_agent._prompt_cache_fork_tag = ( + "review" if write_origin == "background_review" else str(write_origin or "fork") + ) # Same model only: share the warm cached system prompt (~26% cost cut; a rebuilt prompt misses # the byte-exact prefix key) and pin session_start so any re-render (compression, plugin # hooks) stays byte-identical. diff --git a/agent/prompt_cache_scope.py b/agent/prompt_cache_scope.py index 2efec6971e..b3c27e621f 100644 --- a/agent/prompt_cache_scope.py +++ b/agent/prompt_cache_scope.py @@ -146,7 +146,7 @@ def resolve_prompt_cache_scope(agent: Any) -> str: (the physical id without ancestry/DB). Memoized on the agent.""" inherited = getattr(agent, "_inherited_cache_scope", None) if isinstance(inherited, str) and inherited: - return inherited + return _apply_fork_tag(agent, inherited) sid = str(getattr(agent, "session_id", None) or "") if not sid: return "" @@ -155,7 +155,7 @@ def resolve_prompt_cache_scope(agent: Any) -> str: key = (sid, db is not None) memo = getattr(agent, _MEMO_ATTR, None) if isinstance(memo, tuple) and len(memo) == 2 and memo[0] == key: - return memo[1] + return _apply_fork_tag(agent, memo[1]) root = declared_conversation_scope(agent) or _lineage_root(sid, db) scope = root or sid # Memoize on success, with no DB, or when the agent never persists a row. A failed/empty @@ -166,7 +166,52 @@ def resolve_prompt_cache_scope(agent: Any) -> str: setattr(agent, _MEMO_ATTR, (key, scope)) except Exception: pass # frozen/slotted doubles: resolution works, just unmemoized - return scope + return _apply_fork_tag(agent, scope) + + +# ``::`` — double colon so gateway session keys that already contain single +# colons are never mistaken for a fork scope. +FORK_SCOPE_SEPARATOR = "::" + + +def is_slot_keyed_cache_route(provider: Any, model: Any, base_url: Any = "") -> bool: + """True when the cache key selects ONE server-side slot per conversation. + + xAI (``xai``/``xai-oauth``, ``api.x.ai``, or Grok via OpenRouter ``x-ai/grok-*``) pins its + prompt cache to the server picked by ``x-grok-conv-id`` / ``prompt_cache_key``: two divergent + request streams under one key evict each other. Anthropic, DeepSeek and Gemini caches are + content-addressed and OpenAI's ``prompt_cache_key`` only routes over a prefix match, so the + shared scope (#109964) stays a win there. + """ + if str(provider or "").strip().lower() in {"xai", "xai-oauth"}: + return True + if "api.x.ai" in str(base_url or "").lower(): + return True + return str(model or "").strip().lower().startswith(("x-ai/grok-", "xai/grok-")) + + +def is_fork_cache_scope(scope: Any) -> bool: + """True when *scope* is a fork-derived scope (``::``).""" + return isinstance(scope, str) and FORK_SCOPE_SEPARATOR in scope + + +def _apply_fork_tag(agent: Any, scope: str) -> str: + """Derive ``::`` for a tagged cache-parity fork on a slot-keyed provider. + + A same-model fork shares the parent's scope (#109964) so content-addressed caches serve it + warm. On xAI that key makes the fork's divergent stream evict the parent's conversation slot: + measured on grok-4.7 at ~160k context with a same-size fork, the parent's next call read + 1,152 of 162k prompt tokens (2/2) vs 162,176 of 162k (2/2) with a derived key. Evaluated per + call (no DB), so a mid-run provider fallback re-evaluates. + """ + tag = getattr(agent, "_prompt_cache_fork_tag", None) + if not scope or not isinstance(tag, str) or not tag: + return scope + if not is_slot_keyed_cache_route( + getattr(agent, "provider", ""), getattr(agent, "model", ""), getattr(agent, "base_url", ""), + ): + return scope + return f"{scope}{FORK_SCOPE_SEPARATOR}{tag}" def declared_conversation_scope_safe(agent: Any) -> Optional[str]: diff --git a/agent/transports/chat_completions.py b/agent/transports/chat_completions.py index 8f3ae66712..51e0dbdcaf 100644 --- a/agent/transports/chat_completions.py +++ b/agent/transports/chat_completions.py @@ -562,7 +562,7 @@ class ChatCompletionsTransport(ProviderTransport): reasoning_config=reasoning_config, supports_reasoning=params.get("supports_reasoning", False), qwen_session_metadata=params.get("qwen_session_metadata"), model=model, base_url=params.get("base_url"), ollama_num_ctx=params.get("ollama_num_ctx"), - session_id=params.get("session_id"), + session_id=params.get("session_id"), cache_scope_id=params.get("cache_scope_id"), ) api_kwargs.update(top_level_from_profile) diff --git a/plugins/model-providers/openrouter/__init__.py b/plugins/model-providers/openrouter/__init__.py index a81ce3a12c..09716938cf 100644 --- a/plugins/model-providers/openrouter/__init__.py +++ b/plugins/model-providers/openrouter/__init__.py @@ -4,6 +4,7 @@ import logging from typing import Any from agent.portal_tags import get_affinity_scope, get_conversation_context +from agent.prompt_cache_scope import is_fork_cache_scope from agent.transports.codex import _cache_scope_from_session_id from providers import register_provider from providers.base import ProviderProfile @@ -190,6 +191,10 @@ class OpenRouterProfile(ProviderProfile): extra_body["reasoning"] = {"enabled": True, "effort": "medium"} # xAI's prompt cache is pinned per backend server via this header. grok_conv_id = _sticky_key(session_id) + # A cache-parity fork carries the parent's ambient scope; on Grok that key would evict + # the parent's server slot, so honour the fork-derived scope (agent/prompt_cache_scope.py). + if is_fork_cache_scope(context.get("cache_scope_id")): + grok_conv_id = context["cache_scope_id"] if grok_conv_id and model and model.startswith(("x-ai/grok-", "xai/grok-")): top_level["extra_headers"] = {"x-grok-conv-id": grok_conv_id} return extra_body, top_level diff --git a/tests/agent/test_background_review_cache_parity.py b/tests/agent/test_background_review_cache_parity.py index 5d5824b2b0..5218f9f13a 100644 --- a/tests/agent/test_background_review_cache_parity.py +++ b/tests/agent/test_background_review_cache_parity.py @@ -577,3 +577,32 @@ def test_routed_fork_does_not_inherit_cache_scope(): "Routed fork must not inherit the parent's cache scope — its prefix " "is cache-cold on the different model regardless." ) + + +def test_same_model_fork_on_xai_gets_derived_cache_scope(): + """xAI's cache key selects ONE server slot: a fork under the parent's inherited scope + evicts the parent (grok-4.7, ~160k context: the parent's next call read 1,152 of 162k). + The fork is tagged and the resolver derives ``::`` on xAI only; on + content-addressed providers the #109964 inheritance is unchanged.""" + import run_agent + from agent.background_review import build_cache_parity_fork + from agent.prompt_cache_scope import resolve_prompt_cache_scope + + for write_origin, tag in (("background_review", "review"), ("side_question", "side_question")): + agent = _make_agent_stub(run_agent.AIAgent) + agent.provider, agent.model = "xai-oauth", "grok-4.7" + captured = {} + with patch.object(run_agent, "AIAgent", _make_recorder_class(captured)): + fork, _rt, routed = build_cache_parity_fork( + agent, max_iterations=5, write_origin=write_origin + ) + + assert not routed + fork.provider = captured["init_kwargs"].get("provider") + fork.model = captured["init_kwargs"].get("model") + parent_scope = resolve_prompt_cache_scope(agent) + assert fork.session_id == agent.session_id # transcript identity unchanged + assert fork._prompt_cache_fork_tag == tag + assert resolve_prompt_cache_scope(fork) == f"{parent_scope}::{tag}" + fork.provider, fork.model = "anthropic", "claude-opus-4-8" + assert resolve_prompt_cache_scope(fork) == parent_scope diff --git a/tests/agent/test_review_fork_cache_scope.py b/tests/agent/test_review_fork_cache_scope.py new file mode 100644 index 0000000000..36addd4b7e --- /dev/null +++ b/tests/agent/test_review_fork_cache_scope.py @@ -0,0 +1,190 @@ +"""Review/side-question forks get their OWN cache scope on slot-keyed caches. + +A background-review fork shares the parent's ``session_id`` so content- +addressed caches (Anthropic, DeepSeek, Gemini) and OpenAI-style routing keys +serve it warm. xAI is different: ``x-grok-conv-id`` / ``prompt_cache_key`` +select ONE server-side conversation slot, so the fork's divergent request +stream evicts the parent's — measured on xai-oauth, the parent's next call +read 1,152 of ~356k prompt tokens after a fork. The resolver derives +``::`` for tagged forks on xAI routes only. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from agent.prompt_cache_scope import ( + is_fork_cache_scope, + is_slot_keyed_cache_route, + resolve_prompt_cache_scope, +) + +SYS = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, +] + + +def _agent(provider, model="grok-4.3", tag=None, base_url=""): + a = SimpleNamespace( + session_id="parent-sess", + _session_db=None, + provider=provider, + model=model, + base_url=base_url, + ) + if tag is not None: + a._prompt_cache_fork_tag = tag + return a + + +class TestSlotKeyedRoute: + @pytest.mark.parametrize( + "provider,model,base_url", + [ + ("xai-oauth", "grok-4.3", ""), + ("xai", "grok-4", ""), + ("custom", "grok-4", "https://api.x.ai/v1"), + ("openrouter", "x-ai/grok-4", ""), + ], + ) + def test_xai_routes_are_slot_keyed(self, provider, model, base_url): + assert is_slot_keyed_cache_route(provider, model, base_url) + + @pytest.mark.parametrize( + "provider,model", + [ + ("anthropic", "claude-opus-4-8"), + ("claude-bpr", "claude-opus-4-8"), + ("openai-codex", "gpt-5.5"), + ("openrouter", "anthropic/claude-sonnet-4.5"), + ("gemini", "gemini-3-pro"), + ], + ) + def test_content_addressed_and_routing_providers_are_not(self, provider, model): + assert not is_slot_keyed_cache_route(provider, model, "") + + +class TestResolverDerivesForkScope: + def test_parent_on_xai_keeps_its_scope(self): + assert resolve_prompt_cache_scope(_agent("xai-oauth")) == "parent-sess" + + def test_review_fork_on_xai_gets_derived_scope(self): + scope = resolve_prompt_cache_scope(_agent("xai-oauth", tag="review")) + assert scope == "parent-sess::review" + assert is_fork_cache_scope(scope) + + def test_side_question_fork_gets_its_own_tag(self): + assert ( + resolve_prompt_cache_scope(_agent("xai-oauth", tag="side_question")) + == "parent-sess::side_question" + ) + + @pytest.mark.parametrize( + "provider,model", + [("anthropic", "claude-opus-4-8"), ("openai-codex", "gpt-5.5"), + ("openrouter", "anthropic/claude-sonnet-4.5")], + ) + def test_fork_on_shared_key_provider_keeps_parent_scope(self, provider, model): + # #109964-style sharing stays where it is a measured win. + assert ( + resolve_prompt_cache_scope(_agent(provider, model, tag="review")) + == "parent-sess" + ) + + def test_memo_is_not_polluted_and_fallback_is_reevaluated(self): + fork = _agent("xai-oauth", tag="review") + assert resolve_prompt_cache_scope(fork) == "parent-sess::review" + assert resolve_prompt_cache_scope(fork) == "parent-sess::review" + fork.provider, fork.model = "anthropic", "claude-opus-4-8" # fallback hop + assert resolve_prompt_cache_scope(fork) == "parent-sess" + + def test_inherited_parent_scope_is_derived_on_xai_only(self): + """#109964 inheritance stays for content-addressed providers, derives on xAI.""" + fork = _agent("xai-oauth", tag="review") + fork._inherited_cache_scope = "gwk_parentscope" + assert resolve_prompt_cache_scope(fork) == "gwk_parentscope::review" + fork.provider, fork.model = "anthropic", "claude-opus-4-8" + assert resolve_prompt_cache_scope(fork) == "gwk_parentscope" + + def test_gateway_style_single_colon_ids_are_not_fork_scopes(self): + assert not is_fork_cache_scope("agent:main:telegram:dm:42") + + +class TestTransportWire: + """What actually goes on the wire for parent vs fork.""" + + def _xai(self, agent): + from agent.transports.codex import ResponsesApiTransport + + return ResponsesApiTransport().build_kwargs( + model="grok-4.3", messages=SYS, tools=[], + session_id=agent.session_id, + cache_scope_id=resolve_prompt_cache_scope(agent), + is_xai_responses=True, + ) + + def test_xai_fork_sends_distinct_conv_id_and_cache_key(self): + parent = self._xai(_agent("xai-oauth")) + fork = self._xai(_agent("xai-oauth", tag="review")) + assert parent["extra_headers"]["x-grok-conv-id"] == "parent-sess" + assert fork["extra_headers"]["x-grok-conv-id"] == "parent-sess::review" + assert ( + parent["extra_body"]["prompt_cache_key"] + != fork["extra_body"]["prompt_cache_key"] + ) + + def test_negative_control_untagged_fork_collides(self): + """Pre-fix shape: same scope -> same slot -> the eviction.""" + parent = self._xai(_agent("xai-oauth")) + fork = self._xai(_agent("xai-oauth")) + assert parent["extra_headers"]["x-grok-conv-id"] == fork["extra_headers"]["x-grok-conv-id"] + assert parent["extra_body"]["prompt_cache_key"] == fork["extra_body"]["prompt_cache_key"] + + def test_codex_fork_keeps_shared_prompt_cache_key(self): + from agent.transports.codex import ResponsesApiTransport + + def build(agent): + return ResponsesApiTransport().build_kwargs( + model="gpt-5.5", messages=SYS, tools=[], + session_id=agent.session_id, + cache_scope_id=resolve_prompt_cache_scope(agent), + is_codex_backend=True, + ) + + parent = build(_agent("openai-codex", "gpt-5.5")) + fork = build(_agent("openai-codex", "gpt-5.5", tag="review")) + assert parent["prompt_cache_key"] == fork["prompt_cache_key"] + + def test_openrouter_grok_fork_overrides_ambient_conversation(self): + from agent.portal_tags import reset_conversation_context, set_conversation_context + from providers import get_provider_profile + + p = get_provider_profile("openrouter") + token = set_conversation_context("parent-sess") # fork copies parent's Context + try: + _, parent = p.build_api_kwargs_extras( + model="x-ai/grok-4", session_id="parent-sess", + cache_scope_id="parent-sess", + ) + _, fork = p.build_api_kwargs_extras( + model="x-ai/grok-4", session_id="parent-sess", + cache_scope_id="parent-sess::review", + ) + finally: + reset_conversation_context(token) + assert parent["extra_headers"]["x-grok-conv-id"] == "parent-sess" + assert fork["extra_headers"]["x-grok-conv-id"] == "parent-sess::review" + + def test_chat_completions_threads_cache_scope_to_profile(self): + from agent.transports.chat_completions import ChatCompletionsTransport + from providers import get_provider_profile + + kwargs = ChatCompletionsTransport().build_kwargs( + model="x-ai/grok-4", messages=SYS, tools=None, + provider_profile=get_provider_profile("openrouter"), + session_id="parent-sess", cache_scope_id="parent-sess::review", + ) + assert kwargs["extra_headers"]["x-grok-conv-id"] == "parent-sess::review"