diff --git a/desktop/src/__tests__/pages.test.tsx b/desktop/src/__tests__/pages.test.tsx
index 221d760f..74223be0 100644
--- a/desktop/src/__tests__/pages.test.tsx
+++ b/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()
- 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 })
diff --git a/desktop/src/components/chat/ContextUsageIndicator.tsx b/desktop/src/components/chat/ContextUsageIndicator.tsx
index 9bae4ac1..08a17082 100644
--- a/desktop/src/components/chat/ContextUsageIndicator.tsx
+++ b/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)))
)
diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts
index b8b8dda1..5fb1fdca 100644
--- a/src/server/__tests__/conversations.test.ts
+++ b/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(
+ delayMs: number | undefined,
+ callback: () => Promise,
+ ): Promise {
+ 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(
delayMs: number | undefined,
callback: () => Promise,
@@ -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((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'),
diff --git a/src/server/__tests__/fixtures/mock-sdk-cli.ts b/src/server/__tests__/fixtures/mock-sdk-cli.ts
index a58f7664..046a21de 100644
--- a/src/server/__tests__/fixtures/mock-sdk-cli.ts
+++ b/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)
diff --git a/src/server/services/conversationService.ts b/src/server/services/conversationService.ts
index deced161..d94296e3 100644
--- a/src/server/services/conversationService.ts
+++ b/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 {
+ 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,
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)