From 22aa6d617490a0f90e9e03ad97c9e1cd71ceb07c Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Thu, 17 Sep 2026 10:48:18 +0800 Subject: [PATCH 1/2] fix(typescript): bound and cancel background streams --- typescript/src/orchestrator.ts | 247 ++++++++++++++---- typescript/src/utils/helpers.ts | 133 ++++++---- .../tests/streamingCancellation.test.ts | 232 ++++++++++++++++ typescript/tests/utils/helpers.test.ts | 20 +- 4 files changed, 519 insertions(+), 113 deletions(-) create mode 100644 typescript/tests/streamingCancellation.test.ts diff --git a/typescript/src/orchestrator.ts b/typescript/src/orchestrator.ts index aa71912c..ffa6decc 100644 --- a/typescript/src/orchestrator.ts +++ b/typescript/src/orchestrator.ts @@ -85,6 +85,15 @@ export interface AgentSquadConfig { * If MAX_MESSAGE_PAIRS_PER_AGENT is 5, up to 10 messages (5 pairs) will be stored per agent. */ MAX_MESSAGE_PAIRS_PER_AGENT?: number; + + /** Maximum number of chunks accepted from a streaming agent response. */ + MAX_STREAM_CHUNKS?: number; + + /** Maximum number of bytes accepted from a streaming agent response. */ + MAX_STREAM_BYTES?: number; + + /** Maximum number of text bytes retained for conversation storage. */ + MAX_ACCUMULATED_RESPONSE_BYTES?: number; } export const DEFAULT_CONFIG: AgentSquadConfig = { @@ -121,6 +130,15 @@ export const DEFAULT_CONFIG: AgentSquadConfig = { /** Default: Maximum of 100 message pairs (200 individual messages) to retain per agent */ MAX_MESSAGE_PAIRS_PER_AGENT: 100, + + /** Default: bound streaming responses to prevent unbounded resource usage */ + MAX_STREAM_CHUNKS: 10000, + + /** Default: bound streaming responses to 10 MiB */ + MAX_STREAM_BYTES: 10 * 1024 * 1024, + + /** Default: retain at most 10 MiB of response text */ + MAX_ACCUMULATED_RESPONSE_BYTES: 10 * 1024 * 1024, }; export interface DispatchToAgentsParams { @@ -183,7 +201,7 @@ export interface RequestMetadata { export type ThinkingResponse = { content: string; thinking: string; -} +}; export class AgentSquad { private config: AgentSquadConfig; @@ -216,6 +234,13 @@ export class AgentSquad { MAX_MESSAGE_PAIRS_PER_AGENT: options.config?.MAX_MESSAGE_PAIRS_PER_AGENT ?? DEFAULT_CONFIG.MAX_MESSAGE_PAIRS_PER_AGENT, + MAX_STREAM_CHUNKS: + options.config?.MAX_STREAM_CHUNKS ?? DEFAULT_CONFIG.MAX_STREAM_CHUNKS, + MAX_STREAM_BYTES: + options.config?.MAX_STREAM_BYTES ?? DEFAULT_CONFIG.MAX_STREAM_BYTES, + MAX_ACCUMULATED_RESPONSE_BYTES: + options.config?.MAX_ACCUMULATED_RESPONSE_BYTES ?? + DEFAULT_CONFIG.MAX_ACCUMULATED_RESPONSE_BYTES, USE_DEFAULT_AGENT_IF_NONE_IDENTIFIED: options.config?.USE_DEFAULT_AGENT_IF_NONE_IDENTIFIED ?? DEFAULT_CONFIG.USE_DEFAULT_AGENT_IF_NONE_IDENTIFIED, @@ -265,7 +290,7 @@ export class AgentSquad { Object.entries(this.agents).map(([key, { name, description }]) => [ key, { name, description }, - ]) + ]), ); } @@ -274,7 +299,7 @@ export class AgentSquad { } async dispatchToAgent( - params: DispatchToAgentsParams + params: DispatchToAgentsParams, ): Promise | ThinkingResponse> { const { userInput, @@ -292,13 +317,13 @@ export class AgentSquad { const agentChatHistory = await this.storage.fetchChat( userId, sessionId, - selectedAgent.id + selectedAgent.id, ); this.logger.printChatHistory(agentChatHistory, selectedAgent.id); this.logger.info( - `Routing intent "${userInput}" to ${selectedAgent.id} ...` + `Routing intent "${userInput}" to ${selectedAgent.id} ...`, ); const response = await this.measureExecutionTime( @@ -309,8 +334,8 @@ export class AgentSquad { userId, sessionId, agentChatHistory, - additionalParams - ) + additionalParams, + ), ); //if (this.isStream(response)) { @@ -337,8 +362,9 @@ export class AgentSquad { if (thinkingParts.length > 0) { return { - content: contentParts.join(''), thinking: thinkingParts.join('') - } + content: contentParts.join(""), + thinking: thinkingParts.join(""), + }; } responseText = contentParts.join(""); } @@ -353,14 +379,14 @@ export class AgentSquad { async classifyRequest( userInput: string, userId: string, - sessionId: string + sessionId: string, ): Promise { try { const chatHistory = (await this.storage.fetchAllChats(userId, sessionId)) || []; const classifierResult = await this.measureExecutionTime( "Classifying user intent", - () => this.classifier.classify(userInput, chatHistory) + () => this.classifier.classify(userInput, chatHistory), ); this.logger.printIntent(userInput, classifierResult); @@ -387,7 +413,7 @@ export class AgentSquad { userId: string, sessionId: string, classifierResult: ClassifierResult, - additionalParams: Record = {} + additionalParams: Record = {}, ): Promise { try { const agentResponse = await this.dispatchToAgent({ @@ -403,18 +429,21 @@ export class AgentSquad { userInput, userId, sessionId, - additionalParams + additionalParams, ); if (this.isAsyncIterable(agentResponse)) { - const accumulatorTransform = new AccumulatorTransform(); - this.processStreamInBackground( + const accumulatorTransform = new AccumulatorTransform({ + maxAccumulatedResponseBytes: + this.config.MAX_ACCUMULATED_RESPONSE_BYTES, + }); + void this.processStreamInBackground( agentResponse, accumulatorTransform, userInput, userId, sessionId, - classifierResult.selectedAgent + classifierResult.selectedAgent, ); return { metadata, @@ -423,7 +452,9 @@ export class AgentSquad { }; } - const response: string = (agentResponse as ThinkingResponse).content || agentResponse as string; + const response: string = + (agentResponse as ThinkingResponse).content || + (agentResponse as string); if (classifierResult?.selectedAgent.saveChat) { await saveConversationExchange( userInput, @@ -432,7 +463,7 @@ export class AgentSquad { userId, sessionId, classifierResult?.selectedAgent.id, - this.config.MAX_MESSAGE_PAIRS_PER_AGENT + this.config.MAX_MESSAGE_PAIRS_PER_AGENT, ); } @@ -440,7 +471,7 @@ export class AgentSquad { metadata, output: response, streaming: false, - thinking: (agentResponse as ThinkingResponse).thinking + thinking: (agentResponse as ThinkingResponse).thinking, }; } catch (error) { this.logger.error("Error during agent processing:", error); @@ -452,7 +483,7 @@ export class AgentSquad { userInput: string, userId: string, sessionId: string, - additionalParams: Record = {} + additionalParams: Record = {}, ): Promise { this.executionTimes = new Map(); @@ -460,7 +491,7 @@ export class AgentSquad { const classifierResult = await this.classifyRequest( userInput, userId, - sessionId + sessionId, ); if (!classifierResult.selectedAgent) { @@ -470,7 +501,7 @@ export class AgentSquad { userInput, userId, sessionId, - additionalParams + additionalParams, ), output: this.config.NO_SELECTED_AGENT_MESSAGE!, streaming: false, @@ -482,7 +513,7 @@ export class AgentSquad { userId, sessionId, classifierResult, - additionalParams + additionalParams, ); } catch (error) { return { @@ -491,7 +522,7 @@ export class AgentSquad { userInput, userId, sessionId, - additionalParams + additionalParams, ), output: this.config.GENERAL_ROUTING_ERROR_MSG_MESSAGE || String(error), streaming: false, @@ -507,58 +538,174 @@ export class AgentSquad { userInput: string, userId: string, sessionId: string, - agent: Agent + agent: Agent, ): Promise { const streamStartTime = Date.now(); let chunkCount = 0; + let streamBytes = 0; + let completed = false; + let cancelled = false; + let iteratorClosed = false; + let iteratorClosePromise: Promise | undefined; + let resolveCancellation!: () => void; + const cancellation = new Promise((resolve) => { + resolveCancellation = resolve; + }); + const iterator = agentResponse[Symbol.asyncIterator](); + const closeIterator = () => { + if (iteratorClosed) { + return; + } + iteratorClosed = true; + iteratorClosePromise = Promise.resolve() + .then(() => iterator.return?.()) + .catch((error) => { + this.logger.error("Error closing streaming response:", error); + }); + }; + const handleClose = () => { + if (!completed) { + cancelled = true; + resolveCancellation(); + closeIterator(); + } + }; + + accumulatorTransform.once("close", handleClose); try { - for await (const chunk of agentResponse) { + while (!cancelled) { + const nextResult = Promise.resolve().then(() => iterator.next()); + // A provider may reject after cancellation wins the race; consume that rejection. + void nextResult.catch(() => undefined); + const result = await Promise.race([ + nextResult, + cancellation.then(() => ({ done: true, value: undefined })), + ]); + + if (cancelled || result.done) { + break; + } + + const chunkBytes = this.getChunkByteLength(result.value); + if ( + chunkCount >= this.config.MAX_STREAM_CHUNKS! || + streamBytes + chunkBytes > this.config.MAX_STREAM_BYTES! + ) { + throw new Error("Streaming response exceeded configured limits"); + } + if (chunkCount === 0) { const firstChunkTime = Date.now(); const timeToFirstChunk = firstChunkTime - streamStartTime; this.executionTimes.set("Time to first chunk", timeToFirstChunk); this.logger.printExecutionTimes(this.executionTimes); } - accumulatorTransform.write(chunk); + streamBytes += chunkBytes; + const writeAccepted = await new Promise((resolve, reject) => { + let accepted = false; + accepted = accumulatorTransform.write(result.value, (error) => { + if (error) { + reject(error); + } else { + resolve(accepted); + } + }); + }); + if (!writeAccepted && !cancelled) { + await this.waitForStreamDrainOrCancellation( + accumulatorTransform, + cancellation, + ); + } chunkCount++; } + if (cancelled) { + this.logger.debug( + `\nStreaming cancelled after ${chunkCount} chunks received`, + ); + return; + } + + completed = true; accumulatorTransform.end(); this.logger.debug(`\nStreaming completed: ${chunkCount} chunks received`); const fullResponse = accumulatorTransform.getAccumulatedData(); - if (fullResponse) { - if (agent.saveChat) { - await saveConversationExchange( - userInput, - fullResponse, - this.storage, - userId, - sessionId, - agent.id, - this.config.MAX_MESSAGE_PAIRS_PER_AGENT - ); - } - } else { + if (fullResponse && agent.saveChat) { + await saveConversationExchange( + userInput, + fullResponse, + this.storage, + userId, + sessionId, + agent.id, + this.config.MAX_MESSAGE_PAIRS_PER_AGENT, + ); + } else if (!fullResponse) { this.logger.warn("No data accumulated, messages not saved"); } } catch (error) { + if (cancelled) { + return; + } this.logger.error("Error processing stream:", error); - accumulatorTransform.end(); - if (error instanceof Error) { - accumulatorTransform.destroy(error); - } else if (typeof error === "string") { - accumulatorTransform.destroy(new Error(error)); - } else { - accumulatorTransform.destroy(new Error("An unknown error occurred")); + closeIterator(); + accumulatorTransform.destroy( + error instanceof Error + ? error + : new Error( + typeof error === "string" ? error : "An unknown error occurred", + ), + ); + } finally { + accumulatorTransform.removeListener("close", handleClose); + if (iteratorClosed && iteratorClosePromise) { + await iteratorClosePromise; } } } + private waitForStreamDrainOrCancellation( + accumulatorTransform: AccumulatorTransform, + cancellation: Promise, + ): Promise { + return new Promise((resolve) => { + let settled = false; + const onDrain = () => finish(resolve); + const cleanup = () => { + accumulatorTransform.removeListener("drain", onDrain); + }; + const finish = (callback: () => void) => { + if (settled) { + return; + } + settled = true; + cleanup(); + callback(); + }; + + accumulatorTransform.once("drain", onDrain); + cancellation.then(() => finish(resolve)); + }); + } + + private getChunkByteLength(chunk: unknown): number { + if (typeof chunk === "string") { + return Buffer.byteLength(chunk); + } + + try { + return Buffer.byteLength(JSON.stringify(chunk) ?? ""); + } catch { + return Buffer.byteLength(String(chunk)); + } + } + private measureExecutionTime( timerName: string, - fn: () => Promise | T + fn: () => Promise | T, ): Promise { if (!this.config.LOG_EXECUTION_TIMES) { return Promise.resolve(fn()); @@ -579,7 +726,7 @@ export class AgentSquad { const duration = endTime - startTime; this.executionTimes.set(timerName, duration); throw error; - } + }, ); } @@ -588,7 +735,7 @@ export class AgentSquad { userInput: string, userId: string, sessionId: string, - additionalParams: Record + additionalParams: Record, ): RequestMetadata { const baseMetadata = { userInput, diff --git a/typescript/src/utils/helpers.ts b/typescript/src/utils/helpers.ts index e6d954ad..d1f3fac8 100644 --- a/typescript/src/utils/helpers.ts +++ b/typescript/src/utils/helpers.ts @@ -1,75 +1,92 @@ -import { Transform, TransformCallback } from 'stream'; -import { ConversationMessage, ToolInput } from '../types'; +import { Transform, TransformCallback } from "stream"; +import { ConversationMessage, ToolInput } from "../types"; +export interface AccumulatorTransformOptions { + maxAccumulatedResponseBytes?: number; +} export class AccumulatorTransform extends Transform { - private accumulator: string; + private accumulator: string; + private readonly maxAccumulatedResponseBytes: number; + private accumulatedResponseBytes = 0; - constructor() { - super({ - objectMode: true // This allows the transform to handle object chunks - }); - this.accumulator = ''; - } + constructor(options: AccumulatorTransformOptions = {}) { + super({ + objectMode: true, // This allows the transform to handle object chunks + }); + this.accumulator = ""; + this.maxAccumulatedResponseBytes = + options.maxAccumulatedResponseBytes ?? Infinity; + } - _transform(chunk: any, encoding: string, callback: TransformCallback): void { - // A widget chunk is forwarded to the consumer but never folded into the accumulated text - // answer (which is what gets saved to storage). - if (chunk && typeof chunk === 'object' && chunk.ui) { - this.push(chunk); - callback(); - return; - } - const text = this.extractFromChunk(chunk); - if (text) { - this.accumulator += text; - this.push(text); // Push the text, not the original chunk - } + _transform(chunk: any, encoding: string, callback: TransformCallback): void { + // A widget chunk is forwarded to the consumer but never folded into the accumulated text + // answer (which is what gets saved to storage). + if (chunk && typeof chunk === "object" && chunk.ui) { + this.push(chunk); callback(); + return; } - - extractFromChunk(chunk: any): string | null | any { - if (typeof chunk === 'string') { - return chunk; - } else if (chunk.contentBlockDelta?.delta?.text) { - return chunk.contentBlockDelta.delta.text; - } else if (chunk.thinking) { - return chunk; + const text = this.extractFromChunk(chunk); + if (typeof text === "string" && text.length > 0) { + const textBytes = Buffer.byteLength(text, "utf8"); + if ( + this.accumulatedResponseBytes + textBytes > + this.maxAccumulatedResponseBytes + ) { + callback( + new Error("Maximum accumulated streaming response size exceeded"), + ); + return; } - // Add more conditions here if there are other possible structures - return null; - } - - getAccumulatedData(): string { - return this.accumulator; + this.accumulator += text; + this.accumulatedResponseBytes += textBytes; + this.push(text); // Push the text, not the original chunk } + callback(); } - export function extractXML(text: string) { - const xmlRegex = /[\s\S]*?<\/response>/; - const match = text.match(xmlRegex); - return match ? match[0] : null; + extractFromChunk(chunk: any): string | null | any { + if (typeof chunk === "string") { + return chunk; + } else if (chunk.contentBlockDelta?.delta?.text) { + return chunk.contentBlockDelta.delta.text; + } else if (chunk.thinking) { + return chunk; + } + // Add more conditions here if there are other possible structures + return null; } - - export function isClassifierToolInput(input: unknown): input is ToolInput { - return ( - typeof input === 'object' && - input !== null && - 'userinput' in input && - 'selected_agent' in input && - 'confidence' in input - ); + getAccumulatedData(): string { + return this.accumulator; } +} - export function isConversationMessage(result: any): result is ConversationMessage { - return ( - result && - typeof result === "object" && - "role" in result && - "content" in result && - Array.isArray(result.content) - ); - } +export function extractXML(text: string) { + const xmlRegex = /[\s\S]*?<\/response>/; + const match = text.match(xmlRegex); + return match ? match[0] : null; +} +export function isClassifierToolInput(input: unknown): input is ToolInput { + return ( + typeof input === "object" && + input !== null && + "userinput" in input && + "selected_agent" in input && + "confidence" in input + ); +} +export function isConversationMessage( + result: any, +): result is ConversationMessage { + return ( + result && + typeof result === "object" && + "role" in result && + "content" in result && + Array.isArray(result.content) + ); +} diff --git a/typescript/tests/streamingCancellation.test.ts b/typescript/tests/streamingCancellation.test.ts new file mode 100644 index 00000000..caf377eb --- /dev/null +++ b/typescript/tests/streamingCancellation.test.ts @@ -0,0 +1,232 @@ +import { AgentSquad } from "../src/orchestrator"; +import { Agent } from "../src/agents/agent"; +import { Classifier, ClassifierResult } from "../src/classifiers/classifier"; +import { ChatStorage } from "../src/storage/chatStorage"; +import { ConversationMessage } from "../src/types"; +import { AccumulatorTransform } from "../src/utils/helpers"; + +class TestClassifier extends Classifier { + private selectedAgent: Agent | null = null; + + setSelectedAgent(agent: Agent): void { + this.selectedAgent = agent; + } + + async processRequest( + _inputText: string, + _chatHistory: ConversationMessage[], + ): Promise { + return { selectedAgent: this.selectedAgent, confidence: 1 }; + } +} + +class TestStorage extends ChatStorage { + readonly savedMessages: ConversationMessage[] = []; + + async saveChatMessage( + _userId: string, + _sessionId: string, + _agentId: string, + message: ConversationMessage, + ): Promise { + this.savedMessages.push(message); + return this.savedMessages; + } + + async fetchChat(): Promise { + return []; + } + + async fetchAllChats(): Promise { + return []; + } +} + +function createAgent( + processRequest: Agent["processRequest"], + saveChat = true, +): Agent { + return { + id: "stream-agent", + name: "Stream Agent", + description: "A test streaming agent", + saveChat, + processRequest, + } as Agent; +} + +describe("streaming response lifecycle", () => { + it("stops requesting chunks and closes the provider iterator when the output closes", async () => { + let nextCalls = 0; + let iteratorReturned = false; + const pendingNext = new Promise>(() => undefined); + const stream: AsyncIterable = { + [Symbol.asyncIterator]: () => ({ + next: async () => { + nextCalls += 1; + if (nextCalls === 1) { + return { done: false, value: "first" }; + } + return pendingNext; + }, + return: async () => { + iteratorReturned = true; + return { done: true, value: undefined }; + }, + }), + }; + const agent = createAgent(async () => stream, false); + const classifier = new TestClassifier(); + classifier.setSelectedAgent(agent); + const orchestrator = new AgentSquad({ + classifier, + storage: new TestStorage(), + }); + orchestrator.addAgent(agent); + + const response = await orchestrator.routeRequest( + "input", + "user", + "session", + ); + const output = response.output as AccumulatorTransform; + output.on("error", () => undefined); + await new Promise((resolve) => output.once("data", () => resolve())); + output.destroy(); + + await new Promise((resolve) => setImmediate(resolve)); + expect(iteratorReturned).toBe(true); + expect(nextCalls).toBe(1); + }); + + it("saves a finite stream after delivering all chunks", async () => { + const storage = new TestStorage(); + const agent = createAgent(async () => + (async function* () { + yield "hello "; + yield "world"; + })(), + ); + const classifier = new TestClassifier(); + classifier.setSelectedAgent(agent); + const orchestrator = new AgentSquad({ classifier, storage }); + orchestrator.addAgent(agent); + + const response = await orchestrator.routeRequest( + "input", + "user", + "session", + ); + const chunks: string[] = []; + for await (const chunk of response.output as AccumulatorTransform) { + chunks.push(chunk as string); + } + + expect(chunks.join("")).toBe("hello world"); + expect(storage.savedMessages).toHaveLength(2); + expect(storage.savedMessages[1].content[0].text).toBe("hello world"); + }); + + it("rejects a stream that exceeds the configured byte limit", async () => { + const storage = new TestStorage(); + const agent = createAgent(async () => + (async function* () { + yield "12345"; + yield "6"; + })(), + ); + const classifier = new TestClassifier(); + classifier.setSelectedAgent(agent); + const orchestrator = new AgentSquad({ + classifier, + storage, + config: { MAX_STREAM_BYTES: 5 }, + }); + orchestrator.addAgent(agent); + + const response = await orchestrator.routeRequest( + "input", + "user", + "session", + ); + const output = response.output as AccumulatorTransform; + const streamError = new Promise((resolve) => + output.once("error", resolve), + ); + output.resume(); + + await expect(streamError).resolves.toMatchObject({ + message: "Streaming response exceeded configured limits", + }); + expect(storage.savedMessages).toHaveLength(0); + }); + + it("does not save an accumulated response that exceeds its text limit", async () => { + const storage = new TestStorage(); + const agent = createAgent(async () => + (async function* () { + yield "12345"; + yield "6"; + })(), + ); + const classifier = new TestClassifier(); + classifier.setSelectedAgent(agent); + const orchestrator = new AgentSquad({ + classifier, + storage, + config: { + MAX_STREAM_BYTES: 100, + MAX_ACCUMULATED_RESPONSE_BYTES: 5, + }, + }); + orchestrator.addAgent(agent); + + const response = await orchestrator.routeRequest( + "input", + "user", + "session", + ); + const output = response.output as AccumulatorTransform; + const streamError = new Promise((resolve) => + output.once("error", resolve), + ); + output.resume(); + + await expect(streamError).resolves.toMatchObject({ + message: "Maximum accumulated streaming response size exceeded", + }); + expect(storage.savedMessages).toHaveLength(0); + }); + + it("forwards provider errors and closes the iterator", async () => { + const storage = new TestStorage(); + const providerError = new Error("provider failed"); + const iterator = { + next: jest.fn().mockRejectedValue(providerError), + return: jest.fn().mockResolvedValue({ done: true, value: undefined }), + }; + const stream: AsyncIterable = { + [Symbol.asyncIterator]: () => iterator, + }; + const agent = createAgent(async () => stream); + const classifier = new TestClassifier(); + classifier.setSelectedAgent(agent); + const orchestrator = new AgentSquad({ classifier, storage }); + orchestrator.addAgent(agent); + + const response = await orchestrator.routeRequest( + "input", + "user", + "session", + ); + const output = response.output as AccumulatorTransform; + const streamError = new Promise((resolve) => + output.once("error", resolve), + ); + output.resume(); + + await expect(streamError).resolves.toBe(providerError); + expect(iterator.return).toHaveBeenCalledTimes(1); + expect(storage.savedMessages).toHaveLength(0); + }); +}); diff --git a/typescript/tests/utils/helpers.test.ts b/typescript/tests/utils/helpers.test.ts index c12c888b..75829eaf 100644 --- a/typescript/tests/utils/helpers.test.ts +++ b/typescript/tests/utils/helpers.test.ts @@ -5,10 +5,14 @@ describe("AccumulatorTransform", () => { const transform = new AccumulatorTransform(); const out: any[] = []; transform.on("data", (c) => out.push(c)); - const done = new Promise((resolve) => transform.on("end", () => resolve())); + const done = new Promise((resolve) => + transform.on("end", () => resolve()), + ); transform.write("Hello "); - transform.write({ ui: { resourceUri: "ui://x", mimeType: "text/html;profile=mcp-app" } }); + transform.write({ + ui: { resourceUri: "ui://x", mimeType: "text/html;profile=mcp-app" }, + }); transform.write("world"); transform.end(); await done; @@ -19,14 +23,18 @@ describe("AccumulatorTransform", () => { const widget = out.find((c) => c && typeof c === "object" && c.ui); expect(widget.ui.resourceUri).toBe("ui://x"); // ...alongside the text chunks. - expect(out.filter((c) => typeof c === "string").join("")).toBe("Hello world"); + expect(out.filter((c) => typeof c === "string").join("")).toBe( + "Hello world", + ); }); it("accumulates and forwards plain text chunks unchanged", async () => { const transform = new AccumulatorTransform(); const out: string[] = []; transform.on("data", (c) => out.push(c)); - const done = new Promise((resolve) => transform.on("end", () => resolve())); + const done = new Promise((resolve) => + transform.on("end", () => resolve()), + ); transform.write("a"); transform.write("b"); @@ -41,7 +49,9 @@ describe("AccumulatorTransform", () => { const transform = new AccumulatorTransform(); const out: any[] = []; transform.on("data", (c) => out.push(c)); - const done = new Promise((resolve) => transform.on("end", () => resolve())); + const done = new Promise((resolve) => + transform.on("end", () => resolve()), + ); transform.write({ ui: undefined }); // falsy ui → text path → dropped like any unknown chunk transform.write("text"); From 972e6f87d3ca76c1c0656f3a2a3d8d67ae1be4e6 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Fri, 9 Oct 2026 06:50:57 +0800 Subject: [PATCH 2/2] fix(typescript): preserve streamed reasoning chunks --- typescript/src/utils/helpers.ts | 5 ++--- typescript/tests/utils/helpers.test.ts | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/typescript/src/utils/helpers.ts b/typescript/src/utils/helpers.ts index d1f3fac8..00c77f0b 100644 --- a/typescript/src/utils/helpers.ts +++ b/typescript/src/utils/helpers.ts @@ -20,9 +20,8 @@ export class AccumulatorTransform extends Transform { } _transform(chunk: any, encoding: string, callback: TransformCallback): void { - // A widget chunk is forwarded to the consumer but never folded into the accumulated text - // answer (which is what gets saved to storage). - if (chunk && typeof chunk === "object" && chunk.ui) { + // Widget and reasoning chunks are forwarded but not folded into the saved text answer. + if (chunk && typeof chunk === "object" && (chunk.ui || chunk.thinking)) { this.push(chunk); callback(); return; diff --git a/typescript/tests/utils/helpers.test.ts b/typescript/tests/utils/helpers.test.ts index 75829eaf..994e98f4 100644 --- a/typescript/tests/utils/helpers.test.ts +++ b/typescript/tests/utils/helpers.test.ts @@ -28,6 +28,25 @@ describe("AccumulatorTransform", () => { ); }); + it("forwards a thinking chunk without folding it into the accumulated text", async () => { + const transform = new AccumulatorTransform(); + const out: unknown[] = []; + transform.on("data", (chunk) => out.push(chunk)); + const done = new Promise((resolve) => + transform.on("end", () => resolve()), + ); + const thinkingChunk = { thinking: true, content: "reasoning" }; + + transform.write("answer "); + transform.write(thinkingChunk); + transform.write("continues"); + transform.end(); + await done; + + expect(out).toEqual(["answer ", thinkingChunk, "continues"]); + expect(transform.getAccumulatedData()).toBe("answer continues"); + }); + it("accumulates and forwards plain text chunks unchanged", async () => { const transform = new AccumulatorTransform(); const out: string[] = [];