diff --git a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx index ea5cc084ce..866194227b 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-session-actions.test.tsx @@ -6,7 +6,13 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { $terminalTakeover, setTerminalTakeover } from '@/app/right-sidebar/store' import { noteActiveTreeGroup, revealTreePane } from '@/components/pane-shell/tree/store' -import { getAllSessionMessages, getLatestSessionMessages, getSession, type SessionInfo } from '@/hermes' +import { + getAllSessionMessages, + getLatestSessionMessages, + getSession, + type SessionInfo, + type SessionResumeResponse +} from '@/hermes' import { createClientSessionState } from '@/lib/chat-runtime' import { clearSessionDraft, stashSessionDraft, takeSessionDraft } from '@/store/composer' import { $activeGatewayProfile, $newChatProfile, ensureGatewayProfile } from '@/store/profile' @@ -1736,13 +1742,65 @@ describe('resumeSession warm-cache mapping integrity', () => { // resume RPC ran, for the session that was actually requested. const resumeCalls = requestGateway.mock.calls.filter(([method]) => method === 'session.resume') expect(resumeCalls.length).toBe(1) - expect(resumeCalls[0][1]).toMatchObject({ session_id: 'stored-A' }) + expect(resumeCalls[0][1]).toMatchObject({ + defer_history: true, + session_id: 'stored-A' + }) + expect(getLatestSessionMessages).toHaveBeenCalledWith('stored-A', undefined) // The corrupt mapping was purged so it can't mis-resolve again. expect(runtimeIdByStoredSessionIdRef.current.has('stored-A')).toBe(false) expect(sessionStateByRuntimeIdRef.current.has('rt-recycled')).toBe(false) }) + it('paints the bounded latest transcript before deferred resume finishes', async () => { + const latestPage = Array.from({ length: 500 }, (_, index) => ({ + content: `message-${index}`, + role: index % 2 === 0 ? ('user' as const) : ('assistant' as const), + timestamp: index + 1 + })) + + setSessions([storedSession({ id: 'stored-A', message_count: 50_000 })]) + vi.mocked(getLatestSessionMessages).mockReset() + vi.mocked(getLatestSessionMessages).mockResolvedValue({ + messages: latestPage, + pagination: { limit: 500, offset: 0, order: 'latest', returned: 500 }, + session_id: 'stored-A' + }) + + const deferredResume = deferred() + + const requestGatewayMock = vi.fn((method: string, _params?: Record) => { + if (method === 'session.resume') { + return deferredResume.promise + } + + return Promise.resolve({}) + }) + + const requestGateway = (method: string, params?: Record): Promise => + requestGatewayMock(method, params) as Promise + + let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise) | null = null + render( (resume = value)} requestGateway={requestGateway} />) + await waitFor(() => expect(resume).not.toBeNull()) + const resumePromise = resume!('stored-A', true) + + await waitFor(() => expect($messages.get()).toHaveLength(500)) + expect(getLatestSessionMessages).toHaveBeenCalledTimes(1) + expect(getLatestSessionMessages).toHaveBeenCalledWith('stored-A', undefined) + + deferredResume.resolve({ + session_id: 'rt-A', + resumed: 'stored-A', + message_count: 500, + messages: [], + info: {} + }) + await resumePromise + expect($messages.get()).toHaveLength(500) + }) + it('honours a warm cache entry whose stored id matches and refreshes its persisted transcript', async () => { // Correctly-wired mapping: 'rt-A' <-> 'stored-A'. The fast-path should trust // it and never reach session.resume. session.activate refreshes the live diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts index 361a553fd0..57e93046b9 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts @@ -1030,6 +1030,7 @@ export function useSessionActions({ session_id: storedSessionId, cols: 96, source: 'desktop', + defer_history: !watchWindow, // REST is the transcript authority for Desktop. Avoid duplicating a // potentially huge compression lineage in the WebSocket response. // Watch windows attach lazily (live mirror). Every other cold resume @@ -1079,7 +1080,6 @@ export function useSessionActions({ prefetchApplied = true prefetchedStoredSessionId = prefetchedResult.session_id || storedSessionId } - const currentMessages = $messages.get() // Keep the local snapshot when resume would only reshuffle runtime diff --git a/apps/desktop/src/types/hermes.ts b/apps/desktop/src/types/hermes.ts index b6291e8455..c9c28eef8f 100644 --- a/apps/desktop/src/types/hermes.ts +++ b/apps/desktop/src/types/hermes.ts @@ -606,6 +606,7 @@ export interface SessionResumeResponse { attempt: number interrupted_at: number } + hydrating?: boolean inflight?: null | { assistant?: string /** Mid-turn redirect corrections, oldest first. The turn's original prompt diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index fa6ba36e7e..26fa685f18 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -3282,6 +3282,221 @@ def test_lazy_child_watch_resume_serves_candidate_inclusive_display(monkeypatch, assert texts == ["child prompt", "child substantive answer", "child terse reply"] +def test_session_resume_deferred_history_acknowledges_and_reuses(monkeypatch): + history_started = threading.Event() + release_history = threading.Event() + build_started = threading.Event() + history_calls = [] + auto_continue_calls = [] + ancestor = {"role": "assistant", "content": "ancestor"} + loaded = {"role": "user", "content": "loaded"} + + class FakeDB: + def get_session(self, target): + return {"id": target, "message_count": 1200} + + def resolve_resume_session_id(self, target): + return target + + def reopen_session(self, target): + assert target == "large-session" + + def get_resume_conversations(self, target): + history_calls.append(("resume", target)) + history_started.set() + assert release_history.wait(timeout=2.0) + return [loaded], [ancestor, loaded] + + def get_ancestor_display_prefix(self, target): + history_calls.append(("prefix", target)) + return [ancestor] + + monkeypatch.setattr(server, "_get_db", lambda: FakeDB()) + monkeypatch.setattr(server, "_enable_gateway_prompts", lambda: None) + monkeypatch.setattr(server, "_schedule_session_cap_enforcement", lambda: None) + monkeypatch.setattr( + server, + "_start_agent_build", + lambda _sid, _session: build_started.set(), + ) + monkeypatch.setattr( + server, + "_maybe_schedule_auto_continue", + lambda sid, session, stored_id: auto_continue_calls.append( + (sid, session, stored_id) + ), + ) + + try: + first = server._methods["session.resume"]( + "r1", + { + "session_id": "large-session", + "source": "desktop", + "defer_history": True, + }, + ) + + assert first["result"]["hydrating"] is True + assert first["result"]["messages"] == [] + assert first["result"]["message_count"] == 1200 + assert history_started.wait(timeout=1.0) + + second = server._methods["session.resume"]( + "r2", + {"session_id": "large-session", "defer_history": True}, + ) + assert second["result"]["session_id"] == first["result"]["session_id"] + assert second["result"]["hydrating"] is True + assert second["result"]["messages"] == [] + + release_history.set() + sid = first["result"]["session_id"] + assert server._sessions[sid]["resume_history_ready"].wait(timeout=1.0) + assert build_started.wait(timeout=1.0) + assert history_calls == [ + ("resume", "large-session"), + ("prefix", "large-session"), + ] + assert server._sessions[sid]["history"] == [loaded] + assert server._sessions[sid]["display_history_prefix"] == [ancestor] + assert server._sessions[sid]["resume_message_count"] == 2 + assert auto_continue_calls == [(sid, server._sessions[sid], "large-session")] + finally: + release_history.set() + for sid, session in list(server._sessions.items()): + if session.get("session_key") == "large-session": + lease = session.get("active_session_lease") + if lease is not None: + lease.release() + server._sessions.pop(sid, None) + + +def test_session_resume_deferred_history_failure_can_retry(monkeypatch): + first_released = threading.Event() + build_started = threading.Event() + attempts = 0 + + class FakeDB: + def get_session(self, target): + return {"id": target, "message_count": 1} + + def resolve_resume_session_id(self, target): + return target + + def reopen_session(self, _target): + pass + + def get_resume_conversations(self, _target): + nonlocal attempts + attempts += 1 + if attempts == 1: + first_released.set() + raise RuntimeError("sqlite read failed") + loaded = [{"role": "user", "content": "retry loaded"}] + return loaded, loaded + + def get_ancestor_display_prefix(self, _target): + return [] + + monkeypatch.setattr(server, "_get_db", lambda: FakeDB()) + monkeypatch.setattr( + server, + "_claim_active_session_slot", + lambda *_args, **_kwargs: pytest.fail( + "resume must not claim a session slot before the first prompt" + ), + ) + monkeypatch.setattr(server, "_enable_gateway_prompts", lambda: None) + monkeypatch.setattr(server, "_schedule_session_cap_enforcement", lambda: None) + monkeypatch.setattr( + server, + "_start_agent_build", + lambda _sid, _session: build_started.set(), + ) + + try: + first = server._methods["session.resume"]( + "r1", + {"session_id": "retry-session", "defer_history": True}, + ) + first_sid = first["result"]["session_id"] + assert first_released.wait(timeout=1.0) + assert first_sid not in server._sessions + + second = server._methods["session.resume"]( + "r2", + {"session_id": "retry-session", "defer_history": True}, + ) + second_sid = second["result"]["session_id"] + assert second_sid != first_sid + assert server._sessions[second_sid]["resume_history_ready"].wait(timeout=1.0) + assert build_started.wait(timeout=1.0) + finally: + for sid, session in list(server._sessions.items()): + if session.get("session_key") == "retry-session": + lease = session.get("active_session_lease") + if lease is not None: + lease.release() + server._sessions.pop(sid, None) + + +def test_session_resume_deferred_history_close_cancels_build(monkeypatch): + history_started = threading.Event() + release_history = threading.Event() + build_started = threading.Event() + + class FakeDB: + def get_session(self, target): + return {"id": target, "message_count": 1} + + def resolve_resume_session_id(self, target): + return target + + def reopen_session(self, _target): + pass + + def get_resume_conversations(self, _target): + history_started.set() + assert release_history.wait(timeout=2.0) + loaded = [{"role": "user", "content": "late"}] + return loaded, loaded + + def get_ancestor_display_prefix(self, _target): + return [] + + monkeypatch.setattr(server, "_get_db", lambda: FakeDB()) + monkeypatch.setattr(server, "_enable_gateway_prompts", lambda: None) + monkeypatch.setattr(server, "_schedule_session_cap_enforcement", lambda: None) + monkeypatch.setattr( + server, + "_start_agent_build", + lambda _sid, _session: build_started.set(), + ) + + response = {} + try: + response = server._methods["session.resume"]( + "r1", + {"session_id": "cancel-session", "defer_history": True}, + ) + sid = response["result"]["session_id"] + session = server._sessions[sid] + assert history_started.wait(timeout=1.0) + + assert server._close_session_by_id(sid, end_reason="tui_close") is True + assert session["resume_history_ready"].is_set() + assert session["resume_history_error"] == "session resume cancelled" + + release_history.set() + time.sleep(0.05) + assert not build_started.is_set() + assert sid not in server._sessions + finally: + release_history.set() + server._sessions.pop(response.get("result", {}).get("session_id", ""), None) + + def test_session_resume_follows_compression_tip(monkeypatch, tmp_path): """Resuming a rotated-out parent id must load the continuation's messages. diff --git a/tests/tui_gateway/test_session_resume_db_ownership.py b/tests/tui_gateway/test_session_resume_db_ownership.py index 309d1e3998..f31117b7cb 100644 --- a/tests/tui_gateway/test_session_resume_db_ownership.py +++ b/tests/tui_gateway/test_session_resume_db_ownership.py @@ -28,6 +28,7 @@ Pinned here, in both directions: from __future__ import annotations +import threading import types import pytest @@ -196,6 +197,43 @@ def test_resume_closes_profile_db_on_deferred_cold_resume(profile_dbs, monkeypat assert profile_dbs[0].closed == 1 +def test_resume_hands_profile_db_to_deferred_history_worker(profile_dbs, monkeypatch): + """Incremental hydration owns the profile handle until its read completes.""" + history_started = threading.Event() + release_history = threading.Event() + + class _BlockingDB(_RecordingDB): + def get_resume_conversations(self, _target): + history_started.set() + assert release_history.wait(timeout=2.0) + assert self.closed == 0 + return ([], []) + + def _factory(db_path=None, **kwargs): + db = _BlockingDB(db_path=db_path, **kwargs) + db.rows["s1"] = {"id": "s1", "cwd": "", "message_count": 0} + profile_dbs.append(db) + return db + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + monkeypatch.setattr(server, "_stored_session_runtime_overrides", lambda _found: {}) + + try: + resp = _resume( + session_id="s1", profile="work", defer_history=True + ) + sid = resp["result"]["session_id"] + db = profile_dbs[0] + assert history_started.wait(timeout=1.0) + assert db.closed == 0 + + release_history.set() + assert server._sessions[sid]["resume_history_ready"].wait(timeout=1.0) + assert db.closed == 1 + finally: + release_history.set() + + def test_resume_keeps_profile_db_open_after_ownership_transfer(profile_dbs, monkeypatch): """A COMPLETED resume transfers the handle to the agent — do not close it. diff --git a/tui_gateway/methods_session.py b/tui_gateway/methods_session.py index d7e95edcdc..7e357ab030 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -317,6 +317,7 @@ def _(rid, params: dict) -> dict: # local profile's state.db. None/own profile → the launch profile (unchanged). profile = (params.get("profile") or "").strip() or None profile_home = _profile_home(profile) + defer_history = is_truthy_value(params.get("defer_history", False)) # Desktop hydrates persisted transcripts through the authenticated REST # route in parallel. Suppress the duplicate WebSocket transcript only when # the caller explicitly requests it; other clients keep upstream behavior. @@ -424,6 +425,12 @@ def _(rid, params: dict) -> dict: omit_messages=omit_messages, ) payload["resumed"] = target + if defer_history: + payload["messages"] = [] + payload["message_count"] = int( + session.get("resume_message_count") or payload["message_count"] + ) + payload["hydrating"] = bool(session.get("resume_hydrating")) # A lazy watch session never owns a run loop, so its payload's running # flag is always False — overlay the child-run registry so a reconnecting # watch window keeps its busy indicator while the child is still mid-run. @@ -512,6 +519,64 @@ def _(rid, params: dict) -> dict: }, ) + # Desktop can ask for a bounded acknowledgement and hydrate the display + # transcript through the paginated REST endpoint. Register the runtime now, + # then load model history and initialize optional providers in background. + # Repeated requests reuse the record through the live fast path above. + if defer_history and not is_truthy_value(params.get("eager_build", False)): + sid = uuid.uuid4().hex[:8] + source = _resolve_session_source(str(params.get("source") or "").strip() or None) + lease = None # claimed lazily on the first turn (_ensure_active_session_slot) + _enable_gateway_prompts() + overrides = _stored_session_runtime_overrides(found) or {} + model_override = overrides.get("model_override") or {} + cwd = profile_resume_cwd or _default_session_cwd() + record = _deferred_session_record( + target, + cols=cols, + cwd=cwd, + history=[], + lease=lease, + source=source, + close_on_disconnect=is_truthy_value(params.get("close_on_disconnect", False)), + profile_home=profile_home, + model_override=overrides.get("model_override"), + resume_runtime_overrides=overrides or None, + ) + record["resume_history_ready"] = threading.Event() + record["resume_hydrating"] = True + record["resume_message_count"] = int(found.get("message_count") or 0) + if (live := _claim_or_reuse_live(sid, target, record, lease)) is not None: + return _ok(rid, _reuse_live_payload(*live)) + + _schedule_resume_hydration(sid, target, db, close_db=owns_db) + # The hydration worker now owns a profile-scoped handle and closes it + # after the transcript read. The shared launch DB is process-owned. + if owns_db: + owns_db = False + _schedule_session_cap_enforcement() + return _ok( + rid, + { + "session_id": sid, + "resumed": target, + "message_count": record["resume_message_count"], + "messages": [], + "hydrating": True, + "info": _lazy_resume_info( + cwd, + model=model_override.get("model") or "", + provider=overrides.get("provider_override") or "", + profile=profile, + ), + "inflight": None, + "running": False, + "session_key": target, + "started_at": record["created_at"], + "status": "resuming", + }, + ) + # Cold resume default: register the live session and read its stored # transcript, but build the agent OFF the response path. _make_agent can # block for seconds (MCP discovery, prompt/skill build, AIAgent diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 267a6839a5..d8dfcaeda3 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -738,6 +738,10 @@ def _finalize_session(session: dict | None, end_reason: str = "tui_close") -> No if not session or session.get("_finalized"): return session["_finalized"] = True + history_ready = session.get("resume_history_ready") + if history_ready is not None and not history_ready.is_set(): + session["resume_history_error"] = "session resume cancelled" + history_ready.set() _release_active_session_slot(session) stop_event = session.get("_notif_stop") if stop_event is not None: @@ -2277,6 +2281,15 @@ def _start_agent_build(sid: str, session: dict) -> None: owns_db = False profile_home = current.get("profile_home") try: + history_ready = current.get("resume_history_ready") + if history_ready is not None: + if not history_ready.wait(timeout=300.0): + raise TimeoutError("session history hydration timed out") + if history_error := current.get("resume_history_error"): + raise RuntimeError(str(history_error)) + with _sessions_lock: + if _sessions.get(sid) is not current: + return tokens = _set_session_context(key) # Build against the session's profile (global-remote): bind its # HERMES_HOME so config/skills/model resolve to it, and hand the @@ -8352,6 +8365,75 @@ def _schedule_agent_build(sid: str, delay: float = 0.05) -> None: timer.start() +def _schedule_resume_hydration( + sid: str, stored_id: str, db, *, close_db: bool = False +) -> None: + """Load a cold resume's transcript off the JSON-RPC response path.""" + + def _run() -> None: + session = _sessions.get(sid) + try: + if session is None: + return + _emit( + "session.resume_progress", + sid, + {"phase": "history", "status": "loading"}, + ) + db.reopen_session(stored_id) + raw_history, display_history = db.get_resume_conversations(stored_id) + prefix = db.get_ancestor_display_prefix(stored_id) + history = sanitize_replay_history(raw_history) + + if _sessions.get(sid) is not session: + return + with session["history_lock"]: + session["history"] = history + session["display_history_prefix"] = prefix + session["resume_hydrating"] = False + session["resume_message_count"] = len(display_history) + session["resume_history_ready"].set() + _emit( + "session.resume_progress", + sid, + { + "message_count": len(display_history), + "phase": "history", + "status": "complete", + }, + ) + _maybe_schedule_auto_continue(sid, session, stored_id) + _start_agent_build(sid, session) + except Exception as exc: + if _sessions.get(sid) is not session: + return + message = f"resume failed: {exc}" + session["resume_hydrating"] = False + session["resume_history_error"] = message + session["agent_error"] = message + session["resume_history_ready"].set() + session["agent_ready"].set() + _emit( + "session.resume_progress", + sid, + {"message": message, "phase": "history", "status": "failed"}, + ) + _emit("error", sid, {"message": message}) + with _sessions_lock: + discarded = _sessions.pop(sid, None) if _sessions.get(sid) is session else None + lease = (discarded or {}).get("active_session_lease") + if lease is not None: + lease.release() + finally: + if close_db and hasattr(db, "close"): + try: + db.close() + except Exception: + logger.debug("failed to close resume db for %s", sid, exc_info=True) + + threading.Thread(target=_run, daemon=True).start() + + def _session_pending_kind(sid: str) -> str: for rid, (owner_sid, _ev) in list(_pending.items()): if owner_sid != sid: