ai-resumable-stream 1.1.1 → 1.2.1

package/README.md

@@ -38,7 +38,7 @@ sequenceDiagram
         Server->>Server: streamText()
         Server->>Redis: Subscribe to stop channel
         Server->>Redis: Subscribe to stream channel
-        Server->>Redis: Store chunks
+        Server->>Redis: Store chunks (UI-to-SSE conversion)
         Server-->>Client: Stream chunks
     end
 
@@ -47,9 +47,9 @@ sequenceDiagram
         Client->>Server: resumeMessage()
         Server->>Redis: Subscribe to stream channel
         Redis-->>Server: Replay stored chunks
-        Server-->>Client: Stream past chunks
+        Server-->>Client: Stream past chunks (SSE-to-UI conversion)
         Redis-->>Server: Receive new chunks
-        Server-->>Client: Stream new chunks
+        Server-->>Client: Stream new chunks (SSE-to-UI conversion)
     end
 
     rect rgb(255, 248, 240)
@@ -333,7 +333,7 @@ async function startStream(
 ): Promise<AsyncIterableStream<UIMessageChunk>>;
 ```
 
-Starts a new resumable stream. The input stream is tee'd—one branch goes to the client, the other is persisted to Redis.
+Starts a new resumable stream. A single drain loop reads from the source and sends chunks to both the client and Redis simultaneously. If the client disconnects, chunks continue flowing to Redis for resumability.
 
 #### `resumeStream`
 

package/dist/index.d.mts

@@ -1,5 +1,5 @@
-import { AsyncIterableStream } from "ai-stream-utils/utils";
 import { UIMessageChunk } from "ai";
+import { AsyncIterableStream } from "ai-stream-utils";
 import { createClient } from "redis";
 
 //#region src/resumable-ui-message-stream.d.ts
@@ -35,6 +35,9 @@ type CreateResumableUIMessageStream = {
 };
 /**
  * Creates a resumable context for starting, resuming and stopping UI message streams.
+ *
+ * Leverages resumable-stream's internal eager drain pattern which drains the source
+ * eagerly and enqueues directly to the output stream.
  */
 declare function createResumableUIMessageStream(options: CreateResumableUIMessageStream): Promise<{
   startStream: (stream: ReadableStream<UIMessageChunk>) => Promise<AsyncIterableStream<UIMessageChunk>>;

package/dist/index.mjs

@@ -1,10 +1,33 @@
-import { convertSSEToUIMessageStream, convertUIMessageToSSEStream, createAsyncIterableStream } from "ai-stream-utils/utils";
+import { JsonToSseTransformStream, parseJsonEventStream, uiMessageChunkSchema } from "ai";
+import { createAsyncIterableStream } from "ai-stream-utils";
 import { createResumableStreamContext } from "resumable-stream";
 
+//#region src/convert-sse-stream-to-ui-message-stream.ts
+/**
+* Converts an SSE stream to a UI message stream.
+*/
+function convertSSEToUIMessageStream(stream, onComplete) {
+	return parseJsonEventStream({
+		stream: stream.pipeThrough(new TextEncoderStream()),
+		schema: uiMessageChunkSchema
+	}).pipeThrough(new TransformStream({
+		transform(result, controller) {
+			if (result.success) controller.enqueue(result.value);
+		},
+		flush() {
+			onComplete?.();
+		}
+	}));
+}
+
+//#endregion
 //#region src/resumable-ui-message-stream.ts
 const KEY_PREFIX = `ai-resumable-stream`;
 /**
 * Creates a resumable context for starting, resuming and stopping UI message streams.
+*
+* Leverages resumable-stream's internal eager drain pattern which drains the source
+* eagerly and enqueues directly to the output stream.
 */
 async function createResumableUIMessageStream(options) {
 	const { streamId, abortController, publisher, subscriber, waitUntil = null } = options;
@@ -39,25 +62,86 @@ async function createResumableUIMessageStream(options) {
 	}
 	/**
 	* Start a new stream by creating a new resumable stream in Redis and returning a client stream for the UI.
+	*
+	* Uses a single drain loop that:
+	* 1. Reads from source stream
+	* 2. Sends UIMessageChunk directly to client stream (no conversion)
+	* 3. Sends SSE to Redis stream → resumable-stream → Redis
+	* 4. Propagates errors to both streams
 	*/
 	async function startStream(stream) {
 		/**
-		* Tee the stream into two streams: one for the client and one for the resumable stream in Redis.
+		* Track client disconnect to avoid unbounded memory growth
 		*/
-		const [clientStream, resumableStream] = stream.tee();
-		const sseStream = convertUIMessageToSSEStream(resumableStream);
+		let clientCancelled = false;
 		/**
-		* Create a new resumable stream in Redis with the stream ID.
+		* Client stream for sending UI message chunks directly to the client without conversion.
 		*/
-		await context.createNewResumableStream(streamId, () => sseStream);
-		return createAsyncIterableStream(clientStream.pipeThrough(new TransformStream({
-			transform(chunk, controller) {
-				controller.enqueue(chunk);
+		let clientController;
+		const clientStream = new ReadableStream({
+			start(controller) {
+				clientController = controller;
 			},
-			flush() {
-				unsubscribe();
+			cancel() {
+				clientCancelled = true;
+			}
+		});
+		/**
+		* Redis stream with SSE conversion for resumable-stream persistence in Redis.
+		* JsonToSseTransformStream converts UIMessageChunk → SSE string and adds [DONE] on flush.
+		*/
+		let redisController;
+		const redisStream = new ReadableStream({ start(controller) {
+			redisController = controller;
+		} }).pipeThrough(new JsonToSseTransformStream());
+		/**
+		* Get reader synchronously to fail fast if stream is locked
+		*/
+		const reader = stream.getReader();
+		/**
+		* Register Redis stream with resumable-stream for persistence.
+		* Release reader lock if registration fails to avoid locking the source stream.
+		*/
+		try {
+			await context.createNewResumableStream(streamId, () => redisStream);
+		} catch (error) {
+			reader.releaseLock();
+			throw error;
+		}
+		/**
+		* Single drain loop.
+		* Continues draining to Redis even if client disconnects
+		*/
+		(async () => {
+			try {
+				while (true) {
+					const { done, value } = await reader.read();
+					if (done) {
+						redisController.close();
+						if (!clientCancelled) clientController.close();
+						break;
+					}
+					/**
+					* Always enqueue to Redis for persistence
+					*/
+					redisController.enqueue(value);
+					/**
+					* Only enqueue to client if still connected to avoid unbounded memory growth.
+					*
+					*/
+					if (!clientCancelled) clientController.enqueue(value);
+				}
+			} catch (error) {
+				redisController.error(error);
+				if (!clientCancelled) clientController.error(error);
+			} finally {
+				reader.releaseLock();
+				try {
+					await unsubscribe();
+				} catch {}
 			}
-		})));
+		})();
+		return createAsyncIterableStream(clientStream);
 	}
 	/**
 	* Resume an existing stream by fetching the resumable stream from Redis using the stream ID.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-resumable-stream",
-  "version": "1.1.1",
+  "version": "1.2.1",
   "description": "AI SDK: resume and stop UI message streams",
   "keywords": [
     "ai",
@@ -26,7 +26,7 @@
     "./package.json": "./package.json"
   },
   "dependencies": {
-    "ai-stream-utils": "^1.6.0",
+    "ai-stream-utils": "^2.0.0",
     "resumable-stream": "^2.2.10"
   },
   "devDependencies": {
@@ -38,6 +38,7 @@
     "ai": "^6.0.79",
     "husky": "^9.1.7",
     "lint-staged": "^16.2.7",
+    "mitata": "^1.0.34",
     "oxfmt": "^0.31.0",
     "oxlint": "^1.46.0",
     "pkg-pr-new": "^0.0.63",
@@ -47,7 +48,8 @@
     "tsdown": "^0.20.3",
     "tsx": "^4.21.0",
     "typescript": "^5.9.3",
-    "vitest": "^4.0.18"
+    "vitest": "^4.0.18",
+    "zod": "^4.3.6"
   },
   "peerDependencies": {
     "ai": "^6.0.0",