refactor(cache): fold self-review findings on the rotation-scope fix

- prompt_cache_scope: memo key now includes DB presence (a lazily attached
  _session_db re-resolves instead of staying pinned to the physical id);
  _persist_disabled agents (background-review forks that never get a DB row)
  memoize the fallback instead of re-querying the lineage per API call;
  module docstring cross-references get_conversation_root and why the two
  lineage resolvers must not be deduplicated.
- chat_completion_helpers: hoist the triplicated
  _prompt_cache_scope_for_agent(agent) call to a single local above the
  OpenAI-wire dispatch (after the anthropic/bedrock early returns, which
  don't use prompt_cache_key).
- codex transport docstring: x-client-request-id mirrors the derived body
  key, not the raw scope id.
- turn_context comment: acknowledge the first-turn pre-persist fallback.
- tests: +2 (persist-disabled memoization; lazy DB attach re-resolution).
This commit is contained in:
kshitij
2026-08-15 10:55:29 +05:30
parent cee2446222
commit 96cdf19a0b
5 changed files with 75 additions and 22 deletions
+9 -3
View File
@@ -1832,6 +1832,12 @@ def build_api_kwargs(agent, api_messages: list, tools_for_api: list | None = Non
guardrail_config=guardrail,
)
# Rotation-stable logical cache scope, shared by every OpenAI-wire branch
# below (codex + both chat_completions paths). Memoized on the agent —
# cheap after the first call. Resolved after the anthropic/bedrock early
# returns above, which don't use prompt_cache_key.
_cache_scope_id = _prompt_cache_scope_for_agent(agent)
if agent.api_mode == "codex_responses":
_ct = agent._get_transport()
is_github_responses = (
@@ -1897,7 +1903,7 @@ def build_api_kwargs(agent, api_messages: list, tools_for_api: list | None = Non
tools=tools_for_api,
reasoning_config=agent.reasoning_config,
session_id=getattr(agent, "session_id", None),
cache_scope_id=_prompt_cache_scope_for_agent(agent),
cache_scope_id=_cache_scope_id,
base_url=agent.base_url,
max_tokens=agent.max_tokens,
timeout=agent._resolved_api_call_timeout(),
@@ -2007,7 +2013,7 @@ def build_api_kwargs(agent, api_messages: list, tools_for_api: list | None = Non
reasoning_config=agent.reasoning_config,
request_overrides=agent.request_overrides,
session_id=getattr(agent, "session_id", None),
cache_scope_id=_prompt_cache_scope_for_agent(agent),
cache_scope_id=_cache_scope_id,
provider_profile=_profile,
ollama_num_ctx=agent._ollama_num_ctx,
# Context forwarded to profile hooks:
@@ -2040,7 +2046,7 @@ def build_api_kwargs(agent, api_messages: list, tools_for_api: list | None = Non
reasoning_config=agent.reasoning_config,
request_overrides=agent.request_overrides,
session_id=getattr(agent, "session_id", None),
cache_scope_id=_prompt_cache_scope_for_agent(agent),
cache_scope_id=_cache_scope_id,
model_lower=(agent.model or "").lower(),
is_openrouter=_is_or,
is_nous=_is_nous,
+27 -11
View File
@@ -10,7 +10,12 @@ bucket even though it is logically the same conversation continuing
``resolve_prompt_cache_scope()`` maps the physical session id to the ROOT of
its *compression lineage* — the pre-rotation session id — using
``SessionDB.get_compression_lineage()``, whose fork-aware semantics
(hardened in #79193) give exactly the scope boundaries the cache key needs:
(hardened in #79193) give exactly the scope boundaries the cache key needs.
NOT ``SessionDB.get_conversation_root`` / ``run_agent._conversation_root_id``
(the Portal-attribution walk): that one follows ``parent_session_id`` blindly,
collapsing /branch children and whole delegate trees into one id, which would
violate the #79161 isolation this scope must preserve. The two resolvers are
intentionally different — do not "deduplicate" them.
- compression-rotation children walk back to the original segment
(rotation-stable scope — the fix);
@@ -71,20 +76,31 @@ def resolve_prompt_cache_scope(agent: Any) -> str:
sid = str(getattr(agent, "session_id", None) or "")
if not sid:
return ""
memo = getattr(agent, _MEMO_ATTR, None)
if isinstance(memo, tuple) and len(memo) == 2 and memo[0] == sid:
return memo[1]
db = getattr(agent, "_session_db", None)
# Memo key includes DB presence: an agent that starts DB-less and gains a
# handle later (run_agent._get_session_db_for_recall lazily attaches one)
# must re-resolve instead of staying pinned to the physical id.
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]
root = _lineage_root(sid, db) if db is not None else None
scope = root or sid
# Memoize on a successful walk, or when there is no DB to consult at all.
# A failed/empty walk (row not persisted yet, transient DB error) is NOT
# memoized: falling back to the physical id is the correct degraded
# answer right now, but pinning it for the whole segment would keep the
# scope wrong after the session row lands.
if root is not None or db is None:
# Memoize on a successful walk, or when there is no DB to consult at all,
# or when the agent will never persist a row (background-review forks set
# _persist_disabled but still hold a DB handle — without this, every API
# call would re-run the lineage query forever).
# A failed/empty walk on a persisting agent is NOT memoized: falling back
# to the physical id is the correct degraded answer right now (row not
# persisted yet, transient DB error), but pinning it for the whole segment
# would keep the scope wrong after the session row lands.
if (
root is not None
or db is None
or getattr(agent, "_persist_disabled", False)
):
try:
setattr(agent, _MEMO_ATTR, (sid, scope))
setattr(agent, _MEMO_ATTR, (key, scope))
except Exception:
# Frozen/slotted test doubles — resolution still works, just
# unmemoized.
+5 -4
View File
@@ -347,10 +347,11 @@ class ResponsesApiTransport(ProviderTransport):
``cache_scope_id`` is given
cache_scope_id: str | None — rotation-stable logical scope id
(compression-lineage root; see agent/prompt_cache_scope.py).
Preferred over session_id for prompt_cache_key derivation and
the xAI x-grok-conv-id / Codex x-client-request-id routing
headers, so the cache stays warm across context-compression
session rotation (#79017)
Preferred over session_id when deriving the prompt_cache_key
content hash and the xAI x-grok-conv-id header; the Codex
x-client-request-id header mirrors the resulting body key.
Keeps the cache warm across context-compression session
rotation (#79017)
max_tokens: int | None — max_output_tokens
timeout: float | None — per-request timeout forwarded to the SDK
request_overrides: dict | None — extra kwargs merged in
+7 -4
View File
@@ -495,10 +495,13 @@ def build_turn_context(
api_mode=getattr(agent, "api_mode", "") or "",
auth_mode=getattr(agent, "auth_mode", "") or "",
session_id=getattr(agent, "session_id", "") or "",
# Rotation-stable prompt-cache scope, resolved once per turn here
# (memoized per segment — no DB walk on the per-call hot path).
# Stays valid through a mid-turn compression rotation because the
# lineage root is by definition rotation-invariant (#79017).
# Rotation-stable prompt-cache scope. Memoized per segment on the
# agent, so this is a DB walk at most once per segment — except a
# brand-new session whose row lands later in turn setup
# (_ensure_db_session); that first turn falls back to the physical
# id here and the first build_api_kwargs re-resolves. Stays valid
# through a mid-turn compression rotation because the lineage root
# is by definition rotation-invariant (#79017).
cache_scope=resolve_prompt_cache_scope(agent),
)
except Exception:
+27
View File
@@ -157,6 +157,33 @@ class TestResolvePromptCacheScope:
db.create_session("late-row", source="webui", parent_session_id="late-root")
assert resolve_prompt_cache_scope(agent) == "late-root"
def test_persist_disabled_agent_is_memoized_despite_missing_row(self, db):
"""Background-review forks (_persist_disabled) never get a DB row —
they must memoize the fallback instead of re-querying per API call."""
agent = _agent("review-fork", db)
agent._persist_disabled = True
assert resolve_prompt_cache_scope(agent) == "review-fork"
calls = []
original = db.get_compression_lineage
db.get_compression_lineage = lambda sid: calls.append(sid) or original(sid)
try:
assert resolve_prompt_cache_scope(agent) == "review-fork"
assert calls == [] # memoized — no per-call re-query
finally:
db.get_compression_lineage = original
def test_db_attached_later_re_resolves(self, db):
"""A DB-less memo must not survive a lazy _session_db attach."""
db.create_session("root-sess", source="webui")
_rotate(db, "root-sess", "rotated-1")
agent = _agent("rotated-1", None)
# No DB -> physical id, memoized for the DB-less state.
assert resolve_prompt_cache_scope(agent) == "rotated-1"
# Lazy attach (run_agent._get_session_db_for_recall pattern).
agent._session_db = db
assert resolve_prompt_cache_scope(agent) == "root-sess"
def test_bogus_lineage_shape_falls_back(self):
class WeirdDB:
def get_compression_lineage(self, sid):