ai-resumable-stream 1.1.1 → 1.3.0

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)
@@ -309,6 +309,7 @@ await subscriber.quit();
 async function createResumableUIMessageStream(options: CreateResumableUIMessageStream): Promise<{
   startStream: (
     stream: ReadableStream<UIMessageChunk>,
+    options?: { onFlush?: () => void | Promise<void> },
   ) => Promise<AsyncIterableStream<UIMessageChunk>>;
   resumeStream: () => Promise<AsyncIterableStream<UIMessageChunk> | null>;
   stopStream: () => Promise<void>;
@@ -330,10 +331,13 @@ type CreateResumableUIMessageStream = {
 ```typescript
 async function startStream(
   stream: ReadableStream<UIMessageChunk>,
+  options?: { onFlush?: () => void | Promise<void> },
 ): 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.
+
+The optional `onFlush` callback is invoked after the stream finishes draining to Redis and cleanup is complete, regardless of how the stream ended (complete, error, or abort). Use it for cleanup tasks like removing the active stream ID from the database. Errors thrown by `onFlush` are silently caught.
 
 #### `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,9 +35,14 @@ 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>>;
+  startStream: (stream: ReadableStream<UIMessageChunk>, options?: {
+    onFlush?: () => void | Promise<void>;
+  }) => Promise<AsyncIterableStream<UIMessageChunk>>;
   resumeStream: () => Promise<AsyncIterableStream<UIMessageChunk> | null>;
   stopStream: () => Promise<void>;
 }>;

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, onFlush) {
+	return parseJsonEventStream({
+		stream: stream.pipeThrough(new TextEncoderStream()),
+		schema: uiMessageChunkSchema
+	}).pipeThrough(new TransformStream({
+		transform(result, controller) {
+			if (result.success) controller.enqueue(result.value);
+		},
+		flush() {
+			onFlush?.();
+		}
+	}));
+}
+
+//#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,90 @@ 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) {
+	async function startStream(stream, options) {
+		const { onFlush } = options ?? {};
 		/**
-		* 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 {}
+				try {
+					await onFlush?.();
+				} 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.3.0",
   "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",