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:
@@ -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.
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
190
tests/agent/test_review_fork_cache_scope.py
Normal file
190
tests/agent/test_review_fork_cache_scope.py
Normal 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"
|
||||
Reference in New Issue
Block a user