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)