feat (ai): restructure chat transports (vercel#6746)

## Background

Chat resumption on reload needs a separate endpoint to connect to.

## Summary

* restructure chat transports
* move default route for stream reconnection to `/chat/[id]/stream`
* rename `prepareRequest` to `prepareSubmitMessagesRequest``

## Verification
* test `next` example

Author: lgrammel

.changeset/tricky-ravens-kick.md

@@ -0,0 +1,5 @@
+---
+'ai': major
+---
+
+feat (ai): restructure chat transports

examples/next-openai/app/api/use-chat-resume/[id]/stream/route.ts

@@ -0,0 +1,44 @@
+import { loadStreams } from '@/util/chat-store';
+import { createUIMessageStream, JsonToSseTransformStream } from 'ai';
+import { after } from 'next/server';
+import { createResumableStreamContext } from 'resumable-stream';
+
+// Allow streaming responses up to 30 seconds
+export const maxDuration = 30;
+
+export async function GET(
+  request: Request,
+  { params }: { params: Promise<{ id: string }> },
+) {
+  const { id } = await params;
+
+  if (!id) {
+    return new Response('id is required', { status: 400 });
+  }
+
+  const streamIds = await loadStreams(id);
+
+  if (!streamIds.length) {
+    return new Response('No streams found', { status: 204 });
+  }
+
+  const recentStreamId = streamIds.at(-1);
+
+  if (!recentStreamId) {
+    return new Response('No recent stream found', { status: 204 });
+  }
+
+  const emptyDataStream = createUIMessageStream({
+    execute: () => {},
+  });
+
+  const streamContext = createResumableStreamContext({
+    waitUntil: after,
+  });
+
+  return new Response(
+    await streamContext.resumableStream(recentStreamId, () =>
+      emptyDataStream.pipeThrough(new JsonToSseTransformStream()),
+    ),
+  );
+}

examples/next-openai/app/api/use-chat-resume/route.ts

@@ -1,7 +1,6 @@
 import {
   appendMessageToChat,
   appendStreamId,
-  loadStreams,
   saveChat,
 } from '@/util/chat-store';
 import { openai } from '@ai-sdk/openai';
@@ -64,38 +63,3 @@ export async function POST(req: Request) {
     ),
   );
 }
-
-export async function GET(request: Request) {
-  const streamContext = createResumableStreamContext({
-    waitUntil: after,
-  });
-
-  const { searchParams } = new URL(request.url);
-  const chatId = searchParams.get('chatId');
-
-  if (!chatId) {
-    return new Response('id is required', { status: 400 });
-  }
-
-  const streamIds = await loadStreams(chatId);
-
-  if (!streamIds.length) {
-    return new Response('No streams found', { status: 404 });
-  }
-
-  const recentStreamId = streamIds.at(-1);
-
-  if (!recentStreamId) {
-    return new Response('No recent stream found', { status: 404 });
-  }
-
-  const emptyDataStream = createUIMessageStream({
-    execute: () => {},
-  });
-
-  return new Response(
-    await streamContext.resumableStream(recentStreamId, () =>
-      emptyDataStream.pipeThrough(new JsonToSseTransformStream()),
-    ),
-  );
-}

examples/next-openai/app/use-chat-persistence-single-message/[id]/chat.tsx

@@ -14,7 +14,7 @@ export default function Chat({
     transport: new DefaultChatTransport({
       api: '/api/use-chat-persistence-single-message',
       // only send the last message to the server:
-      prepareRequest({ messages, id }) {
+      prepareSubmitMessagesRequest({ messages, id }) {
         return { body: { message: messages[messages.length - 1], id } };
       },
     }),

examples/next/app/api/chat/[id]/stream/route.ts

@@ -0,0 +1,27 @@
+import { readChat } from '@util/chat-store';
+import { UI_MESSAGE_STREAM_HEADERS } from 'ai';
+import { after } from 'next/server';
+import { createResumableStreamContext } from 'resumable-stream';
+
+export async function GET(
+  request: Request,
+  { params }: { params: Promise<{ id: string }> },
+) {
+  const { id } = await params;
+
+  const chat = await readChat(id);
+
+  if (chat.activeStreamId == null) {
+    // no content response when there is no active stream
+    return new Response(null, { status: 204 });
+  }
+
+  const streamContext = createResumableStreamContext({
+    waitUntil: after,
+  });
+
+  return new Response(
+    await streamContext.resumeExistingStream(chat.activeStreamId),
+    { headers: UI_MESSAGE_STREAM_HEADERS },
+  );
+}

examples/next/app/api/chat/route.ts

@@ -1,6 +1,8 @@
 import { MyUIMessage } from '@/util/chat-schema';
 import { readChat, saveChat } from '@util/chat-store';
-import { convertToModelMessages, streamText } from 'ai';
+import { convertToModelMessages, generateId, streamText } from 'ai';
+import { after } from 'next/server';
+import { createResumableStreamContext } from 'resumable-stream';
 
 export async function POST(req: Request) {
   const { message, id }: { message: MyUIMessage; id: string } =
@@ -9,13 +11,14 @@ export async function POST(req: Request) {
   const chat = await readChat(id);
   const messages = [...chat.messages, message];
 
+  // save the user message
+  saveChat({ id, messages, activeStreamId: null });
+
   const result = streamText({
     model: 'openai/gpt-4o-mini',
     messages: convertToModelMessages(messages),
   });
 
-  result.consumeStream(); // TODO always consume the stream even when the client disconnects
-
   return result.toUIMessageStreamResponse({
     originalMessages: messages,
     messageMetadata: ({ part }) => {
@@ -24,8 +27,17 @@ export async function POST(req: Request) {
       }
     },
     onFinish: ({ messages }) => {
-      // TODO fix type safety
-      saveChat({ id, messages: messages as MyUIMessage[] });
+      saveChat({ id, messages, activeStreamId: null });
+    },
+    async consumeSseStream({ stream }) {
+      const streamId = generateId();
+
+      // send the sse stream into a resumable stream sink as well:
+      const streamContext = createResumableStreamContext({ waitUntil: after });
+      await streamContext.createNewResumableStream(streamId, () => stream);
+
+      // update the chat with the streamId
+      saveChat({ id, activeStreamId: streamId });
     },
   });
 }

examples/next/app/chat/[chatId]/chat.tsx

@@ -11,18 +11,21 @@ import Message from './message';
 export default function ChatComponent({
   chatData,
   isNewChat = false,
+  resume = false,
 }: {
   chatData: { id: string; messages: MyUIMessage[] };
   isNewChat?: boolean;
+  resume?: boolean;
 }) {
   const inputRef = useRef<HTMLInputElement>(null);
 
   const { status, sendMessage, messages } = useChat({
     id: chatData.id,
     messages: chatData.messages,
+    resume,
     transport: new DefaultChatTransport({
       // only send the last message to the server to limit the request size:
-      prepareRequest: ({ id, messages }) => ({
+      prepareSubmitMessagesRequest: ({ id, messages }) => ({
         body: { id, message: messages[messages.length - 1] },
       }),
     }),

examples/next/app/chat/[chatId]/page.tsx

@@ -23,7 +23,7 @@ export default async function Page(props: {
           </li>
         ))}
       </ul>
-      <Chat chatData={chatData} />;
+      <Chat chatData={chatData} resume={chatData.activeStreamId !== null} />;
     </div>
   );
 }

examples/next/package.json

@@ -16,8 +16,7 @@
     "react": "^18",
     "react-dom": "^18",
     "react-markdown": "9.0.1",
-    "redis": "^4.7.0",
-    "resumable-stream": "^2.0.0",
+    "resumable-stream": "^2.2.0",
     "zod": "3.25.49"
   },
   "devDependencies": {

examples/next/util/chat-schema.ts

@@ -13,4 +13,5 @@ export type ChatData = {
   id: string;
   messages: MyUIMessage[];
   createdAt: number;
+  activeStreamId: string | null;
 };

examples/next/util/chat-store.ts

@@ -16,13 +16,23 @@ export async function createChat(): Promise<string> {
 
 export async function saveChat({
   id,
+  activeStreamId,
   messages,
 }: {
   id: string;
-  messages: MyUIMessage[];
+  activeStreamId?: string | null;
+  messages?: MyUIMessage[];
 }): Promise<void> {
   const chat = await readChat(id);
-  chat.messages = messages;
+
+  if (messages !== undefined) {
+    chat.messages = messages;
+  }
+
+  if (activeStreamId !== undefined) {
+    chat.activeStreamId = activeStreamId;
+  }
+
   writeChat(chat);
 }
 
@@ -42,6 +52,7 @@ async function writeChat(chat: ChatData) {
   await writeFile(await getChatFile(chat.id), JSON.stringify(chat, null, 2));
 }
 
+// TODO return null if the chat does not exist
 export async function readChat(id: string): Promise<ChatData> {
   return JSON.parse(await readFile(await getChatFile(id), 'utf8'));
 }
@@ -72,28 +83,3 @@ async function getChatFile(id: string): Promise<string> {
 
   return chatFile;
 }
-
-export async function appendStreamId({
-  id,
-  streamId,
-}: {
-  id: string;
-  streamId: string;
-}) {
-  const file = getStreamsFile(id);
-  const streams = await loadStreams(id);
-  streams.push(streamId);
-  await writeFile(file, JSON.stringify(streams, null, 2));
-}
-
-export async function loadStreams(id: string): Promise<string[]> {
-  const file = getStreamsFile(id);
-  if (!existsSync(file)) return [];
-  return JSON.parse(await readFile(file, 'utf8'));
-}
-
-function getStreamsFile(id: string): string {
-  const chatDir = path.join(process.cwd(), '.streams');
-  if (!existsSync(chatDir)) mkdirSync(chatDir, { recursive: true });
-  return path.join(chatDir, `${id}.json`);
-}

packages/ai/src/ui/chat-transport.ts

@@ -9,7 +9,12 @@ export interface ChatTransport<UI_MESSAGE extends UIMessage> {
       chatId: string;
       messages: UI_MESSAGE[];
       abortSignal: AbortSignal | undefined;
-      requestType: 'generate' | 'resume'; // TODO have separate functions
     } & ChatRequestOptions,
   ) => Promise<ReadableStream<UIMessageStreamPart>>;
+
+  reconnectToStream: (
+    options: {
+      chatId: string;
+    } & ChatRequestOptions,
+  ) => Promise<ReadableStream<UIMessageStreamPart> | null>;
 }

packages/ai/src/ui/chat.ts

@@ -29,6 +29,7 @@ import {
   type UIDataTypes,
   type UIMessage,
 } from './ui-messages';
+import { UIMessageStreamPart } from '../ui-message-stream/ui-message-stream-parts';
 
 export type CreateUIMessage<UI_MESSAGE extends UIMessage> = Omit<
   UI_MESSAGE,
@@ -384,15 +385,31 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
 
       this.activeResponse = activeResponse;
 
-      const stream = await this.transport.submitMessages({
-        chatId: this.id,
-        messages: this.state.messages,
-        abortSignal: activeResponse.abortController.signal,
-        metadata,
-        headers,
-        body,
-        requestType,
-      });
+      let stream: ReadableStream<UIMessageStreamPart>;
+
+      if (requestType === 'resume') {
+        const reconnect = await this.transport.reconnectToStream({
+          chatId: this.id,
+          metadata,
+          headers,
+          body,
+        });
+
+        if (reconnect == null) {
+          return; // no active stream found, so we do not resume
+        }
+
+        stream = reconnect;
+      } else {
+        stream = await this.transport.submitMessages({
+          chatId: this.id,
+          messages: this.state.messages,
+          abortSignal: activeResponse.abortController.signal,
+          metadata,
+          headers,
+          body,
+        });
+      }
 
       const runUpdateMessageJob = (
         job: (options: {