diff --git a/agent/prompt_cache_scope.py b/agent/prompt_cache_scope.py index d3d464cea7..46dd0f7e31 100644 --- a/agent/prompt_cache_scope.py +++ b/agent/prompt_cache_scope.py @@ -101,11 +101,17 @@ def _agent_source(agent: Any, session_id: str, session_db: Any) -> str: """The ``sessions.source`` this agent's conversation is recorded under. Read from the agent's own row when it exists, because that is the value - the peer queries below match on. Before the row lands (the first turn - resolves a scope ahead of ``_ensure_db_session``) it falls back to the - platform, which is what the row will be created with — the two can differ - only when ``HERMES_SESSION_SOURCE`` overrides the platform, and the cost of - that is one cold bucket on the first turn, never a crossed identity. + the peer queries below match on. + + Before the row lands — this module resolves the first scope ahead of + ``_ensure_db_session`` — it uses the SAME resolver persistence will use, + ``run_agent._session_source_for_agent``, not ``agent.platform``. The two + diverge whenever ``HERMES_SESSION_SOURCE`` overrides the platform, and the + divergence is not a cosmetic one: the declared scope is non-``None`` + immediately, so ``resolve_prompt_cache_scope`` memoizes it for this session + id and never re-resolves once the authoritative row appears. Both sides of + a ``/new`` would then read the platform domain, miss the boundary recorded + under the override, and hash the same scope. """ if session_id and session_db is not None: try: @@ -117,7 +123,18 @@ def _agent_source(agent: Any, session_id: str, session_db: Any) -> str: source = str(row.get("source") or "").strip() if source: return source - return str(getattr(agent, "platform", "") or "").strip() + platform = getattr(agent, "platform", None) + try: + # Imported lazily: run_agent imports this module, and this is the + # single owner of the source a session row is created with. + from run_agent import _session_source_for_agent + + source = str(_session_source_for_agent(platform) or "").strip() + if source: + return source + except Exception: + logger.debug("declared-scope source authority unavailable", exc_info=True) + return str(platform or "").strip() def _conversation_generation(session_key: str, source: str, session_db: Any) -> str: @@ -141,11 +158,13 @@ def _conversation_generation(session_key: str, source: str, session_db: Any) -> - rotating on every conversation replacement, ``/new`` and the policy auto-resets alike. - The marker pairs the boundary COUNT with the latest ``ended_at`` because - each alone can repeat a previous generation under a different rare - condition — a backwards clock correction defeats the timestamp, retention - pruning defeats the count — and the two do not fail together (see - ``SessionDB.latest_conversation_boundary``). + The marker is a durable counter kept outside prunable session history + (``conversation_generations``, advanced in the same transaction that writes + each boundary). Deriving it from the session rows instead — an aggregate + over ``_RESET_END_REASONS`` boundaries — cannot prove non-reuse: deleting or + pruning an ended row makes the aggregate return a pair it already emitted, + handing a new conversation a retired affinity identity. It is also + wall-clock-free, so a backwards NTP correction cannot reorder it. No counter is introduced anywhere: the marker is read from state the reset paths already write, and it is read on the memoized resolution path, @@ -157,13 +176,10 @@ def _conversation_generation(session_key: str, source: str, session_db: Any) -> reader = getattr(session_db, "latest_conversation_boundary", None) if not callable(reader): return "" - boundary = reader(session_key, source) - if boundary is None: + generation = reader(session_key, source) + if generation is None: return "" - # (crossings, ended_at). Fixed-point on the timestamp so the carrier is - # byte-identical across repr differences between platforms. - crossings, ended_at = boundary - return f"{int(crossings)}:{float(ended_at):.6f}" + return str(int(generation)) def declared_conversation_scope(agent: Any) -> Optional[str]: diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index 3ece6b96c9..8a297354bd 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -7620,11 +7620,10 @@ class APIServerAdapter(BasePlatformAdapter): # (/v1/responses, /v1/runs) record one, so no other # caller's rows change shape. if bind_declared_conversation: - if _declared_selected: - self._bind_declared_conversation( - getattr(agent, "session_id", None) or session_id, - gateway_session_key, - ) + self._bind_declared_conversation( + getattr(agent, "session_id", None) or session_id, + gateway_session_key, + ) clear_session_vars(tokens) self._activate_admitted_request() diff --git a/hermes_state.py b/hermes_state.py index 7314ca3d85..c14a0d7820 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -7860,6 +7860,45 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) self._execute_write(_do) + def _bump_conversation_generation(self, conn, session_id: str, end_reason: str) -> None: + """Advance this peer's conversation generation past a boundary. + + Called inside the transaction that writes the boundary, so the + generation and the ``end_reason`` that caused it commit together. + + Only ``_RESET_END_REASONS`` count: ``compression`` continues one + conversation, and an accidental close is not a replacement. Rows with + no ``session_key`` have no routing peer to advance. + + The counter deliberately does NOT read the session rows. An aggregate + over them (COUNT/MAX of boundaries) can return a pair it already + emitted once ``delete_session()`` or bulk pruning removes an ended row, + which would hand a new conversation a retired affinity identity. This + value only ever increments, so a generation is never reused for a peer + even if every row behind it is gone. + """ + if end_reason not in _RESET_END_REASONS: + return + row = conn.execute( + "SELECT source, session_key FROM sessions WHERE id = ?", + (session_id,), + ).fetchone() + if row is None: + return + source = str(row["source"] or "").strip() + session_key = str(row["session_key"] or "").strip() + if not source or not session_key: + return + conn.execute( + """ + INSERT INTO conversation_generations (source, session_key, generation) + VALUES (?, ?, 1) + ON CONFLICT(source, session_key) DO UPDATE + SET generation = conversation_generations.generation + 1 + """, + (source, session_key), + ) + def end_session(self, session_id: str, end_reason: str) -> None: """Mark a session as ended. @@ -7871,11 +7910,15 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) intentionally need to re-end a closed session with a new reason. """ def _do(conn): - conn.execute( + changed = conn.execute( "UPDATE sessions SET ended_at = ?, end_reason = ? " "WHERE id = ? AND ended_at IS NULL", (time.time(), end_reason, session_id), - ) + ).rowcount + # Only a boundary this call actually wrote advances the generation: + # the first end_reason wins, so a no-op must not rotate the peer. + if changed: + self._bump_conversation_generation(conn, session_id, end_reason) self._execute_write(_do) def reopen_session(self, session_id: str) -> None: @@ -7946,6 +7989,11 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) f"OR end_reason IN ({_RECOVERABLE_END_REASONS_SQL}))", (now, reason, session_id), ) + # /new and the policy auto-resets promote rather than end_session, + # so the generation has to advance here too — in the same + # transaction as the boundary, and only when one was written. + if cursor.rowcount: + self._bump_conversation_generation(conn, session_id, reason) return cursor.rowcount try: @@ -14136,8 +14184,8 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) def latest_conversation_boundary( self, session_key: str, source: str - ) -> Optional[Tuple[int, float]]: - """How many conversation boundaries this peer has crossed, and when. + ) -> Optional[int]: + """How many conversation boundaries this routing peer has crossed. A boundary is a row this peer ended at an intentional conversation break — the ``_RESET_END_REASONS`` set (``/new``, ``/switch``, idle, @@ -14154,41 +14202,34 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) this conversation's affinity identity while recovery correctly refuses to cross the same line. - Returns ``(count, latest_ended_at)``, or ``None`` when the key has - never been reset. BOTH halves are reported because each one alone has - a narrow way to repeat a previous generation: + Returns the count, or ``None`` when this peer has never been reset. - - ``ended_at`` is wall-clock, so a backwards NTP correction between two - resets writes a SMALLER boundary and ``MAX`` keeps returning the older - one — the next conversation would reuse the previous generation; - - ``count`` survives that, but retention pruning of an old ended row - decrements it. + The value comes from ``conversation_generations``, which + :meth:`_bump_conversation_generation` advances inside the transaction + that writes each boundary — NOT from an aggregate over the session + rows. An aggregate cannot prove non-reuse: ``delete_session()`` + orphans children and deletes the row, and bulk prune selects ended + rows, so ``COUNT``/``MAX`` over boundaries can return a pair it already + emitted and hand a new conversation a retired affinity identity. It is + also wall-clock-free, so a backwards NTP correction cannot reorder it. - The two fail under different conditions, so the pair only repeats a - generation if both happen at once. A spurious change is merely one cold - prompt-cache bucket; a repeated one would put two conversations on the - same routing key, so the pair is biased toward changing. Read-only; - ``agent/prompt_cache_scope.py`` uses it to keep a host-declared - conversation key from outliving the conversation it names. + Databases upgraded mid-conversation start at no generation and take + their first one from the next boundary written; a conversation that + reset before the upgrade shares its predecessor's scope once, which + costs a warm prompt-cache bucket and never crosses an identity. """ if not session_key or not source: return None with self._read_ctx() as conn: row = conn.execute( - f""" - SELECT COUNT(*) AS crossings, MAX(ended_at) AS boundary - FROM sessions - WHERE session_key = ? - AND source = ? - AND ended_at IS NOT NULL - AND end_reason IN ({_RESET_END_REASONS_SQL}) - """, - (session_key, source), + "SELECT generation FROM conversation_generations " + "WHERE source = ? AND session_key = ?", + (source, session_key), ).fetchone() - boundary = row["boundary"] if row is not None else None - if boundary is None: + if row is None or row["generation"] is None: return None - return int(row["crossings"] or 0), float(boundary) + generation = int(row["generation"]) + return generation if generation > 0 else None def _is_compression_child_row(self, child: Dict[str, Any]) -> bool: parent_id = child.get("parent_session_id") diff --git a/hermes_state_common.py b/hermes_state_common.py index d7f70bcf54..736ed18ee0 100644 --- a/hermes_state_common.py +++ b/hermes_state_common.py @@ -520,6 +520,27 @@ CREATE TABLE IF NOT EXISTS gateway_hygiene_state ( failure_streak INTEGER NOT NULL DEFAULT 0 ); +-- Monotonic conversation generation per routing peer (#96811). +-- +-- A host-declared conversation key (X-Hermes-Session-Key / build_session_key) +-- is per-CHAT and outlives any single conversation on it, so the prompt-cache +-- affinity scope derived from it must be qualified by which conversation is +-- currently live. Deriving that from the session rows themselves +-- (COUNT/MAX over _RESET_END_REASONS boundaries) cannot prove non-reuse: +-- delete_session() and bulk pruning remove ended rows, so an aggregate can +-- return a pair it already emitted and hand a new conversation a retired +-- affinity identity. +-- +-- This counter lives outside prunable session history and only ever +-- increments, once per boundary actually written, so a generation can never +-- be reused for a peer even if every session row behind it is deleted. +CREATE TABLE IF NOT EXISTS conversation_generations ( + source TEXT NOT NULL, + session_key TEXT NOT NULL, + generation INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY (source, session_key) +); + -- Per-backend liveness heartbeat (#94895). Each serve / tui_gateway process -- registers a row at startup and refreshes ``last_heartbeat`` periodically. -- The startup orphan sweep (sessions.startup_orphan_reap) consults this diff --git a/tests/agent/test_declared_conversation_scope.py b/tests/agent/test_declared_conversation_scope.py index 53326a6c73..e658ec3d32 100644 --- a/tests/agent/test_declared_conversation_scope.py +++ b/tests/agent/test_declared_conversation_scope.py @@ -589,3 +589,153 @@ class TestPeerIdentityIsSourceQualified: assert db.latest_conversation_boundary(self.KEY, "telegram") is not None assert db.latest_conversation_boundary(self.KEY, "api_server") is None assert db.latest_conversation_boundary(self.KEY, "") is None + + +class TestGenerationSurvivesPruning: + """A generation derived from prunable rows cannot prove non-reuse. + + `delete_session()` orphans surviving children and deletes the selected + row, and bulk prune selects ended rows, so an aggregate over + `_RESET_END_REASONS` boundaries can return a pair it already emitted: + `(1, T1) -> (2, T2) -> delete boundary B -> (1, T1)`, handing a new + conversation a retired affinity identity (@andrexibiza on #98811). + + The counter therefore lives in `conversation_generations`, outside session + history, and only ever increments. + """ + + KEY = "agent:main:telegram:dm:777" + SOURCE = "telegram" + + def _keyed(self, db, sid): + db.create_session(session_id=sid, source=self.SOURCE, session_key=self.KEY) + return SimpleNamespace( + session_id=sid, _session_db=db, _gateway_session_key=self.KEY, + platform=self.SOURCE, + ) + + def _gen(self, db): + return db.latest_conversation_boundary(self.KEY, self.SOURCE) + + def test_deleting_the_newest_boundary_does_not_roll_back(self, db): + a = self._keyed(db, "s-a") + scope_a = resolve_prompt_cache_scope(a) + db.end_session("s-a", "session_reset") + + b = self._keyed(db, "s-b") + scope_b = resolve_prompt_cache_scope(b) + db.end_session("s-b", "session_reset") + assert self._gen(db) == 2 + + db.delete_session("s-b") # prune the newest boundary + assert self._gen(db) == 2 # counter is outside that rowset + + c = self._keyed(db, "s-c") + scope_c = resolve_prompt_cache_scope(c) + assert len({scope_a, scope_b, scope_c}) == 3 + + def test_deleting_every_boundary_does_not_roll_back(self, db): + a = self._keyed(db, "s-a") + scope_a = resolve_prompt_cache_scope(a) + db.end_session("s-a", "session_reset") + db.delete_session("s-a") + + b = self._keyed(db, "s-b") + assert resolve_prompt_cache_scope(b) != scope_a + + def test_a_backwards_clock_then_a_prune_still_cannot_repeat(self, db): + """The reviewer's exact shape: (1,T1) -> (2,T1) -> delete -> (1,T1).""" + a = self._keyed(db, "s-a") + scope_a = resolve_prompt_cache_scope(a) + db.end_session("s-a", "session_reset") + + b = self._keyed(db, "s-b") + scope_b = resolve_prompt_cache_scope(b) + db.end_session("s-b", "session_reset") + # Clock went backwards: this boundary lands before the first one. + with db._lock: + db._conn.execute( + "UPDATE sessions SET ended_at = (" + " SELECT MIN(ended_at) FROM sessions WHERE ended_at IS NOT NULL" + ") - 60 WHERE id = ?", + ("s-b",), + ) + db._conn.commit() + db.delete_session("s-b") + + c = self._keyed(db, "s-c") + assert len({scope_a, scope_b, resolve_prompt_cache_scope(c)}) == 3 + + def test_only_real_boundaries_advance_it(self, db): + self._keyed(db, "s-a") + db.end_session("s-a", "compression") + assert self._gen(db) is None + db.create_session(session_id="s-b", source=self.SOURCE, session_key=self.KEY) + db.end_session("s-b", "agent_close") + assert self._gen(db) is None + + def test_a_repeated_end_does_not_double_count(self, db): + """end_session no-ops on an ended row; the first reason wins.""" + self._keyed(db, "s-a") + db.end_session("s-a", "session_reset") + db.end_session("s-a", "session_reset") + db.end_session("s-a", "idle") + assert self._gen(db) == 1 + + def test_promote_advances_it_too(self, db): + """/new and the policy resets promote rather than end_session.""" + self._keyed(db, "s-a") + db.end_session("s-a", "agent_close") + assert self._gen(db) is None + assert db.promote_to_session_reset("s-a", "session_reset") is True + assert self._gen(db) == 1 + + def test_an_unkeyed_row_advances_nothing(self, db): + db.create_session(session_id="s-bare", source=self.SOURCE) + db.end_session("s-bare", "session_reset") + assert self._gen(db) is None + + def test_the_counter_is_peer_scoped(self, db): + self._keyed(db, "s-a") + db.end_session("s-a", "session_reset") + assert self._gen(db) == 1 + assert db.latest_conversation_boundary(self.KEY, "api_server") is None + + +class TestSourceOverrideDomain: + """The scope is memoized immediately, so the source must be right first. + + ``_agent_source`` used ``agent.platform`` before the row landed while + persistence uses ``_session_source_for_agent``, which honors + ``HERMES_SESSION_SOURCE``. Under an override both sides of a ``/new`` + queried the platform domain, missed the boundary stored under the + override, and hashed the same scope (@andrexibiza on #98811). + """ + + KEY = "agent:main:telegram:dm:888" + + def test_the_pre_row_source_matches_persistence(self, db, monkeypatch): + from agent.prompt_cache_scope import _agent_source + + monkeypatch.setenv("HERMES_SESSION_SOURCE", "override-src") + agent = SimpleNamespace( + session_id="s-none", _session_db=db, _gateway_session_key=self.KEY, + platform="telegram", + ) + assert _agent_source(agent, "", db) == "override-src" + + def test_new_rotates_under_a_source_override(self, db, monkeypatch): + monkeypatch.setenv("HERMES_SESSION_SOURCE", "override-src") + + def keyed(sid): + db.create_session( + session_id=sid, source="override-src", session_key=self.KEY + ) + return SimpleNamespace( + session_id=sid, _session_db=db, + _gateway_session_key=self.KEY, platform="telegram", + ) + + before = resolve_prompt_cache_scope(keyed("s-a")) + db.end_session("s-a", "session_reset") + assert resolve_prompt_cache_scope(keyed("s-b")) != before diff --git a/tests/gateway/test_api_server_declared_conversation.py b/tests/gateway/test_api_server_declared_conversation.py index f46489d756..e6d5aa20e3 100644 --- a/tests/gateway/test_api_server_declared_conversation.py +++ b/tests/gateway/test_api_server_declared_conversation.py @@ -14,7 +14,9 @@ the durable conversation boundaries already recorded in ``sessions.end_reason``, and nothing that does not declare a key changes. """ +import asyncio import types +from unittest.mock import MagicMock import pytest from aiohttp import web @@ -586,3 +588,162 @@ class TestRunsHandlerPrecedence: assert seen and seen[0]["session_id"] == "explicit-session" # KEY still resolves to its own conversation, never the explicit one. assert adapter._declared_conversation_session(KEY) == "sess-live" + + +async def _await_run(adapter, run_id, timeout=10.0): + """Wait for a /v1/runs worker to finish, so settlement has happened.""" + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if run_id not in adapter._active_run_agents: + # One more tick so the executor thread's finally can retire. + await asyncio.sleep(0.05) + return True + await asyncio.sleep(0.02) + return False + + +def _stub_agent(adapter, session_id, seen): + """An agent good enough for the REAL _run_agent to drive end to end. + + Replacing `_run_agent` itself cannot exercise its settlement block, which + is where a caller-local name leaked in and raised `NameError` on every + opted-in bind while mocked tests stayed green (@andrexibiza on #98811). + Stubbing one layer lower — `_create_agent` — leaves that block real. + """ + agent = MagicMock() + agent.session_id = session_id + agent.session_prompt_tokens = 0 + agent.session_completion_tokens = 0 + agent.session_total_tokens = 0 + + def _run(**kwargs): + seen.append(kwargs) + db = adapter._ensure_session_db() + # AIAgent._ensure_db_session() creates the row during the turn. + if db is not None and db.get_session(session_id) is None: + db.create_session(session_id=session_id, source=SOURCE, model="m") + return {"final_response": "ok", "messages": [], "api_calls": 1} + + agent.run_conversation.side_effect = _run + return agent + + +class TestRealRunAgentSettlement: + """The real `_run_agent` settlement block, not a stand-in for it.""" + + @pytest.mark.asyncio + async def test_declared_bind_settles_without_raising(self, live): + adapter, db, app = live + seen = [] + created = {} + + def _create(**kw): + created.update(kw) + return _stub_agent(adapter, kw["session_id"], seen) + + adapter._create_agent = _create + + async with TestClient(TestServer(app)) as cli: + resp = await cli.post( + "/v1/responses", + json={"model": "hermes-agent", "input": "hi"}, + headers=_headers(KEY), + ) + assert resp.status == 200 + body = await resp.json() + + assert body["status"] == "completed" + sid = created["session_id"] + # Settlement ran for real: the row is bound and recoverable. + assert db.get_session(sid)["session_key"] == KEY + assert adapter._declared_conversation_session(KEY) == sid + + @pytest.mark.asyncio + async def test_two_replies_settle_on_one_conversation(self, live): + adapter, db, app = live + seen = [] + ids = [] + + def _create(**kw): + ids.append(kw["session_id"]) + return _stub_agent(adapter, kw["session_id"], seen) + + adapter._create_agent = _create + + async with TestClient(TestServer(app)) as cli: + for _ in range(2): + resp = await cli.post( + "/v1/responses", + json={"model": "hermes-agent", "input": "hi"}, + headers=_headers(KEY), + ) + assert resp.status == 200 + + assert len(set(ids)) == 1 + + @pytest.mark.asyncio + async def test_a_chained_turn_never_rebinds_through_real_settlement(self, live): + adapter, db, app = live + seen = [] + + def _create(**kw): + return _stub_agent(adapter, kw["session_id"], seen) + + adapter._create_agent = _create + + _seed(db, "sess-A", key=KEY) + adapter._response_store.put( + "resp_A", + {"conversation_history": [], "session_id": "sess-A", "instructions": None}, + ) + + async with TestClient(TestServer(app)) as cli: + resp = await cli.post( + "/v1/responses", + json={"model": "hermes-agent", "input": "hi", + "previous_response_id": "resp_A"}, + headers=_headers(OTHER_KEY), + ) + assert resp.status == 200 + + assert db.get_session("sess-A")["session_key"] == KEY + assert adapter._declared_conversation_session(OTHER_KEY) is None + + @pytest.mark.asyncio + async def test_runs_explicit_unkeyed_session_stays_unbound(self, live): + """The reviewer's exact case: an unkeyed explicit row must not be taken. + + `_run_sync` bound unconditionally, so an explicit body session_id that + existed with an empty `session_key` was silently adopted by the header + key even though the header lost precedence. + """ + adapter, db, app = live + seen = [] + + def _create(**kw): + return _stub_agent(adapter, kw["session_id"], seen) + + adapter._create_agent = _create + + # Explicit session exists and carries NO routing key. + db.create_session(session_id="explicit-session", source=SOURCE, model="m") + _seed(db, "sess-live", key=KEY) + + async with TestClient(TestServer(app)) as cli: + resp = await cli.post( + "/v1/runs", + json={"model": "hermes-agent", "input": "hi", + "session_id": "explicit-session"}, + headers=_headers(KEY), + ) + assert resp.status in (200, 202) + run_id = (await resp.json()).get("run_id") + assert run_id + # /v1/runs answers before the turn settles, so the assertion must + # wait for the worker's finally to run. Without this the test + # passes on timing rather than on the precedence gate. + assert await _await_run(adapter, run_id), "run never settled" + + row = db.get_session("explicit-session") + assert not (row.get("session_key") or "") + assert adapter._declared_conversation_session(KEY) == "sess-live"