ai-resumable-stream 1.0.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Chris
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,356 @@
+<div align='center'>
+
+# ai-resumable-stream
+
+<p align="center">AI SDK: Resume and stop UI message streams</p>
+<p align="center">
+  <a href="https://www.npmjs.com/package/ai-resumable-stream" alt="ai-resumable-stream"><img src="https://img.shields.io/npm/dt/ai-resumable-stream?label=ai-resumable-stream"></a> <a href="https://github.com/zirkelc/ai-resumable-stream/actions/workflows/ci.yml" alt="CI"><img src="https://img.shields.io/github/actions/workflow/status/zirkelc/ai-resumable-stream/ci.yml?branch=main"></a>
+</p>
+
+</div>
+
+This library provides resumable streaming for UI message streams created by [`streamText()`](https://ai-sdk.dev/docs/reference/ai-sdk-core/stream-text) in the AI SDK. It uses Redis to persist stream data, allowing clients to resume interrupted streams or stop active streams from anywhere.
+
+**Why?**
+
+Streams are ephemeral — once data flows through, it's gone. This creates two hard problems:
+
+**Resume is hard** because the server doesn't track what's been sent. If a client disconnects (network issue, page reload, tab switch), the stream keeps running on the server but the client loses all that data. When reconnecting, there's no way to replay missed chunks without persisting them somewhere.
+
+**Stop is hard** because the client requesting "stop" isn't the same request that started the stream. The user clicks "Stop generating", which fires a new HTTP request, but the original stream is running in a different request/process. Without a central coordination point, you can't signal across requests.
+
+This library implements resumable streams with Redis to support both:
+
+- **Resuming**: Chunks are stored as they arrive, enabling replay on reconnect
+- **Stopping**: Stop signals are broadcast to any process running the stream
+
+## How It Works
+
+```mermaid
+sequenceDiagram
+    participant Client
+    participant Server
+    participant Redis
+
+    rect rgb(240, 248, 255)
+        Note over Client,Redis: Start Stream
+        Client->>Server: sendMessage()
+        Server->>Server: streamText()
+        Server->>Redis: Subscribe to stop channel
+        Server->>Redis: Subscribe to stream channel
+        Server->>Redis: Store chunks
+        Server-->>Client: Stream chunks
+    end
+
+    rect rgb(240, 255, 240)
+        Note over Client,Redis: Resume Stream
+        Client->>Server: resumeMessage()
+        Server->>Redis: Subscribe to stream channel
+        Redis-->>Server: Replay stored chunks
+        Server-->>Client: Stream past chunks
+        Redis-->>Server: Receive new chunks
+        Server-->>Client: Stream new chunks
+    end
+
+    rect rgb(255, 248, 240)
+        Note over Client,Redis: Stop Stream
+        Client->>Server: stopStream()
+        Server->>Redis: Publish to stop channel
+        Redis-->>Server: Notify subscriber
+        Server->>Server: AbortController.abort()
+    end
+```
+
+## Installation
+
+This library requires a [Redis](https://github.com/redis/node-redis) client.
+
+```bash
+npm install ai-resumable-stream redis
+```
+
+## Usage
+
+The library requires two Redis clients (pub/sub needs separate connections). The clients will connect automatically, if not already connected, but the library won't disconnect them afterwards. That means you can manage the connection lifecycle in your application and reuse clients across multiple streams.
+
+```typescript
+import { createClient } from "redis";
+import { createResumableUIMessageStream } from "ai-resumable-stream";
+
+const publisher = createClient({ url: process.env.REDIS_URL });
+const subscriber = createClient({ url: process.env.REDIS_URL });
+
+// Optional: connect clients immediately or let the library connect on demand
+await publisher.connect();
+await subscriber.connect();
+
+const context = await createResumableUIMessageStream({
+  streamId: `stream-123`,
+  publisher,
+  subscriber,
+});
+```
+
+### `startStream`
+
+Start a new stream and persist chunks to Redis. Returns a client stream that can be consumed immediately.
+
+> [!TIP]
+> The stream returned by `startStream` is both a readable stream and an async iterable.
+> That means you can use `return stream` or `yield* stream`.
+
+```typescript
+import { streamText } from "ai";
+
+async function sendMessage() {
+  // Optional: create AbortController to enable stopStream
+  const abortController = new AbortController();
+
+  // Create resumable stream context
+  const context = await createResumableUIMessageStream({
+    streamId: `stream-123`,
+    publisher,
+    subscriber,
+    abortController,
+  });
+
+  const result = streamText({
+    model: `gpt-4o`,
+    prompt: `Tell me a story`,
+    // Optional: pass abort signal to enable stopping
+    abortSignal: abortController.signal,
+  });
+
+  // Start streaming - chunks are stored in Redis as they arrive
+  const stream = await context.startStream(result.toUIMessageStream());
+
+  // Return stream to client
+  return stream;
+}
+```
+
+### `resumeStream`
+
+Resume an existing stream from Redis. Returns all past chunks followed by any remaining chunks, or `null` if no active stream exists.
+
+> [!TIP]
+> The stream returned by `resumeStream` is both a readable stream and an async iterable.
+> That means you can use `return stream` or `yield* stream`.
+
+```typescript
+async function resumeMessage() {
+  // Create resumable stream context
+  const context = await createResumableUIMessageStream({
+    streamId: `stream-123`,
+    publisher,
+    subscriber,
+  });
+
+  // Try to resume an existing stream
+  const stream = await context.resumeStream();
+
+  // If no stream exists, return early
+  if (!stream) {
+    console.log("No active stream to resume");
+    return;
+  }
+
+  // Return resumed stream to client
+  return stream;
+}
+```
+
+### `stopStream`
+
+Stop an active stream from any client. Requires an `AbortController` to be passed when creating the context.
+
+```typescript
+async function stopMessage() {
+  const context = await createResumableUIMessageStream({
+    streamId: `stream-123`,
+    publisher,
+    subscriber,
+  });
+
+  await context.stopStream();
+}
+```
+
+## Examples
+
+### tRPC
+
+Server-side tRPC procedures for sending, resuming, and stopping streams:
+
+```typescript
+// server/router.ts
+import { z } from "zod";
+import { streamText, type UIMessage, type UIMessageChunk } from "ai";
+import { createClient } from "redis";
+import { createResumableUIMessageStream } from "ai-resumable-stream";
+import { publicProcedure, router } from "./trpc";
+
+const publisher = createClient({ url: process.env.REDIS_URL });
+const subscriber = createClient({ url: process.env.REDIS_URL });
+
+export const appRouter = router({
+  sendMessage: publicProcedure
+    .input(z.object({ chatId: z.string(), message: z.custom<UIMessage>() }))
+    .mutation(async function* ({ input }): AsyncGenerator<UIMessageChunk> {
+      const { chatId, message } = input;
+
+      // TODO: Generate and save active stream ID for the chat
+      const activeStreamId = randomUUID();
+      await saveChat({ chatId, activeStreamId });
+
+      const abortController = new AbortController();
+
+      const context = await createResumableUIMessageStream({
+        streamId: activeStreamId,
+        publisher,
+        subscriber,
+        abortController,
+      });
+
+      const result = streamText({
+        model: openai("gpt-4o"),
+        messages: [message],
+        abortSignal: abortController.signal,
+        onFinish: async () => {
+          // TODO: Clear the active stream when finished
+          await saveChat({ chatId, activeStreamId: null });
+        },
+      });
+
+      const stream = await context.startStream(result.toUIMessageStream());
+
+      yield* stream;
+    }),
+
+  resumeMessage: publicProcedure.input(z.object({ chatId: z.string() })).mutation(async function* ({
+    input,
+  }): AsyncGenerator<UIMessageChunk> {
+    const { chatId } = input;
+
+    // TODO: Get active stream ID for the chat
+    const { activeStreamId } = await getChat(chatId);
+
+    const context = await createResumableUIMessageStream({
+      streamId: activeStreamId,
+      publisher,
+      subscriber,
+    });
+
+    const stream = await context.resumeStream();
+    if (!stream) return;
+
+    yield* stream;
+  }),
+
+  stopStream: publicProcedure
+    .input(z.object({ chatId: z.string() }))
+    .mutation(async ({ input }) => {
+      const { chatId } = input;
+
+      // TODO: Get active stream ID for the chat
+      const { activeStreamId } = await getChat(chatId);
+
+      const context = await createResumableUIMessageStream({
+        streamId: activeStreamId,
+        publisher,
+        subscriber,
+      });
+
+      await context.stopStream();
+
+      return { success: true };
+    }),
+});
+```
+
+## Configuration
+
+### Options
+
+| Option            | Type                         | Required | Description                                                    |
+| ----------------- | ---------------------------- | -------- | -------------------------------------------------------------- |
+| `streamId`        | `string`                     | Yes      | Unique identifier for the stream                               |
+| `publisher`       | `Redis`                      | Yes      | Redis client for publishing                                    |
+| `subscriber`      | `Redis`                      | Yes      | Redis client for subscribing (must be separate from publisher) |
+| `abortController` | `AbortController`            | No       | Controller to enable `stopStream` functionality                |
+| `waitUntil`       | `(promise: Promise) => void` | No       | Keep serverless function alive until stream completes          |
+
+### Redis Connection
+
+The library automatically connects Redis clients if they're not already connected. Clients are **not** disconnected after stream completion.
+
+```typescript
+// Clients can be connected or disconnected
+const publisher = createClient({ url: redisUrl });
+const subscriber = createClient({ url: redisUrl });
+
+// Library connects if needed
+const context = await createResumableUIMessageStream({
+  streamId: "stream-123",
+  publisher, // Will connect if not already connected
+  subscriber, // Will connect if not already connected
+});
+
+// Manage disconnection yourself when appropriate
+await publisher.quit();
+await subscriber.quit();
+```
+
+## API Reference
+
+### `createResumableUIMessageStream`
+
+```typescript
+async function createResumableUIMessageStream(options: CreateResumableUIMessageStream): Promise<{
+  startStream: (
+    stream: ReadableStream<UIMessageChunk>,
+  ) => Promise<AsyncIterableStream<UIMessageChunk>>;
+  resumeStream: () => Promise<AsyncIterableStream<UIMessageChunk> | null>;
+  stopStream: () => Promise<void>;
+}>;
+
+type CreateResumableUIMessageStream = {
+  streamId: string;
+  publisher: Redis;
+  subscriber: Redis;
+  abortController?: AbortController;
+  waitUntil?: (promise: Promise<unknown>) => void;
+};
+```
+
+### Return Values
+
+#### `startStream`
+
+```typescript
+async function startStream(
+  stream: ReadableStream<UIMessageChunk>,
+): 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.
+
+#### `resumeStream`
+
+```typescript
+async function resumeStream(): Promise<AsyncIterableStream<UIMessageChunk> | null>;
+```
+
+Resumes an existing stream. Returns `null` if:
+
+- No stream exists for the given `streamId`
+- The stream has already completed
+- The stream TTL has expired
+
+#### `stopStream`
+
+```typescript
+async function stopStream(): Promise<void>;
+```
+
+Publishes a stop message via Redis pub/sub. If an `abortController` was provided during creation, the stream's `AbortController.abort()` is called, ending the stream with an `abort` chunk.

package/dist/index.mjs

@@ -0,0 +1,87 @@
+import { convertSSEToUIMessageStream, convertUIMessageToSSEStream, createAsyncIterableStream } from "ai-stream-utils/utils";
+import { createResumableStreamContext } from "resumable-stream";
+
+//#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.
+*/
+async function createResumableUIMessageStream(options) {
+	const { streamId, abortController, publisher, subscriber, waitUntil = null } = options;
+	const stopChannel = `${KEY_PREFIX}:stop:${streamId}`;
+	const context = createResumableStreamContext({
+		waitUntil,
+		publisher,
+		subscriber,
+		keyPrefix: KEY_PREFIX
+	});
+	await Promise.all([publisher.isOpen ? Promise.resolve() : publisher.connect(), subscriber.isOpen ? Promise.resolve() : subscriber.connect()]);
+	/**
+	* Unsubscribe from stop channel
+	*/
+	async function unsubscribe() {
+		if (!abortController) return;
+		await subscriber.unsubscribe(stopChannel);
+	}
+	/**
+	* Set up stop subscription if abortController provided
+	*/
+	if (abortController) {
+		await subscriber.subscribe(stopChannel, () => {
+			abortController.abort();
+		});
+		/**
+		* Cleanup when abort signal fires
+		*/
+		abortController.signal.addEventListener(`abort`, () => {
+			unsubscribe();
+		}, { once: true });
+	}
+	/**
+	* Start a new stream by creating a new resumable stream in Redis and returning a client stream for the UI.
+	*/
+	async function startStream(stream) {
+		/**
+		* Tee the stream into two streams: one for the client and one for the resumable stream in Redis.
+		*/
+		const [clientStream, resumableStream] = stream.tee();
+		const sseStream = convertUIMessageToSSEStream(resumableStream);
+		/**
+		* Create a new resumable stream in Redis with the stream ID.
+		*/
+		await context.createNewResumableStream(streamId, () => sseStream);
+		return createAsyncIterableStream(clientStream.pipeThrough(new TransformStream({
+			transform(chunk, controller) {
+				controller.enqueue(chunk);
+			},
+			flush() {
+				unsubscribe();
+			}
+		})));
+	}
+	/**
+	* Resume an existing stream by fetching the resumable stream from Redis using the stream ID.
+	*/
+	async function resumeStream() {
+		/**
+		* Resume the existing stream from Redis using the stream ID.
+		*/
+		const resumableStream = await context.resumeExistingStream(streamId);
+		if (!resumableStream) return null;
+		return createAsyncIterableStream(convertSSEToUIMessageStream(resumableStream));
+	}
+	/**
+	* Publish a stop message to the stop channel, which will trigger the abortController to abort the stream.
+	*/
+	async function stopStream() {
+		await publisher.publish(stopChannel, `stop`);
+	}
+	return {
+		startStream,
+		resumeStream,
+		stopStream
+	};
+}
+
+//#endregion
+export { createResumableUIMessageStream };

package/package.json

@@ -0,0 +1,71 @@
+{
+	"name": "ai-resumable-stream",
+	"version": "1.0.0",
+	"description": "AI SDK: resume and stop UI message streams",
+	"keywords": [
+		"ai",
+		"ai-sdk",
+		"redis",
+		"resumable",
+		"resumable-stream",
+		"stream"
+	],
+	"license": "MIT",
+	"author": "Chris Cook",
+	"repository": {
+		"type": "git",
+		"url": "git+https://github.com/zirkelc/ai-resumable-stream.git"
+	},
+	"files": [
+		"dist"
+	],
+	"type": "module",
+	"exports": {
+		".": "./dist/index.mjs",
+		"./package.json": "./package.json"
+	},
+	"scripts": {
+		"prepublishOnly": "pnpm build",
+		"build": "tsdown",
+		"test": "vitest",
+		"lint": "oxlint --fix",
+		"lint:ci": "oxlint",
+		"format": "oxfmt --write",
+		"format:ci": "oxfmt --check",
+		"prepare": "husky"
+	},
+	"dependencies": {
+		"ai-stream-utils": "^1.5.0",
+		"resumable-stream": "^2.2.10"
+	},
+	"devDependencies": {
+		"@ai-sdk/provider": "^3.0.8",
+		"@arethetypeswrong/cli": "^0.18.2",
+		"@total-typescript/tsconfig": "^1.0.4",
+		"@types/node": "^25.2.3",
+		"@vitest/coverage-v8": "^4.0.18",
+		"ai": "^6.0.79",
+		"husky": "^9.1.7",
+		"lint-staged": "^16.2.7",
+		"oxfmt": "^0.31.0",
+		"oxlint": "^1.46.0",
+		"pkg-pr-new": "^0.0.63",
+		"publint": "^0.3.17",
+		"redis": "^5.10.0",
+		"redis-memory-server": "^0.16.0",
+		"tsdown": "^0.20.3",
+		"tsx": "^4.21.0",
+		"typescript": "^5.9.3",
+		"vitest": "^4.0.18"
+	},
+	"lint-staged": {
+		"*.{js,ts,tsx,jsx}": [
+			"pnpm lint",
+			"pnpm format"
+		],
+		"*.{json,jsonc}": [
+			"pnpm format"
+		]
+	},
+	"packageManager": "pnpm@10.0.0"
+}