From fb4c4f768477d4d96392ad56f282dce66085f74f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?= =?UTF-8?q?=29?= Date: Mon, 13 Jul 2026 12:54:13 +0800 Subject: [PATCH] fix(permission): speed up Auto mode switching #978 --- src/server/__tests__/conversations.test.ts | 83 ++++++++++++++++++++++ src/server/services/conversationService.ts | 22 +++++- src/server/ws/handler.ts | 2 +- 3 files changed, 103 insertions(+), 4 deletions(-) diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts index a5c427ae..2587f693 100644 --- a/src/server/__tests__/conversations.test.ts +++ b/src/server/__tests__/conversations.test.ts @@ -131,6 +131,25 @@ describe('ConversationService', () => { await expect(request).resolves.toEqual({ ok: true }) }) + it('should ignore a stale SDK disconnect after a replacement socket attaches', () => { + const svc = new ConversationService() + const sessionId = crypto.randomUUID() + const firstSocket = { send() {} } + const replacementSocket = { send() {} } + const session = { + sdkSocket: null, + pendingOutbound: [], + resolveSdkAttached: null, + } + ;(svc as any).sessions.set(sessionId, session) + + svc.attachSdkConnection(sessionId, firstSocket) + svc.attachSdkConnection(sessionId, replacementSocket) + svc.detachSdkConnection(sessionId, firstSocket) + + expect(session.sdkSocket).toBe(replacementSocket) + }) + it('should forward suggested permission updates for allow-for-session decisions', () => { const svc = new ConversationService() const sent: unknown[] = [] @@ -4106,6 +4125,70 @@ describe('WebSocket Chat Integration', () => { } }, 20_000) + it('should confirm a permission switch as soon as the SDK control channel connects', async () => { + await withMockInitMode('on_first_user', async () => { + const createRes = await fetch(`${baseUrl}/api/sessions`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ workDir: process.cwd(), permissionMode: 'default' }), + }) + expect(createRes.status).toBe(201) + const { sessionId } = await createRes.json() as { sessionId: string } + const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`) + + try { + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + reject(new Error(`Timed out connecting startup permission session ${sessionId}`)) + }, 5_000) + ws.onmessage = (event) => { + const msg = JSON.parse(event.data as string) + if (msg.type === 'connected') { + clearTimeout(timeout) + ws.send(JSON.stringify({ type: 'prewarm_session' })) + resolve() + } + } + ws.onerror = () => { + clearTimeout(timeout) + reject(new Error(`WebSocket error for startup permission session ${sessionId}`)) + } + }) + + await waitUntil( + () => Boolean((conversationService as any).sessions.get(sessionId)?.sdkSocket), + `SDK control channel for startup permission session ${sessionId}`, + ) + expect(conversationService.getSessionInitMessage(sessionId)).toBeNull() + + const switchStartedAt = performance.now() + const confirmation = new Promise<'confirmed'>((resolve, reject) => { + ws.onmessage = (event) => { + const msg = JSON.parse(event.data as string) + if (msg.type === 'permission_mode_changed' && msg.mode === 'auto') { + resolve('confirmed') + } + if (msg.type === 'error') { + reject(new Error(msg.message)) + } + } + }) + ws.send(JSON.stringify({ type: 'set_permission_mode', mode: 'auto' })) + + const outcome = await Promise.race([ + confirmation, + new Promise<'timeout'>((resolve) => setTimeout(() => resolve('timeout'), 1_000)), + ]) + + expect(outcome).toBe('confirmed') + expect(performance.now() - switchStartedAt).toBeLessThan(1_000) + } finally { + ws.close() + conversationService.stopSession(sessionId) + } + }) + }, 10_000) + it('should not persist or broadcast a rejected auto permission switch', async () => { await withMockPermissionModeBehavior('status-before-reject', async () => { const createRes = await fetch(`${baseUrl}/api/sessions`, { diff --git a/src/server/services/conversationService.ts b/src/server/services/conversationService.ts index 3543dece..91dc989d 100644 --- a/src/server/services/conversationService.ts +++ b/src/server/services/conversationService.ts @@ -105,6 +105,8 @@ type SessionProcess = { permissionMode: string sdkToken: string sdkSocket: { send(data: string): void } | null + sdkAttached: Promise + resolveSdkAttached: (() => void) | null pendingOutbound: string[] startupPending: boolean startupExitCode: number | null @@ -346,6 +348,10 @@ export class ConversationService { ) } + let resolveSdkAttached: (() => void) | null = null + const sdkAttached = new Promise((resolve) => { + resolveSdkAttached = resolve + }) const session: SessionProcess = { proc, outputCallbacks: [], @@ -353,6 +359,8 @@ export class ConversationService { permissionMode: options?.permissionMode || 'default', sdkToken: this.getSdkTokenFromUrl(sdkUrl), sdkSocket: null, + sdkAttached, + resolveSdkAttached, pendingOutbound: [], startupPending: true, startupExitCode: null, @@ -377,12 +385,15 @@ export class ConversationService { }) const STARTUP_GRACE_MS = 3000 + let startupGraceTimer: ReturnType | undefined const earlyExitCode = await Promise.race([ proc.exited, + session.sdkAttached.then(() => null), new Promise((resolve) => - setTimeout(() => resolve(null), STARTUP_GRACE_MS), + startupGraceTimer = setTimeout(() => resolve(null), STARTUP_GRACE_MS), ), ]) + if (startupGraceTimer) clearTimeout(startupGraceTimer) const startupExitCode = earlyExitCode ?? session.startupExitCode if (startupExitCode !== null) { @@ -753,6 +764,8 @@ export class ConversationService { if (!session) return false session.sdkSocket = socket + session.resolveSdkAttached?.() + session.resolveSdkAttached = null while (session.pendingOutbound.length > 0) { const line = session.pendingOutbound.shift() if (line) { @@ -762,9 +775,12 @@ export class ConversationService { return true } - detachSdkConnection(sessionId: string): void { + detachSdkConnection( + sessionId: string, + socket: { send(data: string): void }, + ): void { const session = this.sessions.get(sessionId) - if (session) { + if (session?.sdkSocket === socket) { session.sdkSocket = null } } diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 0dcb616c..2d892f9f 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -342,7 +342,7 @@ export const handleWebSocket = { if (channel === 'sdk') { console.log(`[WS] SDK disconnected from session: ${sessionId} (${code}: ${reason})`) - conversationService.detachSdkConnection(sessionId) + conversationService.detachSdkConnection(sessionId, ws) return }