From 93ff129b51d4a764c6fac745ee2c22dfe56b02bf Mon Sep 17 00:00:00 2001 From: Z Date: Tue, 14 Jul 2026 20:30:30 +0300 Subject: [PATCH] fix(gateway): follow async completions across compression --- gateway/run.py | 196 +++++++-- gateway/session.py | 36 ++ .../test_async_delegation_session_binding.py | 373 ++++++++++++++++-- .../test_session_store_runtime_stale_guard.py | 60 +++ 4 files changed, 585 insertions(+), 80 deletions(-) diff --git a/gateway/run.py b/gateway/run.py index 0bbb976258cd..5b6fd4cf4494 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -44,7 +44,7 @@ from collections import OrderedDict from contextvars import copy_context from pathlib import Path from datetime import datetime -from typing import Awaitable, Callable, Dict, Optional, Any, List, Union +from typing import Awaitable, Callable, Dict, Optional, Any, List, Union, cast # account_usage imports the OpenAI SDK chain (~230 ms). Only needed by # /usage; we still import it at module top in the gateway because test @@ -1912,6 +1912,7 @@ from gateway.config import ( ) from gateway.session import ( AsyncSessionStore, + SessionEntry, SessionStore, SessionSource, SessionContext, @@ -10065,6 +10066,156 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew await adapter.send(source.chat_id, content, metadata=metadata) + async def _resolve_async_delegation_session( + self, + session_entry: SessionEntry, + pinned_session_id: str, + ) -> Optional[SessionEntry]: + """Resolve an async completion to its verified owning gateway session. + + A compression rotation ends the physical parent row while continuing + the same logical conversation in a child. Follow that lineage, but + never let a late completion override an unrelated /new or restored + route. Unknown ownership remains fail-closed; the result is still + available in the delegation records. + """ + session_db = cast(Any, self._session_db) + if session_db is None: + logger.warning( + "Async-delegation completion has no session database; " + "dropping injection (#55578 fail-closed)." + ) + return None + + pinned_row = None + try: + pinned_row = await session_db.get_session(pinned_session_id) + except Exception: + logger.debug( + "Async-delegation parent lookup failed for %s", + pinned_session_id, + exc_info=True, + ) + + if pinned_row is None: + logger.warning( + "Async-delegation completion has unknown spawning session %s; " + "dropping injection (#55578 fail-closed).", + pinned_session_id, + ) + return None + + target_session_id = pinned_session_id + follows_compression = False + if pinned_row.get("ended_at"): + if pinned_row.get("end_reason") != "compression": + logger.warning( + "Async-delegation completion pinned to ended session %s " + "(end_reason=%r); dropping injection instead of resurrecting it " + "(#55578 fail-closed).", + pinned_session_id, + pinned_row.get("end_reason"), + ) + return None + + follows_compression = True + try: + target_session_id = await session_db.get_compression_tip( + pinned_session_id + ) + except Exception: + logger.debug( + "Async-delegation compression-tip lookup failed for %s", + pinned_session_id, + exc_info=True, + ) + target_session_id = None + + if not target_session_id or target_session_id == pinned_session_id: + logger.warning( + "Async-delegation completion pinned to compressed session %s " + "without a continuation; dropping injection.", + pinned_session_id, + ) + return None + + try: + tip_row = await session_db.get_session(target_session_id) + except Exception: + tip_row = None + if tip_row is None or tip_row.get("ended_at"): + logger.warning( + "Async-delegation compression continuation %s is %s; " + "dropping injection.", + target_session_id, + "unknown" if tip_row is None else "ended", + ) + return None + + route_owns_lineage = session_entry.session_id in { + pinned_session_id, + target_session_id, + } + if not route_owns_lineage: + # A long-running delegation may survive multiple compression + # rotations. Accept an intermediate stale route only when its + # own verified compression tip is the same live target. + try: + route_row = await session_db.get_session(session_entry.session_id) + route_tip = ( + await session_db.get_compression_tip(session_entry.session_id) + if route_row is not None + and route_row.get("ended_at") + and route_row.get("end_reason") == "compression" + else None + ) + except Exception: + route_tip = None + route_owns_lineage = route_tip == target_session_id + + if not route_owns_lineage: + logger.warning( + "Async-delegation completion for compression lineage %s -> %s " + "does not own current route %s; dropping injection.", + pinned_session_id, + target_session_id, + session_entry.session_id, + ) + return None + + if target_session_id == session_entry.session_id: + return session_entry + + prior_session_id = session_entry.session_id + if follows_compression: + switched = await self.async_session_store.advance_compression_session( + session_entry.session_key, + prior_session_id, + target_session_id, + ) + else: + switched = await self.async_session_store.switch_session( + session_entry.session_key, + target_session_id, + ) + if switched is None: + logger.warning( + "Async-delegation completion could not bind routing key %s to " + "owning session %s; dropping injection.", + session_entry.session_key, + target_session_id, + ) + return None + + logger.info( + "Pinned async-delegation completion to owning session %s " + "(was %s) for routing key %s (#57498)", + target_session_id, + prior_session_id, + session_entry.session_key, + ) + return switched + async def _handle_message(self, event: MessageEvent) -> Optional[str]: """ Handle an incoming message from any platform. @@ -12123,43 +12274,14 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew pinned_session_id = str( (getattr(event, "metadata", None) or {}).get("gateway_session_id") or "" ).strip() - if pinned_session_id and pinned_session_id != session_entry.session_id: - # Fail closed (#55578): the spawning session may have ENDED since - # dispatch (user /new-reset, compression rotation whose parent was - # closed). switch_session() re-opens ended sessions, so pinning - # blindly would RESURRECT a conversation the user explicitly - # ended and inject into it — the same illicit-revival class as - # the ws_orphan_reap loop (#60609). A completion whose spawning - # session is dead is dropped from injection; the subagent's - # output remains in the delegation records. - pinned_row = None - try: - if self._session_db is not None: - # AsyncSessionDB already offloads to a thread. - pinned_row = await self._session_db.get_session(pinned_session_id) - except Exception: - pinned_row = None - if pinned_row is None or pinned_row.get("ended_at"): - logger.warning( - "Async-delegation completion pinned to session %s, which is " - "%s — dropping injection instead of resurrecting it " - "(#55578 fail-closed; result remains in the delegation " - "records).", - pinned_session_id, - "unknown" if pinned_row is None else "ended", - ) + if pinned_session_id: + resolved_entry = await self._resolve_async_delegation_session( + session_entry, + pinned_session_id, + ) + if resolved_entry is None: return - prior_session_id = session_entry.session_id - switched = await self.async_session_store.switch_session(session_key, pinned_session_id) - if switched is not None: - session_entry = switched - logger.info( - "Pinned async-delegation completion to spawning session %s " - "(was %s) for routing key %s (#57498)", - pinned_session_id, - prior_session_id, - session_key, - ) + session_entry = resolved_entry self._cache_session_source(session_key, source) if await asyncio.to_thread(self._is_telegram_topic_lane, source): try: diff --git a/gateway/session.py b/gateway/session.py index 0168d940ffd4..d3020d11822f 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -2401,6 +2401,42 @@ class SessionStore: return new_entry + def advance_compression_session( + self, + session_key: str, + expected_session_id: str, + target_session_id: str, + ) -> Optional[SessionEntry]: + """CAS-advance one route along an already-verified compression lineage. + + Unlike ``switch_session``, this does not end or reopen SQLite rows. The + compression transaction already owns that lifecycle; this method only + repairs the persisted gateway key→session mapping. Returning ``None`` + means the route moved after the caller's snapshot (for example /new), + so the caller must fail closed instead of overwriting the newer route. + """ + if not session_key or not expected_session_id or not target_session_id: + return None + + with self._lock: + self._ensure_loaded_locked() + entry = self._entries.get(session_key) + if entry is None: + return None + if entry.session_id == target_session_id: + return entry + if entry.session_id != expected_session_id: + return None + if not self._heal_compression_tip_locked( + entry, + expected_session_id, + target_session_id, + ): + return None + entry.updated_at = _now() + self._save() + return entry + def switch_session(self, session_key: str, target_session_id: str) -> Optional[SessionEntry]: """Switch a session key to point at an existing session ID. diff --git a/tests/gateway/test_async_delegation_session_binding.py b/tests/gateway/test_async_delegation_session_binding.py index c48b37e47dfa..b5a8835fd8d7 100644 --- a/tests/gateway/test_async_delegation_session_binding.py +++ b/tests/gateway/test_async_delegation_session_binding.py @@ -8,7 +8,6 @@ Three invariants on the messaging-gateway surface, mirroring the TUI rules: 3. /new interrupts the old conversation's in-flight async delegations. """ -import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -62,61 +61,349 @@ class TestInterruptForSessionByParentId: class TestGatewayPinningFailsClosed: - """The gateway injection path must never resurrect an ended session.""" + """The gateway must follow only verified compression continuations.""" - def _make_runner(self, pinned_row): + @staticmethod + def _entry(session_id): + from datetime import datetime + + from gateway.config import Platform + from gateway.session import SessionEntry + + return SessionEntry( + session_key="agent:main:telegram:group:-100:4", + session_id=session_id, + created_at=datetime.now(), + updated_at=datetime.now(), + platform=Platform.TELEGRAM, + chat_type="group", + ) + + def _make_runner( + self, + rows, + *, + compression_tip=None, + compression_error=None, + switched_entry=None, + ): from gateway.run import GatewayRunner + from gateway.session import AsyncSessionStore runner = object.__new__(GatewayRunner) db = MagicMock() - db.get_session = AsyncMock(return_value=pinned_row) + db.get_session = AsyncMock(side_effect=lambda session_id: rows.get(session_id)) + db.get_compression_tip = AsyncMock( + return_value=compression_tip, + side_effect=compression_error, + ) runner._session_db = db - - entry = MagicMock() - entry.session_key = "agent:main:telegram:dm:1" - entry.session_id = "sess_current" runner.session_store = MagicMock() - runner.session_store.get_or_create_session.return_value = entry - runner.session_store.switch_session.return_value = entry - return runner, entry + runner.session_store.switch_session = MagicMock(return_value=switched_entry) + runner.session_store.advance_compression_session = MagicMock( + return_value=switched_entry + ) + runner._async_session_store = AsyncSessionStore(runner.session_store) + return runner - def _run_pinning_prefix(self, runner, pinned_session_id): - """Execute the pinning guard logic exactly as _handle_message does.""" + @staticmethod + def _assert_no_route_change(runner): + getattr(runner.session_store, "switch_session").assert_not_called() + getattr( + runner.session_store, "advance_compression_session" + ).assert_not_called() - async def _go(): - event = MagicMock() - event.metadata = {"gateway_session_id": pinned_session_id} - session_entry = runner.session_store.get_or_create_session(MagicMock()) - pinned = str((getattr(event, "metadata", None) or {}).get("gateway_session_id") or "").strip() - if pinned and pinned != session_entry.session_id: - pinned_row = None - try: - if runner._session_db is not None: - pinned_row = await runner._session_db.get_session(pinned) - except Exception: - pinned_row = None - if pinned_row is None or pinned_row.get("ended_at"): - return "dropped" - switched = runner.session_store.switch_session(session_entry.session_key, pinned) - if switched is not None: - return "pinned" - return "default" + @pytest.mark.asyncio + async def test_live_spawning_session_stays_pinned(self): + current = self._entry("sess_live") + runner = self._make_runner( + {"sess_live": {"id": "sess_live", "ended_at": None}} + ) - return asyncio.run(_go()) + resolved = await runner._resolve_async_delegation_session( + current, "sess_live" + ) - def test_live_spawning_session_pins(self): - runner, _ = self._make_runner({"id": "sess_old", "ended_at": None}) - assert self._run_pinning_prefix(runner, "sess_old") == "pinned" + assert resolved is current + self._assert_no_route_change(runner) - def test_ended_spawning_session_drops(self): - runner, _ = self._make_runner({"id": "sess_old", "ended_at": "2026-07-08T00:00:00"}) - assert self._run_pinning_prefix(runner, "sess_old") == "dropped" - runner.session_store.switch_session.assert_not_called() + @pytest.mark.asyncio + async def test_live_spawning_session_rebinds_from_different_route(self): + current = self._entry("sess_current") + pinned = self._entry("sess_live") + runner = self._make_runner( + {"sess_live": {"id": "sess_live", "ended_at": None}}, + switched_entry=pinned, + ) - def test_unknown_spawning_session_drops(self): - runner, _ = self._make_runner(None) - assert self._run_pinning_prefix(runner, "sess_gone") == "dropped" - runner.session_store.switch_session.assert_not_called() + resolved = await runner._resolve_async_delegation_session( + current, "sess_live" + ) + + assert resolved is pinned + getattr(runner.session_store, "switch_session").assert_called_once_with( + current.session_key, "sess_live" + ) + + @pytest.mark.asyncio + async def test_non_compression_ended_parent_drops(self): + current = self._entry("sess_old") + runner = self._make_runner( + { + "sess_old": { + "id": "sess_old", + "ended_at": "2026-07-08T00:00:00", + "end_reason": "session_reset", + } + } + ) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_old" + ) + + assert resolved is None + self._assert_no_route_change(runner) + + @pytest.mark.asyncio + async def test_compression_parent_advances_stale_route_to_live_tip(self): + current = self._entry("sess_parent") + tip = self._entry("sess_tip") + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-08T00:00:00", + "end_reason": "compression", + }, + "sess_tip": { + "id": "sess_tip", + "ended_at": None, + "parent_session_id": "sess_parent", + }, + }, + compression_tip="sess_tip", + switched_entry=tip, + ) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_parent" + ) + + assert resolved is tip + getattr( + runner.session_store, "advance_compression_session" + ).assert_called_once_with(current.session_key, "sess_parent", "sess_tip") + + @pytest.mark.asyncio + async def test_compression_cas_losing_to_new_drops(self): + current = self._entry("sess_parent") + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-08T00:00:00", + "end_reason": "compression", + }, + "sess_tip": { + "id": "sess_tip", + "ended_at": None, + "parent_session_id": "sess_parent", + }, + }, + compression_tip="sess_tip", + switched_entry=None, + ) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_parent" + ) + + assert resolved is None + getattr( + runner.session_store, "advance_compression_session" + ).assert_called_once_with(current.session_key, "sess_parent", "sess_tip") + + @pytest.mark.asyncio + async def test_intermediate_compression_route_advances_to_same_live_tip(self): + current = self._entry("sess_middle") + tip = self._entry("sess_tip") + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-08T00:00:00", + "end_reason": "compression", + }, + "sess_middle": { + "id": "sess_middle", + "ended_at": "2026-07-08T00:01:00", + "end_reason": "compression", + "parent_session_id": "sess_parent", + }, + "sess_tip": { + "id": "sess_tip", + "ended_at": None, + "parent_session_id": "sess_middle", + }, + }, + compression_tip="sess_tip", + switched_entry=tip, + ) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_parent" + ) + + assert resolved is tip + getattr( + runner.session_store, "advance_compression_session" + ).assert_called_once_with(current.session_key, "sess_middle", "sess_tip") + + @pytest.mark.asyncio + async def test_compression_parent_follows_real_sessiondb_lineage(self, tmp_path): + from gateway.run import GatewayRunner + from gateway.session import AsyncSessionStore + from hermes_state import AsyncSessionDB, SessionDB + + session_db = SessionDB(db_path=tmp_path / "state.db") + session_db.create_session("sess_parent", source="telegram") + session_db.end_session("sess_parent", end_reason="compression") + session_db.create_session( + "sess_tip", + source="telegram", + parent_session_id="sess_parent", + ) + + current = self._entry("sess_parent") + tip = self._entry("sess_tip") + runner = object.__new__(GatewayRunner) + runner._session_db = AsyncSessionDB(session_db) + runner.session_store = MagicMock() + runner.session_store.switch_session = MagicMock(return_value=tip) + runner.session_store.advance_compression_session = MagicMock(return_value=tip) + runner._async_session_store = AsyncSessionStore(runner.session_store) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_parent" + ) + + assert resolved is tip + getattr( + runner.session_store, "advance_compression_session" + ).assert_called_once_with(current.session_key, "sess_parent", "sess_tip") + + @pytest.mark.asyncio + async def test_ended_compression_tip_drops(self): + current = self._entry("sess_parent") + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-08T00:00:00", + "end_reason": "compression", + }, + "sess_tip": { + "id": "sess_tip", + "ended_at": "2026-07-08T00:01:00", + "end_reason": "session_reset", + "parent_session_id": "sess_parent", + }, + }, + compression_tip="sess_tip", + ) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_parent" + ) + + assert resolved is None + self._assert_no_route_change(runner) + + @pytest.mark.asyncio + async def test_compression_lookup_failure_drops(self): + current = self._entry("sess_parent") + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-08T00:00:00", + "end_reason": "compression", + } + }, + compression_error=RuntimeError("db unavailable"), + ) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_parent" + ) + + assert resolved is None + self._assert_no_route_change(runner) + + @pytest.mark.asyncio + async def test_compression_parent_accepts_already_current_tip(self): + current = self._entry("sess_tip") + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-08T00:00:00", + "end_reason": "compression", + }, + "sess_tip": { + "id": "sess_tip", + "ended_at": None, + "parent_session_id": "sess_parent", + }, + }, + compression_tip="sess_tip", + ) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_parent" + ) + + assert resolved is current + self._assert_no_route_change(runner) + + @pytest.mark.asyncio + async def test_compression_parent_does_not_override_new_route(self): + current = self._entry("sess_after_new") + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-08T00:00:00", + "end_reason": "compression", + }, + "sess_tip": { + "id": "sess_tip", + "ended_at": None, + "parent_session_id": "sess_parent", + }, + }, + compression_tip="sess_tip", + ) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_parent" + ) + + assert resolved is None + self._assert_no_route_change(runner) + + @pytest.mark.asyncio + async def test_unknown_spawning_session_drops(self): + current = self._entry("sess_current") + runner = self._make_runner({}) + + resolved = await runner._resolve_async_delegation_session( + current, "sess_gone" + ) + + assert resolved is None + self._assert_no_route_change(runner) class TestResetHandlerInterruptsDelegations: diff --git a/tests/gateway/test_session_store_runtime_stale_guard.py b/tests/gateway/test_session_store_runtime_stale_guard.py index 62c9eff1f4c5..941ae761a973 100644 --- a/tests/gateway/test_session_store_runtime_stale_guard.py +++ b/tests/gateway/test_session_store_runtime_stale_guard.py @@ -267,3 +267,63 @@ class TestRuntimeStaleGuard: db.reopen_session.assert_not_called() # A brand-new session row was created. db.create_session.assert_called_once() + + +class TestAdvanceCompressionSession: + def test_cas_advances_route_without_reopening_rows(self, tmp_path): + db = _db_returning({}) + store = _make_store_with_db(tmp_path, db) + source = _source() + key = store._generate_session_key(source) + original = _make_entry(key, "sid_parent") + store._entries[key] = original + + result = store.advance_compression_session( + key, + "sid_parent", + "sid_tip", + ) + + assert result is not None + assert result is original + assert result.session_id == "sid_tip" + assert store.peek_session_id(key) == "sid_tip" + db.end_session.assert_not_called() + db.reopen_session.assert_not_called() + + def test_cas_is_idempotent_when_another_caller_already_advanced(self, tmp_path): + db = _db_returning({}) + store = _make_store_with_db(tmp_path, db) + source = _source() + key = store._generate_session_key(source) + current = _make_entry(key, "sid_tip") + store._entries[key] = current + + result = store.advance_compression_session( + key, + "sid_parent", + "sid_tip", + ) + + assert result is current + assert store.peek_session_id(key) == "sid_tip" + db.end_session.assert_not_called() + db.reopen_session.assert_not_called() + + def test_cas_refuses_to_overwrite_route_changed_by_new(self, tmp_path): + db = _db_returning({}) + store = _make_store_with_db(tmp_path, db) + source = _source() + key = store._generate_session_key(source) + store._entries[key] = _make_entry(key, "sid_after_new") + + result = store.advance_compression_session( + key, + "sid_parent", + "sid_tip", + ) + + assert result is None + assert store.peek_session_id(key) == "sid_after_new" + db.end_session.assert_not_called() + db.reopen_session.assert_not_called()