"""Regression tests for #29335 — gateway must persist ``session_entry.session_id`` after the agent's compression path mutates it. When ``_compress_context()`` rolls the agent forward into a new session, the agent now returns the new ``session_id`` in its result dict. The gateway updates ``session_entry.session_id`` in memory AND must call ``session_store._save()`` so the new mapping survives a gateway restart. Without ``_save()``, the next turn loads the OLD session's transcript and re-triggers compression forever. Three sites in ``gateway/run.py`` mutate ``session_entry.session_id`` after a compression-induced session split. All three MUST be followed by a ``_save()`` call. This test pins that invariant. ``TestCompressionSessionPropagation`` adds behavioral tests that exercise the actual propagation path inline, verifying that the mock session_entry update and _save() semantics are correct without requiring a live gateway. """ from __future__ import annotations import ast import inspect import textwrap from unittest.mock import MagicMock, call from gateway import run as gateway_run from gateway.session_context import set_current_session_id, get_session_env def _session_id_assignments_followed_by_save(source: str) -> list[tuple[int, bool]]: """For each ``session_entry.session_id = ...`` assignment in *source*, return ``(lineno, saved_within_5_stmts)`` — True iff a ``self.session_store._save()`` call appears in the same block within the next 5 statements (covers normal control flow without false-flagging cleanup that lives 200 lines away). """ tree = ast.parse(textwrap.dedent(source)) results: list[tuple[int, bool]] = [] class _Visitor(ast.NodeVisitor): def _is_session_id_assign(self, node: ast.AST) -> bool: if not isinstance(node, ast.Assign): return False for target in node.targets: if ( isinstance(target, ast.Attribute) and target.attr == "session_id" and isinstance(target.value, ast.Name) and target.value.id == "session_entry" ): return True return False def _block_has_save_after(self, body: list[ast.stmt], idx: int) -> bool: for stmt in body[idx : idx + 6]: for sub in ast.walk(stmt): if ( isinstance(sub, ast.Call) and isinstance(sub.func, ast.Attribute) and sub.func.attr == "_save" ): return True return False def _walk_body(self, body: list[ast.stmt]) -> None: for i, stmt in enumerate(body): if self._is_session_id_assign(stmt): results.append((stmt.lineno, self._block_has_save_after(body, i))) for child in ast.iter_child_nodes(stmt): if isinstance(child, (ast.If, ast.For, ast.While, ast.With, ast.Try, ast.AsyncWith, ast.AsyncFor)): self._walk_node(child) def _walk_node(self, node: ast.AST) -> None: for attr in ("body", "orelse", "finalbody"): inner = getattr(node, attr, None) if isinstance(inner, list): self._walk_body(inner) if hasattr(node, "handlers"): for handler in node.handlers: self._walk_body(handler.body) def visit(self, node: ast.AST) -> None: if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): self._walk_body(node.body) for child in ast.iter_child_nodes(node): self.visit(child) _Visitor().visit(tree) return results def test_every_post_compression_session_id_assignment_persists(): """Every ``session_entry.session_id = ...`` in gateway/run.py must be followed by a ``session_store._save()`` call within the same block. Regression for #29335 — the assignment at the end of ``_handle_message_with_agent`` used to skip ``_save()`` while two sibling sites (hygiene rewrite, manual /compress) already persisted. The agent would compress correctly, the gateway would update its in-memory session_id, then drop it on next gateway restart. """ source = inspect.getsource(gateway_run) assignments = _session_id_assignments_followed_by_save(source) assert assignments, ( "No ``session_entry.session_id = ...`` assignments found in gateway/run.py — " "either the structure changed or the AST walker is broken." ) missing = [lineno for lineno, saved in assignments if not saved] assert not missing, ( f"{len(missing)} ``session_entry.session_id = ...`` site(s) in gateway/run.py " f"are not followed by ``session_store._save()`` within the same block " f"(lines: {missing}). Every post-compression session_id update must persist " f"or the next turn loads the pre-compression transcript and triggers an " f"infinite compression loop. See issue #29335." ) class TestCompressionSessionPropagation: """Behavioral tests for post-compression session_id propagation. The structural AST test above pins that every ``session_entry.session_id`` assignment in gateway/run.py is followed by ``_save()``. These tests exercise the *behavior* of that propagation path inline, using mocks that mirror the objects gateway/run.py works with (``session_entry`` and ``session_store``), verifying the semantics are correct without requiring a live gateway instance. Ordering contract (from the comments added to the source in this PR): 1. The agent thread updates the contextvar in ``conversation_compression.py`` via ``set_current_session_id(agent.session_id)``. 2. After ``run_in_executor`` returns, the gateway propagates the new id to ``session_entry.session_id`` and calls ``session_store._save()``. Both halves must agree for the next turn to route correctly. """ def test_gateway_session_entry_follows_compression_rotation(self) -> None: """The gateway handler must update session_entry and call _save() when the agent result carries a rotated session_id. Simulates the inline propagation block in gateway/run.py: if agent_result.get("session_id") and \\ agent_result["session_id"] != session_entry.session_id: session_entry.session_id = agent_result["session_id"] self.session_store._save() Verifies that session_entry.session_id is mutated and _save is called exactly once — the minimal contract that prevents the restart-loop bug. """ old_sid = "20260101_000000_aaaaaa" new_sid = "20260101_000001_bbbbbb" session_entry = MagicMock() session_entry.session_id = old_sid session_store = MagicMock() agent_result = {"session_id": new_sid, "response": "hello"} # Inline the propagation logic exactly as it appears in gateway/run.py # (around line 9459). This is the behavior we are pinning. if agent_result.get("session_id") and agent_result["session_id"] != session_entry.session_id: session_entry.session_id = agent_result["session_id"] session_store._save() assert session_entry.session_id == new_sid, ( "session_entry.session_id was not updated to the compressed session id. " "The next turn would load the old transcript and re-trigger compression." ) session_store._save.assert_called_once_with(), ( "session_store._save() was not called after session_entry update. " "The new session mapping would not survive a gateway restart." ) def test_no_update_when_session_id_unchanged(self) -> None: """The propagation block must be a no-op when the agent did not compress. If the agent returns the same session_id (normal turn, no compression), session_entry must not be touched and _save must not be called — avoiding spurious writes on every turn. """ same_sid = "20260101_000000_aaaaaa" session_entry = MagicMock() session_entry.session_id = same_sid session_store = MagicMock() # Normal turn: agent returns same session_id (or none at all) agent_result = {"response": "hello"} # no "session_id" key if agent_result.get("session_id") and agent_result["session_id"] != session_entry.session_id: session_entry.session_id = agent_result["session_id"] session_store._save() # session_entry.session_id was set during mock construction; the # propagation block must not have set it again. session_store._save.assert_not_called() def test_contextvar_and_session_entry_agree_after_compression(self) -> None: """After compression, the contextvar and session_entry must carry the same session_id. The agent thread calls ``set_current_session_id(new_sid)`` inside ``conversation_compression.py`` (step 1). The gateway then propagates ``new_sid`` to ``session_entry.session_id`` (step 2). If either step is missing, tool calls and transcript writes will disagree on which session is active. This test simulates both steps and asserts agreement. """ old_sid = "20260101_000000_cccccc" new_sid = "20260101_000002_dddddd" # Step 1: agent thread updates contextvar (mirrors conversation_compression.py # around line 511-513) set_current_session_id(new_sid) # Step 2: gateway propagates to session_entry (mirrors gateway/run.py # around line 9459-9461) session_entry = MagicMock() session_entry.session_id = old_sid agent_result = {"session_id": new_sid} if agent_result.get("session_id") and agent_result["session_id"] != session_entry.session_id: session_entry.session_id = agent_result["session_id"] contextvar_sid = get_session_env("HERMES_SESSION_ID", "") assert contextvar_sid == new_sid, ( f"Contextvar still holds old session_id '{contextvar_sid}' after " f"set_current_session_id('{new_sid}'). Tool calls in the next turn " "will read stale routing state." ) assert session_entry.session_id == new_sid, ( f"session_entry.session_id is '{session_entry.session_id}' but contextvar " f"says '{contextvar_sid}'. The two routing paths disagree after compression." ) assert contextvar_sid == session_entry.session_id, ( "Contextvar and session_entry disagree on the active session_id " "after compression rotation. Exactly one of the two ordering steps " "was skipped." )