From c3139d56602a311b5d106209737b1f52426426c9 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Tue, 26 May 2026 20:40:06 +0800 Subject: [PATCH 1/2] fix(kimi-code): show last-request token count for running subagents Route `agent.status.updated` for subagents to a new `updateSubagentMetrics` on the ToolCallComponent so context-token counts refresh while the agent is still working. Carry `contextTokens` through `subagent.completed` so the final chip reflects the same number. --- .changeset/fix-subagent-token-display.md | 6 ++ .../src/tui/components/messages/tool-call.ts | 79 ++++++++++++++----- apps/kimi-code/src/tui/kimi-tui.ts | 11 ++- packages/agent-core/src/rpc/events.ts | 1 + .../agent-core/src/session/subagent-host.ts | 1 + packages/agent-core/test/session/init.test.ts | 1 + 6 files changed, 77 insertions(+), 22 deletions(-) create mode 100644 .changeset/fix-subagent-token-display.md diff --git a/.changeset/fix-subagent-token-display.md b/.changeset/fix-subagent-token-display.md new file mode 100644 index 0000000..e420de8 --- /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 6e225c8..88367e5 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,13 @@ 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(); + this.subagentContextTokens = payload.contextTokens; this.subagentUsage = payload.usage; this.subagentResultSummary = payload.resultSummary.length > 0 ? payload.resultSummary : undefined; @@ -888,6 +893,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 +1232,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 +1328,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(' · ')}`; } @@ -1675,6 +1706,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 91273ec..00283c6 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -2350,7 +2350,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': @@ -2895,6 +2903,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 40466a8..cd19c71 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 20b744e..95c6209 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 8320daf..c0120b5 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([ From eb68e708918ef2cdc00c575e1ccf2722cae9c60d Mon Sep 17 00:00:00 2001 From: _Kerman Date: Tue, 26 May 2026 21:31:00 +0800 Subject: [PATCH 2/2] fix --- apps/kimi-code/src/tui/components/messages/tool-call.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 88367e5..c4823f6 100644 --- a/apps/kimi-code/src/tui/components/messages/tool-call.ts +++ b/apps/kimi-code/src/tui/components/messages/tool-call.ts @@ -879,7 +879,9 @@ export class ToolCallComponent extends Container { }): void { this.subagentPhase = 'done'; this.subagentEndedAtMs ??= Date.now(); - this.subagentContextTokens = payload.contextTokens; + if (payload.contextTokens !== undefined && payload.contextTokens > 0) { + this.subagentContextTokens = payload.contextTokens; + } this.subagentUsage = payload.usage; this.subagentResultSummary = payload.resultSummary.length > 0 ? payload.resultSummary : undefined;