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)