From a31b3045fd6559ff1b9c2a18f6e5504c33255144 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: Fri, 8 May 2026 13:45:02 +0800 Subject: [PATCH] Preserve explicit desktop runtime selection Desktop sessions can prewarm before the user-selected runtime reaches the server, and an empty-session draft could previously persist an inferred default runtime as if the user had explicitly selected it. The fix keeps new-session defaults implicit while preserving explicit draft selections, and waits for queued runtime restarts before the first turn is delivered to the CLI process. Constraint: Runtime selection is session-scoped and may arrive while prewarm or first-turn startup is already in flight Rejected: Always persist the current default provider/model into new sessions | this masks whether the user actually selected a runtime and can launch stale providers Confidence: high Scope-risk: moderate Directive: Do not send user turns across pending runtime restarts without re-checking the session runtime override Tested: bun test src/server/__tests__/conversations.test.ts Tested: cd desktop && bun run test -- src/pages/EmptySession.test.tsx Tested: git diff --check --- desktop/src/pages/EmptySession.test.tsx | 37 ++++++ desktop/src/pages/EmptySession.tsx | 35 +----- src/server/__tests__/conversations.test.ts | 130 +++++++++++++++++++++ src/server/ws/handler.ts | 74 ++++++++---- 4 files changed, 224 insertions(+), 52 deletions(-) diff --git a/desktop/src/pages/EmptySession.test.tsx b/desktop/src/pages/EmptySession.test.tsx index 52fdcf34..a7cb6217 100644 --- a/desktop/src/pages/EmptySession.test.tsx +++ b/desktop/src/pages/EmptySession.test.tsx @@ -197,6 +197,43 @@ describe('EmptySession', () => { attachments: [], }) expect(mocks.wsConnect).toHaveBeenCalledWith('draft-session') + expect(useSessionRuntimeStore.getState().selections['draft-session']).toBeUndefined() + }) + + it('stores and replays a draft runtime only when the user explicitly selected one', async () => { + useSessionRuntimeStore.getState().setSelection('__draft__', { + providerId: 'provider-explicit', + modelId: 'model-explicit', + }) + + render() + + fireEvent.change(screen.getByRole('textbox'), { + target: { value: 'draft question', selectionStart: 14 }, + }) + + fireEvent.click(screen.getByRole('button', { name: /Run/i })) + + await waitFor(() => { + expect(mocks.createSession).toHaveBeenCalledWith({}) + }) + + expect(useSessionRuntimeStore.getState().selections['draft-session']).toEqual({ + providerId: 'provider-explicit', + modelId: 'model-explicit', + }) + expect(useSessionRuntimeStore.getState().selections['__draft__']).toBeUndefined() + expect(mocks.wsSend.mock.calls.slice(0, 2)).toEqual([ + [ + 'draft-session', + { + type: 'set_runtime_config', + providerId: 'provider-explicit', + modelId: 'model-explicit', + }, + ], + ['draft-session', { type: 'prewarm_session' }], + ]) }) it('shows an actionable repository error when direct branch switching is blocked', async () => { diff --git a/desktop/src/pages/EmptySession.tsx b/desktop/src/pages/EmptySession.tsx index 60e1f179..936e23a8 100644 --- a/desktop/src/pages/EmptySession.tsx +++ b/desktop/src/pages/EmptySession.tsx @@ -4,12 +4,10 @@ import { skillsApi } from '../api/skills' import { useTranslation } from '../i18n' import { useSessionStore } from '../stores/sessionStore' import { useChatStore } from '../stores/chatStore' -import { useProviderStore } from '../stores/providerStore' import { useSessionRuntimeStore, DRAFT_RUNTIME_SELECTION_KEY } from '../stores/sessionRuntimeStore' import { useSettingsStore } from '../stores/settingsStore' import { useUIStore } from '../stores/uiStore' import { SETTINGS_TAB_ID, useTabStore } from '../stores/tabStore' -import { OFFICIAL_DEFAULT_MODEL_ID } from '../constants/modelCatalog' import { RepositoryLaunchControls } from '../components/shared/RepositoryLaunchControls' import { PermissionModeSelector } from '../components/controls/PermissionModeSelector' import { ModelSelector } from '../components/controls/ModelSelector' @@ -207,31 +205,6 @@ export function EmptySession() { [slashCommands], ) - const resolveDraftRuntimeSelection = async () => { - const settings = useSettingsStore.getState() - let providerState = useProviderStore.getState() - if ( - settings.activeProviderName && - providerState.providers.length === 0 && - !providerState.isLoading - ) { - await providerState.fetchProviders() - providerState = useProviderStore.getState() - } - const inferredProviderId = providerState.activeId ?? ( - settings.activeProviderName - ? providerState.providers.find((provider) => provider.name === settings.activeProviderName)?.id ?? null - : null - ) - return ( - useSessionRuntimeStore.getState().selections[DRAFT_RUNTIME_SELECTION_KEY] - ?? { - providerId: inferredProviderId, - modelId: settings.currentModel?.id ?? OFFICIAL_DEFAULT_MODEL_ID, - } - ) - } - const handleWorkDirChange = (newWorkDir: string) => { setWorkDir(newWorkDir) setSelectedBranch(null) @@ -297,15 +270,17 @@ export function EmptySession() { setIsSubmitting(true) try { - const draftSelection = await resolveDraftRuntimeSelection() + const explicitDraftSelection = useSessionRuntimeStore.getState().selections[DRAFT_RUNTIME_SELECTION_KEY] const sessionId = await createSession( workDir || undefined, selectedBranch ? { repository: { branch: selectedBranch, worktree: useWorktree } } : undefined, ) - useSessionRuntimeStore.getState().setSelection(sessionId, draftSelection) - useSessionRuntimeStore.getState().clearSelection(DRAFT_RUNTIME_SELECTION_KEY) + if (explicitDraftSelection) { + useSessionRuntimeStore.getState().setSelection(sessionId, explicitDraftSelection) + useSessionRuntimeStore.getState().clearSelection(DRAFT_RUNTIME_SELECTION_KEY) + } setActiveView('code') useTabStore.getState().openTab(sessionId, 'New Session') connectToSession(sessionId) diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts index 8b97c4cb..03771882 100644 --- a/src/server/__tests__/conversations.test.ts +++ b/src/server/__tests__/conversations.test.ts @@ -1726,6 +1726,136 @@ describe('WebSocket Chat Integration', () => { } }, 20_000) + it('should wait for a runtime restart queued during first-turn startup before sending that turn', async () => { + const providerService = new ProviderService() + const provider = await providerService.addProvider({ + presetId: 'custom', + name: 'Provider First Turn Runtime', + apiKey: 'key-first-turn-runtime', + baseUrl: 'http://127.0.0.1:1/anthropic', + apiFormat: 'anthropic', + models: { + main: 'first-turn-main', + haiku: 'first-turn-haiku', + sonnet: 'first-turn-sonnet', + opus: 'first-turn-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 + }> = [] + const sendCalls: Array<{ content: string; startCallCount: number }> = [] + let releaseFirstStart!: () => void + const firstStartGate = new Promise((resolve) => { + releaseFirstStart = resolve + }) + let markFirstStart!: () => void + const firstStartEntered = new Promise((resolve) => { + markFirstStart = resolve + }) + + 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) { + markFirstStart() + await firstStartGate + } + return originalStartSession(sid, workDir, sdkUrl, options) + }) as typeof conversationService.startSession + + conversationService.sendMessage = (function patchedSendMessage( + sid: string, + content: string, + attachments?: any, + ) { + sendCalls.push({ content, startCallCount: startCalls.length }) + return originalSendMessage(sid, content, attachments) + }) as typeof conversationService.sendMessage + + const ws = new WebSocket(`${wsUrl}/ws/${sessionId}`) + const messages: any[] = [] + try { + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + reject(new Error(`Timed out waiting for first-turn runtime synchronization for session ${sessionId}`)) + }, 15_000) + + ws.onmessage = (event) => { + const msg = JSON.parse(event.data as string) + messages.push(msg) + + if (msg.type === 'connected') { + ws.send(JSON.stringify({ type: 'prewarm_session' })) + void firstStartEntered.then(() => { + ws.send(JSON.stringify({ type: 'user_message', content: 'first turn while runtime changes' })) + ws.send(JSON.stringify({ + type: 'set_runtime_config', + providerId: provider.id, + modelId: 'first-turn-sonnet', + })) + releaseFirstStart() + }) + return + } + + if (msg.type === 'error') { + clearTimeout(timeout) + reject(new Error(msg.message)) + return + } + + if (msg.type === 'message_complete') { + clearTimeout(timeout) + resolve() + } + } + + ws.onerror = () => { + clearTimeout(timeout) + reject(new Error(`WebSocket error for first-turn runtime synchronization session ${sessionId}`)) + } + }) + + expect(startCalls).toHaveLength(2) + expect(startCalls[0]).toMatchObject({ sessionId }) + expect(startCalls[0]?.options?.providerId).toBeNull() + expect(startCalls[1]).toMatchObject({ + sessionId, + options: { + providerId: provider.id, + model: 'first-turn-sonnet', + }, + }) + expect(sendCalls).toEqual([{ + content: 'first turn while runtime changes', + startCallCount: 2, + }]) + expect(messages.some((msg) => msg.type === 'content_delta')).toBe(true) + } finally { + ws.close() + conversationService.startSession = originalStartSession + conversationService.sendMessage = originalSendMessage + conversationService.stopSession(sessionId) + } + }, 20_000) + it('should keep the session idle in the UI while applying a runtime-only model switch', async () => { const providerService = new ProviderService() const provider = await providerService.addProvider({ diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index d5b2afca..936b0461 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -252,28 +252,10 @@ async function handleUserMessage( // Send thinking status sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) - const pendingRuntimeTransition = runtimeTransitionPromises.get(sessionId) - if (pendingRuntimeTransition) { - try { - await pendingRuntimeTransition - } catch (err) { - const errMsg = err instanceof Error ? err.message : String(err) - void diagnosticsService.recordEvent({ - type: 'runtime_transition_failed', - severity: 'error', - sessionId, - summary: errMsg, - details: err, - }) - console.error(`[WS] Runtime transition failed before handling user message for ${sessionId}: ${errMsg}`) - sendMessage(ws, { - type: 'error', - message: `Failed to switch provider/model: ${errMsg}`, - code: 'CLI_RESTART_FAILED', - }) - sendMessage(ws, { type: 'status', state: 'idle' }) - return - } + const initialRuntimeTransition = await waitForRuntimeTransitionBeforeUserTurn(ws, sessionId) + if (!initialRuntimeTransition.ok) return + if (initialRuntimeTransition.waited) { + sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) } // Track and emit the first placeholder title before CLI startup/streaming. @@ -314,6 +296,15 @@ async function handleUserMessage( return } + const startupRuntimeTransition = await waitForRuntimeTransitionBeforeUserTurn(ws, sessionId) + if (startupRuntimeTransition.ok) { + if (startupRuntimeTransition.waited) { + sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) + } + } else { + return + } + // Register the callback before sending the turn so startup errors are not lost. // Keep output muted until the current user turn is enqueued to avoid forwarding // any pre-turn SDK chatter as fresh chat history. @@ -1483,6 +1474,45 @@ function enqueueRuntimeTransition( return next } +async function waitForRuntimeTransitionBeforeUserTurn( + ws: ServerWebSocket, + sessionId: string, +): Promise<{ ok: boolean; waited: boolean }> { + let waited = false + let pendingRuntimeTransition = runtimeTransitionPromises.get(sessionId) + while (pendingRuntimeTransition) { + waited = true + try { + await pendingRuntimeTransition + } catch (err) { + const errMsg = err instanceof Error ? err.message : String(err) + void diagnosticsService.recordEvent({ + type: 'runtime_transition_failed', + severity: 'error', + sessionId, + summary: errMsg, + details: err, + }) + console.error(`[WS] Runtime transition failed before handling user message for ${sessionId}: ${errMsg}`) + sendMessage(ws, { + type: 'error', + message: `Failed to switch provider/model: ${errMsg}`, + code: 'CLI_RESTART_FAILED', + }) + sendMessage(ws, { type: 'status', state: 'idle' }) + return { ok: false, waited } + } + + const nextTransition = runtimeTransitionPromises.get(sessionId) + pendingRuntimeTransition = + nextTransition && nextTransition !== pendingRuntimeTransition + ? nextTransition + : undefined + } + + return { ok: true, waited } +} + /** * Send a message to a specific session's WebSocket (for use by services) */