fix: wait for sdk socket before context inspection

Desktop context inspection could send get_context_usage while a prewarmed CLI process existed but before the SDK socket was connected. That first control request sat in the pending outbound queue and could time out, leaving the composer spinner or a misleading empty-state context display. The server now waits for the SDK socket before sending control requests, and the UI renders the actual initial context snapshot instead of masking empty active-session data.

Constraint: CLI /context and desktop get_context_usage must stay on the shared collectContextData/analyzeContextUsage path

Rejected: Hide zero-token active snapshots in the component | that masks backend timing bugs and prevents showing real initial system/tool context

Confidence: high

Scope-risk: narrow

Directive: Do not queue control requests before the SDK socket is connected; let user messages keep the pending outbound behavior

Tested: bun test src/server/__tests__/conversations.test.ts --test-name-pattern "not queue control|context-only|prewarmed empty session|structured session inspection"

Tested: cd desktop && bun run test -- pages.test.tsx --test-name-pattern "EmptySession shows draft context|first-paint spinner|empty live session|live context usage|runtime model"

Tested: bun run check:server

Tested: cd desktop && bun run lint

Tested: real dev API prewarmed empty session returned 31,190 / 262,144 (12%) on first context-only inspection and 16ms on repeat

Tested: browser at http://127.0.0.1:5174/ showed 上下文用量 12% for an empty active session

desktop/src/__tests__/pages.test.tsx

@@ -833,7 +833,7 @@ describe('Content-only pages render without errors', () => {
     useChatStore.setState({ sessions: {} })
   })
 
-  it('ActiveSession treats an empty live zero-token context as pending instead of 0%', async () => {
+  it('ActiveSession shows initial context usage for an empty live session', async () => {
     const SESSION_ID = 'context-empty-live-session'
     vi.mocked(sessionsApi.getInspection).mockResolvedValueOnce({
       active: true,
@@ -846,12 +846,15 @@ describe('Content-only pages render without errors', () => {
       },
       context: {
         categories: [
-          { name: 'Free space', tokens: 128_000, color: '#9B928C', isDeferred: true },
+          { name: 'System prompt', tokens: 6_800, color: '#8a8a8a' },
+          { name: 'System tools', tokens: 13_200, color: '#9B928C' },
+          { name: 'MCP tools', tokens: 8_000, color: '#06b6d4' },
+          { name: 'Free space', tokens: 100_000, color: '#9B928C', isDeferred: true },
         ],
-        totalTokens: 0,
+        totalTokens: 28_000,
         maxTokens: 128_000,
         rawMaxTokens: 128_000,
-        percentage: 0,
+        percentage: 22,
         gridRows: [],
         model: 'kimi-k2.6',
         memoryFiles: [],
@@ -901,11 +904,10 @@ describe('Content-only pages render without errors', () => {
 
     render(<ActiveSession />)
 
-    const indicator = await screen.findByLabelText('Context usage not calculated')
-    expect(indicator).toHaveTextContent('--')
-    expect(screen.queryByLabelText('Context usage 0%')).not.toBeInTheDocument()
+    const indicator = await screen.findByLabelText('Context usage 22%')
+    expect(indicator).toHaveTextContent('22%')
     expect(screen.getAllByText('kimi-k2.6').length).toBeGreaterThan(0)
-    expect(screen.getByText('Context usage will be calculated after the session starts.')).toBeInTheDocument()
+    expect(screen.queryByText('Context usage will be calculated after the session starts.')).not.toBeInTheDocument()
 
     useTabStore.setState({ tabs: [], activeTabId: null })
     useSessionStore.setState({ sessions: [], activeSessionId: null, isLoading: false, error: null })

desktop/src/components/chat/ContextUsageIndicator.tsx

@@ -53,11 +53,6 @@ function shouldFetchContext(sessionId: string | undefined, draft: boolean) {
   return Boolean(sessionId) && !draft
 }
 
-function isEmptyContextSnapshot(context: SessionContextSnapshot | null, messageCount: number) {
-  if (!context || messageCount > 0) return false
-  return context.totalTokens === 0 && context.percentage === 0
-}
-
 export function ContextUsageIndicator({
   sessionId,
   chatState,
@@ -139,11 +134,11 @@ export function ContextUsageIndicator({
   }, [chatState, messageCount, refresh])
 
   const details = useMemo(() => {
-    if (!context || isEmptyContextSnapshot(context, messageCount)) return []
+    if (!context) return []
     return pickUsedContextCategory(context)
-  }, [context, messageCount])
+  }, [context])
 
-  const displayContext = isEmptyContextSnapshot(context, messageCount) ? null : context
+  const displayContext = context
   const hasPlaceholderContext = !displayContext && (
     draft || (!loading && messageCount === 0 && (!error || isCliNotRunningError(error)))
   )

src/server/__tests__/conversations.test.ts

@@ -75,6 +75,60 @@ describe('ConversationService', () => {
     expect(result).toBe(false)
   })
 
+  it('should not queue control requests before the SDK socket connects', async () => {
+    const svc = new ConversationService()
+    const sid = crypto.randomUUID()
+    const sent: unknown[] = []
+    const session: any = {
+      proc: { kill() {}, exited: Promise.resolve(0) },
+      outputCallbacks: [],
+      workDir: process.cwd(),
+      permissionMode: 'default',
+      sdkToken: 'token',
+      sdkSocket: null,
+      pendingOutbound: [],
+      startupPending: false,
+      startupExitCode: null,
+      stdoutLines: [],
+      stderrLines: [],
+      outputDrain: Promise.resolve(),
+      sdkMessages: [],
+      initMessage: null,
+      pendingPermissionRequests: new Map(),
+    }
+    ;(svc as any).sessions.set(sid, session)
+
+    const request = svc.requestControl(sid, { subtype: 'get_context_usage' }, 1_000)
+    await new Promise((resolve) => setTimeout(resolve, 75))
+
+    expect(session.pendingOutbound).toHaveLength(0)
+    expect(sent).toHaveLength(0)
+
+    session.sdkSocket = {
+      send(data: string) {
+        sent.push(JSON.parse(data))
+      },
+    }
+
+    await new Promise((resolve) => setTimeout(resolve, 75))
+    expect(session.pendingOutbound).toHaveLength(0)
+    expect(sent).toHaveLength(1)
+
+    const requestId = (sent[0] as any).request_id
+    for (const callback of [...session.outputCallbacks]) {
+      callback({
+        type: 'control_response',
+        response: {
+          subtype: 'success',
+          request_id: requestId,
+          response: { ok: true },
+        },
+      })
+    }
+
+    await expect(request).resolves.toEqual({ ok: true })
+  })
+
   it('should forward suggested permission updates for allow-for-session decisions', () => {
     const svc = new ConversationService()
     const sent: unknown[] = []
@@ -425,6 +479,29 @@ describe('WebSocket Chat Integration', () => {
     }
   }
 
+  async function withMockInitDelay<T>(
+    delayMs: number | undefined,
+    callback: () => Promise<T>,
+  ): Promise<T> {
+    const previousDelay = process.env.MOCK_SDK_INIT_DELAY_MS
+
+    if (delayMs && delayMs > 0) {
+      process.env.MOCK_SDK_INIT_DELAY_MS = String(delayMs)
+    } else {
+      delete process.env.MOCK_SDK_INIT_DELAY_MS
+    }
+
+    try {
+      return await callback()
+    } finally {
+      if (previousDelay === undefined) {
+        delete process.env.MOCK_SDK_INIT_DELAY_MS
+      } else {
+        process.env.MOCK_SDK_INIT_DELAY_MS = previousDelay
+      }
+    }
+  }
+
   async function withMockStreamDelay<T>(
     delayMs: number | undefined,
     callback: () => Promise<T>,
@@ -861,6 +938,64 @@ describe('WebSocket Chat Integration', () => {
     })
   })
 
+  it('should return initial context for a prewarmed empty session on the first inspection request', async () => {
+    await withMockInitDelay(500, async () => {
+      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 ws = new WebSocket(`${wsUrl}/ws/${sessionId}`)
+
+      try {
+        await new Promise<void>((resolve, reject) => {
+          const timeout = setTimeout(() => {
+            ws.close()
+            reject(new Error(`Timed out waiting for prewarm connection for ${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)
+            ws.close()
+            reject(new Error(`WebSocket error for prewarm context session ${sessionId}`))
+          }
+        })
+
+        await waitUntil(
+          () => conversationService.hasSession(sessionId),
+          `prewarmed CLI process for ${sessionId}`,
+        )
+
+        const startedAt = performance.now()
+        const res = await fetch(`${baseUrl}/api/sessions/${sessionId}/inspection?includeContext=1&contextOnly=1`)
+        const elapsedMs = performance.now() - startedAt
+        expect(res.status).toBe(200)
+        const body = await res.json() as any
+
+        expect(body.context.model).toBe('mock-opus')
+        expect(body.context.totalTokens).toBeGreaterThan(0)
+        expect(body.context.percentage).toBe(13)
+        expect(body.context.categories.some((category: any) => category.name === 'System prompt')).toBe(true)
+        expect(body.errors).toEqual({})
+        expect(elapsedMs).toBeLessThan(2_000)
+      } finally {
+        ws.close()
+        conversationService.stopSession(sessionId)
+      }
+    })
+  }, 10_000)
+
   it('should complete the client turn when the CLI exits after startup', async () => {
     const messages = await withMockExitAfterFirstUser(50, () =>
       runTurnUntilComplete(`chat-late-exit-${crypto.randomUUID()}`, 'trigger late exit'),

src/server/__tests__/fixtures/mock-sdk-cli.ts

@@ -25,6 +25,7 @@ function extractUserText(message: any): string {
 const sdkUrl = getArg('--sdk-url')
 const sessionId = getArg('--session-id') || crypto.randomUUID()
 const initMode = process.env.MOCK_SDK_INIT_MODE || 'on_open'
+const initDelayMs = Number(process.env.MOCK_SDK_INIT_DELAY_MS || '0')
 const streamDelayMs = Number(process.env.MOCK_SDK_STREAM_DELAY_MS || '0')
 const exitAfterOpenMs = Number(process.env.MOCK_SDK_EXIT_AFTER_OPEN_MS || '0')
 const exitAfterFirstUserMs = Number(process.env.MOCK_SDK_EXIT_AFTER_FIRST_USER_MS || '0')
@@ -63,7 +64,11 @@ function sendInit() {
 
 ws.addEventListener('open', () => {
   if (initMode !== 'on_first_user') {
-    sendInit()
+    if (initDelayMs > 0) {
+      setTimeout(sendInit, initDelayMs)
+    } else {
+      sendInit()
+    }
   }
   if (exitAfterOpenMs > 0) {
     setTimeout(() => process.exit(1), exitAfterOpenMs)

src/server/services/conversationService.ts

@@ -20,6 +20,7 @@ import {
 const MAX_CAPTURED_PROCESS_LINES = 80
 const MAX_CAPTURED_SDK_MESSAGES = 40
 const MAX_CAPTURED_SDK_SUMMARY = 20
+const CONTROL_READY_POLL_MS = 50
 
 type AttachmentRef = {
   type: 'file' | 'image'
@@ -387,7 +388,31 @@ export class ConversationService {
     })
   }
 
-  requestControl(
+  private isControlChannelReady(session: SessionProcess): boolean {
+    return Boolean(session.sdkSocket)
+  }
+
+  private async waitForControlChannelReady(
+    sessionId: string,
+    timeoutMs: number,
+  ): Promise<void> {
+    const startedAt = Date.now()
+
+    while (Date.now() - startedAt < timeoutMs) {
+      const session = this.sessions.get(sessionId)
+      if (!session) {
+        throw new Error('CLI session is not running')
+      }
+      if (this.isControlChannelReady(session)) {
+        return
+      }
+      await new Promise((resolve) => setTimeout(resolve, CONTROL_READY_POLL_MS))
+    }
+
+    throw new Error('Timed out waiting for CLI control channel to become ready')
+  }
+
+  async requestControl(
     sessionId: string,
     request: Record<string, unknown>,
     timeoutMs = 10_000,
@@ -396,12 +421,15 @@ export class ConversationService {
       return Promise.reject(new Error('CLI session is not running'))
     }
 
+    const startedAt = Date.now()
+    await this.waitForControlChannelReady(sessionId, timeoutMs)
+    const responseTimeoutMs = Math.max(1, timeoutMs - (Date.now() - startedAt))
     const requestId = crypto.randomUUID()
     return new Promise((resolve, reject) => {
       const timeout = setTimeout(() => {
         this.removeOutputCallback(sessionId, handleOutput)
         reject(new Error(`Timed out waiting for ${String(request.subtype ?? 'control')} response`))
-      }, timeoutMs)
+      }, responseTimeoutMs)
 
       const finish = (fn: () => void) => {
         clearTimeout(timeout)