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:
@@ -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
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user