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 "<scope>::<tag>" 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)
This commit is contained in:
Kyzcreig
2026-09-25 22:34:33 -07:00
committed by kshitij
parent 3307336404
commit d159b6e0c3
6 changed files with 279 additions and 4 deletions

View File

@@ -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 ``<scope>::<tag>`` 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.

View File

@@ -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)
# ``<scope>::<tag>`` — 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 (``<scope>::<tag>``)."""
return isinstance(scope, str) and FORK_SCOPE_SEPARATOR in scope
def _apply_fork_tag(agent: Any, scope: str) -> str:
"""Derive ``<scope>::<tag>`` 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]:

View File

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

View File

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

View File

@@ -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 ``<parent scope>::<tag>`` 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

View File

@@ -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
``<scope>::<tag>`` 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"