diff --git a/apps/kimi-code/src/tui/controllers/session-event-handler.ts b/apps/kimi-code/src/tui/controllers/session-event-handler.ts index 0e310ba91..ab05e3d9e 100644 --- a/apps/kimi-code/src/tui/controllers/session-event-handler.ts +++ b/apps/kimi-code/src/tui/controllers/session-event-handler.ts @@ -33,6 +33,7 @@ import type { import { MoonLoader } from '../components/chrome/moon-loader'; import { StatusMessageComponent } from '../components/messages/status-message'; +import { workerActivityFromTool } from '../components/messages/swarm-dashboard-model'; import { MAIN_AGENT_ID, OAUTH_LOGIN_REQUIRED_CODE, @@ -232,6 +233,19 @@ export class SessionEventHandler { if (info === undefined || info.parentToolCallId.length === 0) return true; const { parentToolCallId } = info; const sourceName = info.name; + + const swarmDash = streamingUI.getSwarmDashboard(parentToolCallId); + if (swarmDash !== undefined) { + if (event.type === 'tool.call.started') { + swarmDash.apply({ + t: 'worker.toolcall', + id: subagentId, + activity: workerActivityFromTool(event.name, argsRecord(event.args)), + }); + } + return true; // swarm worker events never fall through to appendSubToolCall + } + const toolCall = streamingUI.getToolComponent(parentToolCallId); if (toolCall === undefined) return true; toolCall.setSubagentMeta(subagentId, sourceName); @@ -483,6 +497,15 @@ export class SessionEventHandler { } private handleToolProgress(event: ToolProgressEvent): void { + if (event.update.kind === 'custom' && event.update.customKind === 'swarm') { + const dash = this.host.streamingUI.getSwarmDashboard(event.toolCallId); + if (dash === undefined) return; + const p = event.update.customData as { phase?: string; total?: number }; + if (p.phase === 'planned' && typeof p.total === 'number') dash.apply({ t: 'planned', total: p.total }); + else if (p.phase === 'synthesizing') dash.apply({ t: 'synthesizing' }); + else if (p.phase === 'done') dash.apply({ t: 'done', succeeded: 0, failed: 0 }); + return; + } if (event.update.kind !== 'status') return; const text = event.update.text; if (text === undefined || text.length === 0) return; @@ -695,6 +718,16 @@ export class SessionEventHandler { name: event.subagentName, }); + const swarmDash = streamingUI.getSwarmDashboard(event.parentToolCallId); + if (swarmDash !== undefined) { + swarmDash.apply({ + t: 'worker.spawned', + id: event.subagentId, + role: event.description ?? event.subagentName, + }); + return; + } + if (event.runInBackground) { const meta = this.buildBackgroundAgentMetadata(event); this.backgroundAgentMetadata.set(event.subagentId, meta); @@ -738,6 +771,15 @@ export class SessionEventHandler { this.appendBackgroundAgentEntry('completed', backgroundMeta, extras); return; } + const swarmDashC = streamingUI.getSwarmDashboard(event.parentToolCallId); + if (swarmDashC !== undefined) { + swarmDashC.apply({ + t: 'worker.done', + id: event.subagentId, + ...(event.contextTokens !== undefined ? { tokens: event.contextTokens } : {}), + }); + return; + } const tc = streamingUI.getToolComponent(event.parentToolCallId); if (tc === undefined) return; tc.onSubagentCompleted({ @@ -764,6 +806,11 @@ export class SessionEventHandler { this.appendBackgroundAgentEntry('failed', backgroundMeta, { error: event.error }); return; } + const swarmDashF = streamingUI.getSwarmDashboard(event.parentToolCallId); + if (swarmDashF !== undefined) { + swarmDashF.apply({ t: 'worker.failed', id: event.subagentId, error: event.error }); + return; + } const tc = streamingUI.getToolComponent(event.parentToolCallId); if (tc === undefined) return; tc.onSubagentFailed({ error: event.error }); diff --git a/apps/kimi-code/src/tui/controllers/streaming-ui.ts b/apps/kimi-code/src/tui/controllers/streaming-ui.ts index 769df3767..31e485fb3 100644 --- a/apps/kimi-code/src/tui/controllers/streaming-ui.ts +++ b/apps/kimi-code/src/tui/controllers/streaming-ui.ts @@ -4,6 +4,7 @@ import { AgentGroupComponent } from '../components/messages/agent-group'; import { AssistantMessageComponent } from '../components/messages/assistant-message'; import { CompactionComponent } from '../components/dialogs/compaction'; import { ReadGroupComponent } from '../components/messages/read-group'; +import { SwarmDashboardComponent } from '../components/messages/swarm-dashboard'; import { ThinkingComponent } from '../components/messages/thinking'; import { ToolCallComponent } from '../components/messages/tool-call'; import { STREAMING_UI_FLUSH_MS } from '../constant/streaming'; @@ -60,6 +61,7 @@ export class StreamingUIController { { name?: string; argumentsText: string; startedAtMs: number } >(); private _pendingToolComponents = new Map(); + private readonly _swarmDashboards = new Map(); private _pendingAgentGroup: { readonly turnId: string | undefined; readonly step: number; @@ -156,6 +158,10 @@ export class StreamingUIController { return this._pendingToolComponents.get(id); } + getSwarmDashboard(toolCallId: string): SwarmDashboardComponent | undefined { + return this._swarmDashboards.get(toolCallId); + } + removeToolComponent(id: string): void { this._pendingToolComponents.delete(id); } @@ -503,6 +509,15 @@ export class StreamingUIController { onToolCallStart(toolCall: ToolCallBlockData): void { if (toolCall.name === 'AskUserQuestion') return; + if (toolCall.name === 'Swarm') { + const task = typeof toolCall.args['task'] === 'string' ? toolCall.args['task'] : ''; + const dash = new SwarmDashboardComponent(task, this.host.state.theme.colors, this.host.state.ui); + this._swarmDashboards.set(toolCall.id, dash); + this.host.state.transcriptContainer.addChild(dash); + this.host.state.ui.requestRender(); + return; + } + const { state } = this.host; const tc = new ToolCallComponent( toolCall, diff --git a/apps/kimi-code/test/tui/components/messages/swarm-routing.test.ts b/apps/kimi-code/test/tui/components/messages/swarm-routing.test.ts new file mode 100644 index 000000000..a1e5e2d87 --- /dev/null +++ b/apps/kimi-code/test/tui/components/messages/swarm-routing.test.ts @@ -0,0 +1,96 @@ +import { describe, expect, it } from 'vitest'; + +import type { Event } from '@moonshot-ai/kimi-code-sdk'; + +import { SessionEventHandler, type SessionEventHost } from '#/tui/controllers/session-event-handler'; +import { SwarmDashboardComponent } from '#/tui/components/messages/swarm-dashboard'; +import { workerActivityFromTool } from '#/tui/components/messages/swarm-dashboard-model'; +import { darkColors } from '#/tui/theme/colors'; + +const strip = (t: string): string => t.replaceAll(/\[[0-9;]*m/g, ''); + +describe('swarm dashboard wiring (translation)', () => { + it('produces the expected dashboard from a worker lifecycle sequence', () => { + const dash = new SwarmDashboardComponent('task', darkColors, undefined); + dash.apply({ t: 'planned', total: 2 }); + dash.apply({ t: 'worker.spawned', id: 's1', role: 'Researcher' }); + dash.apply({ t: 'worker.toolcall', id: 's1', activity: workerActivityFromTool('Read', { path: 'a.ts' }) }); + dash.apply({ t: 'worker.done', id: 's1', tokens: 2100 }); + dash.apply({ t: 'worker.spawned', id: 's2', role: 'Analyst' }); + dash.apply({ t: 'worker.failed', id: 's2', error: 'timeout' }); + dash.apply({ t: 'done', succeeded: 1, failed: 1 }); + const out = strip(dash.render(80).join('\n')); + expect(out).toContain('Researcher'); + expect(out).toContain('Analyst'); + expect(out).toContain('timeout'); + expect(out).toMatch(/2 workers/); + }); + + it('routes live swarm events through SessionEventHandler into the dashboard', () => { + const parentToolCallId = 'tc-swarm'; + const dash = new SwarmDashboardComponent('task', darkColors, undefined); + const mockHost = { + streamingUI: { + setTurnId: (): void => {}, + getSwarmDashboard: (id: string): SwarmDashboardComponent | undefined => + id === parentToolCallId ? dash : undefined, + getToolComponent: (): undefined => undefined, + }, + } as unknown as SessionEventHost; + const handler = new SessionEventHandler(mockHost); + const noop = (): void => {}; + + handler.handleEvent( + { + type: 'tool.progress', + agentId: 'main', + sessionId: 's', + turnId: 1, + toolCallId: parentToolCallId, + update: { kind: 'custom', customKind: 'swarm', customData: { phase: 'planned', total: 1 } }, + } as unknown as Event, + noop, + ); + handler.handleEvent( + { + type: 'subagent.spawned', + agentId: 'main', + sessionId: 's', + subagentId: 'w1', + subagentName: 'explore', + parentToolCallId, + description: 'Researcher', + runInBackground: false, + } as unknown as Event, + noop, + ); + handler.handleEvent( + { + type: 'tool.call.started', + agentId: 'w1', + sessionId: 's', + turnId: 1, + toolCallId: 'inner-1', + name: 'Read', + args: { path: 'x.ts' }, + } as unknown as Event, + noop, + ); + handler.handleEvent( + { + type: 'subagent.failed', + agentId: 'main', + sessionId: 's', + subagentId: 'w1', + parentToolCallId, + error: 'boom', + } as unknown as Event, + noop, + ); + + const out = strip(dash.render(80).join('\n')); + expect(out).toContain('Researcher'); + expect(out).toContain('boom'); + expect(out).toMatch(/Workers 1\/1/); + }); +});