continue scaffolding and refine the basic foundation

server/utils/generation.ts

@@ -1,70 +1,63 @@
 import { useDrizzle } from '~~/server/utils/drizzle';
-import { generations, messages as messages_drizzle } from '~~/db/schema';
-import { type GenerationStreamEvent, type ChatMessage, type MessageType } from '~~/server/types/chat';
+import { generations, messages as messages_drizzle, messagesRelations } from '~~/db/schema';
+import { type GenerationStreamEvent, type ChatMessage, type GenerationStatus } from '~~/server/types/chat';
 import { eq } from 'drizzle-orm';
 
-interface ActiveGeneration {
+/**
+ * Streaming generation state - only stores active stream controllers
+ * All persistent state lives in the database
+ */
+interface ActiveGenerationStream {
     userId: string;
     topicId: string;
-    messages: ChatMessage[];
-    content: string;
     clients: Set<ReadableStreamDefaultController<Uint8Array>>;
-    complete: boolean;
+    isGenerating: boolean;
 }
 
-interface PendingGeneration {
-    generationId: string;
-    userId: string;
-    topicId: string;
-    messages: ChatMessage[];
-    timeout: NodeJS.Timeout;
-    expired: boolean;
-}
-
-const activeGenerations = new Map<string, ActiveGeneration>();
-const pendingGenerations = new Map<string, PendingGeneration>();
-
-export const getActiveGeneration = (generationId: string): ActiveGeneration | undefined => {
-    return activeGenerations.get(generationId);
-};
-
-export const getPendingGeneration = (generationId: string): PendingGeneration | undefined => {
-    return pendingGenerations.get(generationId);
-};
-
-export const isGenerationActive = (generationId: string): boolean => {
-    return activeGenerations.has(generationId);
-};
-
-export const isGenerationPending = (generationId: string): boolean => {
-    return pendingGenerations.has(generationId);
-};
+const activeGenerationStreams = new Map<string, ActiveGenerationStream>();
+const db = useDrizzle();
 
+/**
+ * Add a client connection to an active generation stream
+ */
 export const addClientToGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): boolean => {
-    const generation = activeGenerations.get(generationId);
-    if (!generation) {
+    const stream = activeGenerationStreams.get(generationId);
+    if (!stream) {
         return false;
     }
-    generation.clients.add(controller);
+    stream.clients.add(controller);
     return true;
 };
 
+/**
+ * Remove a client connection from an active generation stream
+ */
 export const removeClientFromGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): void => {
-    const generation = activeGenerations.get(generationId);
-    if (generation) {
-        generation.clients.delete(controller);
-        if (generation.clients.size === 0 && generation.complete) {
-            activeGenerations.delete(generationId);
+    const stream = activeGenerationStreams.get(generationId);
+    if (stream) {
+        stream.clients.delete(controller);
+        // Clean up if no clients left and generation is complete
+        if (stream.clients.size === 0 && !stream.isGenerating) {
+            activeGenerationStreams.delete(generationId);
         }
     }
 };
 
-const sendToClients = (generation: ActiveGeneration, event: GenerationStreamEvent): void => {
-    for (const client of generation.clients) {
+/**
+ * Send an event to all connected clients for a generation
+ */
+export const sendToClients = (generationId: string, event: GenerationStreamEvent): void => {
+    const stream = activeGenerationStreams.get(generationId);
+    if (!stream) return;
+
+    for (const client of stream.clients) {
         sendToClient(client, event);
     }
 };
 
+/**
+ * Send an event to a single client
+ */
 export const sendToClient = (client: ReadableStreamDefaultController<Uint8Array>, event: GenerationStreamEvent): void => {
     const data = JSON.stringify(event);
     const encoder = new TextEncoder();
@@ -73,12 +66,15 @@ export const sendToClient = (client: ReadableStreamDefaultController<Uint8Array>
     } catch (error) {
         console.error('Failed to send to client:', error);
     }
-}
+};
 
+/**
+ * Build a prompt from chat messages
+ */
 const buildPrompt = (messages: ChatMessage[]): string => {
     return messages
         .map((msg: ChatMessage) => {
-            const roleMap: Record<MessageType, string> = {
+            const roleMap: Record<typeof msg.type, string> = {
                 system: 'System',
                 user: 'User',
                 agent: 'Assistant'
@@ -88,121 +84,200 @@ const buildPrompt = (messages: ChatMessage[]): string => {
         .join('\n\n');
 };
 
-const db = useDrizzle();
-
-export const startGeneration = async (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): Promise<void> => {
-    const pending = pendingGenerations.get(generationId);
-
-    if (!pending) {
-        console.error(`Generation ${generationId} not found in pending generations`);
-        return;
-    }
-
-    clearTimeout(pending.timeout);
-    pendingGenerations.delete(generationId);
-
-    const { userId, topicId, messages } = pending;
-    const prompt = buildPrompt(messages);
-
-    const generation: ActiveGeneration = {
+/**
+ * Create a new pending generation in the database
+ * Returns the generation ID
+ */
+export const createPendingGeneration = async (userId: string, topicId: string, messages: ChatMessage[], regeneratesFrom?: string): Promise<string> => {
+    const generationValues: any = {
         userId,
         topicId,
-        messages,
-        content: '',
-        clients: new Set([controller]),
-        complete: false
+        status: 'pending' as GenerationStatus,
     };
 
-    await db.insert(generations).values({
-        id: generationId,
-        userId,
-        topicId,
-        messageId: null
-    });
+    if (regeneratesFrom) {
+        generationValues.regeneratesFrom = regeneratesFrom;
+    }
 
-    activeGenerations.set(generationId, generation);
+    const [generation] = await db
+        .insert(generations)
+        .values(generationValues)
+        .returning();
 
+    return generation.id;
+};
+
+/**
+ * Get current generation status and content from database
+ */
+export const getGenerationStatus = async (generationId: string) => {
+    const [generation] = await db
+        .select()
+        .from(generations)
+        .where(eq(generations.id, generationId));
+
+    return generation || null;
+};
+
+/**
+ * Start a generation: update status to active and begin streaming
+ * This is called when a client connects to the stream
+ */
+export const startGeneration = async (
+    generationId: string,
+    userId: string,
+    topicId: string,
+    messages: ChatMessage[],
+    controller: ReadableStreamDefaultController<Uint8Array>
+): Promise<void> => {
     try {
+        // Get current generation from database
+        const generation = await getGenerationStatus(generationId);
+        if (!generation) {
+            throw new Error('Generation not found');
+        }
+
+        // Create active stream tracking
+        activeGenerationStreams.set(generationId, {
+            userId,
+            topicId,
+            clients: new Set([controller]),
+            isGenerating: true,
+        });
+
+        // Update status to active
+        await db
+            .update(generations)
+            .set({
+                status: 'active' as GenerationStatus,
+                startedAt: new Date(),
+            })
+            .where(eq(generations.id, generationId));
+
+        sendToClients(generationId, {
+            type: 'start',
+            data: null,
+        });
+
+        const prompt = buildPrompt(messages);
         const dummyResponse = generateDummyResponse(prompt, messages);
         const tokens = dummyResponse.split(' ');
 
+        // Simulate token streaming
+        let accumulatedContent = '';
         for (const token of tokens) {
             await new Promise(resolve => setTimeout(resolve, 50));
 
-            generation.content += token + ' ';
+            accumulatedContent += token + ' ';
 
-            sendToClients(generation, {
+            sendToClients(generationId, {
                 type: 'token',
-                data: token + ' '
+                data: token + ' ',
             });
         }
 
-        const [message] = await db.insert(messages_drizzle).values({
-            topicId: generation.topicId,
-            userId: generation.userId,
+        const finalContent = accumulatedContent.trim();
+
+        // Create message record for this generation
+        const messageValues: any = {
+            topicId,
+            userId,
+            content: finalContent,
             isUser: false,
-            content: generation.content.trim(),
-            model: 'dummy-model-v1',
-            tokensGenerated: tokens.length,
-            tokensUsedThinking: 0
-        }).returning();
+        };
 
-        generation.complete = true;
+        // If this is a regeneration, set the regeneratedFromId
+        if (generation.regeneratesFrom) {
+            messageValues.regeneratedFromId = generation.regeneratesFrom;
+            messageValues.isRegenerated = true;
+        }
 
-        sendToClients(generation, {
+        let [message] = await db
+            .insert(messages_drizzle)
+            .values(messageValues)
+            .returning();
+
+        // Update generation as completed
+        await db
+            .update(generations)
+            .set({
+                status: 'completed' as GenerationStatus,
+                completedAt: new Date(),
+                messageId: message.id,
+            })
+            .where(eq(generations.id, generationId));
+
+        let fmessage = await db.select().from(messages_drizzle).where(eq(messages_drizzle.userId, userId)).leftJoin(generations, eq(messages_drizzle.id, generations.messageId))
+        console.log(fmessage);
+
+        // Mark stream as no longer generating
+        const stream = activeGenerationStreams.get(generationId);
+        if (stream) {
+            stream.isGenerating = false;
+        }
+
+        sendToClients(generationId, {
             type: 'complete',
-            data: message
+            data: fmessage,
         });
 
-        if (generation.clients.size === 0) {
-            activeGenerations.delete(generationId);
-        } else {
-            generation.clients.forEach((client) => {
+        // Close all client connections
+        const finalStream = activeGenerationStreams.get(generationId);
+        if (finalStream) {
+            for (const client of finalStream.clients) {
                 try {
                     client.close();
-                } catch {
+                } catch (error) {
+                    // Client already closed
                 }
-            });
-            activeGenerations.delete(generationId);
+            }
+            activeGenerationStreams.delete(generationId);
         }
-
-        await db.update(generations).set({
-            messageId: message.id
-        }).where(eq(generations.id, generationId));
-
     } catch (error) {
         console.error('Generation failed:', error);
 
-        sendToClients(generation, {
+        const errorMessage = error instanceof Error ? error.message : 'Unknown error';
+
+        // Update generation as failed
+        await db
+            .update(generations)
+            .set({
+                status: 'failed' as GenerationStatus,
+                error: errorMessage,
+                completedAt: new Date(),
+            })
+            .where(eq(generations.id, generationId));
+
+        sendToClients(generationId, {
             type: 'error',
-            data: error instanceof Error ? error.message : 'Unknown error'
+            data: { error: errorMessage },
         });
 
-        await db.delete(generations).where(eq(generations.id, generationId));
-
-        activeGenerations.delete(generationId);
+        // Close all client connections
+        const stream = activeGenerationStreams.get(generationId);
+        if (stream) {
+            for (const client of stream.clients) {
+                try {
+                    client.close();
+                } catch (error) {
+                    // Client already closed
+                }
+            }
+            activeGenerationStreams.delete(generationId);
+        }
     }
 };
 
-export const registerPendingGeneration = (userId: string, generationId: string, topicId: string, messages: ChatMessage[]): void => {
-    const timeout = setTimeout(() => {
-        const gen = pendingGenerations.get(generationId)
-        if (gen) gen.expired = true;
-        console.log(`Generation ${generationId} expired - no client connected within 60 seconds`);
-    }, 60000);
-
-    pendingGenerations.set(generationId, {
-        generationId,
-        userId,
-        topicId,
-        messages,
-        timeout,
-        expired: false
-    });
-
-    console.log(`Registered pending generation ${generationId}, waiting for client connection...`);
+/**
+ * Check if a generation is currently being streamed
+ */
+export const isGenerationStreaming = (generationId: string): boolean => {
+    return activeGenerationStreams.has(generationId);
 };
 
+/**
+ * Generate a dummy response for testing
+ */
 const generateDummyResponse = (prompt: string, messages: ChatMessage[]): string => {
     const responses = [
         "This is a simulated response to your prompt. In a real implementation, this would be generated by an AI model like GPT-4 or Claude."
@@ -227,4 +302,4 @@ const generateDummyResponse = (prompt: string, messages: ChatMessage[]): string
     }
 
     return responses[Math.floor(Math.random() * responses.length)];
-};
+};