Feat/resumable stream

.env.example

@@ -3,6 +3,9 @@ NEXT_PUBLIC_SUPABASE_URL=your_supabase_project_url
 NEXT_PUBLIC_SUPABASE_ANON_KEY=your_supabase_anon_key
 SUPABASE_SERVICE_ROLE=your_supabase_service_role_key
 
+# Redis Configuration
+REDIS_URL=redis://localhost:6379
+
 # CSRF Protection (required)
 CSRF_SECRET=your_32_character_random_string
 

INSTALL.md

@@ -22,6 +22,9 @@ NEXT_PUBLIC_SUPABASE_URL=your_supabase_url
 NEXT_PUBLIC_SUPABASE_ANON_KEY=your_supabase_anon_key
 SUPABASE_SERVICE_ROLE=your_supabase_service_role_key
 
+# Redis
+REDIS_URL=redis://localhost:6379
+
 # OpenAI
 OPENAI_API_KEY=your_openai_api_key
 
@@ -151,6 +154,14 @@ Here are the detailed steps to set up Google OAuth:
 
 This allows users limited access to try the product before properly creating an account.
 
+### Redis
+
+Start a redis via docker:
+
+```bash
+docker run -p 6379:6379 -d redis:8.0-rc1
+```
+
 ### Database Schema
 
 Create the following tables in your Supabase SQL editor:
@@ -200,6 +211,15 @@ CREATE TABLE chats (
   CONSTRAINT chats_project_id_fkey FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE
 );
 
+-- Stream IDs table
+CREATE TABLE stream_ids (
+  id SERIAL PRIMARY KEY,
+  stream_id UUID NOT NULL,
+  chat_id UUID NOT NULL,
+  created_at TIMESTAMPTZ DEFAULT NOW(),
+  CONSTRAINT stream_ids_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
+);
+
 -- Messages table
 CREATE TABLE messages (
   id SERIAL PRIMARY KEY, -- Using SERIAL for auto-incrementing integer ID

app/api/chat/db.ts

@@ -90,3 +90,24 @@ export async function saveFinalAssistantMessage(
     console.log("Assistant message saved successfully (merged).")
   }
 }
+
+export async function createStreamId(
+  supabase: SupabaseClient<Database>,
+  {
+    streamId,
+    chatId,
+  }: {
+    streamId: string
+    chatId: string
+  }
+) {
+  const { error } = await supabase.from("stream_ids").insert({
+    stream_id: streamId,
+    chat_id: chatId,
+  })
+
+  if (error) {
+    console.error("Error creating stream ID:", error)
+    throw new Error(`Failed to create stream ID: ${error.message}`)
+  }
+}

app/api/chat/route.ts

@@ -1,19 +1,34 @@
+import { getMessagesFromDb } from "@/lib/chat-store/messages/api"
 import { SYSTEM_PROMPT_DEFAULT } from "@/lib/config"
 import { getAllModels } from "@/lib/models"
 import { getProviderForModel } from "@/lib/openproviders/provider-map"
+import { createClient } from "@/lib/supabase/server"
 import type { ProviderWithoutOllama } from "@/lib/user-keys"
+import { generateUUID } from "@/lib/utils"
 import { Attachment } from "@ai-sdk/ui-utils"
-import { Message as MessageAISDK, streamText, ToolSet } from "ai"
+import {
+  createDataStream,
+  Message as MessageAISDK,
+  streamText,
+  ToolSet,
+} from "ai"
+import { after, NextResponse } from "next/server"
+import { createResumableStreamContext } from "resumable-stream"
 import {
   incrementMessageCount,
   logUserMessage,
   storeAssistantMessage,
   validateAndTrackUsage,
 } from "./api"
+import { createStreamId } from "./db"
 import { createErrorResponse, extractErrorMessage } from "./utils"
 
 export const maxDuration = 60
 
+const streamContext = createResumableStreamContext({
+  waitUntil: after,
+})
+
 type ChatRequest = {
   messages: MessageAISDK[]
   chatId: string
@@ -89,39 +104,51 @@ export async function POST(req: Request) {
         undefined
     }
 
-    const result = streamText({
-      model: modelConfig.apiSdk(apiKey, { enableSearch }),
-      system: effectiveSystemPrompt,
-      messages: messages,
-      tools: {} as ToolSet,
-      maxSteps: 10,
-      onError: (err: unknown) => {
-        console.error("Streaming error occurred:", err)
-        // Don't set streamError anymore - let the AI SDK handle it through the stream
-      },
+    const streamId = generateUUID()
+    await createStreamId(supabase!, { streamId, chatId })
 
-      onFinish: async ({ response }) => {
-        if (supabase) {
-          await storeAssistantMessage({
-            supabase,
-            chatId,
-            messages:
-              response.messages as unknown as import("@/app/types/api.types").Message[],
-            message_group_id,
-            model,
-          })
-        }
-      },
-    })
+    const stream = createDataStream({
+      execute: (dataStream) => {
+        const result = streamText({
+          model: modelConfig.apiSdk!(apiKey, { enableSearch }),
+          system: effectiveSystemPrompt,
+          messages: messages,
+          tools: {} as ToolSet,
+          maxSteps: 10,
+          onError: (err: unknown) => {
+            console.error("Streaming error occurred:", err)
+            // Don't set streamError anymore - let the AI SDK handle it through the stream
+          },
 
-    return result.toDataStreamResponse({
-      sendReasoning: true,
-      sendSources: true,
-      getErrorMessage: (error: unknown) => {
-        console.error("Error forwarded to client:", error)
-        return extractErrorMessage(error)
+          onFinish: async ({ response }) => {
+            if (supabase) {
+              await storeAssistantMessage({
+                supabase,
+                chatId,
+                messages:
+                  response.messages as unknown as import("@/app/types/api.types").Message[],
+                message_group_id,
+                model,
+              })
+            }
+          },
+        })
+        result.consumeStream()
+
+        result.mergeIntoDataStream(dataStream, {
+          sendReasoning: true,
+          sendSources: true,
+        })
+      },
+      onError: (err: unknown) => {
+        console.error("Error forwarded to client:", err)
+        return extractErrorMessage(err)
       },
     })
+
+    return new Response(
+      await streamContext.resumableStream(streamId, () => stream)
+    )
   } catch (err: unknown) {
     console.error("Error in /api/chat:", err)
     const error = err as {
@@ -133,3 +160,105 @@ export async function POST(req: Request) {
     return createErrorResponse(error)
   }
 }
+
+export async function GET(req: Request) {
+  try {
+    const { searchParams } = new URL(req.url)
+    const chatId = searchParams.get("chatId")
+
+    if (!chatId) {
+      return new Response("id is required", { status: 400 })
+    }
+    const supabase = await createClient()
+
+    if (!supabase) {
+      return NextResponse.json(
+        { error: "Database connection failed" },
+        { status: 500 }
+      )
+    }
+
+    // Get the current user
+    const {
+      data: { user },
+      error: authError,
+    } = await supabase.auth.getUser()
+
+    if (authError || !user) {
+      return NextResponse.json({ error: "Unauthorized" }, { status: 401 })
+    }
+
+    // Get the chat by id
+    const chat = await supabase
+      .from("chats")
+      .select("*")
+      .eq("id", chatId)
+      .single()
+
+    if (!chat.data) {
+      return new Response("Not found", { status: 404 })
+    }
+
+    if (chat.data.user_id !== user.id) {
+      return new Response("Forbidden", { status: 403 })
+    }
+
+    // Get the stream ids by chat id
+    const streamIds = await supabase
+      .from("stream_ids")
+      .select("*")
+      .eq("chat_id", chatId)
+
+    if (!streamIds.data) {
+      return new Response("No streams found", { status: 404 })
+    }
+
+    const recentStreamId = streamIds.data.at(-1)?.stream_id
+
+    if (!recentStreamId) {
+      return new Response("No recent stream found", { status: 404 })
+    }
+
+    const emptyDataStream = createDataStream({
+      execute: () => {},
+    })
+
+    const stream = await streamContext.resumableStream(
+      recentStreamId,
+      () => emptyDataStream
+    )
+
+    if (stream) {
+      return new Response(stream, { status: 200 })
+    }
+
+    /*
+     * For when the generation is "active" during SSR but the
+     * resumable stream has concluded after reaching this point.
+     */
+
+    const messages = await getMessagesFromDb(chatId)
+    const mostRecentMessage = messages.at(-1)
+
+    if (!mostRecentMessage || mostRecentMessage.role !== "assistant") {
+      return new Response(emptyDataStream, { status: 200 })
+    }
+
+    const streamWithMessage = createDataStream({
+      execute: (buffer) => {
+        buffer.writeData({
+          type: "append-message",
+          message: JSON.stringify(mostRecentMessage),
+        })
+      },
+    })
+
+    return new Response(streamWithMessage, { status: 200 })
+  } catch (error) {
+    console.error("Error in chat GET API:", error)
+    return NextResponse.json(
+      { error: "Internal server error" },
+      { status: 500 }
+    )
+  }
+}

app/c/[chatId]/page.tsx

@@ -1,25 +1,34 @@
 import { ChatContainer } from "@/app/components/chat/chat-container"
 import { LayoutApp } from "@/app/components/layout/layout-app"
+import { getMessagesFromDb } from "@/lib/chat-store/messages/api"
 import { MessagesProvider } from "@/lib/chat-store/messages/provider"
 import { isSupabaseEnabled } from "@/lib/supabase/config"
 import { createClient } from "@/lib/supabase/server"
+import { Message as MessageAISDK } from "ai"
 import { redirect } from "next/navigation"
 
-export default async function Page() {
+export default async function Page({
+  params,
+}: {
+  params: Promise<{ chatId: string }>
+}) {
+  let initialMessages: MessageAISDK[] = []
+  const { chatId } = await params
   if (isSupabaseEnabled) {
     const supabase = await createClient()
     if (supabase) {
       const { data: userData, error: userError } = await supabase.auth.getUser()
       if (userError || !userData?.user) {
         redirect("/")
       }
+      initialMessages = await getMessagesFromDb(chatId)
     }
   }
 
   return (
     <MessagesProvider>
       <LayoutApp>
-        <ChatContainer />
+        <ChatContainer autoResume={true} initialMessages={initialMessages} />
       </LayoutApp>
     </MessagesProvider>
   )

app/components/chat/chat-container.tsx

@@ -2,15 +2,22 @@
 
 import { MultiChat } from "@/app/components/multi-chat/multi-chat"
 import { useUserPreferences } from "@/lib/user-preference-store/provider"
+import type { Message as MessageAISDK } from "ai"
 import { Chat } from "./chat"
 
-export function ChatContainer() {
+export function ChatContainer({
+  initialMessages,
+  autoResume,
+}: {
+  initialMessages?: MessageAISDK[]
+  autoResume?: boolean
+}) {
   const { preferences } = useUserPreferences()
   const multiModelEnabled = preferences.multiModelEnabled
 
   if (multiModelEnabled) {
     return <MultiChat />
   }
 
-  return <Chat />
+  return <Chat autoResume={autoResume} initialMessages={initialMessages} />
 }