diff --git a/apps/desktop/e2e/approval-reconnect.spec.ts b/apps/desktop/e2e/approval-reconnect.spec.ts new file mode 100644 index 0000000000000..5f7c660e057c7 --- /dev/null +++ b/apps/desktop/e2e/approval-reconnect.spec.ts @@ -0,0 +1,129 @@ +import { writeEnvFile, writeMockProviderConfig } from '../../../tests-js/scripts/mock-provider-config' +import { MOCK_REPLY, receivedUserTexts, startMockServer } from '../../../tests-js/scripts/mock-server' + +import { buildAppEnv, createSandbox, launchDesktop, type MockBackendFixture, waitForAppReady } from './fixtures' +import { expect, test } from './test' + +interface ReconnectProbe { + blocked: boolean + sockets: Set + requests: Map + resumed: string[] + rejectedApprovals: string[] + submitted: string[] +} + +declare global { + interface Window { + __approvalReconnectProbe: ReconnectProbe + } +} + +let fixture: MockBackendFixture | undefined + +test.afterEach(async () => { + await fixture?.cleanup() +}) + +test('reconnects after runtime reclamation without querying approvals on the retired runtime', async () => { + test.setTimeout(180_000) + const mock = await startMockServer() + const sandbox = createSandbox('approval-reconnect') + + writeMockProviderConfig(sandbox.hermesHome, mock.url) + writeEnvFile(sandbox.hermesHome, 'e2e-mock-key', mock.url) + const { app, page } = await launchDesktop(buildAppEnv(sandbox)) + + fixture = { + app, page, mock, mockUrl: mock.url, sandbox, + cleanup: async () => { + await app.close().catch(() => undefined) + await mock.close() + sandbox.cleanup() + } + } + await waitForAppReady(fixture, 90_000) + + await page.evaluate(() => { + const NativeWebSocket = window.WebSocket + const send = NativeWebSocket.prototype.send + + const probe: ReconnectProbe = { + blocked: false, sockets: new Set(), requests: new Map(), resumed: [], rejectedApprovals: [], submitted: [] + } + + window.__approvalReconnectProbe = probe + + NativeWebSocket.prototype.send = function (data) { + if (!probe.sockets.has(this)) { + probe.sockets.add(this) + this.addEventListener('message', event => { + const frame = JSON.parse(String(event.data)) + const request = probe.requests.get(String(frame.id)) + + if (request?.method === 'approval.pending' && frame.error?.code === 4001) { + probe.rejectedApprovals.push(request.sessionId ?? '') + } + + if (request?.method === 'session.resume' && frame.result?.session_id) { + probe.resumed.push(frame.result.session_id) + } + }) + } + + const frame = JSON.parse(String(data)) + + if (frame.method) { + probe.requests.set(String(frame.id), { method: frame.method, sessionId: frame.params?.session_id }) + + if (frame.method === 'prompt.submit') { + probe.submitted.push(frame.params.session_id) + } + } + + return send.call(this, data) + } + + window.WebSocket = class extends NativeWebSocket { + constructor(url: string | URL, protocols?: string | string[]) { + super(url, protocols) + + if (probe.blocked) { + queueMicrotask(() => this.close()) + } + } + } + }) + + const composer = page.locator('[contenteditable="true"]').first() + await composer.fill('E2E_APPROVAL_RECONNECT_BEFORE') + await composer.press('Enter') + await expect(page.locator('[data-slot="aui_thread-viewport"]')).toContainText(MOCK_REPLY, { timeout: 60_000 }) + const oldRuntime = await page.evaluate(() => window.__approvalReconnectProbe.submitted.at(-1)) + + expect(oldRuntime).toBeTruthy() + await page.evaluate(() => { + const probe = window.__approvalReconnectProbe + + probe.blocked = true + probe.resumed.length = 0 + + for (const socket of probe.sockets) { + socket.close() + } + }) + // Exceed the real backend's default 20-second orphan grace. + await page.waitForTimeout(25_000) + await page.evaluate(() => { window.__approvalReconnectProbe.blocked = false }) + await expect.poll( + () => page.evaluate(() => window.__approvalReconnectProbe.resumed.at(-1)), + { timeout: 45_000 } + ).toEqual(expect.any(String)) + expect(await page.evaluate(() => window.__approvalReconnectProbe.resumed.at(-1))).not.toBe(oldRuntime) + + await composer.fill('E2E_APPROVAL_RECONNECT_AFTER') + await composer.press('Enter') + await expect.poll(() => receivedUserTexts().includes('E2E_APPROVAL_RECONNECT_AFTER'), { timeout: 30_000 }).toBe(true) + await expect(page.locator('[data-slot="aui_thread-viewport"]')).toContainText('E2E_APPROVAL_RECONNECT_BEFORE') + expect(await page.evaluate(() => window.__approvalReconnectProbe.rejectedApprovals)).toEqual([]) +}) diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/index.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/index.ts index 901003ab81c91..5f8147dd0ab2a 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/index.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event/index.ts @@ -199,7 +199,7 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { const isActiveEvent = !!sessionId && sessionId === activeSessionIdRef.current - const replaySessionId = approvalReplaySessionId(event.type, activeSessionIdRef.current, sessionId, { + const replaySessionId = approvalReplaySessionId(event.type, sessionId, { explicit: Boolean(explicitSid), isGone: isSessionGone }) diff --git a/apps/desktop/src/lib/gateway-events.test.ts b/apps/desktop/src/lib/gateway-events.test.ts index 11eabbefe80e4..e7e8828460f69 100644 --- a/apps/desktop/src/lib/gateway-events.test.ts +++ b/apps/desktop/src/lib/gateway-events.test.ts @@ -3,22 +3,22 @@ import { describe, expect, it } from 'vitest' import { approvalReplaySessionId, gatewayEventRequiresSessionId, resolveGatewayEventSessionId } from './gateway-events' describe('gateway event routing', () => { - it('rehydrates pending approvals on reconnect ready and resumed session info', () => { - expect(approvalReplaySessionId('gateway.ready', 'active-1', null)).toBe('active-1') - expect(approvalReplaySessionId('session.info', 'active-1', 'routed-1')).toBe('routed-1') - expect(approvalReplaySessionId('message.delta', 'active-1', 'routed-1')).toBeNull() + it('rehydrates pending approvals after session rebind, not socket readiness', () => { + expect(approvalReplaySessionId('gateway.ready', 'active-1')).toBeNull() + expect(approvalReplaySessionId('session.info', 'routed-1')).toBe('routed-1') + expect(approvalReplaySessionId('message.delta', 'routed-1')).toBeNull() }) it('does not replay against an active runtime the gateway already reported gone', () => { const isGone = (sid: string) => sid === 'dead-1' // Unscoped fan-out attributed to a dead active session: skip. - expect(approvalReplaySessionId('session.info', 'dead-1', 'dead-1', { explicit: false, isGone })).toBeNull() - expect(approvalReplaySessionId('gateway.ready', 'dead-1', null, { explicit: false, isGone })).toBeNull() + expect(approvalReplaySessionId('session.info', 'dead-1', { explicit: false, isGone })).toBeNull() + expect(approvalReplaySessionId('gateway.ready', 'dead-1', { explicit: false, isGone })).toBeNull() // A live active session still replays. - expect(approvalReplaySessionId('session.info', 'live-1', 'live-1', { explicit: false, isGone })).toBe('live-1') + expect(approvalReplaySessionId('session.info', 'live-1', { explicit: false, isGone })).toBe('live-1') // An explicitly scoped frame is the runtime speaking for itself — never skipped. - expect(approvalReplaySessionId('session.info', 'dead-1', 'dead-1', { explicit: true, isGone })).toBe('dead-1') + expect(approvalReplaySessionId('session.info', 'dead-1', { explicit: true, isGone })).toBe('dead-1') }) it('drops only unscoped subagent events (genuinely background work)', () => { diff --git a/apps/desktop/src/lib/gateway-events.ts b/apps/desktop/src/lib/gateway-events.ts index d775343ec3918..2f5d0c7443e8a 100644 --- a/apps/desktop/src/lib/gateway-events.ts +++ b/apps/desktop/src/lib/gateway-events.ts @@ -94,7 +94,9 @@ export interface GatewayEventSessionRoute { /** Which session (if any) to re-pull `approval.pending` for after `eventType`. * - * `gateway.ready` and `session.info` are the two rehydration points. An + * `gateway.ready` proves only the socket is open; its cached runtime may have + * been reaped while disconnected. Resume restores pending approvals, and a + * subsequent `session.info` can safely refresh them. An * UNSCOPED `session.info` (the approvals-loop / broadcast fan-out, no * `session_id` on the frame) reaches here attributed to the active session by * the routing fallback; when `isGone(activeSessionId)` — the gateway already @@ -103,15 +105,12 @@ export interface GatewayEventSessionRoute { * session explicitly is the runtime speaking for itself and is never gone. */ export function approvalReplaySessionId( eventType: string | undefined, - activeSessionId: null | string, routedSessionId: null | string, options?: { explicit?: boolean; isGone?: (sessionId: string) => boolean } ): null | string { let target: null | string = null - if (eventType === 'gateway.ready') { - target = activeSessionId - } else if (eventType === 'session.info') { + if (eventType === 'session.info') { target = routedSessionId } diff --git a/apps/desktop/src/store/prompts-reconnect.test.ts b/apps/desktop/src/store/prompts-reconnect.test.ts new file mode 100644 index 0000000000000..c97b99d8ea33d --- /dev/null +++ b/apps/desktop/src/store/prompts-reconnect.test.ts @@ -0,0 +1,56 @@ +import { JsonRpcGatewayError } from '@hermes/shared' +import { afterEach, expect, it } from 'vitest' + +import { approvalReplaySessionId } from '@/lib/gateway-events' + +import { clearAllPrompts, replayPendingApproval, sessionApprovalRequests } from './prompts' +import { isSessionGone, resetBackgroundPollingGuard, resetBackgroundPollingGuardAfterRebind } from './session-gone-latch' + +afterEach(() => { + clearAllPrompts() + resetBackgroundPollingGuard() +}) + +it('waits for session rebind before replaying approvals across repeated reconnects', async () => { + const queried: string[] = [] + let liveRuntime = 'runtime-0' + + const gateway = { + request: async (method: string, params: Record) => { + if (method === 'approval.pending') { + queried.push(String(params.session_id)) + + if (params.session_id !== liveRuntime) { + throw new JsonRpcGatewayError('session not found', { code: 4001 }) + } + + return { approvals: [{ request_id: `approval-${liveRuntime}`, command: 'echo test', description: 'test' }] } + } + + return { acknowledged: true } + } + } + + for (let reconnect = 1; reconnect <= 3; reconnect++) { + const cachedRuntime = liveRuntime + // The disconnected runtime was reaped before the next socket became ready. + liveRuntime = '' + resetBackgroundPollingGuard() + await replayPendingApproval( + gateway, + approvalReplaySessionId('gateway.ready', cachedRuntime, { explicit: false, isGone: isSessionGone }) + ) + + liveRuntime = `runtime-${reconnect}` + resetBackgroundPollingGuardAfterRebind('session.resume', { session_id: 'stored-session' }, { session_id: liveRuntime }) + await replayPendingApproval( + gateway, + approvalReplaySessionId('session.info', liveRuntime, { explicit: true, isGone: isSessionGone }) + ) + expect(sessionApprovalRequests(liveRuntime).get().map(request => request.requestId)).toEqual([ + `approval-${liveRuntime}` + ]) + } + + expect(queried).toEqual(['runtime-1', 'runtime-2', 'runtime-3']) +})