fix(permission): speed up Auto mode switching #978

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<void>((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`, {

src/server/services/conversationService.ts

@@ -105,6 +105,8 @@ type SessionProcess = {
   permissionMode: string
   sdkToken: string
   sdkSocket: { send(data: string): void } | null
+  sdkAttached: Promise<void>
+  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<void>((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<typeof setTimeout> | undefined
     const earlyExitCode = await Promise.race([
       proc.exited,
+      session.sdkAttached.then(() => null),
       new Promise<null>((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
     }
   }

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
     }