diff --git a/.changeset/fix-subagent-token-display.md b/.changeset/fix-subagent-token-display.md new file mode 100644 index 00000000..e420de82 --- /dev/null +++ b/.changeset/fix-subagent-token-display.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/agent-core": patch +"@moonshot-ai/kimi-code": patch +--- + +Restore real-time token display for running subagents in the TUI. diff --git a/apps/kimi-code/src/tui/components/messages/tool-call.ts b/apps/kimi-code/src/tui/components/messages/tool-call.ts index fae50c64..2c84ed3f 100644 --- a/apps/kimi-code/src/tui/components/messages/tool-call.ts +++ b/apps/kimi-code/src/tui/components/messages/tool-call.ts @@ -19,6 +19,7 @@ import { import { STATUS_BULLET } from '#/tui/constant/symbols'; import type { ColorPalette } from '#/tui/theme/colors'; import type { ToolCallBlockData, ToolResultBlockData } from '#/tui/types'; +import type { TokenUsage } from '@moonshot-ai/kimi-code-sdk'; import { appendStreamingArgsPreview } from '#/tui/utils/event-payload'; import { decodeMcpToolName } from '#/tui/utils/mcp-tool-name'; @@ -37,13 +38,6 @@ const PROGRESS_URL_RE = /https?:\/\/\S+/g; type SubagentTextKind = 'thinking' | 'text'; -interface SubagentTokenUsage { - readonly input?: number | undefined; - readonly inputOther?: number | undefined; - readonly inputCacheRead?: number | undefined; - readonly inputCacheCreation?: number | undefined; - readonly output: number; -} interface FinishedSubCall { readonly name: string; @@ -106,16 +100,23 @@ function str(v: unknown): string { return typeof v === 'string' ? v : ''; } -function usageInputTotal(usage: SubagentTokenUsage): number { - return ( - usage.input ?? - (usage.inputOther ?? 0) + (usage.inputCacheRead ?? 0) + (usage.inputCacheCreation ?? 0) - ); +function formatSubagentContextTokens(contextTokens: number | undefined): string | undefined { + if (contextTokens === undefined || contextTokens <= 0) return undefined; + const formatted = contextTokens >= 1000 ? `${(contextTokens / 1000).toFixed(1)}k` : String(contextTokens); + return `${formatted} tok`; } -function formatSubagentTokens(usage: SubagentTokenUsage | undefined): string | undefined { - if (usage === undefined) return undefined; - const total = usageInputTotal(usage) + usage.output; +function usageInputTotal(usage: TokenUsage): number { + return (usage.inputOther ?? 0) + (usage.inputCacheRead ?? 0) + (usage.inputCacheCreation ?? 0); +} + +function usageTotal(usage: TokenUsage | undefined): number { + if (usage === undefined) return 0; + return usageInputTotal(usage) + usage.output; +} + +function formatSubagentTokens(usage: TokenUsage | undefined): string | undefined { + const total = usageTotal(usage); if (total <= 0) return undefined; const formatted = total >= 1000 ? `${(total / 1000).toFixed(1)}k` : String(total); return `${formatted} tok`; @@ -463,7 +464,8 @@ export class ToolCallComponent extends Container { private subagentThinkingText = ''; // ── Subagent lifecycle state from subagent.spawned/completed/failed ── private subagentPhase: 'spawning' | 'running' | 'done' | 'failed' | 'backgrounded' | undefined; - private subagentUsage: SubagentTokenUsage | undefined; + private subagentContextTokens: number | undefined; + private subagentUsage: TokenUsage | undefined; private subagentResultSummary: string | undefined; private subagentError: string | undefined; private streamingProgressTimer: ReturnType | undefined; @@ -672,10 +674,11 @@ export class ToolCallComponent extends Container { getSubagentSnapshot(): ToolCallSubagentSnapshot { const finished = this.finishedSubCalls.length + this.hiddenSubCallCount; + const contextTokens = this.subagentContextTokens; const tokens = - this.subagentUsage === undefined - ? 0 - : usageInputTotal(this.subagentUsage) + this.subagentUsage.output; + contextTokens && contextTokens > 0 + ? contextTokens + : (this.subagentUsage === undefined ? 0 : usageTotal(this.subagentUsage)); const latestActivity = computeLatestActivity( this.ongoingSubCalls, this.finishedSubCalls, @@ -870,11 +873,15 @@ export class ToolCallComponent extends Container { * token usage plus the result summary for the header chip and tail summary. */ onSubagentCompleted(payload: { - usage?: SubagentTokenUsage | undefined; + contextTokens?: number | undefined; + usage?: TokenUsage | undefined; resultSummary: string; }): void { this.subagentPhase = 'done'; this.subagentEndedAtMs ??= Date.now(); + if (payload.contextTokens !== undefined && payload.contextTokens > 0) { + this.subagentContextTokens = payload.contextTokens; + } this.subagentUsage = payload.usage; this.subagentResultSummary = payload.resultSummary.length > 0 ? payload.resultSummary : undefined; @@ -888,6 +895,23 @@ export class ToolCallComponent extends Container { this.ui?.requestRender(); } + /** Handles SDK `agent.status.updated` from the child agent. */ + updateSubagentMetrics(payload: { + contextTokens?: number | undefined; + usage?: TokenUsage | undefined; + }): void { + if (payload.contextTokens !== undefined && payload.contextTokens > 0) { + this.subagentContextTokens = payload.contextTokens; + } + if (payload.usage !== undefined) { + this.subagentUsage = payload.usage; + } + this.headerText.setText(this.buildHeader()); + this.invalidate(); + this.notifySnapshotChange(); + this.ui?.requestRender(); + } + /** Handles SDK `subagent.failed`. */ onSubagentFailed(payload: { error: string }): void { this.subagentPhase = 'failed'; @@ -1210,7 +1234,9 @@ export class ToolCallComponent extends Container { parts.push(chalk.hex(this.colors.success)('✓ done')); const toolCount = this.finishedSubCalls.length + this.hiddenSubCallCount; if (toolCount > 0) parts.push(`${String(toolCount)} tool${toolCount > 1 ? 's' : ''}`); - const tokens = formatSubagentTokens(this.subagentUsage); + const tokens = + formatSubagentContextTokens(this.subagentContextTokens) ?? + formatSubagentTokens(this.subagentUsage); if (tokens !== undefined) parts.push(tokens); break; } @@ -1304,6 +1330,13 @@ export class ToolCallComponent extends Container { ]; const elapsed = this.getSubagentElapsedSeconds(); if (elapsed !== undefined) parts.push(formatElapsed(elapsed)); + const tokens = + this.subagentContextTokens && this.subagentContextTokens > 0 + ? this.subagentContextTokens + : this.subagentUsage === undefined + ? 0 + : usageTotal(this.subagentUsage); + if (tokens > 0) parts.push(formatTokens(tokens)); return ` · ${parts.join(' · ')}`; } @@ -1685,6 +1718,12 @@ function computeLatestActivity( return undefined; } +function formatTokens(n: number): string { + if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M tok`; + if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k tok`; + return `${String(n)} tok`; +} + function formatActivityLine( verb: string, toolName: string, diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index dfe4762c..a185e386 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -2462,7 +2462,15 @@ export class KimiTUI { is_error: event.isError, }); return true; - case 'agent.status.updated': + case 'agent.status.updated': { + const usageObj = event.usage; + const totalUsage = usageObj?.total ?? usageObj?.currentTurn; + toolCall.updateSubagentMetrics({ + contextTokens: event.contextTokens, + usage: totalUsage, + }); + return true; + } case 'background.task.started': case 'background.task.updated': case 'background.task.terminated': @@ -3007,6 +3015,7 @@ export class KimiTUI { const tc = this.state.pendingToolComponents.get(event.parentToolCallId); if (tc === undefined) return; tc.onSubagentCompleted({ + contextTokens: event.contextTokens, usage: event.usage, resultSummary: event.resultSummary, }); diff --git a/packages/agent-core/src/rpc/events.ts b/packages/agent-core/src/rpc/events.ts index 40466a80..cd19c712 100644 --- a/packages/agent-core/src/rpc/events.ts +++ b/packages/agent-core/src/rpc/events.ts @@ -202,6 +202,7 @@ export interface SubagentCompletedEvent { readonly parentToolCallId: string; readonly resultSummary: string; readonly usage?: TokenUsage | undefined; + readonly contextTokens?: number | undefined; } export interface SubagentFailedEvent { diff --git a/packages/agent-core/src/session/subagent-host.ts b/packages/agent-core/src/session/subagent-host.ts index 20b744ec..95c6209c 100644 --- a/packages/agent-core/src/session/subagent-host.ts +++ b/packages/agent-core/src/session/subagent-host.ts @@ -255,6 +255,7 @@ export class SessionSubagentHost { parentToolCallId: options.parentToolCallId, resultSummary: result, usage, + contextTokens: child.context.tokenCount, }); this.triggerSubagentStop(parent, profileName, result); return { result, usage }; diff --git a/packages/agent-core/test/session/init.test.ts b/packages/agent-core/test/session/init.test.ts index 8320daf4..c0120b5c 100644 --- a/packages/agent-core/test/session/init.test.ts +++ b/packages/agent-core/test/session/init.test.ts @@ -99,6 +99,7 @@ describe('Session.init', () => { agentId: 'main', subagentId: 'agent-0', parentToolCallId: 'generate-agents-md', + contextTokens: expect.any(Number), }), ); expect(scripted.calls[0]?.history).toMatchObject([