diff --git a/tests/tui_gateway/test_prompt_submit_durable_refresh.py b/tests/tui_gateway/test_prompt_submit_durable_refresh.py new file mode 100644 index 0000000000000..d903bc0fcc109 --- /dev/null +++ b/tests/tui_gateway/test_prompt_submit_durable_refresh.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import contextlib +import logging +import threading + +from tui_gateway.prompt_history_sync import ( + refresh_resumed_history_before_submit, + wrap_prompt_submit, +) + + +class _FakeDB: + def __init__(self, history, *, on_read=None): + self.history = list(history) + self.on_read = on_read + self.reads = 0 + + def get_resume_conversations(self, session_key): + assert session_key == "cron_session" + self.reads += 1 + if self.on_read is not None: + self.on_read() + return list(self.history), list(self.history) + + +class _FakeServer: + def __init__(self, session, db): + self._sessions = {"live-sid": session} + self._db = db + self.logger = logging.getLogger("test_prompt_submit_durable_refresh") + + def _session_db(self, _session): + return contextlib.nullcontext(self._db) + + @staticmethod + def sanitize_replay_history(history): + return list(history) + + +def _messages(count): + return [ + {"role": "user" if i % 2 == 0 else "assistant", "content": f"m{i}"} + for i in range(count) + ] + + +def _resumed_session(history): + ready = threading.Event() + ready.set() + return { + "history": list(history), + "history_lock": threading.Lock(), + "history_version": 0, + "resume_history_ready": ready, + "resume_session_id": "cron_session", + "running": False, + "session_key": "cron_session", + } + + +def test_wrap_refreshes_open_time_prefix_before_prompt_handler_runs(): + """Regression for #91508: the handler must see all externally persisted rows.""" + opened_at = _messages(12) + durable_after_cron_finished = _messages(37) + session = _resumed_session(opened_at) + db = _FakeDB(durable_after_cron_finished) + server = _FakeServer(session, db) + seen = {} + + def handler(_rid, _params): + seen["history"] = list(session["history"]) + return {"ok": True} + + result = wrap_prompt_submit(server, handler)( + "1", {"session_id": "live-sid", "text": "follow-up"} + ) + + assert result == {"ok": True} + assert len(seen["history"]) == 37 + assert seen["history"][-1]["content"] == "m36" + assert session["history_version"] == 1 + assert db.reads == 1 + + +def test_refresh_does_not_touch_new_non_resumed_sessions(): + session = _resumed_session(_messages(2)) + session.pop("resume_session_id") + db = _FakeDB(_messages(8)) + server = _FakeServer(session, db) + + assert refresh_resumed_history_before_submit( + server, {"session_id": "live-sid", "text": "hello"} + ) is False + assert len(session["history"]) == 2 + assert session["history_version"] == 0 + assert db.reads == 0 + + +def test_refresh_never_replaces_memory_with_a_shorter_durable_projection(): + session = _resumed_session(_messages(8)) + db = _FakeDB(_messages(6)) + server = _FakeServer(session, db) + + assert refresh_resumed_history_before_submit( + server, {"session_id": "live-sid", "text": "hello"} + ) is False + assert len(session["history"]) == 8 + assert session["history_version"] == 0 + + +def test_local_history_mutation_wins_if_it_races_the_durable_read(): + session = _resumed_session(_messages(12)) + + def mutate_locally(): + with session["history_lock"]: + session["history"] = _messages(13) + session["history_version"] += 1 + + db = _FakeDB(_messages(37), on_read=mutate_locally) + server = _FakeServer(session, db) + + assert refresh_resumed_history_before_submit( + server, {"session_id": "live-sid", "text": "hello"} + ) is False + assert len(session["history"]) == 13 + assert session["history_version"] == 1 + + +def test_busy_session_keeps_the_active_turn_as_history_owner(): + session = _resumed_session(_messages(12)) + session["running"] = True + db = _FakeDB(_messages(37)) + server = _FakeServer(session, db) + + assert refresh_resumed_history_before_submit( + server, {"session_id": "live-sid", "text": "queued follow-up"} + ) is False + assert len(session["history"]) == 12 + assert db.reads == 0 diff --git a/tui_gateway/method_ctx.py b/tui_gateway/method_ctx.py index a72f802bfaec8..7b5f17a3bc8bc 100644 --- a/tui_gateway/method_ctx.py +++ b/tui_gateway/method_ctx.py @@ -48,6 +48,16 @@ def install(self, server) -> None: real.__kwdefaults__ = fn.__kwdefaults__ real.__doc__ = fn.__doc__ real.__dict__.update(fn.__dict__) + if name == "prompt.submit": + # A resumed session's model-history snapshot can be advanced by + # an external writer (notably cron) without touching this + # gateway process. Refresh at the dispatch boundary so the + # handler claims the turn only after adopting newer durable + # context. Kept here to preserve methods_prompt.py's mechanical + # handler split and avoid a second copy of its large body. + from .prompt_history_sync import wrap_prompt_submit + + real = wrap_prompt_submit(server, real) if getattr(fn, "_hermes_profile_scoped", False): real = server._profile_scoped(real) server._methods[name] = real diff --git a/tui_gateway/prompt_history_sync.py b/tui_gateway/prompt_history_sync.py new file mode 100644 index 0000000000000..6e01cb8e08dd9 --- /dev/null +++ b/tui_gateway/prompt_history_sync.py @@ -0,0 +1,160 @@ +"""Durable-history refresh for resumed ``prompt.submit`` calls. + +A resumed Desktop/TUI session keeps an in-memory model-history snapshot. Other +processes can continue writing the same durable session (cron is the important +case), so the display transcript may advance while the model snapshot stays at +its open-time prefix. Before a new turn is claimed, refresh an idle resumed +session from its owning state.db when the durable model projection has grown. +""" + +from __future__ import annotations + +import functools +import logging +from typing import Any, Callable + + +_LOG = logging.getLogger(__name__) + + +def _server_logger(server): + return getattr(server, "logger", _LOG) + + +def refresh_resumed_history_before_submit(server, params: dict[str, Any]) -> bool: + """Adopt newer durable model history for an idle resumed session. + + Returns ``True`` only when ``session['history']`` was replaced. The update + is deliberately monotonic: a shorter/equal durable projection never + replaces memory, and a local history mutation that races the DB read wins. + This keeps the guard safe for queued/live turns while repairing the stale + open-time snapshot produced by cross-process writers. + """ + if not isinstance(params, dict): + return False + + sid = str(params.get("session_id") or "") + sessions = getattr(server, "_sessions", None) + session = sessions.get(sid) if isinstance(sessions, dict) else None + if not isinstance(session, dict) or not session.get("resume_session_id"): + return False + + # Deferred Desktop resumes hydrate once in a background worker. Never race + # that assignment: wait for its completion/error signal before taking a + # fresh durable snapshot. A bounded wait avoids making prompt dispatch + # permanently dependent on a stuck hydration worker. + ready = session.get("resume_history_ready") + if ready is not None and callable(getattr(ready, "wait", None)): + try: + if not ready.wait(timeout=30.0): + _server_logger(server).warning( + "prompt.submit: resume history still hydrating for session %s; " + "skipping durable refresh", + sid, + ) + return False + except Exception: + _server_logger(server).debug( + "prompt.submit: failed waiting for resume history for session %s", + sid, + exc_info=True, + ) + return False + + if session.get("resume_history_error"): + return False + + lock = session.get("history_lock") + if lock is None: + return False + + with lock: + # A busy submit is handled by the normal queue/interrupt path. Its + # active turn owns the in-memory transcript and must not be overwritten + # from a concurrent durable read. + if session.get("running"): + return False + start_version = int(session.get("history_version", 0)) + start_history = list(session.get("history") or []) + + session_key = str( + session.get("session_key") or session.get("resume_session_id") or sid + ) + if not session_key: + return False + + try: + session_db = getattr(server, "_session_db") + with session_db(session) as db: + if db is None: + return False + get_resume = getattr(db, "get_resume_conversations", None) + if callable(get_resume): + raw_history, _display_history = get_resume(session_key) + else: + get_history = getattr(db, "get_messages_as_conversation", None) + if not callable(get_history): + return False + raw_history = get_history( + session_key, + repair_alternation=True, + include_row_ids=True, + ) + except Exception: + _server_logger(server).debug( + "prompt.submit: failed refreshing durable history for resumed session %s", + session_key, + exc_info=True, + ) + return False + + if not isinstance(raw_history, list): + return False + sanitize = getattr(server, "sanitize_replay_history", None) + try: + durable_history = sanitize(raw_history) if callable(sanitize) else raw_history + except Exception: + _server_logger(server).debug( + "prompt.submit: failed sanitizing refreshed history for session %s", + session_key, + exc_info=True, + ) + return False + if not isinstance(durable_history, list): + return False + + with lock: + # A local mutation after the snapshot wins. Checking both the explicit + # version and the value catches legacy/in-place mutations that forgot to + # bump history_version. + if session.get("running"): + return False + if int(session.get("history_version", 0)) != start_version: + return False + if list(session.get("history") or []) != start_history: + return False + if len(durable_history) <= len(start_history): + return False + + session["history"] = list(durable_history) + session["history_version"] = start_version + 1 + + _server_logger(server).info( + "prompt.submit: refreshed resumed session %s history %d -> %d messages " + "from durable state", + session_key, + len(start_history), + len(durable_history), + ) + return True + + +def wrap_prompt_submit(server, handler: Callable) -> Callable: + """Wrap the registered ``prompt.submit`` handler with the refresh guard.""" + + @functools.wraps(handler) + def wrapped(rid, params): + refresh_resumed_history_before_submit(server, params) + return handler(rid, params) + + return wrapped