From d72ce5e43443714dd99d55d3fa85d522fc56ad72 Mon Sep 17 00:00:00 2001 From: Brooklyn Nicholson Date: Mon, 31 Aug 2026 20:05:43 -0500 Subject: [PATCH] fix(desktop): rebind the visible chat after a reaped runtime session.reclaimed now heals through markRuntimeGone before dropping cache. A 4001 on the visible session's dispatcher request asks for a durable resume. prompt.submit uses the window dispatcher so the turn lease survives the ACK, and a primary route no longer registers a phantom turn lease. Co-authored-by: Lester Liang <153183032+lesterlxt@users.noreply.github.com> Co-authored-by: wz-heng <68931789+wz-heng@users.noreply.github.com> --- .../contrib/session-rpc-dispatcher.test.ts | 69 ++++++++++++++++++- .../src/app/contrib/session-rpc-dispatcher.ts | 25 ++++++- .../gateway-event/lifecycle.ts | 3 + .../session-reclaimed.test.tsx | 29 ++++++++ .../hooks/use-prompt-actions/index.test.tsx | 6 +- .../session/hooks/use-prompt-actions/index.ts | 31 +-------- apps/desktop/src/store/gateway.ts | 7 ++ 7 files changed, 134 insertions(+), 36 deletions(-) diff --git a/apps/desktop/src/app/contrib/session-rpc-dispatcher.test.ts b/apps/desktop/src/app/contrib/session-rpc-dispatcher.test.ts index 9e16eb1697..8c6e83dcca 100644 --- a/apps/desktop/src/app/contrib/session-rpc-dispatcher.test.ts +++ b/apps/desktop/src/app/contrib/session-rpc-dispatcher.test.ts @@ -23,12 +23,18 @@ vi.mock('@/store/gateway', async importActual => ({ })) const probe = vi.hoisted(() => ({ resolveSessionOwner: vi.fn(async () => undefined as unknown) })) +const sessionMocks = vi.hoisted(() => ({ requestSessionResume: vi.fn() })) vi.mock('@/app/session/hooks/use-session-actions/utils', async importActual => ({ ...(await importActual>()), resolveSessionOwner: probe.resolveSessionOwner })) +vi.mock('@/store/session', async importActual => ({ + ...(await importActual>()), + requestSessionResume: sessionMocks.requestSessionResume +})) + const { createSessionRpcDispatcher } = await import('./session-rpc-dispatcher') const { $connectionsRegistry } = await import('@/store/connection-registry-state') const { $profiles } = await import('@/store/profile') @@ -40,13 +46,16 @@ const { isSessionOwnerResolutionError } = await import('@/store/session-owner-re const { $sessionTiles } = await import('@/store/session-states') const { makeSessionInfo } = await import('@/test/session-info') -function dispatcher(ambientRequest = vi.fn(async () => ({ ambient: true }))) { +function dispatcher( + ambientRequest = vi.fn(async () => ({ ambient: true })), + selectedStoredSessionId: null | string = null +) { return { ambientRequest, request: createSessionRpcDispatcher({ ambientRequest: ambientRequest as never, runtimeIdByStoredSessionIdRef: { current: new Map([['stored-omar', 'rt-omar']]) }, - selectedStoredSessionIdRef: { current: null }, + selectedStoredSessionIdRef: { current: selectedStoredSessionId }, sessionStateByRuntimeIdRef: { current: new Map() } }) } @@ -67,6 +76,7 @@ afterEach(() => { $sessionTiles.set([]) $profiles.set([]) _resetSessionOwnerHintsForTests({ storage: true }) + sessionMocks.requestSessionResume.mockReset() vi.clearAllMocks() }) @@ -194,3 +204,58 @@ describe('createSessionRpcDispatcher: exact owner rungs', () => { }) }) }) + +describe('createSessionRpcDispatcher: stale runtime recovery', () => { + it('requests a durable rebind for the visible session after a structured 4001', async () => { + setSessions([makeSessionInfo({ connection_id: 'local', id: 'stored-omar', profile: 'omar' })]) + gatewayMocks.requestGatewayForAgent.mockRejectedValueOnce( + Object.assign(new Error('runtime was reaped'), { code: 4001 }) + ) + const { request } = dispatcher(undefined, 'stored-omar') + + await expect(request('process.list', { session_id: 'rt-omar' })).rejects.toThrow('runtime was reaped') + + expect(sessionMocks.requestSessionResume).toHaveBeenCalledWith('stored-omar', { + connectionId: 'local', + profile: 'omar' + }) + }) + + it('does not let a background 4001 pull a different session into the foreground', async () => { + setSessions([makeSessionInfo({ connection_id: 'local', id: 'stored-omar', profile: 'omar' })]) + gatewayMocks.requestGatewayForAgent.mockRejectedValueOnce( + Object.assign(new Error('session not found'), { code: 4001 }) + ) + const { request } = dispatcher(undefined, 'stored-other') + + await expect(request('process.list', { session_id: 'rt-omar' })).rejects.toThrow('session not found') + + expect(sessionMocks.requestSessionResume).not.toHaveBeenCalled() + }) + + it('does not interpret an unrelated coded RPC failure as a stale runtime', async () => { + setSessions([makeSessionInfo({ connection_id: 'local', id: 'stored-omar', profile: 'omar' })]) + gatewayMocks.requestGatewayForAgent.mockRejectedValueOnce( + Object.assign(new Error('tool output says session not found'), { code: 5007 }) + ) + const { request } = dispatcher(undefined, 'stored-omar') + + await expect(request('process.list', { session_id: 'rt-omar' })).rejects.toThrow( + 'tool output says session not found' + ) + + expect(sessionMocks.requestSessionResume).not.toHaveBeenCalled() + }) + + it('leaves the warm resume lifecycle to recover its own session.activate failure', async () => { + setSessions([makeSessionInfo({ connection_id: 'local', id: 'stored-omar', profile: 'omar' })]) + gatewayMocks.requestGatewayForAgent.mockRejectedValueOnce( + Object.assign(new Error('session not found'), { code: 4001 }) + ) + const { request } = dispatcher(undefined, 'stored-omar') + + await expect(request('session.activate', { session_id: 'rt-omar' })).rejects.toThrow('session not found') + + expect(sessionMocks.requestSessionResume).not.toHaveBeenCalled() + }) +}) diff --git a/apps/desktop/src/app/contrib/session-rpc-dispatcher.ts b/apps/desktop/src/app/contrib/session-rpc-dispatcher.ts index 152fedd45c..cf9bfd94b4 100644 --- a/apps/desktop/src/app/contrib/session-rpc-dispatcher.ts +++ b/apps/desktop/src/app/contrib/session-rpc-dispatcher.ts @@ -35,7 +35,8 @@ import type { MutableRefObject } from 'react' import { resolveSessionOwner } from '@/app/session/hooks/use-session-actions/utils' import type { ClientSessionState } from '@/app/types' -import { getSessionOwnerHint, knownSessionOwner, ownerLookupSessionRows } from '@/store/session' +import { isSessionGoneForBackgroundPolling } from '@/store/runtime-gone' +import { getSessionOwnerHint, knownSessionOwner, ownerLookupSessionRows, requestSessionResume } from '@/store/session' import { assertSessionOwnerResolved } from '@/store/session-owner-resolution' import { requestForSessionProfile, type SessionOwnerScope } from '@/store/session-request-router' import { $focusedStoredSessionId, sessionTileOwnerRoute, storedSessionIdForRuntimeId } from '@/store/session-states' @@ -97,6 +98,26 @@ export function createSessionRpcDispatcher(deps: SessionRpcDispatcherDeps): Ambi // backend "session not found" on a backend that never held the runtime. assertSessionOwnerResolved(owner, { method, sessionId: paramSessionId ? routingSessionId : null }) - return requestForSessionProfile(owner, ambientRequest, method, params ?? {}, timeoutMs, signal) + try { + return await requestForSessionProfile(owner, ambientRequest, method, params ?? {}, timeoutMs, signal) + } catch (error) { + // A missed session.reclaimed leaves later RPCs answering 4001 against a + // still-resumable stored row. Prompt actions already retry their own + // calls; this seam covers the other session-scoped callers and wakes + // route-resume for the visible main session only. Do not retry the + // failing RPC — it may be destructive, and a fresh binding is async. + if ( + method !== 'session.resume' && + method !== 'session.activate' && + paramSessionId && + routingSessionId && + routingSessionId === selectedStoredSessionIdRef.current && + isSessionGoneForBackgroundPolling(error) + ) { + requestSessionResume(routingSessionId, typeof owner === 'object' && owner ? owner : undefined) + } + + throw error + } } } diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/lifecycle.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/lifecycle.ts index a5df27ced1..2395498d94 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/lifecycle.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/lifecycle.ts @@ -9,6 +9,7 @@ import { type PetChangeMeta, setChangeEventsAvailable } from '@/store/live-sync' +import { markRuntimeGone } from '@/store/runtime-gone' import { dropSessionState, unbindTileRuntime } from '@/store/session-states' // Leaf import (not the `@/themes` barrel) to avoid pulling the ThemeProvider // module graph into the gateway event hot path. @@ -80,6 +81,8 @@ export function handleLifecycleEvent(ctx: GatewayEventContext): boolean { const reclaimedRuntimeId = String((payload as { session_id?: string } | undefined)?.session_id ?? '') if (reclaimedRuntimeId) { + // Heal while the cached stored-id mapping is still intact, then drop. + markRuntimeGone(reclaimedRuntimeId) dropSessionState(reclaimedRuntimeId) // A tile bound to the reclaimed runtime would otherwise render an // empty transcript forever: its view reads $sessionStates[runtime] 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 index 7fd8473e16..0dcface159 100644 --- 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 @@ -4,6 +4,8 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import type { ClientSessionState } from '@/app/types' import { createClientSessionState } from '@/lib/chat-runtime' +import { resetRuntimeGoneHealing } from '@/store/runtime-gone' +import { $activeSessionId, $sessionResumeRequest } from '@/store/session' import { $sessionStates, $sessionTiles, publishSessionState } from '@/store/session-states' import type { RpcEvent } from '@/types/hermes' @@ -36,14 +38,20 @@ const reclaim = (sessionId: string, reason = 'ws_orphan_reap') => beforeEach(() => { queryClient = new QueryClient() + resetRuntimeGoneHealing() $sessionStates.set({}) $sessionTiles.set([]) + $activeSessionId.set(null) + $sessionResumeRequest.set(null) }) afterEach(() => { cleanup() + resetRuntimeGoneHealing() $sessionStates.set({}) $sessionTiles.set([]) + $activeSessionId.set(null) + $sessionResumeRequest.set(null) vi.restoreAllMocks() }) @@ -129,4 +137,25 @@ describe('session.reclaimed', () => { expect(wiringCache.has('live-gone')).toBe(false) expect(wiringCache.has('live-kept')).toBe(true) }) + + it('requests a durable resume when the reclaimed runtime is the active chat', () => { + mountStream() + $activeSessionId.set('live-gone') + publishSessionState('live-gone', createClientSessionState('stored-1')) + + reclaim('live-gone') + + expect($sessionResumeRequest.get()?.sessionId).toBe('stored-1') + }) + + it('does not navigate the primary chat when a background runtime is reclaimed', () => { + mountStream() + $activeSessionId.set('live-kept') + publishSessionState('live-gone', createClientSessionState('stored-1')) + publishSessionState('live-kept', createClientSessionState('stored-2')) + + reclaim('live-gone') + + expect($sessionResumeRequest.get()).toBeNull() + }) }) diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx index 89a873a8ca..e0347ce0b8 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx @@ -1891,14 +1891,12 @@ describe('usePromptActions submit / queue drain semantics', () => { ) expect(await handle!.submitText('continue remotely')).toBe(true) - expect(requestGatewayForAgent).toHaveBeenCalledWith( - 'hermes01', - 'default', + expect(ambientRequest).toHaveBeenCalledWith( 'prompt.submit', { session_id: 'runtime-remote', text: 'continue remotely' }, 1_800_000 ) - expect(ambientRequest).not.toHaveBeenCalled() + expect(requestGatewayForAgent).not.toHaveBeenCalled() }) it('clears a leftover interrupted flag on a fresh submit (so the new turn streams)', async () => { diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts index 75712af235..3a4cb0a888 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts @@ -22,17 +22,14 @@ import { updateComposerAttachment } from '@/store/composer' import { resetSessionBackground } from '@/store/composer-status' -import { requestGatewayForAgent } from '@/store/gateway' import { clearNotifications, notify, notifyError } from '@/store/notifications' import { clearPreviewArtifacts } from '@/store/preview-status' import { clearAllPrompts } from '@/store/prompts' import { $busy, - $connection, $currentCwd, $messages, $terminalBackend, - getSessionOwnerHint, setActiveSessionId, setAwaitingResponse, setBusy, @@ -486,30 +483,6 @@ export function usePromptActions({ } }, [activeSessionId, composerAttachments, eagerlyUploadAttachment]) - // Session resume can be routed through a registry connection while the - // render-time requestGateway still points at the local default socket. Keep - // every follow-up session RPC on the same composite owner; otherwise resume - // succeeds on HERMES01 and prompt.submit immediately fails locally with - // "session not found". - const requestForPromptSession = useCallback( - (method, params = {}, timeoutMs) => { - const storedSessionId = selectedStoredSessionIdRef.current - const owner = storedSessionId ? getSessionOwnerHint(storedSessionId) : undefined - const ambientConnection = $connection.get() - - const connectionId = - owner?.connectionId || - (ambientConnection?.mode === 'remote' ? ambientConnection.connectionId?.trim() || '' : '') - - if (connectionId) { - return requestGatewayForAgent(connectionId, owner?.profile || 'default', method, params, timeoutMs) - } - - return timeoutMs === undefined ? requestGateway(method, params) : requestGateway(method, params, timeoutMs) - }, - [requestGateway, selectedStoredSessionIdRef] - ) - const submitPromptText = useSubmitPrompt({ activeSessionIdRef, busyRef, @@ -518,7 +491,9 @@ export function usePromptActions({ getRoutedStoredSessionId, getRuntimeIdForStoredSession, getRouteToken, - requestGateway: requestForPromptSession, + // Window dispatcher: exact owner + turn lease through the terminal event. + // A private requestGatewayForAgent wrapper released the only client at ACK. + requestGateway, runtimeIdByStoredSessionIdRef, resumeStoredSession, selectedStoredSessionIdRef, diff --git a/apps/desktop/src/store/gateway.ts b/apps/desktop/src/store/gateway.ts index 60c5fada83..7601ea258e 100644 --- a/apps/desktop/src/store/gateway.ts +++ b/apps/desktop/src/store/gateway.ts @@ -1225,6 +1225,13 @@ export async function retainGatewayForSessionTurn( profile: string, sessionId: string ): Promise<() => void> { + // Primary events do not flow through a Secondary's terminal-event listener. + // Registering a no-op lease here would leave a phantom key that can suppress + // the real hold if this route is later re-homed as a secondary. + if (isPrimaryRegistryRoute(connectionId, normKey(profile))) { + return () => undefined + } + const scope = registryBackendScopeKey(connectionId, normKey(profile)) const key = turnLeaseKey(scope, sessionId)