diff --git a/hermes_state.py b/hermes_state.py index 1a3a4ff4e542d..1d5fb355a454f 100644 --- a/hermes_state.py +++ b/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. diff --git a/tests/test_hermes_state.py b/tests/test_hermes_state.py index f5e4f69ae66cb..b13ed4700c385 100644 --- a/tests/test_hermes_state.py +++ b/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 # ========================================================================= diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 338218cd8f90e..291e34ead00d7 100644 --- a/tui_gateway/server.py +++ b/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: