fix(kimi-code): show last-request token count for running subagents

.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.

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<typeof setInterval> | 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,

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,
     });

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 {

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 };

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([