From 11c556f20504e53d995c9bd07edabc0d2d231bc3 Mon Sep 17 00:00:00 2001 From: liruifengv Date: Mon, 25 May 2026 18:59:22 +0800 Subject: [PATCH 1/2] fix: prevent streaming tool argument CPU spikes --- .changeset/steady-streaming-previews.md | 6 + .../src/tui/components/messages/tool-call.ts | 43 +++-- apps/kimi-code/src/tui/constant/streaming.ts | 6 + apps/kimi-code/src/tui/kimi-tui.ts | 182 ++++++++++++++---- apps/kimi-code/src/tui/utils/event-payload.ts | 28 ++- apps/kimi-code/src/tui/utils/object-patch.ts | 8 + .../test/tui/kimi-tui-message-flow.test.ts | 119 ++++++++++++ .../test/tui/utils/event-payload.test.ts | 30 +++ packages/kosong/src/providers/anthropic.ts | 9 +- packages/kosong/test/anthropic-errors.test.ts | 20 +- packages/kosong/test/anthropic.test.ts | 25 +-- 11 files changed, 396 insertions(+), 80 deletions(-) create mode 100644 .changeset/steady-streaming-previews.md create mode 100644 apps/kimi-code/src/tui/utils/object-patch.ts create mode 100644 apps/kimi-code/test/tui/utils/event-payload.test.ts diff --git a/.changeset/steady-streaming-previews.md b/.changeset/steady-streaming-previews.md new file mode 100644 index 0000000000..67b500e08f --- /dev/null +++ b/.changeset/steady-streaming-previews.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/kosong": patch +"@moonshot-ai/kimi-code": patch +--- + +Avoid CPU spikes from large streamed tool arguments and coalesce high-frequency streaming UI updates. 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 dc831e8bef..6e225c84ff 100644 --- a/apps/kimi-code/src/tui/components/messages/tool-call.ts +++ b/apps/kimi-code/src/tui/components/messages/tool-call.ts @@ -12,10 +12,14 @@ import chalk from 'chalk'; import { highlightLines, langFromPath } from '#/tui/components/media/code-highlight'; import { renderDiffLinesClustered } from '#/tui/components/media/diff-preview'; import { COMMAND_PREVIEW_LINES } from '#/tui/constant/rendering'; -import { STREAMING_ARGS_FIELD_RE } from '#/tui/constant/streaming'; +import { + STREAMING_ARGS_FIELD_RE, + STREAMING_ARGS_PREVIEW_MAX_CHARS, +} from '#/tui/constant/streaming'; import { STATUS_BULLET } from '#/tui/constant/symbols'; import type { ColorPalette } from '#/tui/theme/colors'; import type { ToolCallBlockData, ToolResultBlockData } from '#/tui/types'; +import { appendStreamingArgsPreview } from '#/tui/utils/event-payload'; import { decodeMcpToolName } from '#/tui/utils/mcp-tool-name'; import { PlanBoxComponent } from './plan-box'; @@ -278,10 +282,14 @@ function extractPartialStringField(text: string, key: string): string | undefine } function parseArgsPreview(value: string): Record { - if (value.trim().length === 0) return {}; - if (value.trimEnd().endsWith('}')) { + const previewText = value.slice(0, STREAMING_ARGS_PREVIEW_MAX_CHARS); + if (previewText.trim().length === 0) return {}; + if ( + value.length <= STREAMING_ARGS_PREVIEW_MAX_CHARS && + previewText.trimEnd().endsWith('}') + ) { try { - const parsed = JSON.parse(value) as unknown; + const parsed = JSON.parse(previewText) as unknown; if (typeof parsed === 'object' && parsed !== null && !Array.isArray(parsed)) { return parsed as Record; } @@ -290,7 +298,7 @@ function parseArgsPreview(value: string): Record { } } const result: Record = {}; - for (const match of value.matchAll(STREAMING_ARGS_FIELD_RE)) { + for (const match of previewText.matchAll(STREAMING_ARGS_FIELD_RE)) { const key = match[1]; const rawValue = match[2]; if (key === undefined || rawValue === undefined) continue; @@ -933,7 +941,10 @@ export class ToolCallComponent extends Container { argumentsPart: string | null; }): void { const existing = this.ongoingSubCalls.get(delta.id); - const nextArgsText = `${existing?.streamingArguments ?? ''}${delta.argumentsPart ?? ''}`; + const nextArgsText = appendStreamingArgsPreview( + existing?.streamingArguments, + delta.argumentsPart, + ); const parsed = parseArgsPreview(nextArgsText); this.ongoingSubCalls.set(delta.id, { name: delta.name ?? existing?.name ?? 'Tool', @@ -1429,18 +1440,18 @@ export class ToolCallComponent extends Container { * `extractPartialStringField`) and render a stable high-signal * preview: Write's `content` as highlighted code, Edit's argument * receive progress, Bash's `$ command`, etc. While args are still - * streaming we render the body in full (no 10-line cap) — once the - * result lands, the preview snaps to the collapsed cap unless the user - * has expanded. + * streaming we render from a bounded preview buffer; once the result lands, + * the preview snaps to the collapsed cap unless the user has expanded. */ private buildStreamingPreview(streamText: string): void { const name = this.toolCall.name; + const previewText = streamText.slice(0, STREAMING_ARGS_PREVIEW_MAX_CHARS); if (name === 'Write') { - const content = extractPartialStringField(streamText, 'content'); + const content = extractPartialStringField(previewText, 'content'); if (content === undefined || content.length === 0) return; const filePath = - extractPartialStringField(streamText, 'file_path') ?? - extractPartialStringField(streamText, 'path') ?? + extractPartialStringField(previewText, 'file_path') ?? + extractPartialStringField(previewText, 'path') ?? ''; const lang = langFromPath(filePath); const allLines = highlightLines(content, lang); @@ -1452,10 +1463,10 @@ export class ToolCallComponent extends Container { } if (name === 'Edit') { const filePath = - extractPartialStringField(streamText, 'file_path') ?? - extractPartialStringField(streamText, 'path') ?? + extractPartialStringField(previewText, 'file_path') ?? + extractPartialStringField(previewText, 'path') ?? ''; - const bytes = Buffer.byteLength(streamText, 'utf8'); + const bytes = Buffer.byteLength(previewText, 'utf8'); const startedAtMs = this.toolCall.streamingStartedAtMs; const elapsedSeconds = startedAtMs === undefined ? 0 : Math.max(0, Math.floor((Date.now() - startedAtMs) / 1000)); @@ -1467,7 +1478,7 @@ export class ToolCallComponent extends Container { return; } if (name === 'Bash') { - const cmd = extractPartialStringField(streamText, 'command'); + const cmd = extractPartialStringField(previewText, 'command'); if (cmd === undefined || cmd.length === 0) return; this.addChild( new ShellExecutionComponent({ diff --git a/apps/kimi-code/src/tui/constant/streaming.ts b/apps/kimi-code/src/tui/constant/streaming.ts index e52117249b..b8f88a014d 100644 --- a/apps/kimi-code/src/tui/constant/streaming.ts +++ b/apps/kimi-code/src/tui/constant/streaming.ts @@ -2,3 +2,9 @@ // This is intentionally a preview parser, not a full JSON parser. export const STREAMING_ARGS_FIELD_RE = /"(path|file_path|command|pattern|query|url|description|title|name)"\s*:\s*"((?:\\.|[^"\\])*)"/g; + +// Bounds live tool-argument previews; final tool.call payloads remain complete. +export const STREAMING_ARGS_PREVIEW_MAX_CHARS = 64 * 1024; + +// Coalesces high-frequency model/tool deltas before rebuilding TUI components. +export const STREAMING_UI_FLUSH_MS = 50; diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index ada55d5062..e393086b5a 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -184,6 +184,7 @@ import { OAUTH_LOGIN_REQUIRED_CODE, OAUTH_LOGIN_REQUIRED_STARTUP_NOTICE, } from './constant/kimi-tui'; +import { STREAMING_UI_FLUSH_MS } from './constant/streaming'; import { adaptPanelResponse } from './reverse-rpc/approval/adapter'; import { ApprovalController } from './reverse-rpc/approval/controller'; import { createApprovalRequestHandler } from './reverse-rpc/approval/handler'; @@ -208,6 +209,7 @@ import { formatBackgroundAgentTranscript } from './utils/background-agent-status import { formatBackgroundTaskTranscript } from './utils/background-task-status'; import { hasDispose, isExpandable, isPlanExpandable } from './utils/component-capabilities'; import { + appendStreamingArgsPreview, argsRecord, formatErrorMessage, isTodoItemShape, @@ -225,6 +227,7 @@ import { type McpServerStatusSnapshot, selectMcpStartupStatusRows, } from './utils/mcp-server-status'; +import { hasPatchChanges } from './utils/object-patch'; import { openUrl } from './utils/open-url'; import { setProcessTitle } from './utils/proctitle'; import { sessionRowsForPicker } from './utils/session-picker-rows'; @@ -563,6 +566,13 @@ export class KimiTUI { private readonly migrationPlan: MigrationPlan | null; // When true, the migration screen is the whole session: run it, then exit. private readonly migrateOnly: boolean; + // High-frequency model/tool deltas update draft state immediately, then use + // these flags to coalesce expensive component rebuilds into periodic flushes. + private streamingUiFlushTimer: ReturnType | undefined; + private lastStreamingUiFlushAt: number | undefined; + private pendingAssistantFlush = false; + private pendingThinkingFlush = false; + private readonly pendingToolCallFlushIds = new Set(); public onExit?: (exitCode?: number) => Promise; @@ -855,6 +865,7 @@ export class KimiTUI { // Stops UI resources, active sessions, reverse-RPC handlers, and the harness. async stop(): Promise { this.aborted = true; + this.discardPendingStreamingUiUpdates(); if (this.pendingExit) { clearTimeout(this.pendingExit.timer); this.pendingExit = null; @@ -1650,8 +1661,107 @@ export class KimiTUI { }); } + private hasPendingStreamingUiUpdates(): boolean { + return ( + this.pendingAssistantFlush || + this.pendingThinkingFlush || + this.pendingToolCallFlushIds.size > 0 + ); + } + + private clearStreamingUiFlushTimer(): void { + if (this.streamingUiFlushTimer === undefined) return; + clearTimeout(this.streamingUiFlushTimer); + this.streamingUiFlushTimer = undefined; + } + + private clearStreamingUiFlushTimerIfIdle(): void { + if (this.hasPendingStreamingUiUpdates()) return; + this.clearStreamingUiFlushTimer(); + } + + private discardPendingStreamingUiUpdates(): void { + this.clearStreamingUiFlushTimer(); + this.pendingAssistantFlush = false; + this.pendingThinkingFlush = false; + this.pendingToolCallFlushIds.clear(); + } + + // Schedule trailing UI work for streaming deltas. Terminal drawing is already + // coalesced by pi-tui; this avoids doing our own markdown/tool preview rebuild + // work on every chunk before pi-tui even gets a chance to render. + private scheduleStreamingUiFlush(): void { + if (!this.hasPendingStreamingUiUpdates()) return; + if (this.streamingUiFlushTimer !== undefined) return; + const delay = + this.lastStreamingUiFlushAt === undefined + ? 0 + : Math.max(0, STREAMING_UI_FLUSH_MS - (Date.now() - this.lastStreamingUiFlushAt)); + this.streamingUiFlushTimer = setTimeout(() => { + this.streamingUiFlushTimer = undefined; + this.flushStreamingUiUpdates(); + }, delay); + } + + // Final events such as tool.result or turn.ended must observe all streamed + // draft content, so they bypass the timer and drain pending UI work first. + private flushStreamingUiUpdatesNow(): void { + this.clearStreamingUiFlushTimer(); + this.flushStreamingUiUpdates(); + } + + private flushStreamingUiUpdates(): void { + if (!this.hasPendingStreamingUiUpdates()) return; + this.lastStreamingUiFlushAt = Date.now(); + const shouldFlushThinking = this.pendingThinkingFlush; + const shouldFlushAssistant = this.pendingAssistantFlush; + const toolCallIds = [...this.pendingToolCallFlushIds]; + this.pendingThinkingFlush = false; + this.pendingAssistantFlush = false; + this.pendingToolCallFlushIds.clear(); + + if (shouldFlushThinking && this.state.thinkingDraft.length > 0) { + this.onThinkingUpdate(this.state.thinkingDraft); + } + if (shouldFlushAssistant) { + this.onStreamingTextUpdate(this.state.assistantDraft); + } + for (const id of toolCallIds) { + this.flushStreamingToolCallPreview(id); + } + } + + // Materializes the latest bounded argument preview for one in-flight tool + // call. The final tool.call event still replaces this with authoritative args. + private flushStreamingToolCallPreview(id: string): void { + const streaming = this.state.streamingToolCallArguments.get(id); + if (streaming === undefined) return; + const toolCall: ToolCallBlockData = { + id, + name: streaming.name ?? this.state.activeToolCalls.get(id)?.name ?? 'Tool', + args: parseStreamingArgs(streaming.argumentsText), + streamingArguments: streaming.argumentsText, + streamingStartedAtMs: streaming.startedAtMs, + step: this.state.currentStep, + turnId: this.state.currentTurnId, + }; + this.state.activeToolCalls.set(id, toolCall); + + if (this.state.thinkingDraft.length > 0 || this.state.assistantStreamActive) { + this.finalizeLiveTextBuffers('tool'); + } + + const existingComponent = this.state.pendingToolComponents.get(id); + if (existingComponent !== undefined) { + existingComponent.updateToolCall(toolCall); + } else if (toolCall.name !== 'Agent') { + this.onToolCallStart(toolCall); + } + } + // Finalizes live thinking output and moves the live pane to the next mode. private flushThinkingToTranscript(nextMode: LivePaneState['mode'] = 'idle'): void { + this.flushStreamingUiUpdatesNow(); if (this.state.thinkingDraft.length === 0) { this.patchLivePane({ mode: nextMode }); return; @@ -1663,6 +1773,7 @@ export class KimiTUI { // Finalizes live assistant text and clears streaming component state. private finalizeAssistantStream(): void { + this.flushStreamingUiUpdatesNow(); if (this.state.assistantStreamActive) { this.onStreamingTextEnd(); this.state.assistantStreamActive = false; @@ -1674,6 +1785,9 @@ export class KimiTUI { // Discards live thinking and assistant text state without finalizing transcript output. private resetLiveTextRuntime(): void { + this.pendingAssistantFlush = false; + this.pendingThinkingFlush = false; + this.clearStreamingUiFlushTimerIfIdle(); this.state.assistantDraft = ''; this.state.assistantStreamActive = false; this.state.streamingComponent = undefined; @@ -1684,6 +1798,8 @@ export class KimiTUI { // Clears live tool UI state while preserving active tool-call tracking. private resetLiveToolUiState(): void { + this.pendingToolCallFlushIds.clear(); + this.clearStreamingUiFlushTimerIfIdle(); this.state.streamingToolCallArguments.clear(); this.disposeAndClearPendingToolComponents(); this.state.pendingAgentGroup = null; @@ -1738,6 +1854,7 @@ export class KimiTUI { // Applies app-state changes and refreshes dependent UI surfaces. private setAppState(patch: Partial): void { + if (!hasPatchChanges(this.state.appState, patch)) return; const busyChanged = 'isStreaming' in patch || 'isCompacting' in patch; Object.assign(this.state.appState, patch); if ('planMode' in patch) this.updateEditorBorderHighlight(); @@ -1749,6 +1866,7 @@ export class KimiTUI { // Applies live-pane changes and refreshes activity presentation. private patchLivePane(patch: Partial): void { + if (!hasPatchChanges(this.state.livePane, patch)) return; Object.assign(this.state.livePane, patch); this.updateActivityPane(); this.state.ui.requestRender(); @@ -1887,6 +2005,7 @@ export class KimiTUI { // Resets turn, tool, queue, and background-agent state for a session switch. private resetSessionRuntime(): void { this.aborted = false; + this.discardPendingStreamingUiUpdates(); this.state.queuedMessages = []; this.harness.interactiveAgentId = MAIN_AGENT_ID; this.resetToolCallState(); @@ -2242,6 +2361,7 @@ export class KimiTUI { // Finalizes turn-scoped state when the SDK completes a turn. private handleTurnEnd(_event: TurnEndedEvent, sendQueued: (item: QueuedMessage) => void): void { void _event; + this.flushStreamingUiUpdatesNow(); const todos = this.state.todoPanel.getTodos(); if (todos.length > 0 && todos.every((t) => t.status === 'done')) { this.setTodoList([]); @@ -2252,6 +2372,7 @@ export class KimiTUI { // Resets live render state for a new turn step. private handleStepBegin(event: TurnStepStartedEvent): void { + this.flushStreamingUiUpdatesNow(); this.state.currentStep = event.step; this.resetLiveToolUiState(); this.finalizeLiveTextBuffers('waiting'); @@ -2275,6 +2396,7 @@ export class KimiTUI { // wrong. Flip those into a visible 'Truncated' state and append a // notice pointing at the config knob. private handleStepCompleted(event: TurnStepCompletedEvent): void { + this.flushStreamingUiUpdatesNow(); if (event.finishReason !== 'max_tokens') return; // Scope the truncation marking to tool calls that belong to the @@ -2320,6 +2442,7 @@ export class KimiTUI { // Renders user-facing status for an interrupted turn step. private handleStepInterrupted(event: TurnStepInterruptedEvent): void { + this.flushStreamingUiUpdatesNow(); this.resetLiveToolUiState(); this.finalizeLiveTextBuffers('idle'); const reason = event.reason; @@ -2338,9 +2461,12 @@ export class KimiTUI { // Appends a thinking delta to the live thinking block. private handleThinkingDelta(event: ThinkingDeltaEvent): void { this.state.thinkingDraft += event.delta; - this.onThinkingUpdate(this.state.thinkingDraft); + this.pendingThinkingFlush = true; this.patchLivePane({ mode: 'idle' }); - this.setAppState({ streamingPhase: 'thinking' }); + if (this.state.appState.streamingPhase !== 'thinking') { + this.setAppState({ streamingPhase: 'thinking', streamingStartTime: Date.now() }); + } + this.scheduleStreamingUiFlush(); } // Appends an assistant text delta to the live assistant block. @@ -2355,20 +2481,21 @@ export class KimiTUI { } this.state.assistantDraft += event.delta; - this.onStreamingTextUpdate(this.state.assistantDraft); + this.pendingAssistantFlush = true; this.patchLivePane({ mode: 'idle', pendingApproval: null, pendingQuestion: null, }); - this.setAppState({ - streamingPhase: 'composing', - streamingStartTime: Date.now(), - }); + if (this.state.appState.streamingPhase !== 'composing') { + this.setAppState({ streamingPhase: 'composing', streamingStartTime: Date.now() }); + } + this.scheduleStreamingUiFlush(); } private handleHookResult(event: HookResultEvent): void { + this.flushStreamingUiUpdatesNow(); if (this.state.thinkingDraft.length > 0) { this.flushThinkingToTranscript('idle'); } @@ -2389,6 +2516,7 @@ export class KimiTUI { // Starts or updates a rendered tool call from a tool-call start event. private handleToolCall(event: ToolCallStartedEvent): void { + this.flushStreamingUiUpdatesNow(); const toolCall: ToolCallBlockData = { id: event.toolCallId, name: event.name, @@ -2400,6 +2528,7 @@ export class KimiTUI { }; const existing = this.state.activeToolCalls.get(event.toolCallId); this.state.activeToolCalls.set(event.toolCallId, toolCall); + this.pendingToolCallFlushIds.delete(event.toolCallId); this.state.streamingToolCallArguments.delete(event.toolCallId); const existingComponent = this.state.pendingToolComponents.get(event.toolCallId); if (existingComponent !== undefined) { @@ -2422,42 +2551,24 @@ export class KimiTUI { if (event.toolCallId.length === 0) return; const id = event.toolCallId; const existing = this.state.streamingToolCallArguments.get(id); - const argumentsText = `${existing?.argumentsText ?? ''}${event.argumentsPart ?? ''}`; + const argumentsText = appendStreamingArgsPreview( + existing?.argumentsText, + event.argumentsPart, + ); const name = event.name ?? existing?.name ?? this.state.activeToolCalls.get(id)?.name ?? 'Tool'; const startedAtMs = existing?.startedAtMs ?? Date.now(); this.state.streamingToolCallArguments.set(id, { name, argumentsText, startedAtMs }); - - const toolCall: ToolCallBlockData = { - id, - name, - args: parseStreamingArgs(argumentsText), - streamingArguments: argumentsText, - streamingStartedAtMs: startedAtMs, - step: this.state.currentStep, - turnId: this.state.currentTurnId, - }; - this.state.activeToolCalls.set(id, toolCall); - - if (this.state.thinkingDraft.length > 0 || this.state.assistantStreamActive) { - this.finalizeLiveTextBuffers('tool'); - } - - const existingComponent = this.state.pendingToolComponents.get(id); - if (existingComponent !== undefined) { - existingComponent.updateToolCall(toolCall); - } else if (name !== 'Agent') { - this.onToolCallStart(toolCall); - } + this.pendingToolCallFlushIds.add(id); this.patchLivePane({ mode: 'tool', pendingApproval: null, pendingQuestion: null, }); - this.setAppState({ - streamingPhase: 'composing', - streamingStartTime: Date.now(), - }); + if (this.state.appState.streamingPhase !== 'composing') { + this.setAppState({ streamingPhase: 'composing', streamingStartTime: Date.now() }); + } + this.scheduleStreamingUiFlush(); } // Streams a `{kind:'status'}` progress text into the live tool box so @@ -2476,6 +2587,7 @@ export class KimiTUI { // Completes a tool call and applies any tool-specific UI side effects. private handleToolResult(event: ToolResultEvent): void { + this.flushStreamingUiUpdatesNow(); const matchedCall = this.state.activeToolCalls.get(event.toolCallId); const resultData: ToolResultBlockData = { tool_call_id: event.toolCallId, @@ -2528,6 +2640,7 @@ export class KimiTUI { // Finalizes live buffers and renders a session error. private handleSessionError(event: ErrorEvent): void { + this.flushStreamingUiUpdatesNow(); this.resetLiveToolUiState(); this.finalizeLiveTextBuffers('idle'); if (event.code === OAUTH_LOGIN_REQUIRED_CODE) { @@ -3398,6 +3511,7 @@ export class KimiTUI { // Clears transcript-related state and redraws the welcome view. private clearTranscriptAndRedraw(): void { + this.discardPendingStreamingUiUpdates(); this.state.transcriptEntries = []; this.disposeActiveCompactionBlock(); this.resetLiveTextRuntime(); diff --git a/apps/kimi-code/src/tui/utils/event-payload.ts b/apps/kimi-code/src/tui/utils/event-payload.ts index 195f7d3925..51de71f801 100644 --- a/apps/kimi-code/src/tui/utils/event-payload.ts +++ b/apps/kimi-code/src/tui/utils/event-payload.ts @@ -1,6 +1,20 @@ import { isKimiError } from '@moonshot-ai/kimi-code-sdk'; -import { STREAMING_ARGS_FIELD_RE } from '#/tui/constant/streaming'; +import { + STREAMING_ARGS_FIELD_RE, + STREAMING_ARGS_PREVIEW_MAX_CHARS, +} from '#/tui/constant/streaming'; + +export function appendStreamingArgsPreview( + current: string | undefined, + next: string | null | undefined, +): string { + const existing = (current ?? '').slice(0, STREAMING_ARGS_PREVIEW_MAX_CHARS); + if (next === null || next === undefined || next.length === 0) return existing; + const remaining = STREAMING_ARGS_PREVIEW_MAX_CHARS - existing.length; + if (remaining <= 0) return existing; + return `${existing}${next.slice(0, remaining)}`; +} function unescapeJsonString(s: string): string { return s.replaceAll(/\\(["\\/bfnrt])/g, (_, ch: string) => { @@ -28,10 +42,14 @@ function unescapeJsonString(s: string): string { } export function parseStreamingArgs(argumentsText: string): Record { - if (argumentsText.trim().length === 0) return {}; - if (argumentsText.trimEnd().endsWith('}')) { + const previewText = argumentsText.slice(0, STREAMING_ARGS_PREVIEW_MAX_CHARS); + if (previewText.trim().length === 0) return {}; + if ( + argumentsText.length <= STREAMING_ARGS_PREVIEW_MAX_CHARS && + previewText.trimEnd().endsWith('}') + ) { try { - const parsed = JSON.parse(argumentsText) as unknown; + const parsed = JSON.parse(previewText) as unknown; if (typeof parsed === 'object' && parsed !== null && !Array.isArray(parsed)) { return parsed as Record; } @@ -40,7 +58,7 @@ export function parseStreamingArgs(argumentsText: string): Record = {}; - for (const match of argumentsText.matchAll(STREAMING_ARGS_FIELD_RE)) { + for (const match of previewText.matchAll(STREAMING_ARGS_FIELD_RE)) { const key = match[1]; const rawValue = match[2]; if (key === undefined || rawValue === undefined) continue; diff --git a/apps/kimi-code/src/tui/utils/object-patch.ts b/apps/kimi-code/src/tui/utils/object-patch.ts new file mode 100644 index 0000000000..a44f6f5661 --- /dev/null +++ b/apps/kimi-code/src/tui/utils/object-patch.ts @@ -0,0 +1,8 @@ +// Returns true when applying `patch` would change at least one own property. +// Used before UI refresh paths so repeated equivalent state patches are cheap. +export function hasPatchChanges(target: T, patch: Partial): boolean { + for (const key of Object.keys(patch) as Array) { + if (!Object.is(target[key], patch[key])) return true; + } + return false; +} diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index 4edeb7780c..243b133da9 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -616,6 +616,125 @@ describe('KimiTUI message flow', () => { } }); + it('coalesces assistant delta component updates', async () => { + vi.useFakeTimers(); + try { + const { driver } = await makeDriver(); + vi.mocked(driver.state.ui.requestRender).mockClear(); + + driver.handleEvent( + { + type: 'assistant.delta', + agentId: 'main', + sessionId: 'ses-1', + turnId: 1, + delta: 'a', + } as Event, + vi.fn(), + ); + const component = driver.state.streamingComponent; + if (component === undefined) throw new Error('expected streaming component'); + const updateSpy = vi.spyOn(component, 'updateContent'); + + driver.handleEvent( + { + type: 'assistant.delta', + agentId: 'main', + sessionId: 'ses-1', + turnId: 1, + delta: 'b', + } as Event, + vi.fn(), + ); + driver.handleEvent( + { + type: 'assistant.delta', + agentId: 'main', + sessionId: 'ses-1', + turnId: 1, + delta: 'c', + } as Event, + vi.fn(), + ); + + expect(updateSpy).not.toHaveBeenCalled(); + await vi.runOnlyPendingTimersAsync(); + + expect(updateSpy).toHaveBeenCalledTimes(1); + expect(updateSpy).toHaveBeenLastCalledWith('abc'); + } finally { + vi.useRealTimers(); + } + }); + + it('flushes pending assistant deltas before turn completion', async () => { + vi.useFakeTimers(); + try { + const { driver } = await makeDriver(); + const sendQueued = vi.fn(); + driver.state.appState.isStreaming = true; + + driver.handleEvent( + { + type: 'assistant.delta', + agentId: 'main', + sessionId: 'ses-1', + turnId: 1, + delta: 'done', + } as Event, + sendQueued, + ); + driver.handleEvent( + { + type: 'turn.ended', + agentId: 'main', + sessionId: 'ses-1', + turnId: 1, + reason: 'completed', + } as Event, + sendQueued, + ); + + expect(stripSgr(renderTranscript(driver))).toContain('done'); + } finally { + vi.useRealTimers(); + } + }); + + it('coalesces streaming tool-call argument preview updates', async () => { + vi.useFakeTimers(); + try { + const { driver } = await makeDriver(); + driver.state.currentTurnId = '1'; + driver.state.currentStep = 1; + + driver.handleEvent( + { + type: 'tool.call.delta', + agentId: 'main', + sessionId: 'ses-1', + turnId: 1, + toolCallId: 'call_bash', + name: 'Bash', + argumentsPart: '{"command":"echo hi"}', + } as Event, + vi.fn(), + ); + + expect(driver.state.pendingToolComponents.has('call_bash')).toBe(false); + expect(driver.state.activeToolCalls.has('call_bash')).toBe(false); + + await vi.runOnlyPendingTimersAsync(); + + expect(driver.state.pendingToolComponents.has('call_bash')).toBe(true); + expect(driver.state.activeToolCalls.get('call_bash')?.args).toMatchObject({ + command: 'echo hi', + }); + } finally { + vi.useRealTimers(); + } + }); + it('cancels manual compaction from the editor', async () => { const { driver, session } = await makeDriver(); driver.handleEvent( diff --git a/apps/kimi-code/test/tui/utils/event-payload.test.ts b/apps/kimi-code/test/tui/utils/event-payload.test.ts new file mode 100644 index 0000000000..f5d9f04c28 --- /dev/null +++ b/apps/kimi-code/test/tui/utils/event-payload.test.ts @@ -0,0 +1,30 @@ +import { describe, expect, it } from 'vitest'; + +import { STREAMING_ARGS_PREVIEW_MAX_CHARS } from '#/tui/constant/streaming'; +import { + appendStreamingArgsPreview, + parseStreamingArgs, +} from '#/tui/utils/event-payload'; + +describe('streaming tool argument payload helpers', () => { + it('parses complete JSON arguments for finalized small previews', () => { + expect(parseStreamingArgs('{"command":"echo hi","path":"/tmp/a"}')).toEqual({ + command: 'echo hi', + path: '/tmp/a', + }); + }); + + it('caps accumulated streaming preview text', () => { + const current = 'a'.repeat(STREAMING_ARGS_PREVIEW_MAX_CHARS - 2); + + expect(appendStreamingArgsPreview(current, 'bcdef')).toBe(`${current}bc`); + }); + + it('parses only bounded preview fields from oversized streaming arguments', () => { + const oversized = `{"command":"echo ok","description":"${'x'.repeat( + STREAMING_ARGS_PREVIEW_MAX_CHARS + 100, + )}"}`; + + expect(parseStreamingArgs(oversized)).toEqual({ command: 'echo ok' }); + }); +}); diff --git a/packages/kosong/src/providers/anthropic.ts b/packages/kosong/src/providers/anthropic.ts index d70dff794c..f028de28b3 100644 --- a/packages/kosong/src/providers/anthropic.ts +++ b/packages/kosong/src/providers/anthropic.ts @@ -26,6 +26,7 @@ import type { Tool as AnthropicTool, ContentBlockParam, MessageCreateParams, + MessageCreateParamsStreaming, MessageParam, MessageStreamEvent, RawContentBlockDeltaEvent, @@ -968,10 +969,12 @@ export class AnthropicChatProvider implements ChatProvider { const client = this._createClient(options?.auth); if (this._stream) { - // Streaming mode: use client.messages.stream() which returns an AsyncIterable + // Use the raw Messages stream instead of the SDK MessageStream helper. + // The helper reparses accumulated input_json_delta buffers on every chunk, + // which becomes synchronous O(n^2) work for large streamed tool arguments. try { - const stream = client.messages.stream( - createParams as unknown as MessageCreateParams, + const stream = await client.messages.create( + { ...createParams, stream: true } as unknown as MessageCreateParamsStreaming, finalRequestOptions, ); return new AnthropicStreamedMessage(stream, true); diff --git a/packages/kosong/test/anthropic-errors.test.ts b/packages/kosong/test/anthropic-errors.test.ts index 0ddf92874c..fd36fcb486 100644 --- a/packages/kosong/test/anthropic-errors.test.ts +++ b/packages/kosong/test/anthropic-errors.test.ts @@ -232,9 +232,9 @@ describe('stream error propagation', () => { it('APIConnectionTimeoutError during stream iteration is converted', async () => { const provider = createStreamProvider(); const sdkError = new AnthropicTimeoutError({ message: 'stream timed out' }); - (provider as any)._client.messages.stream = vi + (provider as any)._client.messages.create = vi .fn() - .mockReturnValue(makeErrorStream(sdkError)) as never; + .mockResolvedValue(makeErrorStream(sdkError)) as never; const result = await provider.generate( '', @@ -254,9 +254,9 @@ describe('stream error propagation', () => { it('APIConnectionError during stream iteration is converted', async () => { const provider = createStreamProvider(); const sdkError = new AnthropicConnectionError({ message: 'connection reset' }); - (provider as any)._client.messages.stream = vi + (provider as any)._client.messages.create = vi .fn() - .mockReturnValue(makeErrorStream(sdkError)) as never; + .mockResolvedValue(makeErrorStream(sdkError)) as never; const result = await provider.generate( '', @@ -280,9 +280,9 @@ describe('stream error propagation', () => { 'internal error', new Headers(), ); - (provider as any)._client.messages.stream = vi + (provider as any)._client.messages.create = vi .fn() - .mockReturnValue(makeErrorStream(sdkError)) as never; + .mockResolvedValue(makeErrorStream(sdkError)) as never; const result = await provider.generate( '', @@ -306,9 +306,9 @@ describe('stream error propagation', () => { 'too many requests', new Headers(), ); - (provider as any)._client.messages.stream = vi + (provider as any)._client.messages.create = vi .fn() - .mockReturnValue(makeErrorStream(sdkError)) as never; + .mockResolvedValue(makeErrorStream(sdkError)) as never; const result = await provider.generate( '', @@ -334,9 +334,9 @@ describe('stream error propagation', () => { 'invalid', new Headers(), ); - (provider as any)._client.messages.stream = vi + (provider as any)._client.messages.create = vi .fn() - .mockReturnValue(makeErrorStream(sdkError)) as never; + .mockResolvedValue(makeErrorStream(sdkError)) as never; const result = await provider.generate( '', diff --git a/packages/kosong/test/anthropic.test.ts b/packages/kosong/test/anthropic.test.ts index 5cfe4292c6..2dfd0520a9 100644 --- a/packages/kosong/test/anthropic.test.ts +++ b/packages/kosong/test/anthropic.test.ts @@ -1416,7 +1416,7 @@ describe('AnthropicChatProvider', () => { { type: 'message_stop' }, ]); - (provider as any)._client.messages.stream = vi.fn().mockReturnValue(stream) as never; + (provider as any)._client.messages.create = vi.fn().mockResolvedValue(stream) as never; const result = await provider.generate( '', @@ -1477,7 +1477,7 @@ describe('AnthropicChatProvider', () => { { type: 'message_stop' }, ]); - (provider as any)._client.messages.stream = vi.fn().mockReturnValue(stream) as never; + (provider as any)._client.messages.create = vi.fn().mockResolvedValue(stream) as never; const result = await provider.generate( '', @@ -1536,7 +1536,7 @@ describe('AnthropicChatProvider', () => { { type: 'message_stop' }, ]); - (provider as any)._client.messages.stream = vi.fn().mockReturnValue(stream) as never; + (provider as any)._client.messages.create = vi.fn().mockResolvedValue(stream) as never; const result = await provider.generate( '', @@ -1610,7 +1610,7 @@ describe('AnthropicChatProvider', () => { { type: 'message_stop' }, ]); - (provider as any)._client.messages.stream = vi.fn().mockReturnValue(stream) as never; + (provider as any)._client.messages.create = vi.fn().mockResolvedValue(stream) as never; const result = await provider.generate( '', @@ -1694,7 +1694,7 @@ describe('AnthropicChatProvider', () => { { type: 'message_stop' }, ]); - (provider as any)._client.messages.stream = vi.fn().mockReturnValue(stream) as never; + (provider as any)._client.messages.create = vi.fn().mockResolvedValue(stream) as never; const { message } = await generate( provider, @@ -1731,7 +1731,7 @@ describe('AnthropicChatProvider', () => { }, }; - (provider as any)._client.messages.stream = vi.fn().mockReturnValue(errorStream) as never; + (provider as any)._client.messages.create = vi.fn().mockResolvedValue(errorStream) as never; const result = await provider.generate( '', @@ -1771,7 +1771,7 @@ describe('AnthropicChatProvider', () => { { type: 'message_stop' }, ]); - (provider as any)._client.messages.stream = vi.fn().mockReturnValue(stream) as never; + (provider as any)._client.messages.create = vi.fn().mockResolvedValue(stream) as never; const result = await provider.generate( '', @@ -1807,7 +1807,7 @@ describe('AnthropicChatProvider', () => { { type: 'message_stop' }, ]); - (provider as any)._client.messages.stream = vi.fn().mockReturnValue(stream) as never; + (provider as any)._client.messages.create = vi.fn().mockResolvedValue(stream) as never; const result = await provider.generate( '', @@ -1826,7 +1826,7 @@ describe('AnthropicChatProvider', () => { }); describe('stream option', () => { - it('defaults to stream: true and calls messages.stream', async () => { + it('defaults to stream: true and calls messages.create with stream enabled', async () => { const provider = new AnthropicChatProvider({ model: 'k25', apiKey: 'test-key', @@ -1843,8 +1843,8 @@ describe('AnthropicChatProvider', () => { { type: 'message_stop' }, ]); - const streamFn = vi.fn().mockReturnValue(stream); - (provider as any)._client.messages.stream = streamFn as never; + const createFn = vi.fn().mockResolvedValue(stream); + (provider as any)._client.messages.create = createFn as never; const result = await provider.generate( '', @@ -1853,7 +1853,8 @@ describe('AnthropicChatProvider', () => { ); await collectParts(result); - expect(streamFn).toHaveBeenCalledTimes(1); + expect(createFn).toHaveBeenCalledTimes(1); + expect(createFn.mock.calls[0]?.[0]).toMatchObject({ stream: true }); }); it('stream: false calls messages.create', async () => { From b7e0750f196a79237b70538c54008d554719ccdd Mon Sep 17 00:00:00 2001 From: liruifengv Date: Mon, 25 May 2026 19:30:09 +0800 Subject: [PATCH 2/2] test: update Anthropic streaming mock --- packages/kosong/test/finish-reason.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kosong/test/finish-reason.test.ts b/packages/kosong/test/finish-reason.test.ts index a9ed250aec..5119198415 100644 --- a/packages/kosong/test/finish-reason.test.ts +++ b/packages/kosong/test/finish-reason.test.ts @@ -424,7 +424,7 @@ function createAnthropicStreamProvider(events: unknown[]): AnthropicChatProvider clientFactory: () => ({ messages: { - stream: vi.fn().mockReturnValue(anthropicMockStream(events)), + create: vi.fn().mockResolvedValue(anthropicMockStream(events)), }, }) as never, });