Merge branch 'main' into bee/streaming-chunk

Move streaming message processing logic from Task class into a new StreamingChunkProcessor class to improve code organization and maintainability. This change reduces complexity in the main task execution flow while maintaining the same functionality.

src/core/task/index.ts

@@ -2,7 +2,6 @@ import { setTimeout as setTimeoutPromise } from "node:timers/promises"
 import { Anthropic } from "@anthropic-ai/sdk"
 import { ApiHandler, ApiProviderInfo, buildApiHandler } from "@core/api"
 import { ApiStream } from "@core/api/transform/stream"
-import { AssistantMessageContent, parseAssistantMessageV2 } from "@core/assistant-message"
 import { ContextManager } from "@core/context/context-management/ContextManager"
 import { checkContextWindowExceededError } from "@core/context/context-management/context-error-handling"
 import { getContextWindowInfo } from "@core/context/context-management/context-window-utils"
@@ -63,7 +62,6 @@ import {
 } from "@shared/ExtensionMessage"
 import { HistoryItem } from "@shared/HistoryItem"
 import { DEFAULT_LANGUAGE_SETTINGS, getLanguageKey, LanguageDisplay } from "@shared/Languages"
-import { CLINE_MCP_TOOL_IDENTIFIER } from "@shared/mcp"
 import { convertClineMessageToProto } from "@shared/proto-conversions/cline-message"
 import { ClineDefaultTool } from "@shared/tools"
 import { ClineAskResponse } from "@shared/WebviewMessage"
@@ -76,7 +74,6 @@ import pWaitFor from "p-wait-for"
 import * as path from "path"
 import { ulid } from "ulid"
 import * as vscode from "vscode"
-import { ToolUseHandler } from "@/core/api/transform/tool-use-handler"
 import type { SystemPromptContext } from "@/core/prompts/system-prompt"
 import { getSystemPrompt } from "@/core/prompts/system-prompt"
 import { HostProvider } from "@/hosts/host-provider"
@@ -92,6 +89,7 @@ import { Controller } from "../controller"
 import { StateManager } from "../storage/StateManager"
 import { FocusChainManager } from "./focus-chain"
 import { MessageStateHandler } from "./message-state"
+import { StreamingChunkProcessor, StreamingChunkState } from "./streaming-chunk-processor"
 import { TaskState } from "./TaskState"
 import { ToolExecutor } from "./ToolExecutor"
 import { detectAvailableCliTools, extractProviderDomainFromUrl, updateApiReqMsg } from "./utils"
@@ -204,7 +202,6 @@ export class Task {
 	 * because of the expected format from the tool calls is different.
 	 */
 	private useNativeToolCalls: boolean = false
-	private toolUseHandler: ToolUseHandler
 
 	private terminalExecutionMode: "vscodeTerminal" | "backgroundExec"
 	private activeBackgroundCommand?: {
@@ -318,7 +315,6 @@ export class Task {
 		this.browserSession = new BrowserSession(stateManager)
 		this.contextManager = new ContextManager()
 		this.diffViewProvider = HostProvider.get().createDiffViewProvider()
-		this.toolUseHandler = new ToolUseHandler()
 		this.cwd = cwd
 		this.stateManager = stateManager
 		this.workspaceManager = workspaceManager
@@ -2622,6 +2618,18 @@ export class Task {
 			let outputTokens = 0
 			let totalCost: number | undefined
 
+			const streamingState: StreamingChunkState = {
+				// For UI display (includes XML)
+				assistantMessage: "",
+				// For API history (text only, no tool XML)
+				assistantTextOnly: "",
+				reasoningMessage: "",
+				reasoningDetails: [],
+				antThinkingContent: [],
+			}
+
+			let chunkProcessor: StreamingChunkProcessor | undefined
+
 			const abortStream = async (cancelReason: ClineApiReqCancelReason, streamingFailedMessage?: string) => {
 				if (this.diffViewProvider.isEditing) {
 					await this.diffViewProvider.revertChanges() // closes diff view
@@ -2644,7 +2652,7 @@ export class Task {
 						{
 							type: "text",
 							text:
-								assistantMessage +
+								streamingState.assistantMessage +
 								`\n\n[${
 									cancelReason === "streaming_failed"
 										? "Response interrupted by API Error"
@@ -2654,6 +2662,20 @@ export class Task {
 					],
 				})
 
+				const usageSnapshot = chunkProcessor?.getUsageSnapshot() ?? {
+					inputTokens,
+					outputTokens,
+					cacheWriteTokens,
+					cacheReadTokens,
+					totalCost,
+				}
+
+				inputTokens = usageSnapshot.inputTokens
+				outputTokens = usageSnapshot.outputTokens
+				cacheWriteTokens = usageSnapshot.cacheWriteTokens
+				cacheReadTokens = usageSnapshot.cacheReadTokens
+				totalCost = usageSnapshot.totalCost
+
 				// update api_req_started to have cancelled and cost, so that we can display the cost of the partial stream
 				await updateApiReqMsg({
 					messageStateHandler: this.messageStateHandler,
@@ -2701,173 +2723,40 @@ export class Task {
 			this.taskState.presentAssistantMessageHasPendingUpdates = false
 			this.taskState.didAutomaticallyRetryFailedApiRequest = false
 			await this.diffViewProvider.reset()
-			this.toolUseHandler.reset()
 			this.taskState.toolUseIdMap.clear()
 
+			chunkProcessor = new StreamingChunkProcessor({
+				taskState: this.taskState,
+				say: this.say.bind(this),
+				presentAssistantMessage: this.presentAssistantMessage.bind(this),
+				useNativeToolCalls: this.useNativeToolCalls,
+				abortStream,
+				streamingState,
+			})
+
 			const stream = this.attemptApiRequest(previousApiReqIndex) // yields only if the first chunk is successful, otherwise will allow the user to retry the request (most likely due to rate limit error, which gets thrown on the first chunk)
-			let assistantMessage = "" // For UI display (includes XML)
-			let assistantTextOnly = "" // For API history (text only, no tool XML)
-			let reasoningMessage = ""
-			const reasoningDetails = []
-			const antThinkingContent: (Anthropic.Messages.RedactedThinkingBlock | Anthropic.Messages.ThinkingBlock)[] = []
+
 			this.taskState.isStreaming = true
 			let didReceiveUsageChunk = false
+
+			const syncUsageFromProcessor = () => {
+				if (!chunkProcessor) {
+					return
+				}
+				const usageSnapshot = chunkProcessor.getUsageSnapshot()
+				inputTokens = usageSnapshot.inputTokens
+				outputTokens = usageSnapshot.outputTokens
+				cacheWriteTokens = usageSnapshot.cacheWriteTokens
+				cacheReadTokens = usageSnapshot.cacheReadTokens
+				totalCost = usageSnapshot.totalCost
+			}
+
 			try {
-				for await (const chunk of stream) {
-					if (!chunk) {
-						continue
-					}
-					switch (chunk.type) {
-						case "usage":
-							didReceiveUsageChunk = true
-							inputTokens += chunk.inputTokens
-							outputTokens += chunk.outputTokens
-							cacheWriteTokens += chunk.cacheWriteTokens ?? 0
-							cacheReadTokens += chunk.cacheReadTokens ?? 0
-							totalCost = chunk.totalCost
-							break
-						case "reasoning":
-							// reasoning will always come before assistant message
-							reasoningMessage += chunk.reasoning
-							// fixes bug where cancelling task > aborts task > for loop may be in middle of streaming reasoning > say function throws error before we get a chance to properly clean up and cancel the task.
-							if (!this.taskState.abort) {
-								await this.say("reasoning", reasoningMessage, undefined, undefined, true)
-							}
-							break
-						// for cline/openrouter providers
-						case "reasoning_details":
-							// reasoning_details may be an array of 0 or 1 items depending on how openrouter returns it
-							if (Array.isArray(chunk.reasoning_details)) {
-								reasoningDetails.push(...chunk.reasoning_details)
-							} else {
-								reasoningDetails.push(chunk.reasoning_details)
-							}
-							break
-						// for anthropic providers
-						case "ant_thinking":
-							antThinkingContent.push({
-								type: "thinking",
-								thinking: chunk.thinking,
-								signature: chunk.signature,
-							})
-							break
-						case "ant_redacted_thinking":
-							antThinkingContent.push({
-								type: "redacted_thinking",
-								data: chunk.data,
-							})
-							break
-						case "tool_calls": {
-							if (!chunk.tool_call) {
-								console.log("no tool call in chunk, skipping...", chunk)
-								break
-							}
-							// Accumulate tool use blocks in proper Anthropic format
-							this.toolUseHandler.processToolUseDelta({
-								id: chunk.tool_call.function?.id,
-								type: "tool_use",
-								name: chunk.tool_call.function?.name,
-								input: chunk.tool_call.function?.arguments,
-							})
-
-							// Extract and store tool_use_id for creating proper ToolResultBlockParam
-							if (chunk.tool_call.function?.id && chunk.tool_call.function?.name) {
-								this.taskState.toolUseIdMap.set(chunk.tool_call.function.name, chunk.tool_call.function.id)
-
-								// For MCP tools, also store the mapping with the transformed name
-								// since getPartialToolUsesAsContent() will transform the name to "use_mcp_tool"
-								if (chunk.tool_call.function.name.includes(CLINE_MCP_TOOL_IDENTIFIER)) {
-									this.taskState.toolUseIdMap.set(ClineDefaultTool.MCP_USE, chunk.tool_call.function.id)
-								}
-							}
-
-							const prevLength = this.taskState.assistantMessageContent.length
-
-							// Combine any text content with tool uses
-							const textContent = assistantTextOnly.trim()
-							const textBlocks: AssistantMessageContent[] = textContent
-								? [{ type: "text", content: textContent, partial: false }]
-								: []
-							const toolBlocks = this.toolUseHandler.getPartialToolUsesAsContent()
-							assistantMessage += toolBlocks.map((block) => JSON.stringify(block)).join("\n")
-							this.taskState.assistantMessageContent = [...textBlocks, ...toolBlocks]
-
-							if (this.taskState.assistantMessageContent.length > prevLength) {
-								this.taskState.userMessageContentReady = false
-							}
-							this.presentAssistantMessage()
-							break
-						}
-						case "text": {
-							if (reasoningMessage && assistantMessage.length === 0) {
-								// complete reasoning message
-								await this.say("reasoning", reasoningMessage, undefined, undefined, false)
-							}
-							assistantMessage += chunk.text
-							assistantTextOnly += chunk.text // Accumulate text separately
-							// parse raw assistant message into content blocks
-							const prevLength = this.taskState.assistantMessageContent.length
-
-							this.taskState.assistantMessageContent = parseAssistantMessageV2(assistantMessage)
-
-							if (this.taskState.assistantMessageContent.length > prevLength) {
-								this.taskState.userMessageContentReady = false // new content we need to present, reset to false in case previous content set this to true
-							}
-							// present content to user
-							this.presentAssistantMessage()
-							break
-						}
-					}
-
-					if (this.taskState.abort) {
-						console.log("aborting stream...")
-						if (!this.taskState.abandoned) {
-							// only need to gracefully abort if this instance isn't abandoned (sometimes openrouter stream hangs, in which case this would affect future instances of cline)
-							await abortStream("user_cancelled")
-						}
-						break // aborts the stream
-					}
-
-					if (this.taskState.didRejectTool) {
-						// userContent has a tool rejection, so interrupt the assistant's response to present the user's feedback
-						assistantMessage += "\n\n[Response interrupted by user feedback]"
-						// this.userMessageContentReady = true // instead of setting this preemptively, we allow the present iterator to finish and set userMessageContentReady when its ready
-						break
-					}
-
-					// PREV: we need to let the request finish for openrouter to get generation details
-					// UPDATE: it's better UX to interrupt the request at the cost of the api cost not being retrieved
-					if (this.taskState.didAlreadyUseTool) {
-						assistantMessage +=
-							"\n\n[Response interrupted by a tool use result. Only one tool may be used at a time and should be placed at the end of the message.]"
-						break
-					}
-				}
-
-				// Finalize any remaining tool calls at the end of the stream
-				if (this.useNativeToolCalls) {
-					// For native tool calls, mark all pending tool uses as complete
-					const prevLength = this.taskState.assistantMessageContent.length
-
-					// Get finalized tool uses and mark them as complete
-					const textContent = assistantTextOnly.trim()
-					const textBlocks: AssistantMessageContent[] = textContent
-						? [{ type: "text", content: textContent, partial: false }]
-						: []
-
-					// Get all finalized tool uses and mark as complete
-					const toolBlocks = this.toolUseHandler
-						.getPartialToolUsesAsContent()
-						.map((block) => ({ ...block, partial: false }))
-
-					this.taskState.assistantMessageContent = [...textBlocks, ...toolBlocks]
-
-					if (this.taskState.assistantMessageContent.length > prevLength) {
-						this.taskState.userMessageContentReady = false
-					}
-					this.presentAssistantMessage()
-				}
+				const { didReceiveUsageChunk: usageFlag } = await chunkProcessor.processStream(stream)
+				didReceiveUsageChunk = usageFlag
+				syncUsageFromProcessor()
 			} catch (error) {
+				syncUsageFromProcessor()
 				// abandoned happens when extension is no longer waiting for the cline instance to finish aborting (error is thrown here when any function in the for loop throws due to this.abort)
 				if (!this.taskState.abandoned) {
 					const clineError = ErrorService.get().toClineError(error, this.api.getModel().id)
@@ -2981,8 +2870,7 @@ export class Task {
 			// now add to apiconversationhistory
 			// need to save assistant responses to file before proceeding to tool use since user can exit at any moment and we wouldn't be able to save the assistant's response
 			let didEndLoop = false
-			if (assistantMessage.length > 0 || this.useNativeToolCalls) {
-				const currentMode = this.stateManager.getGlobalSettingsKey("mode")
+			if (streamingState.assistantMessage.length > 0 || this.useNativeToolCalls) {
 				telemetryService.captureConversationTurnEvent(this.ulid, providerId, model.id, "assistant", currentMode, {
 					tokensIn: inputTokens,
 					tokensOut: outputTokens,
@@ -2992,7 +2880,7 @@ export class Task {
 				})
 
 				// Get finalized tool use blocks from the handler
-				const toolUseBlocks = this.toolUseHandler.getAllFinalizedToolUses()
+				const toolUseBlocks = chunkProcessor.getFinalizedToolCalls()
 
 				// Build content array with thinking blocks, text (if any), and tool use blocks
 				const assistantContent: Array<
@@ -3004,17 +2892,18 @@ export class Task {
 					// This is critical for maintaining the model's reasoning flow and conversation integrity.
 					// "When providing thinking blocks, the entire sequence of consecutive thinking blocks must match the outputs generated by the model during the original request; you cannot rearrange or modify the sequence of these blocks." The signature_delta is used to verify that the thinking was generated by Claude, and the thinking blocks will be ignored if it's incorrect or missing.
 					// https://docs.claude.com/en/docs/build-with-claude/extended-thinking#preserving-thinking-blocks
-					...antThinkingContent,
+					...streamingState.antThinkingContent,
 				]
 
 				// Only add text block if there's actual text (not just tool XML)
-				if (assistantTextOnly.trim().length > 0) {
+				if (streamingState.assistantTextOnly.trim().length > 0) {
 					assistantContent.push({
 						type: "text",
-						text: assistantTextOnly,
+						text: streamingState.assistantTextOnly,
 						// reasoning_details only exists for cline/openrouter providers
 						// @ts-ignore-next-line
-						reasoning_details: reasoningDetails.length > 0 ? reasoningDetails : undefined,
+						reasoning_details:
+							streamingState.reasoningDetails.length > 0 ? streamingState.reasoningDetails : undefined,
 					})
 				}
 

src/core/task/streaming-chunk-processor.test.ts

@@ -0,0 +1,340 @@
+import type {
+	ApiStream,
+	ApiStreamAnthropicRedactedThinkingChunk,
+	ApiStreamAnthropicThinkingChunk,
+	ApiStreamChunk,
+	ApiStreamReasoningChunk,
+	ApiStreamReasoningDetailsChunk,
+	ApiStreamTextChunk,
+	ApiStreamToolCallsChunk,
+	ApiStreamUsageChunk,
+} from "@core/api/transform/stream"
+import { ClineDefaultTool } from "@shared/tools"
+import { expect } from "chai"
+import { afterEach, beforeEach, describe, it } from "mocha"
+import sinon from "sinon"
+import type { StreamingChunkState } from "./streaming-chunk-processor"
+import { StreamingChunkProcessor } from "./streaming-chunk-processor"
+import { TaskState } from "./TaskState"
+
+type TestContext = {
+	taskState: TaskState
+	streamingState: StreamingChunkState
+	sayStub: sinon.SinonStub
+	presentStub: sinon.SinonStub
+	abortStub: sinon.SinonStub
+	processor: StreamingChunkProcessor
+}
+
+const createStreamingState = (): StreamingChunkState => ({
+	assistantMessage: "",
+	assistantTextOnly: "",
+	reasoningMessage: "",
+	reasoningDetails: [],
+	antThinkingContent: [],
+})
+
+const createProcessorContext = (
+	sandbox: sinon.SinonSandbox,
+	overrides?: Partial<{ useNativeToolCalls: boolean }>,
+): TestContext => {
+	const taskState = new TaskState()
+	const streamingState = createStreamingState()
+	const sayStub = sandbox.stub().resolves(undefined)
+	const presentStub = sandbox.stub()
+	const abortStub = sandbox.stub().resolves(undefined)
+
+	const processor = new StreamingChunkProcessor({
+		taskState,
+		streamingState,
+		say: sayStub,
+		presentAssistantMessage: presentStub,
+		useNativeToolCalls: overrides?.useNativeToolCalls ?? false,
+		abortStream: abortStub,
+	})
+
+	return {
+		taskState,
+		streamingState,
+		sayStub,
+		presentStub,
+		abortStub,
+		processor,
+	}
+}
+
+const createStream = (chunks: ApiStreamChunk[]): ApiStream =>
+	(async function* () {
+		for (const chunk of chunks) {
+			yield chunk
+		}
+	})()
+
+describe("StreamingChunkProcessor", () => {
+	let sandbox: sinon.SinonSandbox
+
+	beforeEach(() => {
+		sandbox = sinon.createSandbox()
+	})
+
+	afterEach(() => {
+		sandbox.restore()
+	})
+
+	describe("handleChunk", () => {
+		type ChunkCase = {
+			name: string
+			chunk:
+				| ApiStreamUsageChunk
+				| ApiStreamReasoningChunk
+				| ApiStreamReasoningDetailsChunk
+				| ApiStreamAnthropicThinkingChunk
+				| ApiStreamAnthropicRedactedThinkingChunk
+				| ApiStreamTextChunk
+				| ApiStreamToolCallsChunk
+			setup?: (ctx: TestContext) => void
+			verify: (ctx: TestContext) => void | Promise<void>
+		}
+
+		const chunkCases: ChunkCase[] = [
+			{
+				name: "usage chunk updates token snapshot",
+				chunk: {
+					type: "usage",
+					inputTokens: 10,
+					outputTokens: 5,
+					cacheWriteTokens: 3,
+					cacheReadTokens: 7,
+					totalCost: 0.42,
+				},
+				verify: (ctx) => {
+					const snapshot = ctx.processor.getUsageSnapshot()
+					expect(ctx.processor.didReceiveUsage).to.be.true
+					expect(snapshot).to.deep.equal({
+						inputTokens: 10,
+						outputTokens: 5,
+						cacheWriteTokens: 3,
+						cacheReadTokens: 7,
+						totalCost: 0.42,
+					})
+				},
+			},
+			{
+				name: "reasoning chunk streams partial reasoning",
+				chunk: {
+					type: "reasoning",
+					reasoning: "Thinking aloud...",
+				},
+				verify: (ctx) => {
+					expect(ctx.streamingState.reasoningMessage).to.equal("Thinking aloud...")
+					expect(ctx.sayStub.calledWithExactly("reasoning", "Thinking aloud...", undefined, undefined, true)).to.be.true
+				},
+			},
+			{
+				name: "reasoning details chunk aggregates traces",
+				chunk: {
+					type: "reasoning_details",
+					reasoning_details: [{ foo: "bar" }],
+				},
+				verify: (ctx) => {
+					expect(ctx.streamingState.reasoningDetails).to.deep.equal([{ foo: "bar" }])
+				},
+			},
+			{
+				name: "anthropic thinking chunk records blocks",
+				chunk: {
+					type: "ant_thinking",
+					thinking: "Calculating step",
+					signature: "sig-123",
+				},
+				verify: (ctx) => {
+					expect(ctx.streamingState.antThinkingContent).to.deep.equal([
+						{
+							type: "thinking",
+							thinking: "Calculating step",
+							signature: "sig-123",
+						},
+					])
+				},
+			},
+			{
+				name: "anthropic redacted thinking chunk records data",
+				chunk: {
+					type: "ant_redacted_thinking",
+					data: "redacted",
+				},
+				verify: (ctx) => {
+					expect(ctx.streamingState.antThinkingContent).to.deep.equal([
+						{
+							type: "redacted_thinking",
+							data: "redacted",
+						},
+					])
+				},
+			},
+			{
+				name: "text chunk updates assistant content and presentation",
+				chunk: {
+					type: "text",
+					text: "Here is the answer.",
+				},
+				setup: (ctx) => {
+					ctx.streamingState.reasoningMessage = "pre-reasoning"
+				},
+				verify: (ctx) => {
+					expect(ctx.sayStub.calledWithExactly("reasoning", "pre-reasoning", undefined, undefined, false)).to.be.true
+					expect(ctx.streamingState.assistantMessage).to.equal("Here is the answer.")
+					expect(ctx.streamingState.assistantTextOnly).to.equal("Here is the answer.")
+					expect(ctx.taskState.assistantMessageContent).to.not.be.empty
+					expect(ctx.presentStub.calledOnce).to.be.true
+				},
+			},
+			{
+				name: "tool call chunk captures partial tool use",
+				chunk: {
+					type: "tool_calls",
+					tool_call: {
+						function: {
+							id: "tool-1",
+							name: ClineDefaultTool.BASH,
+							arguments: '{"command":"ls"}',
+						},
+					},
+				},
+				verify: (ctx) => {
+					expect(ctx.taskState.toolUseIdMap.get(ClineDefaultTool.BASH)).to.equal("tool-1")
+					expect(ctx.taskState.assistantMessageContent).to.have.lengthOf(1)
+					const toolBlock = ctx.taskState.assistantMessageContent[0] as any
+					expect(toolBlock).to.include({
+						type: "tool_use",
+						name: ClineDefaultTool.BASH,
+						partial: true,
+					})
+					expect(toolBlock.params).to.deep.equal({ command: "ls" })
+					const serialized = ctx.streamingState.assistantMessage.trim()
+					expect(serialized).to.not.equal("")
+					expect(JSON.parse(serialized)).to.deep.equal(toolBlock)
+					expect(ctx.presentStub.calledOnce).to.be.true
+				},
+			},
+		]
+
+		for (const chunkCase of chunkCases) {
+			it(`handleChunk ${chunkCase.name}`, async () => {
+				const ctx = createProcessorContext(sandbox)
+
+				chunkCase.setup?.(ctx)
+				await ctx.processor.handleChunk(chunkCase.chunk)
+				await chunkCase.verify(ctx)
+			})
+		}
+	})
+
+	describe("processStream control flow", () => {
+		type InterruptionCase = {
+			name: string
+			setup: (ctx: TestContext) => void
+			verify: (ctx: TestContext, result: Awaited<ReturnType<StreamingChunkProcessor["processStream"]>>) => void
+		}
+
+		const interruptionChunks: ApiStreamChunk[] = [
+			{
+				type: "text",
+				text: "Partial reply",
+			},
+			{
+				type: "usage",
+				inputTokens: 2,
+				outputTokens: 1,
+			},
+		]
+
+		const interruptionCases: InterruptionCase[] = [
+			{
+				name: "aborts stream when task is cancelled",
+				setup: (ctx) => {
+					ctx.taskState.abort = true
+					ctx.taskState.abandoned = false
+				},
+				verify: (ctx, result) => {
+					expect(result.didReceiveUsageChunk).to.be.false
+					expect(ctx.abortStub.calledWithExactly("user_cancelled")).to.be.true
+				},
+			},
+			{
+				name: "interrupts when tool rejection occurs",
+				setup: (ctx) => {
+					ctx.taskState.didRejectTool = true
+				},
+				verify: (ctx, result) => {
+					expect(result.didReceiveUsageChunk).to.be.false
+					expect(ctx.streamingState.assistantMessage).to.include("[Response interrupted by user feedback]")
+					expect(ctx.abortStub.notCalled).to.be.true
+				},
+			},
+			{
+				name: "interrupts when another tool was already used",
+				setup: (ctx) => {
+					ctx.taskState.didAlreadyUseTool = true
+				},
+				verify: (ctx, result) => {
+					expect(result.didReceiveUsageChunk).to.be.false
+					expect(ctx.streamingState.assistantMessage).to.include(
+						"[Response interrupted by a tool use result. Only one tool may be used at a time and should be placed at the end of the message.]",
+					)
+					expect(ctx.abortStub.notCalled).to.be.true
+				},
+			},
+		]
+
+		for (const interruptionCase of interruptionCases) {
+			it(`processStream ${interruptionCase.name}`, async () => {
+				const ctx = createProcessorContext(sandbox)
+				interruptionCase.setup(ctx)
+
+				const result = await ctx.processor.processStream(createStream(interruptionChunks))
+
+				interruptionCase.verify(ctx, result)
+			})
+		}
+	})
+
+	it("processStream finalizes tool calls and aggregates usage when native tools enabled", async () => {
+		const ctx = createProcessorContext(sandbox, { useNativeToolCalls: true })
+
+		const chunks: ApiStreamChunk[] = [
+			{
+				type: "tool_calls",
+				tool_call: {
+					function: {
+						id: "tool-2",
+						name: ClineDefaultTool.FILE_READ,
+						arguments: JSON.stringify({ path: "README.md" }),
+					},
+				},
+			},
+			{
+				type: "usage",
+				inputTokens: 4,
+				outputTokens: 6,
+				cacheReadTokens: 1,
+			},
+		]
+
+		const result = await ctx.processor.processStream(createStream(chunks))
+
+		expect(result.didReceiveUsageChunk).to.be.true
+		expect(ctx.processor.getUsageSnapshot()).to.deep.equal({
+			inputTokens: 4,
+			outputTokens: 6,
+			cacheWriteTokens: 0,
+			cacheReadTokens: 1,
+			totalCost: undefined,
+		})
+		expect(ctx.taskState.assistantMessageContent).to.have.lengthOf(1)
+		const finalizedBlock = ctx.taskState.assistantMessageContent[0] as any
+		expect(finalizedBlock.partial).to.be.false
+		expect(finalizedBlock.name).to.equal(ClineDefaultTool.FILE_READ)
+		expect(ctx.presentStub.callCount).to.equal(2)
+	})
+})

src/core/task/streaming-chunk-processor.ts

@@ -0,0 +1,293 @@
+import type { Anthropic } from "@anthropic-ai/sdk"
+import type { ApiStream, ApiStreamChunk, ApiStreamUsageChunk } from "@core/api/transform/stream"
+import type { AssistantMessageContent } from "@core/assistant-message"
+import { parseAssistantMessageV2 } from "@core/assistant-message"
+import type { ClineApiReqCancelReason, ClineSay } from "@shared/ExtensionMessage"
+import { CLINE_MCP_TOOL_IDENTIFIER } from "@shared/mcp"
+import { ClineDefaultTool } from "@shared/tools"
+import { ToolUseHandler } from "@/core/api/transform/tool-use-handler"
+import type { TaskState } from "./TaskState"
+
+export type StreamingChunkState = {
+	// For UI display (includes XML)
+	assistantMessage: string
+	// For API history (text only, no tool XML)
+	assistantTextOnly: string
+	reasoningMessage: string
+	reasoningDetails: Array<unknown>
+	antThinkingContent: Array<Anthropic.Messages.RedactedThinkingBlock | Anthropic.Messages.ThinkingBlock>
+}
+
+export type TokenUsageSnapshot = {
+	inputTokens: number
+	outputTokens: number
+	cacheWriteTokens: number
+	cacheReadTokens: number
+	totalCost?: number
+}
+
+type SayFn = (
+	type: ClineSay,
+	text?: string,
+	images?: string[],
+	files?: string[],
+	partial?: boolean,
+) => Promise<number | undefined>
+
+type AbortStreamHandler = (cancelReason: ClineApiReqCancelReason, streamingFailedMessage?: string) => Promise<void>
+
+type StreamingChunkProcessorOptions = {
+	taskState: TaskState
+	say: SayFn
+	presentAssistantMessage: () => void
+	useNativeToolCalls: boolean
+	abortStream: AbortStreamHandler
+	streamingState: StreamingChunkState
+}
+
+export type StreamingChunkProcessingResult = {
+	didReceiveUsageChunk: boolean
+}
+
+export class StreamingChunkProcessor {
+	private readonly taskState: TaskState
+	private readonly toolUseHandler: ToolUseHandler
+	private readonly say: SayFn
+	private readonly presentAssistantMessage: () => void
+	private readonly useNativeToolCalls: boolean
+	private readonly abortStream: AbortStreamHandler
+	private readonly streamingState: StreamingChunkState
+
+	private shouldStopStreaming = false
+	private didReceiveUsageChunk = false
+
+	private inputTokens = 0
+	private outputTokens = 0
+	private cacheWriteTokens = 0
+	private cacheReadTokens = 0
+	private totalCost: number | undefined
+
+	constructor(options: StreamingChunkProcessorOptions) {
+		this.taskState = options.taskState
+		this.toolUseHandler = new ToolUseHandler()
+		this.say = options.say
+		this.presentAssistantMessage = options.presentAssistantMessage
+		this.useNativeToolCalls = options.useNativeToolCalls
+		this.abortStream = options.abortStream
+		this.streamingState = options.streamingState
+	}
+
+	public async processStream(stream: ApiStream): Promise<StreamingChunkProcessingResult> {
+		for await (const chunk of stream) {
+			if (!chunk) {
+				continue
+			}
+
+			await this.handleChunk(chunk)
+
+			if (this.shouldStopStreaming) {
+				break
+			}
+		}
+
+		// Finalize any remaining tool calls at the end of the stream
+		this.finalizeNativeToolCalls()
+
+		return {
+			didReceiveUsageChunk: this.didReceiveUsageChunk,
+		}
+	}
+
+	public getUsageSnapshot(): TokenUsageSnapshot {
+		return {
+			inputTokens: this.inputTokens,
+			outputTokens: this.outputTokens,
+			cacheWriteTokens: this.cacheWriteTokens,
+			cacheReadTokens: this.cacheReadTokens,
+			totalCost: this.totalCost,
+		}
+	}
+
+	private updateTokenUsage(usage: ApiStreamUsageChunk): void {
+		this.inputTokens += usage.inputTokens
+		this.outputTokens += usage.outputTokens
+		this.cacheWriteTokens += usage.cacheWriteTokens ?? 0
+		this.cacheReadTokens += usage.cacheReadTokens ?? 0
+		this.totalCost = usage.totalCost
+	}
+
+	async handleChunk(chunk: ApiStreamChunk): Promise<void> {
+		switch (chunk.type) {
+			case "usage":
+				this.didReceiveUsageChunk = true
+				this.updateTokenUsage(chunk)
+				break
+			case "reasoning":
+			case "reasoning_details":
+			case "ant_thinking":
+			case "ant_redacted_thinking":
+				await this.handleReasoningChunk(chunk)
+				break
+			case "tool_calls":
+				this.handleToolCallChunk(chunk)
+				break
+			case "text":
+				await this.handleTextChunk(chunk)
+				break
+		}
+
+		await this.handlePostChunkActions()
+	}
+
+	private async handleReasoningChunk(
+		chunk: Extract<ApiStreamChunk, { type: "reasoning" | "reasoning_details" | "ant_thinking" | "ant_redacted_thinking" }>,
+	): Promise<void> {
+		switch (chunk.type) {
+			case "reasoning":
+				// reasoning will always come before assistant message
+				this.streamingState.reasoningMessage += chunk.reasoning
+				// fixes bug where cancelling task > aborts task > for loop may be in middle of streaming reasoning > say function throws error before we get a chance to properly clean up and cancel the task.
+				if (!this.taskState.abort) {
+					await this.say("reasoning", this.streamingState.reasoningMessage, undefined, undefined, true)
+				}
+				break
+			// for cline/openrouter providers
+			case "reasoning_details":
+				// reasoning_details may be an array of 0 or 1 items depending on how openrouter returns it
+				const details = Array.isArray(chunk.reasoning_details) ? chunk.reasoning_details : [chunk.reasoning_details]
+				this.streamingState.reasoningDetails = [...this.streamingState.reasoningDetails, ...details]
+				break
+			// for anthropic providers
+			case "ant_thinking":
+				this.streamingState.antThinkingContent.push({
+					type: "thinking",
+					thinking: chunk.thinking,
+					signature: chunk.signature,
+				})
+				break
+			case "ant_redacted_thinking":
+				this.streamingState.antThinkingContent.push({
+					type: "redacted_thinking",
+					data: chunk.data,
+				})
+				break
+		}
+	}
+
+	private handleToolCallChunk(chunk: Extract<ApiStreamChunk, { type: "tool_calls" }>): void {
+		if (!chunk.tool_call) {
+			console.log("no tool call in chunk, skipping...", chunk)
+			return
+		}
+
+		// Accumulate tool use blocks in proper Anthropic format
+		this.toolUseHandler.processToolUseDelta({
+			id: chunk.tool_call.function?.id,
+			type: "tool_use",
+			name: chunk.tool_call.function?.name,
+			input: chunk.tool_call.function?.arguments,
+		})
+
+		// Extract and store tool_use_id for creating proper ToolResultBlockParam
+		if (chunk.tool_call.function?.id && chunk.tool_call.function?.name) {
+			this.taskState.toolUseIdMap.set(chunk.tool_call.function.name, chunk.tool_call.function.id)
+
+			// For MCP tools, also store the mapping with the transformed name
+			// since getPartialToolUsesAsContent() will transform the name to "use_mcp_tool"
+			if (chunk.tool_call.function.name.includes(CLINE_MCP_TOOL_IDENTIFIER)) {
+				this.taskState.toolUseIdMap.set(ClineDefaultTool.MCP_USE, chunk.tool_call.function.id)
+			}
+		}
+
+		const { textBlocks, toolBlocks, prevLength } = this.getUserMessageContent()
+
+		// Combine any text content with tool uses
+		this.streamingState.assistantMessage += toolBlocks.map((block) => JSON.stringify(block)).join("\n")
+		this.taskState.assistantMessageContent = [...textBlocks, ...toolBlocks]
+
+		if (this.taskState.assistantMessageContent.length > prevLength) {
+			this.taskState.userMessageContentReady = false
+		}
+		this.presentAssistantMessage()
+	}
+
+	private async handleTextChunk(chunk: Extract<ApiStreamChunk, { type: "text" }>): Promise<void> {
+		if (this.streamingState.reasoningMessage && this.streamingState.assistantMessage.length === 0) {
+			await this.say("reasoning", this.streamingState.reasoningMessage, undefined, undefined, false)
+		}
+
+		this.streamingState.assistantMessage += chunk.text
+		this.streamingState.assistantTextOnly += chunk.text
+
+		const { prevLength } = this.getUserMessageContent()
+
+		this.taskState.assistantMessageContent = parseAssistantMessageV2(this.streamingState.assistantMessage)
+
+		if (this.taskState.assistantMessageContent.length > prevLength) {
+			this.taskState.userMessageContentReady = false
+		}
+
+		this.presentAssistantMessage()
+	}
+
+	private async handlePostChunkActions(): Promise<void> {
+		if (this.taskState.abort) {
+			console.log("aborting stream...")
+			if (!this.taskState.abandoned) {
+				// only need to gracefully abort if this instance isn't abandoned (sometimes openrouter stream hangs, in which case this would affect future instances of cline)
+				await this.abortStream("user_cancelled")
+			}
+			this.shouldStopStreaming = true
+			return // aborts the stream
+		}
+
+		if (this.taskState.didRejectTool) {
+			// userContent has a tool rejection, so interrupt the assistant's response to present the user's feedback
+			this.streamingState.assistantMessage += "\n\n[Response interrupted by user feedback]"
+			this.shouldStopStreaming = true
+			return
+		}
+		// PREV: we need to let the request finish for openrouter to get generation details
+		// UPDATE: it's better UX to interrupt the request at the cost of the api cost not being retrieved
+		if (this.taskState.didAlreadyUseTool) {
+			this.streamingState.assistantMessage +=
+				"\n\n[Response interrupted by a tool use result. Only one tool may be used at a time and should be placed at the end of the message.]"
+			this.shouldStopStreaming = true
+		}
+	}
+
+	private finalizeNativeToolCalls(): void {
+		if (!this.useNativeToolCalls) {
+			return
+		}
+
+		// For native tool calls, mark all pending tool uses as complete
+		const { textBlocks, toolBlocks, prevLength } = this.getUserMessageContent()
+		// Get all finalized tool uses and mark as complete
+		const finalizedToolBlocks = toolBlocks.map((block) => ({ ...block, partial: false }))
+
+		this.taskState.assistantMessageContent = [...textBlocks, ...finalizedToolBlocks]
+
+		if (this.taskState.assistantMessageContent.length > prevLength) {
+			this.taskState.userMessageContentReady = false
+		}
+		this.presentAssistantMessage()
+	}
+
+	get didReceiveUsage(): boolean {
+		return this.didReceiveUsageChunk
+	}
+
+	private getUserMessageContent() {
+		const prevLength = this.taskState.assistantMessageContent.length
+		const textContent = this.streamingState.assistantTextOnly.trim()
+		const textBlocks: AssistantMessageContent[] = textContent ? [{ type: "text", content: textContent, partial: false }] : []
+		const toolBlocks = this.toolUseHandler.getPartialToolUsesAsContent()
+
+		return { textBlocks, toolBlocks, prevLength }
+	}
+
+	public getFinalizedToolCalls() {
+		return this.toolUseHandler.getAllFinalizedToolUses()
+	}
+}