fix(tui): resolve in-memory session UUID in session.resume for crash recovery #39013

hermes_state.py

@@ -265,6 +265,7 @@ def _log_wal_fallback_once(db_label: str, exc: Exception) -> None:
     handoff_error TEXT,
     rewind_count INTEGER NOT NULL DEFAULT 0,
     archived INTEGER NOT NULL DEFAULT 0,
+    in_memory_id TEXT,
     FOREIGN KEY (parent_session_id) REFERENCES sessions(id)
 );
 
@@ -1316,6 +1317,24 @@ def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
             row = cursor.fetchone()
         return dict(row) if row else None
 
+    def get_session_by_in_memory_id(self, in_memory_id: str) -> Optional[Dict[str, Any]]:
+        """Get a session by its in-memory UUID (set by TUI gateway)."""
+        with self._lock:
+            cursor = self._conn.execute(
+                "SELECT * FROM sessions WHERE in_memory_id = ?", (in_memory_id,)
+            )
+            row = cursor.fetchone()
+        return dict(row) if row else None
+
+    def set_session_in_memory_id(self, session_id: str, in_memory_id: str) -> None:
+        """Store the in-memory UUID for a session so it can be resolved later."""
+        def _do(conn):
+            conn.execute(
+                "UPDATE sessions SET in_memory_id = ? WHERE id = ?",
+                (in_memory_id, session_id),
+            )
+        self._execute_write(_do)
+
     def resolve_session_id(self, session_id_or_prefix: str) -> Optional[str]:
         """Resolve an exact or uniquely prefixed session ID to the full ID.
 

tests/test_hermes_state.py

@@ -331,6 +331,51 @@ def connect_without_fts(*args, **kwargs):
             restored.close()
 
 
+# =========================================================================
+# In-memory ID resolution (TUI crash recovery)
+# =========================================================================
+
+class TestInMemoryIdResolution:
+    """Verify that sessions can be looked up by their in-memory UUID,
+    which the TUI stores as ``sid`` and sends during gateway crash recovery."""
+
+    def test_set_and_get_by_in_memory_id(self, db):
+        db.create_session(session_id="stored_key_123", source="tui")
+        db.set_session_in_memory_id("stored_key_123", "inmem_abc123")
+
+        session = db.get_session_by_in_memory_id("inmem_abc123")
+        assert session is not None
+        assert session["id"] == "stored_key_123"
+
+    def test_in_memory_id_none_when_not_set(self, db):
+        db.create_session(session_id="s1", source="tui")
+        session = db.get_session_by_in_memory_id("nonexistent")
+        assert session is None
+
+    def test_in_memory_id_overwrite(self, db):
+        db.create_session(session_id="s1", source="tui")
+        db.set_session_in_memory_id("s1", "old_uuid")
+        db.set_session_in_memory_id("s1", "new_uuid")
+
+        assert db.get_session_by_in_memory_id("old_uuid") is None
+        assert db.get_session_by_in_memory_id("new_uuid") is not None
+
+    def test_in_memory_id_survives_end_and_reopen(self, db):
+        """After a gateway crash, the session is ended then reopened.
+        The in_memory_id column must still resolve correctly."""
+        db.create_session(session_id="s1", source="tui")
+        db.set_session_in_memory_id("s1", "inmem_xyz")
+        db.end_session("s1", end_reason="tui_close")
+
+        session = db.get_session_by_in_memory_id("inmem_xyz")
+        assert session is not None
+        assert session["id"] == "s1"
+
+        db.reopen_session("s1")
+        session = db.get_session_by_in_memory_id("inmem_xyz")
+        assert session is not None
+
+
 # =========================================================================
 # Message storage
 # =========================================================================

tui_gateway/server.py

@@ -702,7 +702,7 @@ def _register_session_cwd(session: dict | None) -> None:
         pass
 
 
-def _ensure_session_db_row(session: dict) -> None:
+def _ensure_session_db_row(session: dict, in_memory_id: str | None = None) -> None:
     """Idempotently persist the session's DB row on first real activity.
 
     Called from prompt.submit so a row only exists once the user actually sends
@@ -716,6 +716,10 @@ def _ensure_session_db_row(session: dict) -> None:
     picked a folder for gets grouped under whatever directory the desktop
     happened to launch in (e.g. "desktop"). Leaving it null groups them under
     "No workspace", which is the desired default.
+
+    When *in_memory_id* is provided (the TUI's in-process session UUID), it is
+    written to the ``in_memory_id`` column so that ``session.resume`` can
+    resolve the in-memory UUID back to the DB row after a gateway respawn.
     """
     key = session.get("session_key")
     if not key:
@@ -730,6 +734,11 @@ def _ensure_session_db_row(session: dict) -> None:
             model=_resolve_model(),
             cwd=_session_cwd(session) if session.get("explicit_cwd") else None,
         )
+        if in_memory_id:
+            try:
+                db.set_session_in_memory_id(key, in_memory_id)
+            except Exception:
+                logger.debug("failed to set in_memory_id on session row", exc_info=True)
     except Exception:
         logger.debug("failed to persist desktop session row", exc_info=True)
 
@@ -2988,7 +2997,16 @@ def _(rid, params: dict) -> dict:
         if found:
             target = found["id"]
         else:
-            return _err(rid, 4007, "session not found")
+            # Fallback: the TUI stores an in-memory session UUID that is
+            # different from the DB primary key.  After a gateway respawn
+            # (crash recovery), the TUI sends this in-memory UUID.  Resolve
+            # it via the ``in_memory_id`` column stored by
+            # ``_ensure_session_db_row``.
+            found = db.get_session_by_in_memory_id(target)
+            if found:
+                target = found["id"]
+            else:
+                return _err(rid, 4007, "session not found")
     sid = uuid.uuid4().hex[:8]
     _enable_gateway_prompts()
     try:
@@ -3959,7 +3977,7 @@ def _(rid, params: dict) -> dict:
         _start_inflight_turn(session, text)
 
     # Persist the DB row lazily, now that the user has actually sent a message.
-    _ensure_session_db_row(session)
+    _ensure_session_db_row(session, in_memory_id=sid)
     _start_agent_build(sid, session)
 
     def run_after_agent_ready() -> None: