From 80b6b46ca967bfc13cda89df8368f124158910e1 Mon Sep 17 00:00:00 2001
From: mikemikimike <13286568797@163.com>
Date: Tue, 18 Aug 2026 07:10:20 +0800
Subject: [PATCH] fix(workers-ai-provider): honor abort signals for binding
 streams

---
 .changeset/calm-streams-abort.md              |  5 ++
 packages/workers-ai-provider/src/streaming.ts | 46 ++++++++++++++++++-
 .../src/workersai-chat-language-model.ts      |  2 +-
 .../test/streaming.test.ts                    | 25 ++++++++++
 4 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 .changeset/calm-streams-abort.md
 create mode 100644 packages/workers-ai-provider/test/streaming.test.ts

diff --git a/.changeset/calm-streams-abort.md b/.changeset/calm-streams-abort.md
new file mode 100644
index 0000000000..9fc0b725e5
--- /dev/null
+++ b/.changeset/calm-streams-abort.md
@@ -0,0 +1,5 @@
+---
+"workers-ai-provider": patch
+---
+
+Honor abort signals while reading streaming responses returned by Workers AI bindings.
diff --git a/packages/workers-ai-provider/src/streaming.ts b/packages/workers-ai-provider/src/streaming.ts
index 114d54f70a..7501b87e60 100644
--- a/packages/workers-ai-provider/src/streaming.ts
+++ b/packages/workers-ai-provider/src/streaming.ts
@@ -74,6 +74,7 @@ export function getMappedStream(
 		tools: Array<{ function: { name?: string } }> | undefined;
 		toolChoice: unknown;
 	},
+	signal?: AbortSignal,
 ): ReadableStream<LanguageModelV4StreamPart> {
 	const rawStream =
 		response instanceof ReadableStream
@@ -84,6 +85,8 @@ export function getMappedStream(
 		throw new Error("No readable stream available for SSE parsing.");
 	}
 
+	const stream = signal ? raceAbort(rawStream, signal) : rawStream;
+
 	// gpt-oss harmony quirk: a forced tool call can be streamed as `content`
 	// text deltas instead of structured tool calls. When a tool was forced,
 	// buffer the text content (rather than emitting it incrementally) so we can
@@ -123,7 +126,7 @@ export function getMappedStream(
 	let lastActiveToolIndex: number | null = null;
 
 	// Step 1: Decode bytes into SSE lines
-	const sseStream = rawStream.pipeThrough(new SSEDecoder());
+	const sseStream = stream.pipeThrough(new SSEDecoder());
 
 	// Step 2: Transform SSE events into LanguageModelV4StreamPart
 	return sseStream.pipeThrough(
@@ -427,3 +430,44 @@ export function getMappedStream(
 		}
 	}
 }
+
+/** Make pending reads from binding streams observe the SDK abort signal. */
+function raceAbort(
+	stream: ReadableStream<Uint8Array>,
+	signal: AbortSignal,
+): ReadableStream<Uint8Array> {
+	const reader = stream.getReader();
+	let abortHandler: (() => void) | undefined;
+
+	const abortPromise = new Promise<never>((_, reject) => {
+		abortHandler = () => {
+			reject(signal.reason ?? new DOMException("Aborted", "AbortError"));
+		};
+		if (signal.aborted) abortHandler();
+		else signal.addEventListener("abort", abortHandler, { once: true });
+	});
+	abortPromise.catch(() => {});
+	const cleanup = () => {
+		if (abortHandler) signal.removeEventListener("abort", abortHandler);
+	};
+
+	return new ReadableStream<Uint8Array>({
+		async pull(controller) {
+			try {
+				const result = await Promise.race([reader.read(), abortPromise]);
+				if (result.done) {
+					cleanup();
+					controller.close();
+				} else controller.enqueue(result.value);
+			} catch (error) {
+				cleanup();
+				await reader.cancel(error).catch(() => {});
+				controller.error(error);
+			}
+		},
+		cancel(reason) {
+			cleanup();
+			return reader.cancel(reason);
+		},
+	});
+}
diff --git a/packages/workers-ai-provider/src/workersai-chat-language-model.ts b/packages/workers-ai-provider/src/workersai-chat-language-model.ts
index 38f7ce8f74..1a0c225878 100644
--- a/packages/workers-ai-provider/src/workersai-chat-language-model.ts
+++ b/packages/workers-ai-provider/src/workersai-chat-language-model.ts
@@ -412,7 +412,7 @@ export class WorkersAIChatLanguageModel implements LanguageModelV4 {
 					getMappedStream(response, {
 						tools: args.tools,
 						toolChoice: args.tool_choice,
-					}),
+					}, options.abortSignal),
 					warnings,
 				),
 			};
diff --git a/packages/workers-ai-provider/test/streaming.test.ts b/packages/workers-ai-provider/test/streaming.test.ts
new file mode 100644
index 0000000000..a55b59d672
--- /dev/null
+++ b/packages/workers-ai-provider/test/streaming.test.ts
@@ -0,0 +1,25 @@
+import { describe, expect, it } from "vitest";
+import { getMappedStream } from "../src/streaming";
+
+describe("getMappedStream abort handling", () => {
+	it("rejects a pending binding read with the abort reason", async () => {
+		const abortController = new AbortController();
+		const abortReason = new Error("timed out");
+		let cancelled = false;
+		const source = new ReadableStream<Uint8Array>({
+			pull() {
+				return new Promise(() => {});
+			},
+			cancel() {
+				cancelled = true;
+			},
+		});
+
+		const reader = getMappedStream(source, undefined, abortController.signal).getReader();
+		const pendingRead = reader.read();
+		abortController.abort(abortReason);
+
+		await expect(pendingRead).rejects.toBe(abortReason);
+		expect(cancelled).toBe(true);
+	});
+});