diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 4dd5dbed95..6c645c17dd 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -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, diff --git a/agent/prompt_cache_scope.py b/agent/prompt_cache_scope.py index b9c8e86bc9..6d23e3d954 100644 --- a/agent/prompt_cache_scope.py +++ b/agent/prompt_cache_scope.py @@ -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. diff --git a/agent/transports/codex.py b/agent/transports/codex.py index 1370bdabcc..bf73ba971f 100644 --- a/agent/transports/codex.py +++ b/agent/transports/codex.py @@ -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 diff --git a/agent/turn_context.py b/agent/turn_context.py index 60c9279a79..0998ae2213 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -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: diff --git a/tests/agent/test_prompt_cache_scope.py b/tests/agent/test_prompt_cache_scope.py index 4088769b70..c7147fa9b5 100644 --- a/tests/agent/test_prompt_cache_scope.py +++ b/tests/agent/test_prompt_cache_scope.py @@ -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):