diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts index 5970b8a072..7897c437d4 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts @@ -59,6 +59,7 @@ import { setTurnStartedAt, setYoloActive } from '@/store/session' +import { dropSessionState } from '@/store/session-states' import { pruneDelegateFallbackSubagents, pruneFinishedSessionSubagents, upsertSubagent } from '@/store/subagents' import { clearActiveSessionTodos } from '@/store/todos' import { recordToolDiff } from '@/store/tool-diffs' @@ -330,6 +331,23 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { } } + return + } else if (event.type === 'session.reclaimed') { + // The backend reclaimed a live session we may still be holding (idle + // TTL, LRU cap, or the WS-orphan reap). Without this the runtime id + // stays cached until something fails against it, which reads as the + // session vanishing rather than being reclaimed. Drop the cached state + // now — the stored row is untouched, so the sidebar keeps the + // conversation and reopening it resumes from the DB. + const reclaimedRuntimeId = String((payload as { session_id?: string } | undefined)?.session_id ?? '') + + if (reclaimedRuntimeId) { + dropSessionState(reclaimedRuntimeId) + } + + // The row's ended_at moved, so refresh the lists that render it. + notifySessionsChanged() + return } else if (event.type === 'session.info') { // Apply session-scoped fields when the event targets the active diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/session-reclaimed.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream/session-reclaimed.test.tsx new file mode 100644 index 0000000000..f08556ddd8 --- /dev/null +++ b/apps/desktop/src/app/session/hooks/use-message-stream/session-reclaimed.test.tsx @@ -0,0 +1,125 @@ +import { QueryClient } from '@tanstack/react-query' +import { act, cleanup, render, waitFor } from '@testing-library/react' +import { useEffect, useRef } from 'react' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import type { ClientSessionState } from '@/app/types' +import { createClientSessionState } from '@/lib/chat-runtime' +import { $sessionStates, publishSessionState } from '@/store/session-states' +import type { RpcEvent } from '@/types/hermes' + +import { useMessageStream } from './index' + +// `session.reclaimed`: the backend tore down a live session we're still +// holding (idle TTL, LRU cap, WS-orphan reap). Before this event the runtime id +// stayed cached until something failed against it, which read to the user as +// the session vanishing rather than being reclaimed. + +const ACTIVE_SID = 'session-active' +const ACTIVE_PROFILE = 'compass' +let handleEvent: ((event: RpcEvent) => void) | null = null +let queryClient: QueryClient + +function Harness() { + const activeSessionIdRef = useRef(ACTIVE_SID) + const sessionStateByRuntimeIdRef = useRef(new Map()) + + const stream = useMessageStream({ + activeGatewayProfile: ACTIVE_PROFILE, + activeSessionIdRef, + hydrateFromStoredSession: vi.fn(async () => undefined), + queryClient, + refreshHermesConfig: vi.fn<() => Promise>(async () => undefined), + refreshSessions: vi.fn<() => Promise>(async () => undefined), + sessionStateByRuntimeIdRef, + updateSessionState: (sessionId, updater) => { + const current = sessionStateByRuntimeIdRef.current.get(sessionId) ?? createClientSessionState() + const next = updater(current) + sessionStateByRuntimeIdRef.current.set(sessionId, next) + + return next + } + }) + + useEffect(() => { + handleEvent = stream.handleGatewayEvent + }, [stream.handleGatewayEvent]) + + return null +} + +async function mountStream() { + render() + await waitFor(() => expect(handleEvent).not.toBeNull()) +} + +const reclaim = (sessionId: string, reason = 'ws_orphan_reap') => + act(() => + handleEvent!({ + payload: { reason, session_id: sessionId, stored_session_id: 'stored-1' }, + session_id: '', + type: 'session.reclaimed' + } as RpcEvent) + ) + +beforeEach(() => { + handleEvent = null + queryClient = new QueryClient() + $sessionStates.set({}) +}) + +afterEach(() => { + cleanup() + $sessionStates.set({}) + vi.restoreAllMocks() +}) + +describe('session.reclaimed', () => { + it('drops the cached state for the reclaimed runtime', async () => { + await mountStream() + publishSessionState('live-gone', createClientSessionState()) + expect($sessionStates.get()['live-gone']).toBeDefined() + + reclaim('live-gone') + + expect($sessionStates.get()['live-gone']).toBeUndefined() + }) + + it('leaves every other live session alone', async () => { + await mountStream() + publishSessionState('live-gone', createClientSessionState()) + publishSessionState('live-kept', createClientSessionState()) + + reclaim('live-gone') + + // Both halves matter: the target went, the bystander stayed. Asserting + // only the survivor would pass with no handler at all. + expect($sessionStates.get()['live-gone']).toBeUndefined() + expect($sessionStates.get()['live-kept']).toBeDefined() + }) + + it('ignores a payload with no runtime id instead of clearing everything', async () => { + await mountStream() + publishSessionState('live-a', createClientSessionState()) + publishSessionState('live-b', createClientSessionState()) + + reclaim('') + + // A malformed/empty id must be a no-op, never a blanket wipe. + expect(Object.keys($sessionStates.get()).sort()).toEqual(['live-a', 'live-b']) + }) + + it('drops the runtime regardless of which reclaim reason fired', async () => { + for (const reason of ['idle_timeout', 'lru_evict', 'ws_orphan_reap']) { + $sessionStates.set({}) + cleanup() + handleEvent = null + await mountStream() + publishSessionState('live-gone', createClientSessionState()) + + reclaim('live-gone', reason) + + expect($sessionStates.get()['live-gone'], reason).toBeUndefined() + } + }) +}) diff --git a/tests/tui_gateway/test_session_reclaim_notify.py b/tests/tui_gateway/test_session_reclaim_notify.py new file mode 100644 index 0000000000..8f7e8be5b0 --- /dev/null +++ b/tests/tui_gateway/test_session_reclaim_notify.py @@ -0,0 +1,87 @@ +"""A backend-reclaimed session must tell the clients still holding it. + +The idle-TTL reaper, the LRU cap, and the WS-orphan reap all tear a live +session down without the client asking. Before ``session.reclaimed`` the client +kept a runtime id the backend had already forgotten, and only discovered it by +failing a later prompt — the "sessions suddenly lost in the backend" report. + +Contracts here: a reclaim broadcasts with its reason, a user-initiated close +stays silent, and a notify failure never breaks teardown. +""" + +import pytest + +from tui_gateway import server + + +@pytest.fixture() +def captured(monkeypatch): + events = [] + monkeypatch.setattr( + server, "_broadcast_global_event", lambda ev, payload=None: events.append((ev, payload)) + ) + # Teardown's real work (finalize, agent close, notifier unregister) is out + # of scope — this is about what reaches the client. + monkeypatch.setattr(server, "_finalize_session", lambda *a, **k: None) + return events + + +def _session(): + return {"_sid": "live-abc", "session_key": "20260731_120000_aaaaaa"} + + +@pytest.mark.parametrize("reason", ["idle_timeout", "lru_evict", "ws_orphan_reap"]) +def test_reclaim_reasons_announce_to_clients(captured, reason): + server._teardown_session(_session(), end_reason=reason) + + assert captured == [ + ( + "session.reclaimed", + { + "session_id": "live-abc", + "stored_session_id": "20260731_120000_aaaaaa", + "reason": reason, + }, + ) + ] + + +@pytest.mark.parametrize("reason", ["tui_close", "tui_shutdown", "ws_disconnect", "branched"]) +def test_client_initiated_closes_stay_silent(captured, reason): + """The client asked for these, so announcing them would be noise.""" + server._teardown_session(_session(), end_reason=reason) + + assert captured == [] + + +def test_broadcast_failure_does_not_break_teardown(monkeypatch): + """A wedged peer must not leave a session half torn down.""" + finalized = [] + + def _boom(*_a, **_k): + raise RuntimeError("transport gone") + + monkeypatch.setattr(server, "_broadcast_global_event", _boom) + monkeypatch.setattr( + server, "_finalize_session", lambda s, **k: finalized.append(k.get("end_reason")) + ) + + server._teardown_session(_session(), end_reason="ws_orphan_reap") + + assert finalized == ["ws_orphan_reap"] + + +def test_reap_paths_stamp_the_runtime_id_the_client_holds(): + """_pop_session_by_id stamps ``_sid``; the payload is useless without it. + + All three reclaim paths pop before tearing down, so this is the invariant + that makes the broadcast addressable on the client. + """ + server._sessions["live-xyz"] = {"session_key": "k"} + try: + popped = server._pop_session_by_id("live-xyz") + finally: + server._sessions.pop("live-xyz", None) + + assert popped is not None + assert popped["_sid"] == "live-xyz" diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 327ae0f717..b2bf7e0a5e 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -793,6 +793,39 @@ def _finalize_session(session: dict | None, end_reason: str = "tui_close") -> No pass +# End reasons where the BACKEND reclaimed a session the client never asked to +# close: the idle-TTL reaper, the LRU cap, and the WS-orphan reap. A client +# holding that live session id gets no signal today — its next prompt fails +# against an id the backend has already forgotten, which reads as the session +# silently vanishing rather than being reclaimed. ``tui_close`` and friends are +# deliberately absent: the client initiated those and already knows. +_RECLAIM_END_REASONS = frozenset({"idle_timeout", "lru_evict", "ws_orphan_reap"}) + + +def _announce_session_reclaimed(session: dict, end_reason: str) -> None: + """Tell connected clients a session was reclaimed out from under them. + + Broadcast rather than session-targeted: the reap paths run on background + timer threads with no contextvar binding, and the WS-orphan case has by + definition lost its own transport — ``_emit`` would bottom out on stdio and + the peer that owns the session would never see it. Best-effort; a failed + notify must never break teardown. + """ + if end_reason not in _RECLAIM_END_REASONS: + return + try: + _broadcast_global_event( + "session.reclaimed", + { + "session_id": str(session.get("_sid") or ""), + "stored_session_id": str(session.get("session_key") or ""), + "reason": end_reason, + }, + ) + except Exception: + logger.debug("session.reclaimed broadcast failed", exc_info=True) + + def _teardown_session(session: dict | None, *, end_reason: str = "tui_close") -> None: """Fully tear down a session: finalize, unregister, close agent + worker. @@ -806,6 +839,7 @@ def _teardown_session(session: dict | None, *, end_reason: str = "tui_close") -> if not session: return _finalize_session(session, end_reason=end_reason) + _announce_session_reclaimed(session, end_reason) try: from tools.approval import unregister_gateway_notify