From fc568b9a34fcf1db9bf4ae317a2128ac0d27fc5d Mon Sep 17 00:00:00 2001 From: JATMN Date: Wed, 8 Jul 2026 18:12:39 -0700 Subject: [PATCH] refactor(messages): extract streaming helpers (2 of 8) (#1899) * refactor(messages): extract streaming helpers * fix(messages): remove streaming EOF whitespace * chore(messages): clean streaming extraction imports --- src/utils/messages.ts | 192 +---------------- .../streaming.test.ts} | 4 +- src/utils/messages/streaming.ts | 198 ++++++++++++++++++ 3 files changed, 202 insertions(+), 192 deletions(-) rename src/utils/{messages.streamingToolUses.test.ts => messages/streaming.test.ts} (97%) create mode 100644 src/utils/messages/streaming.ts diff --git a/src/utils/messages.ts b/src/utils/messages.ts index b215c0bcc..6e9073ae9 100644 --- a/src/utils/messages.ts +++ b/src/utils/messages.ts @@ -50,9 +50,7 @@ import type { NormalizedUserMessage, PartialCompactDirection, ProgressMessage, - RequestStartEvent, StopHookInfo, - StreamEvent, SystemAgentsKilledMessage, SystemAPIErrorMessage, SystemApiMetricsMessage, @@ -69,7 +67,6 @@ import type { SystemScheduledTaskFireMessage, SystemStopHookSummaryMessage, SystemTurnDurationMessage, - TombstoneMessage, ToolUseSummaryMessage, UserMessage, } from '../types/message.js' @@ -123,7 +120,6 @@ import { GLOB_TOOL_NAME } from 'src/tools/GlobTool/prompt.js' import { GREP_TOOL_NAME } from 'src/tools/GrepTool/prompt.js' import type { DeepImmutable } from 'src/types/utils.js' import { getStrictToolResultPairing } from '../bootstrap/state.js' -import type { SpinnerMode } from '../components/Spinner.js' import { COMMAND_ARGS_TAG, COMMAND_MESSAGE_TAG, @@ -3191,192 +3187,8 @@ export function getContentText( return null } -export type StreamingToolUse = { - index: number - contentBlock: BetaToolUseBlock - unparsedToolInput: string -} - -export type StreamingThinking = { - thinking: string - isStreaming: boolean - streamingEndedAt?: number -} - -/** - * Handles messages from a stream, updating response length for deltas and appending completed messages - */ -export function handleMessageFromStream( - message: - | Message - | TombstoneMessage - | StreamEvent - | RequestStartEvent - | ToolUseSummaryMessage, - onMessage: (message: Message) => void, - onUpdateLength: (newContent: string) => void, - onSetStreamMode: (mode: SpinnerMode) => void, - onStreamingToolUses: ( - f: (streamingToolUse: StreamingToolUse[]) => StreamingToolUse[], - ) => void, - onTombstone?: (message: Message) => void, - onStreamingThinking?: ( - f: (current: StreamingThinking | null) => StreamingThinking | null, - ) => void, - onApiMetrics?: (metrics: { ttftMs: number }) => void, - onStreamingText?: (f: (current: string | null) => string | null) => void, -): void { - if ( - message.type !== 'stream_event' && - message.type !== 'stream_request_start' - ) { - // Handle tombstone messages - remove the targeted message instead of adding - if (message.type === 'tombstone') { - onTombstone?.(message.message) - return - } - // Tool use summary messages are SDK-only, ignore them in stream handling - if (message.type === 'tool_use_summary') { - return - } - // Capture complete thinking blocks for real-time display in transcript mode - if (message.type === 'assistant') { - const thinkingBlock = message.message.content.find( - block => block.type === 'thinking', - ) - if (thinkingBlock && thinkingBlock.type === 'thinking') { - onStreamingThinking?.(() => ({ - thinking: thinkingBlock.thinking, - isStreaming: false, - streamingEndedAt: Date.now(), - })) - } - } - // Clear streaming text NOW so the render can switch displayedMessages - // from deferredMessages to messages in the same batch, making the - // transition from streaming text → final message atomic (no gap, no duplication). - onStreamingText?.(() => null) - onMessage(message) - return - } - - if (message.type === 'stream_request_start') { - onSetStreamMode('requesting') - return - } - - if (message.event.type === 'message_start') { - if (message.ttftMs != null) { - onApiMetrics?.({ ttftMs: message.ttftMs }) - } - } - - if (message.event.type === 'message_stop') { - onSetStreamMode('tool-use') - onStreamingToolUses(() => []) - return - } - - switch (message.event.type) { - case 'content_block_start': - onStreamingText?.(() => null) - if ( - feature('CONNECTOR_TEXT') && - isConnectorTextBlock(message.event.content_block) - ) { - onSetStreamMode('responding') - return - } - switch (message.event.content_block.type) { - case 'thinking': - case 'redacted_thinking': - onSetStreamMode('thinking') - return - case 'text': - onSetStreamMode('responding') - return - case 'tool_use': { - onSetStreamMode('tool-input') - const contentBlock = message.event.content_block - const index = message.event.index - onStreamingToolUses(_ => [ - ..._, - { - index, - contentBlock, - unparsedToolInput: '', - }, - ]) - return - } - case 'server_tool_use': - case 'web_search_tool_result': - case 'code_execution_tool_result': - case 'mcp_tool_use': - case 'mcp_tool_result': - case 'container_upload': - case 'web_fetch_tool_result': - case 'bash_code_execution_tool_result': - case 'text_editor_code_execution_tool_result': - case 'tool_search_tool_result': - case 'compaction': - onSetStreamMode('tool-input') - return - } - return - case 'content_block_delta': - switch (message.event.delta.type) { - case 'text_delta': { - const deltaText = message.event.delta.text - onUpdateLength(deltaText) - onStreamingText?.(text => (text ?? '') + deltaText) - return - } - case 'input_json_delta': { - const delta = message.event.delta.partial_json - const index = message.event.index - onUpdateLength(delta) - onStreamingToolUses(_ => { - // Update in place (preserve array order). The previous - // filter-then-append moved the updated tool to the end, which - // shuffled concurrently-streaming tools and broke the index-aligned - // contentBlock check in the Messages memo comparator. - let found = false - const next = _.map(element => { - if (element.index !== index) { - return element - } - found = true - return { - ...element, - unparsedToolInput: element.unparsedToolInput + delta, - } - }) - return found ? next : _ - }) - return - } - case 'thinking_delta': - onUpdateLength(message.event.delta.thinking) - return - case 'signature_delta': - // Signatures are cryptographic authentication strings, not model - // output. Excluding them from onUpdateLength prevents them from - // inflating the OTPS metric and the animated token counter. - return - default: - return - } - case 'content_block_stop': - return - case 'message_delta': - onSetStreamMode('responding') - return - default: - onSetStreamMode('responding') - return - } -} +export { handleMessageFromStream } from './messages/streaming.js' +export type { StreamingThinking, StreamingToolUse } from './messages/streaming.js' export function wrapInSystemReminder(content: string): string { return `\n${content}\n` diff --git a/src/utils/messages.streamingToolUses.test.ts b/src/utils/messages/streaming.test.ts similarity index 97% rename from src/utils/messages.streamingToolUses.test.ts rename to src/utils/messages/streaming.test.ts index 76726f81f..4e6f9703c 100644 --- a/src/utils/messages.streamingToolUses.test.ts +++ b/src/utils/messages/streaming.test.ts @@ -1,7 +1,7 @@ import { expect, test } from 'bun:test' -import { handleMessageFromStream, type StreamingToolUse } from './messages.js' -import type { StreamEvent } from '../types/message.js' +import { handleMessageFromStream, type StreamingToolUse } from './streaming.js' +import type { StreamEvent } from '../../types/message.js' // Regression for the PR #1744 change that switched input_json_delta handling // from filter-then-append to an in-place update. Concurrently-streaming tool diff --git a/src/utils/messages/streaming.ts b/src/utils/messages/streaming.ts new file mode 100644 index 000000000..fba97e5a8 --- /dev/null +++ b/src/utils/messages/streaming.ts @@ -0,0 +1,198 @@ +import { feature } from 'bun:bundle' +import type { BetaToolUseBlock } from '@anthropic-ai/sdk/resources/beta/messages/messages.mjs' +import type { SpinnerMode } from '../../components/Spinner.js' +import { isConnectorTextBlock } from '../../types/connectorText.js' +import type { + Message, + RequestStartEvent, + StreamEvent, + TombstoneMessage, + ToolUseSummaryMessage, +} from '../../types/message.js' + +export type StreamingToolUse = { + index: number + contentBlock: BetaToolUseBlock + unparsedToolInput: string +} + +export type StreamingThinking = { + thinking: string + isStreaming: boolean + streamingEndedAt?: number +} + +/** + * Handles messages from a stream, updating response length for deltas and appending completed messages + */ +export function handleMessageFromStream( + message: + | Message + | TombstoneMessage + | StreamEvent + | RequestStartEvent + | ToolUseSummaryMessage, + onMessage: (message: Message) => void, + onUpdateLength: (newContent: string) => void, + onSetStreamMode: (mode: SpinnerMode) => void, + onStreamingToolUses: ( + f: (streamingToolUse: StreamingToolUse[]) => StreamingToolUse[], + ) => void, + onTombstone?: (message: Message) => void, + onStreamingThinking?: ( + f: (current: StreamingThinking | null) => StreamingThinking | null, + ) => void, + onApiMetrics?: (metrics: { ttftMs: number }) => void, + onStreamingText?: (f: (current: string | null) => string | null) => void, +): void { + if ( + message.type !== 'stream_event' && + message.type !== 'stream_request_start' + ) { + // Handle tombstone messages - remove the targeted message instead of adding + if (message.type === 'tombstone') { + onTombstone?.(message.message) + return + } + // Tool use summary messages are SDK-only, ignore them in stream handling + if (message.type === 'tool_use_summary') { + return + } + // Capture complete thinking blocks for real-time display in transcript mode + if (message.type === 'assistant') { + const thinkingBlock = message.message.content.find( + block => block.type === 'thinking', + ) + if (thinkingBlock && thinkingBlock.type === 'thinking') { + onStreamingThinking?.(() => ({ + thinking: thinkingBlock.thinking, + isStreaming: false, + streamingEndedAt: Date.now(), + })) + } + } + // Clear streaming text NOW so the render can switch displayedMessages + // from deferredMessages to messages in the same batch, making the + // transition from streaming text → final message atomic (no gap, no duplication). + onStreamingText?.(() => null) + onMessage(message) + return + } + + if (message.type === 'stream_request_start') { + onSetStreamMode('requesting') + return + } + + if (message.event.type === 'message_start') { + if (message.ttftMs != null) { + onApiMetrics?.({ ttftMs: message.ttftMs }) + } + } + + if (message.event.type === 'message_stop') { + onSetStreamMode('tool-use') + onStreamingToolUses(() => []) + return + } + + switch (message.event.type) { + case 'content_block_start': + onStreamingText?.(() => null) + if ( + feature('CONNECTOR_TEXT') && + isConnectorTextBlock(message.event.content_block) + ) { + onSetStreamMode('responding') + return + } + switch (message.event.content_block.type) { + case 'thinking': + case 'redacted_thinking': + onSetStreamMode('thinking') + return + case 'text': + onSetStreamMode('responding') + return + case 'tool_use': { + onSetStreamMode('tool-input') + const contentBlock = message.event.content_block + const index = message.event.index + onStreamingToolUses(_ => [ + ..._, + { + index, + contentBlock, + unparsedToolInput: '', + }, + ]) + return + } + case 'server_tool_use': + case 'web_search_tool_result': + case 'code_execution_tool_result': + case 'mcp_tool_use': + case 'mcp_tool_result': + case 'container_upload': + case 'web_fetch_tool_result': + case 'bash_code_execution_tool_result': + case 'text_editor_code_execution_tool_result': + case 'tool_search_tool_result': + case 'compaction': + onSetStreamMode('tool-input') + return + } + return + case 'content_block_delta': + switch (message.event.delta.type) { + case 'text_delta': { + const deltaText = message.event.delta.text + onUpdateLength(deltaText) + onStreamingText?.(text => (text ?? '') + deltaText) + return + } + case 'input_json_delta': { + const delta = message.event.delta.partial_json + const index = message.event.index + onUpdateLength(delta) + onStreamingToolUses(_ => { + // Update in place (preserve array order). The previous + // filter-then-append moved the updated tool to the end, which + // shuffled concurrently-streaming tools and broke the index-aligned + // contentBlock check in the Messages memo comparator. + let found = false + const next = _.map(element => { + if (element.index !== index) { + return element + } + found = true + return { + ...element, + unparsedToolInput: element.unparsedToolInput + delta, + } + }) + return found ? next : _ + }) + return + } + case 'thinking_delta': + onUpdateLength(message.event.delta.thinking) + return + case 'signature_delta': + // Signatures are cryptographic authentication strings, not model + // output. Excluding them from onUpdateLength prevents them from + // inflating the OTPS metric and the animated token counter. + return + default: + return + } + case 'content_block_stop': + return + case 'message_delta': + onSetStreamMode('responding') + return + default: + onSetStreamMode('responding') + return + } +}