From 1a2f3aea9a62ecb728b602d37aa050b6e0b7c004 Mon Sep 17 00:00:00 2001 From: Wilson Kinyua Date: Fri, 10 Jul 2026 13:59:07 +0200 Subject: [PATCH] fix tui finalize persist drop conversation_history so disconnect saves chat finalize passed conversation_history=history aliasing the snapshot so flush skipped every message and wrote nothing. now flush _session_messages via marker dedup like gateway shutdown. add real db e2e tests. --- .../test_finalize_session_persist.py | 182 ++++++++++++++++-- tui_gateway/server.py | 27 ++- 2 files changed, 176 insertions(+), 33 deletions(-) diff --git a/tests/tui_gateway/test_finalize_session_persist.py b/tests/tui_gateway/test_finalize_session_persist.py index e1fe7ea5372..c927488de57 100644 --- a/tests/tui_gateway/test_finalize_session_persist.py +++ b/tests/tui_gateway/test_finalize_session_persist.py @@ -54,11 +54,11 @@ def _make_session(agent=None, history=None, session_key="test_key_001"): class TestFinalizeSessionPersist: """Verify _finalize_session flushes messages via _persist_session.""" - def test_persist_called_with_history(self): - """History from session is passed to agent._persist_session. - - When _session_messages is None (not yet set by any turn), - the session["history"] is used as the snapshot. + def test_no_session_messages_skips_persist(self): + """When _session_messages is empty/None the agent processed nothing + this session, so there is nothing new to flush. Falling back to + session["history"] here re-appended already-durable resumed rows as + duplicates, so finalize must NOT write in that case. """ from tui_gateway.server import _finalize_session @@ -66,20 +66,16 @@ class TestFinalizeSessionPersist: {"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi there"}, ] - agent = _make_agent() + agent = _make_agent() # _session_messages is None session = _make_session(agent=agent, history=history) _finalize_session(session, end_reason="test") - agent._persist_session.assert_called_once() - # snapshot = history (since _session_messages is None) - called_with = agent._persist_session.call_args[0][0] - assert called_with == history - # conversation_history kwarg passed for correct flush indexing - assert agent._persist_session.call_args[1].get("conversation_history") == history + agent._persist_session.assert_not_called() - def test_persist_uses_session_messages_when_available(self): - """agent._session_messages takes priority over session['history'].""" + def test_persist_uses_session_messages(self): + """agent._session_messages is flushed via the marker-based dedup path + (no conversation_history — passing the same list neutered the write).""" from tui_gateway.server import _finalize_session history = [{"role": "user", "content": "old"}] @@ -93,10 +89,10 @@ class TestFinalizeSessionPersist: _finalize_session(session) - agent._persist_session.assert_called_once() - called_with = agent._persist_session.call_args[0][0] - assert called_with == session_msgs # _session_messages wins - assert agent._persist_session.call_args[1].get("conversation_history") == history + agent._persist_session.assert_called_once_with(session_msgs) + # conversation_history must NOT be passed — it aliases the snapshot and + # makes _flush_messages_to_session_db skip every message. + assert "conversation_history" not in agent._persist_session.call_args[1] def test_commit_memory_still_called(self): """Existing memory commit path is preserved.""" @@ -158,6 +154,7 @@ class TestFinalizeSessionPersist: from tui_gateway.server import _finalize_session agent = _make_agent() + agent._session_messages = [{"role": "user", "content": "x"}] agent._persist_session.side_effect = RuntimeError("db is down") session = _make_session( agent=agent, @@ -165,6 +162,7 @@ class TestFinalizeSessionPersist: ) _finalize_session(session) # must not raise + agent._persist_session.assert_called_once() # commit_memory_session should still be called agent.commit_memory_session.assert_called_once() @@ -184,6 +182,154 @@ class TestFinalizeSessionPersist: mock_db.end_session.assert_called_once_with("sess_123", "test") +class TestFinalizeSessionPersistE2E: + """End-to-end: _finalize_session must actually land unflushed turns in + state.db on disconnect/restart. + + The mock-based tests above assert that _persist_session is *called*, but a + call whose arguments neuter the underlying flush persists nothing. These + tests drive the REAL AIAgent flush against a REAL SessionDB, reproducing + the "conversation contains many events yet is absent from state.db across + disconnect/restart" symptom. + """ + + @staticmethod + def _real_agent(db, session_id, session_messages): + from run_agent import AIAgent + + agent = object.__new__(AIAgent) + agent._session_db = db + agent._session_db_created = True + agent.session_id = session_id + agent.platform = "tui" + agent.model = "test-model" + agent._session_messages = session_messages + agent._last_flushed_db_idx = 0 + agent._flushed_db_message_ids = set() + agent._flushed_db_message_session_id = None + agent._persist_disabled = False + agent._cached_system_prompt = None + agent._session_init_model_config = None + agent._parent_session_id = None + agent._session_json_enabled = False + agent.quiet_mode = True + # commit_memory_session runs heavy machinery we don't exercise here. + agent.commit_memory_session = lambda *a, **k: None + return agent + + def test_unflushed_turn_survives_disconnect(self, tmp_path, monkeypatch): + """A completed turn whose transcript flush did NOT durably persist + (messages live only in agent._session_messages / session['history'], + never written to the DB) must be flushed to state.db when the WS + disconnect tears the session down.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes")) + from hermes_state import SessionDB + import tui_gateway.server as srv + + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "sess-unflushed" + db.create_session(session_id=session_id, source="tui") + monkeypatch.setattr(srv, "_get_db", lambda: db) + + # The live turn list that became session["history"] AND + # agent._session_messages (same object), but was never persisted. + turn = [ + {"role": "user", "content": "scan the repo and summarise"}, + {"role": "assistant", "content": "Here is the summary…"}, + {"role": "user", "content": "now open a PR"}, + {"role": "assistant", "content": "PR opened."}, + ] + agent = self._real_agent(db, session_id, turn) + session = _make_session(agent=agent, history=turn, session_key=session_id) + + assert db.get_messages_as_conversation(session_id) == [] + + srv._finalize_session(session, end_reason="ws_disconnect") + + after = db.get_messages_as_conversation(session_id) + contents = [m.get("content") for m in after] + assert len(after) == 4, after + assert any("scan the repo" in (c or "") for c in contents), contents + assert any("PR opened" in (c or "") for c in contents), contents + + def test_resumed_session_not_reflushed_as_duplicates(self, tmp_path, monkeypatch): + """A resumed session torn down before any new turn (its transcript is + already durable in the DB) must NOT re-append duplicate rows.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes")) + from hermes_state import SessionDB + import tui_gateway.server as srv + + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "sess-resumed" + db.create_session(session_id=session_id, source="tui") + loaded = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi there"}, + ] + for m in loaded: + db.append_message(session_id=session_id, role=m["role"], content=m["content"]) + monkeypatch.setattr(srv, "_get_db", lambda: db) + + # Resumed session: history hydrated from the DB, no turn ran, so the + # agent processed nothing this session. + agent = self._real_agent(db, session_id, []) + session = _make_session(agent=agent, history=loaded, session_key=session_id) + + srv._finalize_session(session, end_reason="ws_disconnect") + + after = db.get_messages_as_conversation(session_id) + assert len(after) == 2, after + + def test_resumed_then_run_turn_not_duplicated(self, tmp_path, monkeypatch): + """A resumed session that RUNS a turn must not have its loaded (durable) + prefix re-appended by finalize. + + This exercises the exact path the ``conversation_history`` argument used + to guard: the in-turn flush stamps the loaded prefix with + ``_DB_PERSISTED_MARKER`` (recognising it as durable), so the marker-only + finalize flush skips it. Without that stamping — or if finalize wrote a + markerless copy — the durable prefix would double. The + ``_session_messages``-empty test above skips the flush entirely, so it + can't catch a duplicate-write regression; this one drives a real flush. + """ + monkeypatch.setenv("HERMES_HOME", str(tmp_path / ".hermes")) + from hermes_state import SessionDB + import tui_gateway.server as srv + + db = SessionDB(db_path=tmp_path / "state.db") + session_id = "sess-resume-run" + db.create_session(session_id=session_id, source="tui") + loaded = [ + {"role": "user", "content": "remember: my cat is Mochi"}, + {"role": "assistant", "content": "Noted — Mochi."}, + ] + for m in loaded: + db.append_message(session_id=session_id, role=m["role"], content=m["content"]) + monkeypatch.setattr(srv, "_get_db", lambda: db) + + # Live turn list = loaded prefix (same dicts, as run_conversation copies + # conversation_history) + the new turn. + new_turn = [ + {"role": "user", "content": "what's the cat's name?"}, + {"role": "assistant", "content": "Mochi."}, + ] + messages = list(loaded) + new_turn + agent = self._real_agent(db, session_id, messages) + + # Drive the in-turn flush the way run_conversation does — the loaded + # prefix rides in as conversation_history, so it is recognised durable + # (and marker-stamped) while only the new turn is written. + agent._flush_messages_to_session_db(messages, conversation_history=loaded) + assert len(db.get_messages_as_conversation(session_id)) == 4 + + # WS disconnect → finalize. Must re-append nothing. + session = _make_session(agent=agent, history=messages, session_key=session_id) + srv._finalize_session(session, end_reason="ws_disconnect") + + after = db.get_messages_as_conversation(session_id) + assert len(after) == 4, after + + class TestOnSessionEndHook: """Verify on_session_end plugin hook fires on finalize.""" diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 8e825a22669..64e37d3b540 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -568,24 +568,21 @@ def _finalize_session(session: dict | None, end_reason: str = "tui_close") -> No history = list(session.get("history", [])) # ── Persist unflushed messages to SQLite ────────────────────────── - # Two sources, tried in order of freshness: - # 1. agent._session_messages — set by the last _persist_session() - # call inside run_conversation(). This is the most recent - # snapshot the agent thread wrote, and may include partial - # turn data that hasn't reached session["history"] yet. - # 2. session["history"] — updated after run_conversation() - # returns. Stale when the agent is mid‑turn, but correct - # when the turn completed before finalize. - # Best‑effort — the agent thread may still be mid‑turn, so only - # previously completed messages are guaranteed. + # Flush ``agent._session_messages`` via ``_persist_session``'s marker-based + # dedup (same contract as the gateway-shutdown flush, #13121). Do NOT pass + # ``conversation_history``: ``session["history"]`` and ``_session_messages`` + # alias the SAME list once a turn completes, so passing it made + # ``_flush_messages_to_session_db`` treat every message as already-durable + # and skip it — a data-loss bug when finalize is the sole persist path after + # a WS disconnect/restart (e.g. the in-turn flush hit a transient SQLite + # failure). Markers persist the genuinely-unflushed tail without duplicating + # durable rows (including a resumed-but-not-run session's already-in-DB + # transcript, which stays in ``session["history"]`` only). if agent is not None and hasattr(agent, "_persist_session"): - snapshot = ( - getattr(agent, "_session_messages", None) - or history - ) + snapshot = getattr(agent, "_session_messages", None) if snapshot: try: - agent._persist_session(snapshot, conversation_history=history) + agent._persist_session(snapshot) except Exception: pass