chore: update to ai sdk v5 beta (#1074)

app/(chat)/api/chat/[id]/stream/route.ts

@@ -0,0 +1,112 @@
+import { auth } from '@/app/(auth)/auth';
+import {
+  getChatById,
+  getMessagesByChatId,
+  getStreamIdsByChatId,
+} from '@/lib/db/queries';
+import type { Chat } from '@/lib/db/schema';
+import { ChatSDKError } from '@/lib/errors';
+import type { ChatMessage } from '@/lib/types';
+import { createUIMessageStream, JsonToSseTransformStream } from 'ai';
+import { getStreamContext } from '../../route';
+import { differenceInSeconds } from 'date-fns';
+
+export async function GET(
+  _: Request,
+  { params }: { params: Promise<{ id: string }> },
+) {
+  const { id: chatId } = await params;
+
+  const streamContext = getStreamContext();
+  const resumeRequestedAt = new Date();
+
+  if (!streamContext) {
+    return new Response(null, { status: 204 });
+  }
+
+  if (!chatId) {
+    return new ChatSDKError('bad_request:api').toResponse();
+  }
+
+  const session = await auth();
+
+  if (!session?.user) {
+    return new ChatSDKError('unauthorized:chat').toResponse();
+  }
+
+  let chat: Chat;
+
+  try {
+    chat = await getChatById({ id: chatId });
+  } catch {
+    return new ChatSDKError('not_found:chat').toResponse();
+  }
+
+  if (!chat) {
+    return new ChatSDKError('not_found:chat').toResponse();
+  }
+
+  if (chat.visibility === 'private' && chat.userId !== session.user.id) {
+    return new ChatSDKError('forbidden:chat').toResponse();
+  }
+
+  const streamIds = await getStreamIdsByChatId({ chatId });
+
+  if (!streamIds.length) {
+    return new ChatSDKError('not_found:stream').toResponse();
+  }
+
+  const recentStreamId = streamIds.at(-1);
+
+  if (!recentStreamId) {
+    return new ChatSDKError('not_found:stream').toResponse();
+  }
+
+  const emptyDataStream = createUIMessageStream<ChatMessage>({
+    execute: () => {},
+  });
+
+  const stream = await streamContext.resumableStream(recentStreamId, () =>
+    emptyDataStream.pipeThrough(new JsonToSseTransformStream()),
+  );
+
+  /*
+   * For when the generation is streaming during SSR
+   * but the resumable stream has concluded at this point.
+   */
+  if (!stream) {
+    const messages = await getMessagesByChatId({ id: chatId });
+    const mostRecentMessage = messages.at(-1);
+
+    if (!mostRecentMessage) {
+      return new Response(emptyDataStream, { status: 200 });
+    }
+
+    if (mostRecentMessage.role !== 'assistant') {
+      return new Response(emptyDataStream, { status: 200 });
+    }
+
+    const messageCreatedAt = new Date(mostRecentMessage.createdAt);
+
+    if (differenceInSeconds(resumeRequestedAt, messageCreatedAt) > 15) {
+      return new Response(emptyDataStream, { status: 200 });
+    }
+
+    const restoredStream = createUIMessageStream<ChatMessage>({
+      execute: ({ writer }) => {
+        writer.write({
+          type: 'data-appendMessage',
+          data: JSON.stringify(mostRecentMessage),
+          transient: true,
+        });
+      },
+    });
+
+    return new Response(
+      restoredStream.pipeThrough(new JsonToSseTransformStream()),
+      { status: 200 },
+    );
+  }
+
+  return new Response(stream, { status: 200 });
+}