ai-resumable-stream 1.2.1 → 1.4.0

package/README.md

@@ -301,6 +301,42 @@ await publisher.quit();
 await subscriber.quit();
 ```
 
+### Keep Alive
+
+By default, the resumable-stream producer tears down immediately when the source stream ends. This means resume requests arriving during post-stream work (e.g. saving the assistant message to the database) will get `null` back, even though the full response is still in memory.
+
+Pass a `keepAlive` promise to `startStream` to defer teardown until your post-stream work is complete:
+
+```typescript
+const { promise, resolve } = Promise.withResolvers<void>();
+
+const stream = await context.startStream(result.toUIMessageStream(), {
+  keepAlive: promise,
+});
+
+yield * stream;
+
+// Producer stays alive — resume requests are served from the in-memory buffer
+await saveAssistantMessage(result);
+
+// Done — producer tears down
+resolve();
+```
+
+### Flush
+
+The `onFlush` callback is invoked after the producer has torn down (Redis stream closed, sentinel set to DONE). Use it for cleanup tasks like removing the active stream ID from the database. Errors thrown by `onFlush` are silently caught.
+
+```typescript
+const stream = await context.startStream(result.toUIMessageStream(), {
+  onFlush: async () => {
+    await saveChat({ chatId, activeStreamId: null });
+  },
+});
+```
+
+When used together with `keepAlive`, `onFlush` fires after the `keepAlive` promise resolves and the producer tears down.
+
 ## API Reference
 
 ### `createResumableUIMessageStream`
@@ -309,6 +345,7 @@ await subscriber.quit();
 async function createResumableUIMessageStream(options: CreateResumableUIMessageStream): Promise<{
   startStream: (
     stream: ReadableStream<UIMessageChunk>,
+    options?: StartStreamOptions,
   ) => Promise<AsyncIterableStream<UIMessageChunk>>;
   resumeStream: () => Promise<AsyncIterableStream<UIMessageChunk> | null>;
   stopStream: () => Promise<void>;
@@ -321,6 +358,11 @@ type CreateResumableUIMessageStream = {
   abortController?: AbortController;
   waitUntil?: (promise: Promise<unknown>) => void;
 };
+
+type StartStreamOptions = {
+  keepAlive?: Promise<void>;
+  onFlush?: () => void | Promise<void>;
+};
 ```
 
 ### Return Values
@@ -330,11 +372,29 @@ type CreateResumableUIMessageStream = {
 ```typescript
 async function startStream(
   stream: ReadableStream<UIMessageChunk>,
+  options?: StartStreamOptions,
 ): Promise<AsyncIterableStream<UIMessageChunk>>;
+
+type StartStreamOptions = {
+  keepAlive?: Promise<void>;
+  onFlush?: () => void | Promise<void>;
+};
 ```
 
 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.
 
+##### `keepAlive`
+
+A promise that defers closing the Redis stream until it resolves. This keeps the resumable-stream producer alive after the source stream ends, so late resume requests (e.g. during post-stream DB writes) can still be served from the in-memory chunk buffer.
+
+If the promise rejects, the producer tears down normally. If the source stream errors, the `keepAlive` promise is not awaited.
+
+##### `onFlush`
+
+A callback invoked after the Redis stream is closed and the producer has torn down, 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.
+
+When used together with `keepAlive`, `onFlush` fires after the `keepAlive` promise resolves and the producer tears down.
+
 #### `resumeStream`
 
 ```typescript

package/dist/index.d.mts

@@ -33,6 +33,29 @@ type CreateResumableUIMessageStream = {
    */
   waitUntil?: (promise: Promise<unknown>) => void;
 };
+type StartStreamOptions = {
+  /**
+   * A promise that is awaited before closing the Redis stream, keeping the
+   * resumable-stream producer alive so late resume requests can be served
+   * from the in-memory buffer. The producer tears down when the promise resolves.
+   *
+   * Use `Promise.withResolvers()` to create a deferred promise and resolve it
+   * after post-stream work (e.g. DB writes) is complete:
+   * ```ts
+   * const { promise, resolve } = Promise.withResolvers<void>();
+   * const stream = await ctx.startStream(input, { keepAlive: promise });
+   * for await (const chunk of stream) { ... }
+   * await saveToDb();
+   * resolve();
+   * ```
+   */
+  keepAlive?: Promise<void>;
+  /**
+   * Called after the Redis stream is closed and the producer has torn down.
+   * Use for post-teardown cleanup.
+   */
+  onFlush?: () => void | Promise<void>;
+};
 /**
  * Creates a resumable context for starting, resuming and stopping UI message streams.
  *
@@ -40,7 +63,7 @@ type CreateResumableUIMessageStream = {
  * 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?: StartStreamOptions) => Promise<AsyncIterableStream<UIMessageChunk>>;
   resumeStream: () => Promise<AsyncIterableStream<UIMessageChunk> | null>;
   stopStream: () => Promise<void>;
 }>;

package/dist/index.mjs

@@ -6,7 +6,7 @@ import { createResumableStreamContext } from "resumable-stream";
 /**
 * Converts an SSE stream to a UI message stream.
 */
-function convertSSEToUIMessageStream(stream, onComplete) {
+function convertSSEToUIMessageStream(stream, onFlush) {
 	return parseJsonEventStream({
 		stream: stream.pipeThrough(new TextEncoderStream()),
 		schema: uiMessageChunkSchema
@@ -15,7 +15,7 @@ function convertSSEToUIMessageStream(stream, onComplete) {
 			if (result.success) controller.enqueue(result.value);
 		},
 		flush() {
-			onComplete?.();
+			onFlush?.();
 		}
 	}));
 }
@@ -67,9 +67,13 @@ async function createResumableUIMessageStream(options) {
 	* 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
+	* 4. If `keepAlive` is provided, awaits it before closing the Redis stream
+	* 5. Closes the Redis stream, triggering resumable-stream teardown
+	* 6. Calls `onFlush` after teardown
+	* 7. Propagates errors to both streams
 	*/
-	async function startStream(stream) {
+	async function startStream(stream, options) {
+		const { keepAlive = Promise.resolve(), onFlush } = options ?? {};
 		/**
 		* Track client disconnect to avoid unbounded memory growth
 		*/
@@ -113,12 +117,18 @@ async function createResumableUIMessageStream(options) {
 		* Continues draining to Redis even if client disconnects
 		*/
 		(async () => {
+			/**
+			* Tracks whether the source stream completed successfully (reader.read() returned done: true).
+			* Used to guard the keepAlive await in finally — on error, the caller may never resolve
+			* the keepAlive promise, so awaiting it would hang the drain loop forever.
+			*/
+			let sourceDone = false;
 			try {
 				while (true) {
 					const { done, value } = await reader.read();
 					if (done) {
-						redisController.close();
 						if (!clientCancelled) clientController.close();
+						sourceDone = true;
 						break;
 					}
 					/**
@@ -139,6 +149,23 @@ async function createResumableUIMessageStream(options) {
 				try {
 					await unsubscribe();
 				} catch {}
+				/**
+				* Await keepAlive promise before closing the Redis stream.
+				* This keeps the resumable-stream producer alive so late resume requests
+				* can be served from the in-memory chunk buffer.
+				* Defaults to a resolved promise (no-op) when not provided.
+				* Only await on successful source completion — on error the caller
+				* may never resolve the promise, which would hang the drain loop.
+				*/
+				if (sourceDone) try {
+					await keepAlive;
+				} catch {}
+				try {
+					redisController.close();
+				} catch {}
+				try {
+					await onFlush?.();
+				} catch {}
 			}
 		})();
 		return createAsyncIterableStream(clientStream);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-resumable-stream",
-  "version": "1.2.1",
+  "version": "1.4.0",
   "description": "AI SDK: resume and stop UI message streams",
   "keywords": [
     "ai",