fix: defer runtime restarts during active turns (#626)

Avoid killing an in-flight desktop SDK turn when runtime config changes during tool execution. Persist the requested runtime immediately, then restart after the turn emits its terminal result. Tested: - bun run check:server - real provider tool-turn model switch smoke - bun run check:coverage (changed-line coverage passed; unrelated global lanes still fail) Scope-risk: narrow
2026-07-16 13:03:31 +08:00 · 2026-06-09 21:05:51 +08:00 · 2026-06-09 21:05:51 +08:00 · 385b996736
commit 385b996736
parent 3ff6a79e62
2 changed files with 494 additions and 16 deletions
--- a/src/server/tests/conversations.test.ts
+++ b/src/server/tests/conversations.test.ts
@ -2514,6 +2514,417 @@ describe('WebSocket Chat Integration', () => {
    }
  }, 20_000)

+  it('should clear active turn tracking when sending a user message fails after startup', async () => {
+    const providerService = new ProviderService()
+    const provider = await providerService.addProvider({
+      presetId: 'custom',
+      name: 'Provider Send Failure Runtime',
+      apiKey: 'key-send-failure-runtime',
+      baseUrl: 'http://127.0.0.1:1/anthropic',
+      apiFormat: 'anthropic',
+      models: {
+        main: 'send-failure-main',
+        haiku: 'send-failure-haiku',
+        sonnet: 'send-failure-sonnet',
+        opus: 'send-failure-opus',
+      },
+    })
+
+    const createRes = await fetch(`${baseUrl}/api/sessions`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({ workDir: process.cwd() }),
+    })
+    expect(createRes.status).toBe(201)
+    const { sessionId } = await createRes.json() as { sessionId: string }
+
+    const originalStartSession = conversationService.startSession.bind(conversationService)
+    const originalSendMessage = conversationService.sendMessage.bind(conversationService)
+    const startCalls: Array<{
+      sessionId: string
+      options: { permissionMode?: string; model?: string; effort?: string; providerId?: string | null } | undefined
+    }> = []
+
+    conversationService.startSession = (async function patchedStartSession(
+      sid: string,
+      workDir: string,
+      sdkUrl: string,
+      options?: { permissionMode?: string; model?: string; effort?: string; thinking?: 'enabled' | 'adaptive' | 'disabled'; providerId?: string | null },
+    ) {
+      startCalls.push({ sessionId: sid, options })
+      return originalStartSession(sid, workDir, sdkUrl, options)
+    }) as typeof conversationService.startSession
+    conversationService.sendMessage = (async () => false) as typeof conversationService.sendMessage
+
+    const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`)
+    const messages: any[] = []
+    let sendFailureIdle = false
+    try {
+      await new Promise<void>((resolve, reject) => {
+        const timeout = setTimeout(() => {
+          ws.close()
+          reject(new Error(`Timed out waiting for send-failure runtime switch for session ${sessionId}`))
+        }, 10_000)
+
+        ws.onmessage = (event) => {
+          const msg = JSON.parse(event.data as string)
+          messages.push(msg)
+
+          if (msg.type === 'connected') {
+            ws.send(JSON.stringify({ type: 'user_message', content: 'send failure active turn cleanup' }))
+            return
+          }
+
+          if (msg.type === 'error' && msg.code !== 'CLI_NOT_RUNNING') {
+            clearTimeout(timeout)
+            ws.close()
+            reject(new Error(msg.message))
+            return
+          }
+
+          if (msg.type === 'status' && msg.state === 'idle' && !sendFailureIdle) {
+            if (!messages.some((item) => item.type === 'error' && item.code === 'CLI_NOT_RUNNING')) {
+              return
+            }
+            sendFailureIdle = true
+            ws.send(JSON.stringify({
+              type: 'set_runtime_config',
+              providerId: provider.id,
+              modelId: 'send-failure-sonnet',
+            }))
+            return
+          }
+
+          if (msg.type === 'status' && msg.state === 'idle' && sendFailureIdle && startCalls.length > 1) {
+            clearTimeout(timeout)
+            ws.close()
+            resolve()
+          }
+        }
+
+        ws.onerror = () => {
+          clearTimeout(timeout)
+          reject(new Error(`WebSocket error for send-failure runtime switch session ${sessionId}`))
+        }
+      })
+
+      expect(startCalls).toHaveLength(2)
+      expect(startCalls[1]).toMatchObject({
+        sessionId,
+        options: {
+          providerId: provider.id,
+          model: 'send-failure-sonnet',
+        },
+      })
+    } finally {
+      ws.close()
+      conversationService.startSession = originalStartSession
+      conversationService.sendMessage = originalSendMessage
+      conversationService.stopSession(sessionId)
+    }
+  }, 20_000)
+
+  it('should defer runtime model switches until the active turn completes', async () => {
+    await withMockStreamDelay(350, async () => {
+      const providerService = new ProviderService()
+      const providerA = await providerService.addProvider({
+        presetId: 'custom',
+        name: 'Provider Active Runtime A',
+        apiKey: 'key-active-runtime-a',
+        baseUrl: 'http://127.0.0.1:1/anthropic',
+        apiFormat: 'anthropic',
+        models: {
+          main: 'active-a-main',
+          haiku: 'active-a-haiku',
+          sonnet: 'active-a-sonnet',
+          opus: 'active-a-opus',
+        },
+      })
+      const providerB = await providerService.addProvider({
+        presetId: 'custom',
+        name: 'Provider Active Runtime B',
+        apiKey: 'key-active-runtime-b',
+        baseUrl: 'http://127.0.0.1:1/anthropic',
+        apiFormat: 'anthropic',
+        models: {
+          main: 'active-b-main',
+          haiku: 'active-b-haiku',
+          sonnet: 'active-b-sonnet',
+          opus: 'active-b-opus',
+        },
+      })
+
+      const createRes = await fetch(`${baseUrl}/api/sessions`, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ workDir: process.cwd() }),
+      })
+      expect(createRes.status).toBe(201)
+      const { sessionId } = await createRes.json() as { sessionId: string }
+
+      const originalStartSession = conversationService.startSession.bind(conversationService)
+      const startCalls: Array<{
+        sessionId: string
+        options: { permissionMode?: string; model?: string; effort?: string; providerId?: string | null } | undefined
+      }> = []
+
+      conversationService.startSession = (async function patchedStartSession(
+        sid: string,
+        workDir: string,
+        sdkUrl: string,
+        options?: { permissionMode?: string; model?: string; effort?: string; thinking?: 'enabled' | 'adaptive' | 'disabled'; providerId?: string | null },
+      ) {
+        startCalls.push({ sessionId: sid, options })
+        return originalStartSession(sid, workDir, sdkUrl, options)
+      }) as typeof conversationService.startSession
+
+      const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`)
+      let switchTriggered = false
+      let turnComplete = false
+      try {
+        await new Promise<void>((resolve, reject) => {
+          const timeout = setTimeout(() => {
+            ws.close()
+            reject(new Error(`Timed out waiting for active-turn runtime switch for session ${sessionId}`))
+          }, 10_000)
+
+          ws.onmessage = (event) => {
+            const msg = JSON.parse(event.data as string)
+
+            if (msg.type === 'connected') {
+              ws.send(JSON.stringify({
+                type: 'set_runtime_config',
+                providerId: providerA.id,
+                modelId: 'active-a-sonnet',
+              }))
+              ws.send(JSON.stringify({ type: 'user_message', content: 'active turn runtime switch' }))
+              return
+            }
+
+            if (msg.type === 'error') {
+              clearTimeout(timeout)
+              ws.close()
+              reject(new Error(msg.message))
+              return
+            }
+
+            if (
+              msg.type === 'content_delta' &&
+              typeof msg.text === 'string' &&
+              msg.text.includes('active turn runtime switch') &&
+              !switchTriggered
+            ) {
+              switchTriggered = true
+              ws.send(JSON.stringify({
+                type: 'set_runtime_config',
+                providerId: providerB.id,
+                modelId: 'active-b-opus',
+              }))
+              return
+            }
+
+            if (
+              msg.type === 'status' &&
+              msg.state === 'idle' &&
+              switchTriggered &&
+              !turnComplete &&
+              startCalls.length > 1
+            ) {
+              clearTimeout(timeout)
+              ws.close()
+              reject(new Error('Runtime restarted before the active turn completed'))
+              return
+            }
+
+            if (msg.type === 'message_complete' && switchTriggered && !turnComplete) {
+              turnComplete = true
+              expect(startCalls).toHaveLength(1)
+              return
+            }
+
+            if (msg.type === 'status' && msg.state === 'idle' && turnComplete) {
+              clearTimeout(timeout)
+              ws.close()
+              resolve()
+            }
+          }
+
+          ws.onerror = () => {
+            clearTimeout(timeout)
+            reject(new Error(`WebSocket error for active-turn runtime switch session ${sessionId}`))
+          }
+        })
+
+        expect(startCalls).toHaveLength(2)
+        expect(startCalls[0]).toMatchObject({
+          sessionId,
+          options: {
+            providerId: providerA.id,
+            model: 'active-a-sonnet',
+          },
+        })
+        expect(startCalls[1]).toMatchObject({
+          sessionId,
+          options: {
+            providerId: providerB.id,
+            model: 'active-b-opus',
+          },
+        })
+      } finally {
+        ws.close()
+        conversationService.startSession = originalStartSession
+        conversationService.stopSession(sessionId)
+      }
+    })
+  }, 20_000)
+
+  it('should surface deferred runtime restart failures after the active turn completes', async () => {
+    await withMockStreamDelay(350, async () => {
+      const providerService = new ProviderService()
+      const providerA = await providerService.addProvider({
+        presetId: 'custom',
+        name: 'Provider Deferred Failure A',
+        apiKey: 'key-deferred-failure-a',
+        baseUrl: 'http://127.0.0.1:1/anthropic',
+        apiFormat: 'anthropic',
+        models: {
+          main: 'deferred-failure-a-main',
+          haiku: 'deferred-failure-a-haiku',
+          sonnet: 'deferred-failure-a-sonnet',
+          opus: 'deferred-failure-a-opus',
+        },
+      })
+      const providerB = await providerService.addProvider({
+        presetId: 'custom',
+        name: 'Provider Deferred Failure B',
+        apiKey: 'key-deferred-failure-b',
+        baseUrl: 'http://127.0.0.1:1/anthropic',
+        apiFormat: 'anthropic',
+        models: {
+          main: 'deferred-failure-b-main',
+          haiku: 'deferred-failure-b-haiku',
+          sonnet: 'deferred-failure-b-sonnet',
+          opus: 'deferred-failure-b-opus',
+        },
+      })
+
+      const createRes = await fetch(`${baseUrl}/api/sessions`, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ workDir: process.cwd() }),
+      })
+      expect(createRes.status).toBe(201)
+      const { sessionId } = await createRes.json() as { sessionId: string }
+
+      const originalStartSession = conversationService.startSession.bind(conversationService)
+      const startCalls: Array<{
+        sessionId: string
+        options: { permissionMode?: string; model?: string; effort?: string; providerId?: string | null } | undefined
+      }> = []
+
+      conversationService.startSession = (async function patchedStartSession(
+        sid: string,
+        workDir: string,
+        sdkUrl: string,
+        options?: { permissionMode?: string; model?: string; effort?: string; thinking?: 'enabled' | 'adaptive' | 'disabled'; providerId?: string | null },
+      ) {
+        startCalls.push({ sessionId: sid, options })
+        if (startCalls.length > 1) {
+          throw new Error('deferred restart failed')
+        }
+        return originalStartSession(sid, workDir, sdkUrl, options)
+      }) as typeof conversationService.startSession
+
+      const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`)
+      let switchTriggered = false
+      let turnComplete = false
+      let restartError = false
+      try {
+        await new Promise<void>((resolve, reject) => {
+          const timeout = setTimeout(() => {
+            ws.close()
+            reject(new Error(`Timed out waiting for deferred restart failure for session ${sessionId}`))
+          }, 10_000)
+
+          ws.onmessage = (event) => {
+            const msg = JSON.parse(event.data as string)
+
+            if (msg.type === 'connected') {
+              ws.send(JSON.stringify({
+                type: 'set_runtime_config',
+                providerId: providerA.id,
+                modelId: 'deferred-failure-a-sonnet',
+              }))
+              ws.send(JSON.stringify({ type: 'user_message', content: 'deferred runtime restart failure' }))
+              return
+            }
+
+            if (
+              msg.type === 'content_delta' &&
+              typeof msg.text === 'string' &&
+              msg.text.includes('deferred runtime restart failure') &&
+              !switchTriggered
+            ) {
+              switchTriggered = true
+              ws.send(JSON.stringify({
+                type: 'set_runtime_config',
+                providerId: providerB.id,
+                modelId: 'deferred-failure-b-opus',
+              }))
+              return
+            }
+
+            if (msg.type === 'message_complete' && switchTriggered && !turnComplete) {
+              turnComplete = true
+              expect(startCalls).toHaveLength(1)
+              return
+            }
+
+            if (msg.type === 'error') {
+              if (!turnComplete) {
+                clearTimeout(timeout)
+                ws.close()
+                reject(new Error(`Deferred restart failed before turn completion: ${msg.message}`))
+                return
+              }
+              restartError = msg.code === 'CLI_RESTART_FAILED' &&
+                typeof msg.message === 'string' &&
+                msg.message.includes('deferred restart failed')
+              return
+            }
+
+            if (msg.type === 'status' && msg.state === 'idle' && restartError) {
+              clearTimeout(timeout)
+              ws.close()
+              resolve()
+            }
+          }
+
+          ws.onerror = () => {
+            clearTimeout(timeout)
+            reject(new Error(`WebSocket error for deferred restart failure session ${sessionId}`))
+          }
+        })
+
+        expect(switchTriggered).toBe(true)
+        expect(turnComplete).toBe(true)
+        expect(restartError).toBe(true)
+        expect(startCalls).toHaveLength(2)
+        expect(startCalls[1]).toMatchObject({
+          sessionId,
+          options: {
+            providerId: providerB.id,
+            model: 'deferred-failure-b-opus',
+          },
+        })
+      } finally {
+        ws.close()
+        conversationService.startSession = originalStartSession
+        conversationService.stopSession(sessionId)
+      }
+    })
+  }, 20_000)
+
  it('should keep the session idle in the UI while restarting for a bypass permission switch', async () => {
    await fetch(`${baseUrl}/api/permissions/mode`, {
      method: 'PUT',
@ -2721,7 +3132,7 @@ describe('WebSocket Chat Integration', () => {
    }
  }, 20_000)

-  it('should restart when switching from bypass permissions back to default', async () => {
+  it('should switch from bypass permissions back to default without restarting', async () => {
    const createRes = await fetch(`${baseUrl}/api/sessions`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
@ -2789,24 +3200,13 @@ describe('WebSocket Chat Integration', () => {
        mode: 'default',
      }))

-      await waitUntil(
-        async () => messages.slice(switchStartIndex).some((msg) => msg.type === 'status' && msg.state === 'idle'),
-        `bypass-to-default permission switch completion for ${sessionId}`,
-      )
-
-      expect(startCalls).toHaveLength(2)
-      expect(startCalls[1]).toMatchObject({
-        sessionId,
-        options: {
-          permissionMode: 'default',
-        },
-      })
      await waitUntil(async () => {
        const res = await fetch(`${baseUrl}/api/sessions/${sessionId}/inspection?includeContext=0`)
        if (!res.ok) return false
        const body = await res.json() as { status?: { permissionMode?: string } }
        return body.status?.permissionMode === 'default'
      }, `persisted bypass-to-default permission switch for ${sessionId}`)
+      expect(startCalls).toHaveLength(1)
      expect(messages.slice(switchStartIndex).some((msg) => msg.type === 'error')).toBe(false)
    } finally {
      ws.close()
--- a/src/server/ws/handler.ts
+++ b/src/server/ws/handler.ts
@ -76,11 +76,19 @@ const sessionTitleState = new Map<string, {
  generationSeq: number
 }>()

-const runtimeOverrides = new Map<string, {
+type RuntimeOverride = {
  providerId: string | null
  modelId: string
  effort?: string
-}>()
+}
+
+type ActiveUserTurnState = {
+  messageSent: boolean
+}
+
+const runtimeOverrides = new Map<string, RuntimeOverride>()
+const activeUserTurns = new Map<string, ActiveUserTurnState>()
+const deferredRuntimeRestarts = new Map<string, RuntimeOverride>()

 const runtimeTransitionPromises = new Map<string, Promise<void>>()
 const sessionStartupPromises = new Map<string, Promise<void>>()
@ -306,8 +314,14 @@ async function handleUserMessage(
  // Send thinking status
  sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' })

+  const activeTurn: ActiveUserTurnState = { messageSent: false }
+  activeUserTurns.set(sessionId, activeTurn)
+
  const initialRuntimeTransition = await waitForRuntimeTransitionBeforeUserTurn(ws, sessionId)
-  if (!initialRuntimeTransition.ok) return
+  if (!initialRuntimeTransition.ok) {
+    clearActiveUserTurn(sessionId, activeTurn)
+    return
+  }
  if (initialRuntimeTransition.waited) {
    sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' })
  }
@ -357,6 +371,7 @@ async function handleUserMessage(
        err instanceof ConversationStartupError ? err.retryable : false,
    })
    sendMessage(ws, { type: 'status', state: 'idle' })
+    clearActiveUserTurn(sessionId, activeTurn)
    return
  }

@ -366,6 +381,7 @@ async function handleUserMessage(
      sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' })
    }
  } else {
+    clearActiveUserTurn(sessionId, activeTurn)
    return
  }

@ -387,6 +403,7 @@ async function handleUserMessage(
      return shouldForwardCurrentTurnLocalCommand(cliMsg)
    },
  })
+  const removeActiveTurnOutputCallback = bindActiveUserTurnCompletion(ws, sessionId, activeTurn)

  const sent = await conversationService.sendMessage(
    sessionId,
@ -394,6 +411,8 @@ async function handleUserMessage(
    message.attachments
  )
  if (!sent) {
+    removeActiveTurnOutputCallback()
+    clearActiveUserTurn(sessionId, activeTurn)
    removeTitleOutputCallback?.()
    discardActiveTitleTurn(sessionId, titleTurnNumber)
    sendMessage(ws, {
@ -406,6 +425,57 @@ async function handleUserMessage(
  }

  userMessageSent = true
+  activeTurn.messageSent = true
+}
+
+function clearActiveUserTurn(sessionId: string, activeTurn: ActiveUserTurnState): void {
+  if (activeUserTurns.get(sessionId) === activeTurn) {
+    activeUserTurns.delete(sessionId)
+  }
+}
+
+function bindActiveUserTurnCompletion(
+  ws: ServerWebSocket<WebSocketData>,
+  sessionId: string,
+  activeTurn: ActiveUserTurnState,
+): () => void {
+  const callback = (cliMsg: any) => {
+    if (!activeTurn.messageSent || cliMsg?.type !== 'result') return
+
+    conversationService.removeOutputCallback(sessionId, callback)
+    clearActiveUserTurn(sessionId, activeTurn)
+    applyDeferredRuntimeRestartAfterActiveTurn(ws, sessionId)
+  }
+
+  conversationService.onOutput(sessionId, callback)
+  return () => conversationService.removeOutputCallback(sessionId, callback)
+}
+
+function shouldDeferRuntimeRestartForActiveTurn(sessionId: string): boolean {
+  return activeUserTurns.get(sessionId)?.messageSent === true
+}
+
+function applyDeferredRuntimeRestartAfterActiveTurn(
+  ws: ServerWebSocket<WebSocketData>,
+  sessionId: string,
+): void {
+  const deferred = deferredRuntimeRestarts.get(sessionId)
+  if (!deferred) return
+
+  deferredRuntimeRestarts.delete(sessionId)
+  void enqueueRuntimeTransition(sessionId, async () => {
+    const currentOverride = runtimeOverrides.get(sessionId)
+    if (
+      !currentOverride ||
+      currentOverride.providerId !== deferred.providerId ||
+      currentOverride.modelId !== deferred.modelId ||
+      currentOverride.effort !== deferred.effort ||
+      !conversationService.hasSession(sessionId)
+    ) {
+      return
+    }
+    await restartSessionWithRuntimeConfig(ws, sessionId)
+  })
 }

 async function handleDesktopClearCommand(
@ -621,6 +691,12 @@ async function handleSetRuntimeConfig(
    (runtimeOverrideVersions.get(sessionId) ?? 0) + 1,
  )

+  if (shouldDeferRuntimeRestartForActiveTurn(sessionId)) {
+    deferredRuntimeRestarts.set(sessionId, nextOverride)
+    await persistSessionRuntimeConfig(sessionId, nextOverride)
+    return
+  }
+
  if (conversationService.hasSession(sessionId)) {
    await enqueueRuntimeTransition(sessionId, async () => {
      await persistSessionRuntimeConfig(sessionId, nextOverride)
@ -1081,6 +1157,8 @@ function cleanupSessionRuntimeState(sessionId: string) {
  sessionSlashCommands.delete(sessionId)
  sessionTitleState.delete(sessionId)
  runtimeOverrides.delete(sessionId)
+  activeUserTurns.delete(sessionId)
+  deferredRuntimeRestarts.delete(sessionId)
  runtimeTransitionPromises.delete(sessionId)
  sessionStartupPromises.delete(sessionId)
  lastResolvedStartupWorkDirs.delete(sessionId)