assistant-stream 0.3.25 → 0.3.27

package/src/core/serialization/ui-message-stream/UIMessageStream.test.ts

@@ -35,6 +35,34 @@ function createUIMessageStream(events: string[]): ReadableStream<Uint8Array> {
 
 describe("UIMessageStreamDecoder", () => {
   it("should decode text deltas", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "text-start", id: "text_1" }),
+      JSON.stringify({ type: "text-delta", id: "text_1", delta: "Hello" }),
+      JSON.stringify({ type: "text-delta", id: "text_1", delta: " world" }),
+      JSON.stringify({ type: "text-end" }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { inputTokens: 10, outputTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const textDeltas = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "text-delta" } =>
+        c.type === "text-delta",
+    );
+    expect(textDeltas).toHaveLength(2);
+    expect(textDeltas[0]?.textDelta).toBe("Hello");
+    expect(textDeltas[1]?.textDelta).toBe(" world");
+  });
+
+  it("should decode legacy textDelta text deltas", async () => {
     const events = [
       JSON.stringify({ type: "start", messageId: "msg_123" }),
       JSON.stringify({ type: "text-start", id: "text_1" }),
@@ -138,7 +166,50 @@ describe("UIMessageStreamDecoder", () => {
     expect(result?.result).toEqual({ temp: 72 });
   });
 
-  it("should decode source parts", async () => {
+  it("should decode source-url parts", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "source-url",
+        sourceId: "src_1",
+        url: "https://example.com",
+        title: "Example",
+      }),
+      JSON.stringify({
+        type: "source-url",
+        sourceId: "src_2",
+        url: "https://example.org",
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { inputTokens: 10, outputTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const sourceStarts = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "part-start" } =>
+        c.type === "part-start" && c.part.type === "source",
+    );
+    expect(sourceStarts).toHaveLength(2);
+    if (sourceStarts[0]?.part.type === "source") {
+      expect(sourceStarts[0].part.id).toBe("src_1");
+      expect(sourceStarts[0].part.url).toBe("https://example.com");
+      expect(sourceStarts[0].part.title).toBe("Example");
+    }
+    if (sourceStarts[1]?.part.type === "source") {
+      expect(sourceStarts[1].part.id).toBe("src_2");
+      expect(sourceStarts[1].part.url).toBe("https://example.org");
+      expect(sourceStarts[1].part.title).toBeUndefined();
+    }
+  });
+
+  it("should decode legacy source parts", async () => {
     const events = [
       JSON.stringify({ type: "start", messageId: "msg_123" }),
       JSON.stringify({
@@ -168,12 +239,44 @@ describe("UIMessageStreamDecoder", () => {
     );
     expect(sourceStart).toBeDefined();
     if (sourceStart?.part.type === "source") {
+      expect(sourceStart.part.id).toBe("src_1");
       expect(sourceStart.part.url).toBe("https://example.com");
       expect(sourceStart.part.title).toBe("Example");
     }
   });
 
   it("should decode file parts", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "file",
+        url: "https://example.com/image.png",
+        mediaType: "image/png",
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { inputTokens: 10, outputTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const fileStart = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "part-start" } =>
+        c.type === "part-start" && c.part.type === "file",
+    );
+    expect(fileStart).toBeDefined();
+    if (fileStart?.part.type === "file") {
+      expect(fileStart.part.mimeType).toBe("image/png");
+      expect(fileStart.part.data).toBe("https://example.com/image.png");
+    }
+  });
+
+  it("should decode legacy file parts", async () => {
     const events = [
       JSON.stringify({ type: "start", messageId: "msg_123" }),
       JSON.stringify({
@@ -277,7 +380,7 @@ describe("UIMessageStreamDecoder", () => {
       JSON.stringify({ type: "start", messageId: "msg_123" }),
       JSON.stringify({ type: "start-step", messageId: "step_1" }),
       JSON.stringify({ type: "text-start", id: "text_1" }),
-      JSON.stringify({ type: "text-delta", textDelta: "Hello" }),
+      JSON.stringify({ type: "text-delta", id: "text_1", delta: "Hello" }),
       JSON.stringify({ type: "text-end" }),
       JSON.stringify({
         type: "finish-step",
@@ -308,6 +411,96 @@ describe("UIMessageStreamDecoder", () => {
     );
     expect(stepFinish).toBeDefined();
     expect(stepFinish?.finishReason).toBe("stop");
+    expect(stepFinish?.usage).toEqual({ inputTokens: 10, outputTokens: 5 });
+    expect(stepFinish?.isContinued).toBe(false);
+  });
+
+  it("should handle the stock v6 lifecycle with bare step chunks", async () => {
+    const events = [
+      JSON.stringify({ type: "start" }),
+      JSON.stringify({ type: "start-step" }),
+      JSON.stringify({ type: "text-start", id: "text_1" }),
+      JSON.stringify({ type: "text-delta", id: "text_1", delta: "Hello" }),
+      JSON.stringify({ type: "text-end" }),
+      JSON.stringify({ type: "finish-step" }),
+      JSON.stringify({ type: "finish", finishReason: "stop" }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const stepStarts = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "step-start" } =>
+        c.type === "step-start",
+    );
+    expect(stepStarts).toHaveLength(2);
+    expect(stepStarts[0]?.messageId).toBeTruthy();
+    expect(stepStarts[1]?.messageId).toBe(stepStarts[0]?.messageId);
+
+    const stepFinish = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "step-finish" } =>
+        c.type === "step-finish",
+    );
+    expect(stepFinish?.finishReason).toBe("unknown");
+    expect(stepFinish?.usage).toEqual({ inputTokens: 0, outputTokens: 0 });
+    expect(stepFinish?.isContinued).toBe(false);
+
+    const messageFinish = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "message-finish" } =>
+        c.type === "message-finish",
+    );
+    expect(messageFinish?.finishReason).toBe("stop");
+  });
+
+  it("should default finishReason on a bare finish chunk", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "text-start", id: "text_1" }),
+      JSON.stringify({ type: "text-delta", id: "text_1", delta: "Hello" }),
+      JSON.stringify({ type: "text-end" }),
+      JSON.stringify({ type: "finish" }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const messageFinish = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "message-finish" } =>
+        c.type === "message-finish",
+    );
+    expect(messageFinish?.finishReason).toBe("unknown");
+    expect(messageFinish?.usage).toEqual({ inputTokens: 0, outputTokens: 0 });
+  });
+
+  it("should ignore malformed legacy source and file chunks", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "source" }),
+      JSON.stringify({ type: "source", source: null }),
+      JSON.stringify({ type: "file" }),
+      JSON.stringify({ type: "file", file: null }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { inputTokens: 10, outputTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const partStarts = chunks.filter(
+      (c) =>
+        c.type === "part-start" &&
+        (c.part.type === "source" || c.part.type === "file"),
+    );
+    expect(partStarts).toHaveLength(0);
   });
 
   it("should handle errors", async () => {
@@ -331,8 +524,8 @@ describe("UIMessageStreamDecoder", () => {
   it("should throw when stream ends without [DONE]", async () => {
     const encoder = new TextEncoder();
     const sseText =
-      'data: {"type":"text-delta","textDelta":"Hello"}\n\n' +
-      'data: {"type":"text-delta","textDelta":" world"}\n\n';
+      'data: {"type":"text-delta","id":"text_1","delta":"Hello"}\n\n' +
+      'data: {"type":"text-delta","id":"text_1","delta":" world"}\n\n';
 
     const stream = new ReadableStream<Uint8Array>({
       start(controller) {
@@ -348,6 +541,21 @@ describe("UIMessageStreamDecoder", () => {
     );
   });
 
+  it("should discard an unterminated [DONE] event", async () => {
+    const stream = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(new TextEncoder().encode("data: [DONE]\n"));
+        controller.close();
+      },
+    });
+
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+
+    await expect(collectChunks(decodedStream)).rejects.toThrow(
+      "Stream ended abruptly without receiving [DONE] marker",
+    );
+  });
+
   it("should ignore unknown chunk types for forward compatibility", async () => {
     const events = [
       JSON.stringify({ type: "start", messageId: "msg_123" }),
@@ -367,4 +575,115 @@ describe("UIMessageStreamDecoder", () => {
     // Should not throw, should complete successfully
     expect(chunks.some((c) => c.type === "message-finish")).toBe(true);
   });
+
+  it("drops frames that are not objects with a string type", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const invalidFrames = [
+      "null",
+      "5",
+      '"text"',
+      "[1,2]",
+      '{"foo":1}',
+      '{"type":5}',
+    ];
+    const events = [
+      ...invalidFrames,
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "text-delta", id: "t", delta: "ok" }),
+      "[DONE]",
+    ];
+
+    const chunks = await collectChunks(
+      createUIMessageStream(events).pipeThrough(new UIMessageStreamDecoder()),
+    );
+
+    expect(warn.mock.calls.map((c) => c[0])).toEqual(
+      invalidFrames.map((f) => `Dropped invalid UIMessageStream chunk: ${f}`),
+    );
+    expect(chunks.some((c) => c.type === "text-delta")).toBe(true);
+    warn.mockRestore();
+  });
+
+  it("drops frames that are not valid JSON", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const invalidFrames = ['{"type":"te', "not json"];
+    const events = [
+      ...invalidFrames,
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "text-delta", id: "t", delta: "ok" }),
+      "[DONE]",
+    ];
+
+    const chunks = await collectChunks(
+      createUIMessageStream(events).pipeThrough(new UIMessageStreamDecoder()),
+    );
+
+    expect(warn.mock.calls.map((c) => c[0])).toEqual(
+      invalidFrames.map((f) => `Dropped invalid UIMessageStream chunk: ${f}`),
+    );
+    expect(chunks.some((c) => c.type === "text-delta")).toBe(true);
+    warn.mockRestore();
+  });
+
+  it("ignores a tool-call-delta arriving after tool-result", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "tool-call-start",
+        toolCallId: "call_abc",
+        toolName: "weather",
+      }),
+      JSON.stringify({ type: "tool-call-delta", argsText: '{"city":"NYC"}' }),
+      JSON.stringify({
+        type: "tool-result",
+        toolCallId: "call_abc",
+        result: { temp: 72 },
+      }),
+      JSON.stringify({ type: "tool-call-delta", argsText: '{"late":true}' }),
+      "[DONE]",
+    ];
+
+    const chunks = await collectChunks(
+      createUIMessageStream(events).pipeThrough(new UIMessageStreamDecoder()),
+    );
+
+    const argDeltas = chunks.filter((c) => c.type === "text-delta");
+    expect(argDeltas).toHaveLength(1);
+    expect(chunks.some((c) => c.type === "result")).toBe(true);
+  });
+
+  it("keeps the active tool call writable when another call receives its result", async () => {
+    const events = [
+      JSON.stringify({
+        type: "tool-call-start",
+        toolCallId: "call_a",
+        toolName: "first",
+      }),
+      JSON.stringify({
+        type: "tool-call-start",
+        toolCallId: "call_b",
+        toolName: "second",
+      }),
+      JSON.stringify({
+        type: "tool-result",
+        toolCallId: "call_a",
+        result: { ok: true },
+      }),
+      JSON.stringify({ type: "tool-call-delta", argsText: '{"x":1}' }),
+      "[DONE]",
+    ];
+
+    const chunks = await collectChunks(
+      createUIMessageStream(events).pipeThrough(new UIMessageStreamDecoder()),
+    );
+
+    const argDeltas = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "text-delta" } =>
+        c.type === "text-delta",
+    );
+    expect(argDeltas).toHaveLength(2);
+    const byPath = new Map(argDeltas.map((c) => [c.path[0], c.textDelta]));
+    expect(byPath.get(0)).toBe("{}");
+    expect(byPath.get(1)).toBe('{"x":1}');
+  });
 });

package/src/core/serialization/ui-message-stream/UIMessageStream.ts

@@ -3,7 +3,10 @@ import type { ToolCallStreamController } from "../../modules/tool-call";
 import type { TextStreamController } from "../../modules/text";
 import { AssistantTransformStream } from "../../utils/stream/AssistantTransformStream";
 import { PipeableTransformStream } from "../../utils/stream/PipeableTransformStream";
-import { LineDecoderStream } from "../../utils/stream/LineDecoderStream";
+import {
+  SSEEventDecoderStream,
+  type PipelineSSEEvent,
+} from "../../utils/stream/SSEEventDecoderStream";
 import type {
   UIMessageStreamChunk,
   UIMessageStreamDataChunk,
@@ -21,72 +24,6 @@ export type UIMessageStreamDecoderOptions = {
   }) => void;
 };
 
-type SSEEvent = {
-  event: string;
-  data: string;
-  id?: string | undefined;
-  retry?: number | undefined;
-};
-
-class SSEEventStream extends TransformStream<string, SSEEvent> {
-  constructor() {
-    let eventBuffer: Partial<SSEEvent> = {};
-    let dataLines: string[] = [];
-
-    super({
-      start() {
-        eventBuffer = {};
-        dataLines = [];
-      },
-      transform(line, controller) {
-        if (line.startsWith(":")) return;
-
-        if (line === "") {
-          if (dataLines.length > 0) {
-            controller.enqueue({
-              event: eventBuffer.event || "message",
-              data: dataLines.join("\n"),
-              id: eventBuffer.id,
-              retry: eventBuffer.retry,
-            });
-          }
-          eventBuffer = {};
-          dataLines = [];
-          return;
-        }
-
-        const [field, ...rest] = line.split(":");
-        const value = rest.join(":").trimStart();
-
-        switch (field) {
-          case "event":
-            eventBuffer.event = value;
-            break;
-          case "data":
-            dataLines.push(value);
-            break;
-          case "id":
-            eventBuffer.id = value;
-            break;
-          case "retry":
-            eventBuffer.retry = Number(value);
-            break;
-        }
-      },
-      flush(controller) {
-        if (dataLines.length > 0) {
-          controller.enqueue({
-            event: eventBuffer.event || "message",
-            data: dataLines.join("\n"),
-            id: eventBuffer.id,
-            retry: eventBuffer.retry,
-          });
-        }
-      },
-    });
-  }
-}
-
 const isDataChunk = (
   chunk: UIMessageStreamChunk,
 ): chunk is UIMessageStreamDataChunk => chunk.type.startsWith("data-");
@@ -212,6 +149,9 @@ export class UIMessageStreamDecoder extends PipeableTransformStream<
                   `Encountered tool result with unknown id: ${chunk.toolCallId}`,
                 );
               }
+              if (toolCallController.argsText === activeToolCallArgsText) {
+                activeToolCallArgsText = undefined;
+              }
               toolCallController.setResponse({
                 result: chunk.result,
                 isError: chunk.isError ?? false,
@@ -271,10 +211,9 @@ export class UIMessageStreamDecoder extends PipeableTransformStream<
 
       return readable
         .pipeThrough(new TextDecoderStream())
-        .pipeThrough(new LineDecoderStream())
-        .pipeThrough(new SSEEventStream())
+        .pipeThrough(new SSEEventDecoderStream())
         .pipeThrough(
-          new TransformStream<SSEEvent, UIMessageStreamChunk>({
+          new TransformStream<PipelineSSEEvent, UIMessageStreamChunk>({
             transform(event, controller) {
               if (event.event !== "message") {
                 throw new Error(`Unknown SSE event type: ${event.event}`);
@@ -286,7 +225,78 @@ export class UIMessageStreamDecoder extends PipeableTransformStream<
                 return;
               }
 
-              controller.enqueue(JSON.parse(event.data));
+              let chunk;
+              try {
+                chunk = JSON.parse(event.data);
+              } catch {
+                chunk = undefined;
+              }
+              if (
+                typeof chunk !== "object" ||
+                chunk === null ||
+                Array.isArray(chunk) ||
+                typeof chunk.type !== "string"
+              ) {
+                console.warn(
+                  `Dropped invalid UIMessageStream chunk: ${event.data.slice(0, 200)}`,
+                );
+                return;
+              }
+              if (
+                chunk.type === "text-delta" &&
+                chunk.textDelta === undefined
+              ) {
+                const { delta, ...rest } = chunk;
+                controller.enqueue({ ...rest, textDelta: delta ?? "" });
+                return;
+              }
+              if (chunk.type === "start") {
+                controller.enqueue({
+                  ...chunk,
+                  messageId: chunk.messageId ?? generateId(),
+                });
+                return;
+              }
+              if (chunk.type === "source-url") {
+                controller.enqueue({
+                  type: "source",
+                  source: {
+                    sourceType: "url",
+                    id: chunk.sourceId,
+                    url: chunk.url,
+                    ...(chunk.title && { title: chunk.title }),
+                  },
+                });
+                return;
+              }
+              if (chunk.type === "source" && chunk.source == null) return;
+              if (chunk.type === "file" && chunk.file == null) {
+                if (chunk.url === undefined) return;
+                controller.enqueue({
+                  type: "file",
+                  file: { mimeType: chunk.mediaType, data: chunk.url },
+                });
+                return;
+              }
+              if (chunk.type === "finish-step") {
+                controller.enqueue({
+                  ...chunk,
+                  finishReason: chunk.finishReason ?? "unknown",
+                  usage: chunk.usage ?? { inputTokens: 0, outputTokens: 0 },
+                  isContinued: chunk.isContinued ?? false,
+                });
+                return;
+              }
+              if (chunk.type === "finish") {
+                controller.enqueue({
+                  ...chunk,
+                  finishReason: chunk.finishReason ?? "unknown",
+                  usage: chunk.usage ?? { inputTokens: 0, outputTokens: 0 },
+                });
+                return;
+              }
+
+              controller.enqueue(chunk);
             },
             flush() {
               if (!receivedDone) {

package/src/core/tool/ToolExecutionStream.ts

@@ -51,6 +51,7 @@ export class ToolExecutionStream extends PipeableTransformStream<
       string,
       ToolCallReaderImpl<ReadonlyJSONObject, ReadonlyJSONValue>
     >();
+    const toolCallIdsWithBackendResult = new Set<string>();
 
     super((readable) => {
       const transform = new TransformStream<
@@ -109,6 +110,7 @@ export class ToolExecutionStream extends PipeableTransformStream<
                   modelContent: chunk.modelContent,
                 }),
               );
+              toolCallIdsWithBackendResult.add(toolCallId);
               break;
             }
             case "tool-call-args-text-finish": {
@@ -123,6 +125,11 @@ export class ToolExecutionStream extends PipeableTransformStream<
               // resolve. Awaited so the close settles before the writer is reused.
               await streamController.finishArgsText();
 
+              // A backend result is authoritative. Closing the args stream still
+              // emits this finish chunk, but must not parse stale/incomplete args,
+              // execute the frontend tool, or enqueue a second result.
+              if (toolCallIdsWithBackendResult.has(toolCallId)) break;
+
               let isExecuting = false;
               const promise = withPromiseOrValue(
                 () => {
@@ -206,10 +213,13 @@ export class ToolExecutionStream extends PipeableTransformStream<
                 toolCallPromise.then(() => {
                   toolCallPromises.delete(toolCallId);
                   toolCallControllers.delete(toolCallId);
+                  toolCallIdsWithBackendResult.delete(toolCallId);
 
                   controller.enqueue(chunk);
                 });
               } else {
+                toolCallControllers.delete(toolCallId);
+                toolCallIdsWithBackendResult.delete(toolCallId);
                 controller.enqueue(chunk);
               }
             }

package/src/core/tool/toolResultStream.test.ts

@@ -17,6 +17,56 @@ const createDelayedTool = (delay: number, result?: string): Tool => ({
 });
 
 describe("unstable_runPendingTools", () => {
+  it("removes the abort listener after tool execution settles", async () => {
+    const abortController = new AbortController();
+    const addEventListener = vi.spyOn(
+      abortController.signal,
+      "addEventListener",
+    );
+    const removeEventListener = vi.spyOn(
+      abortController.signal,
+      "removeEventListener",
+    );
+    const message: AssistantMessage = {
+      role: "assistant",
+      status: { type: "requires-action", reason: "tool-calls" },
+      parts: [
+        {
+          type: "tool-call",
+          toolCallId: "1",
+          toolName: "tool",
+          args: {},
+        } as ToolCallPart,
+      ],
+      content: [],
+      metadata: {
+        unstable_state: {},
+        unstable_data: [],
+        unstable_annotations: [],
+        steps: [],
+        custom: {},
+      },
+    };
+
+    await unstable_runPendingTools(
+      message,
+      {
+        tool: {
+          parameters: { type: "object", properties: {} },
+          execute: async () => "done",
+        },
+      },
+      abortController.signal,
+      async () => {},
+    );
+
+    expect(addEventListener).toHaveBeenCalledOnce();
+    expect(removeEventListener).toHaveBeenCalledWith(
+      "abort",
+      addEventListener.mock.calls[0]![1],
+    );
+  });
+
   describe("parallel execution", () => {
     it("should run tool calls in parallel", async () => {
       const tool1 = createDelayedTool(100, "Tool 1");

package/src/core/tool/toolResultStream.ts

@@ -59,9 +59,10 @@ function getToolResponse(
 
     // Create abort promise that resolves after 2 microtasks
     // This gives tools that handle abort a chance to win the race
+    let onAbort!: () => void;
     const abortPromise = new Promise<ToolResponse<ReadonlyJSONValue>>(
       (resolve) => {
-        const onAbort = () => {
+        onAbort = () => {
           queueMicrotask(() => {
             queueMicrotask(() => {
               resolve(
@@ -116,7 +117,11 @@ function getToolResponse(
       return response;
     })();
 
-    return Promise.race([executePromise, abortPromise]);
+    try {
+      return await Promise.race([executePromise, abortPromise]);
+    } finally {
+      abortSignal.removeEventListener("abort", onAbort);
+    }
   };
 
   return getResult(tool.execute);