From 0508277998a72b482ea1bed0fc09264b1546b82f 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: Tue, 28 Apr 2026 23:05:52 +0800 Subject: [PATCH] fix: keep inspector context responsive The session inspector now asks the resumed CLI for a fast structural context estimate instead of forcing the full token-counting API path. This preserves live CLI state for system prompt, tools, MCP tools, skills, and messages while avoiding the 20s timeout that made historical sessions appear stuck or fall back to transcript-only estimates. Constraint: Inspector requests must return quickly for third-party providers and historical resumed sessions. Rejected: Increase the server timeout | the slow path can still block on provider token counting and keeps the UI feeling broken. Confidence: high Scope-risk: narrow Directive: Keep interactive inspector context on the estimateOnly control path unless the UI explicitly supports a slow precise refresh. Tested: bun test src/server/__tests__/conversations.test.ts -t 'structured session inspection|Sonnet 4.6 transcript usage' Tested: cd desktop && bun run test -- --run src/i18n/index.test.tsx Tested: cd desktop && bun run build Tested: direct /api/sessions/:id/inspection?includeContext=1 returned live context in 0.046867s with System prompt, System tools, MCP tools, Messages Tested: agent-browser automation verified no loading/error and visible System prompt/System tools/MCP tools/Messages Not-tested: bare root tsc --noEmit, because current tsconfig scans existing desktop/src-tauri/target generated binary assets unrelated to this change --- src/cli/print.ts | 1 + .../context/context-noninteractive.ts | 3 + src/entrypoints/sdk/controlSchemas.ts | 1 + src/server/__tests__/conversations.test.ts | 1 + src/server/__tests__/fixtures/mock-sdk-cli.ts | 1 + src/server/api/sessions.ts | 6 +- src/utils/analyzeContext.ts | 56 +++++++++++++++---- 7 files changed, 57 insertions(+), 12 deletions(-) diff --git a/src/cli/print.ts b/src/cli/print.ts index a71d3245..9f320b90 100644 --- a/src/cli/print.ts +++ b/src/cli/print.ts @@ -2973,6 +2973,7 @@ function runHeadlessStreaming( agentDefinitions: appState.agentDefinitions, customSystemPrompt: options.systemPrompt, appendSystemPrompt: options.appendSystemPrompt, + estimateOnly: Boolean(message.request.estimateOnly), }, }) sendControlResponseSuccess(message, { ...data }) diff --git a/src/commands/context/context-noninteractive.ts b/src/commands/context/context-noninteractive.ts index 6a366d5f..1b12b63c 100644 --- a/src/commands/context/context-noninteractive.ts +++ b/src/commands/context/context-noninteractive.ts @@ -28,6 +28,7 @@ type CollectContextDataInput = { agentDefinitions: AgentDefinitionsResult customSystemPrompt?: string appendSystemPrompt?: string + estimateOnly?: boolean } } @@ -43,6 +44,7 @@ export async function collectContextData( agentDefinitions, customSystemPrompt, appendSystemPrompt, + estimateOnly, }, } = context @@ -73,6 +75,7 @@ export async function collectContextData( >, undefined, // mainThreadAgentDefinition apiView, // original messages for API usage extraction + { estimateOnly }, ) } diff --git a/src/entrypoints/sdk/controlSchemas.ts b/src/entrypoints/sdk/controlSchemas.ts index 74b5792f..29f26ab5 100644 --- a/src/entrypoints/sdk/controlSchemas.ts +++ b/src/entrypoints/sdk/controlSchemas.ts @@ -176,6 +176,7 @@ export const SDKControlGetContextUsageRequestSchema = lazySchema(() => z .object({ subtype: z.literal('get_context_usage'), + estimateOnly: z.boolean().optional(), }) .describe( 'Requests a breakdown of current context window usage by category.', diff --git a/src/server/__tests__/conversations.test.ts b/src/server/__tests__/conversations.test.ts index 7c428020..48505ba6 100644 --- a/src/server/__tests__/conversations.test.ts +++ b/src/server/__tests__/conversations.test.ts @@ -690,6 +690,7 @@ describe('WebSocket Chat Integration', () => { expect(body.usage.costDisplay).toBe('$0.1234') expect(body.usage.source).toBe('current_process') expect(body.context.model).toBe('mock-opus') + expect(body.context.estimateOnly).toBe(true) expect(body.status.mcpServers).toEqual([{ name: 'mock', status: 'connected' }]) const basicRes = await fetch(`${baseUrl}/api/sessions/${sessionId}/inspection?includeContext=0`) diff --git a/src/server/__tests__/fixtures/mock-sdk-cli.ts b/src/server/__tests__/fixtures/mock-sdk-cli.ts index 582999c5..605aaaed 100644 --- a/src/server/__tests__/fixtures/mock-sdk-cli.ts +++ b/src/server/__tests__/fixtures/mock-sdk-cli.ts @@ -239,6 +239,7 @@ ws.addEventListener('message', (event) => { }), ), model: 'mock-opus', + estimateOnly: parsed.request.estimateOnly === true, memoryFiles: [], mcpTools: [{ name: 'mock_tool', serverName: 'mock', tokens: 144, isLoaded: true }], agents: [], diff --git a/src/server/api/sessions.ts b/src/server/api/sessions.ts index d7dae359..155289b4 100644 --- a/src/server/api/sessions.ts +++ b/src/server/api/sessions.ts @@ -285,7 +285,11 @@ async function getSessionInspection(sessionId: string, url: URL): Promise { + if (estimateOnly) { + return roughTokenCountEstimation(jsonStringify({ messages, tools })) + } + try { const result = await countMessagesTokensWithAPI(messages, tools) if (result !== null) { @@ -236,6 +241,7 @@ export async function countToolDefinitionTokens( getToolPermissionContext: () => Promise, agentInfo: AgentDefinitionsResult | null, model?: string, + estimateOnly = false, ): Promise { const toolSchemas = await Promise.all( tools.map(tool => @@ -247,7 +253,7 @@ export async function countToolDefinitionTokens( }), ), ) - const result = await countTokensWithFallback([], toolSchemas) + const result = await countTokensWithFallback([], toolSchemas, estimateOnly) if (result === null || result === 0) { const toolNames = tools.map(t => t.name).join(', ') logForDebugging( @@ -271,6 +277,7 @@ function extractSectionName(content: string): string { async function countSystemTokens( effectiveSystemPrompt: readonly string[], + estimateOnly = false, ): Promise<{ systemPromptTokens: number systemPromptSections: SystemPromptSectionDetail[] @@ -298,7 +305,7 @@ async function countSystemTokens( const systemTokenCounts = await Promise.all( namedEntries.map(({ content }) => - countTokensWithFallback([{ role: 'user', content }], []), + countTokensWithFallback([{ role: 'user', content }], [], estimateOnly), ), ) @@ -317,7 +324,7 @@ async function countSystemTokens( return { systemPromptTokens, systemPromptSections } } -async function countMemoryFileTokens(): Promise<{ +async function countMemoryFileTokens(estimateOnly = false): Promise<{ memoryFileDetails: MemoryFile[] claudeMdTokens: number }> { @@ -342,6 +349,7 @@ async function countMemoryFileTokens(): Promise<{ const tokens = await countTokensWithFallback( [{ role: 'user', content: file.content }], [], + estimateOnly, ) return { file, tokens: tokens || 0 } @@ -366,6 +374,7 @@ async function countBuiltInToolTokens( agentInfo: AgentDefinitionsResult | null, model?: string, messages?: Message[], + estimateOnly = false, ): Promise<{ builtInToolTokens: number deferredBuiltinDetails: DeferredBuiltinTool[] @@ -405,6 +414,7 @@ async function countBuiltInToolTokens( getToolPermissionContext, agentInfo, model, + estimateOnly, ) : 0 @@ -469,6 +479,7 @@ async function countBuiltInToolTokens( getToolPermissionContext, agentInfo, model, + estimateOnly, ), ), ) @@ -496,6 +507,7 @@ async function countBuiltInToolTokens( getToolPermissionContext, agentInfo, model, + estimateOnly, ) return { builtInToolTokens: alwaysLoadedTokens + deferredTokens, @@ -522,6 +534,7 @@ async function countSlashCommandTokens( tools: Tools, getToolPermissionContext: () => Promise, agentInfo: AgentDefinitionsResult | null, + estimateOnly = false, ): Promise<{ slashCommandTokens: number commandInfo: { totalCommands: number; includedCommands: number } @@ -540,6 +553,8 @@ async function countSlashCommandTokens( [slashCommandTool], getToolPermissionContext, agentInfo, + undefined, + estimateOnly, ) return { @@ -555,6 +570,7 @@ async function countSkillTokens( tools: Tools, getToolPermissionContext: () => Promise, agentInfo: AgentDefinitionsResult | null, + estimateOnly = false, ): Promise<{ skillTokens: number skillInfo: { @@ -582,6 +598,8 @@ async function countSkillTokens( [slashCommandTool], getToolPermissionContext, agentInfo, + undefined, + estimateOnly, ) // Calculate per-skill token estimates based on frontmatter only @@ -619,6 +637,7 @@ export async function countMcpToolTokens( agentInfo: AgentDefinitionsResult | null, model: string, messages?: Message[], + estimateOnly = false, ): Promise<{ mcpToolTokens: number mcpToolDetails: McpTool[] @@ -633,6 +652,7 @@ export async function countMcpToolTokens( getToolPermissionContext, agentInfo, model, + estimateOnly, ) // Subtract the single overhead since we made one bulk call const totalTokens = Math.max( @@ -729,9 +749,10 @@ export async function countMcpToolTokens( } } -async function countCustomAgentTokens(agentDefinitions: { - activeAgents: AgentDefinition[] -}): Promise<{ +async function countCustomAgentTokens( + agentDefinitions: { activeAgents: AgentDefinition[] }, + estimateOnly = false, +): Promise<{ agentTokens: number agentDetails: Agent[] }> { @@ -751,6 +772,7 @@ async function countCustomAgentTokens(agentDefinitions: { }, ], [], + estimateOnly, ), ), ) @@ -852,6 +874,7 @@ function processAttachment( async function approximateMessageTokens( messages: Message[], + estimateOnly = false, ): Promise { const microcompactResult = await microcompactMessages(messages) @@ -909,6 +932,7 @@ async function approximateMessageTokens( return _.message }), [], + estimateOnly, ) breakdown.totalTokens = approximateMessageTokens ?? 0 @@ -926,7 +950,9 @@ export async function analyzeContextUsage( mainThreadAgentDefinition?: AgentDefinition, /** Original messages before microcompact, used to extract API usage */ originalMessages?: Message[], + analysisOptions?: { estimateOnly?: boolean }, ): Promise { + const estimateOnly = analysisOptions?.estimateOnly ?? false const runtimeModel = getRuntimeMainLoopModel({ permissionMode: (await getToolPermissionContext()).mode, mainLoopModel: model, @@ -961,14 +987,15 @@ export async function analyzeContextUsage( { slashCommandTokens, commandInfo }, messageBreakdown, ] = await Promise.all([ - countSystemTokens(effectiveSystemPrompt), - countMemoryFileTokens(), + countSystemTokens(effectiveSystemPrompt, estimateOnly), + countMemoryFileTokens(estimateOnly), countBuiltInToolTokens( tools, getToolPermissionContext, agentDefinitions, runtimeModel, messages, + estimateOnly, ), countMcpToolTokens( tools, @@ -976,10 +1003,16 @@ export async function analyzeContextUsage( agentDefinitions, runtimeModel, messages, + estimateOnly, ), - countCustomAgentTokens(agentDefinitions), - countSlashCommandTokens(tools, getToolPermissionContext, agentDefinitions), - approximateMessageTokens(messages), + countCustomAgentTokens(agentDefinitions, estimateOnly), + countSlashCommandTokens( + tools, + getToolPermissionContext, + agentDefinitions, + estimateOnly, + ), + approximateMessageTokens(messages, estimateOnly), ]) // Count skills separately with error isolation @@ -987,6 +1020,7 @@ export async function analyzeContextUsage( tools, getToolPermissionContext, agentDefinitions, + estimateOnly, ) const skillInfo = skillResult.skillInfo // Use sum of individual skill token estimates (matches what's shown in details)