From 9824e9585cfe8697395c60f4b37518e37a0bcf10 Mon Sep 17 00:00:00 2001 From: Brooklyn Nicholson Date: Fri, 31 Jul 2026 23:04:45 -0500 Subject: [PATCH] fix(desktop): drop cached state for a reclaimed session Evict the runtime the backend just reclaimed instead of waiting for a resume to 404, and refresh the lists whose ended_at moved. The stored row is untouched, so reopening resumes from the DB. --- .../hooks/use-message-stream/gateway-event.ts | 18 +++ .../session-reclaimed.test.tsx | 125 ++++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 apps/desktop/src/app/session/hooks/use-message-stream/session-reclaimed.test.tsx 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 5970b8a07203f..7897c437d49ee 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 0000000000000..f08556ddd8789 --- /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() + } + }) +})