feat: stream persistent

frontend/src/components/chatManager.ts

@@ -6,6 +6,7 @@ import {
   getSettings,
   getModels,
   sendMessage,
+  streamChatUpdates,
   getChatMessages,
   listChats,
   deleteChat,
@@ -41,6 +42,7 @@ Alpine.data('chatManager', () => ({
   selectedChatID: null as string | null,
   chatListOpen: false,
   loading: false,
+  activeStreamChatID: null as string | null,
 
   async init() {
     // Acquire Data
@@ -109,66 +111,27 @@ Alpine.data('chatManager', () => ({
       this.selectedChatID = IN_PROGRESS_UUID;
     }
 
-    // New User Message
-    let userMessage: Message = {
+    // Add Optimistic User Message
+    const currentChat: Chat = this.chats.find(
+      (c) => c.id === this.selectedChatID,
+    )!;
+    currentChat.messages.push({
       id: IN_PROGRESS_UUID,
       chat_id: this.selectedChatID,
       role: 'user',
       thinking: '',
       content: message,
       created_at: new Date().toISOString(),
-    };
-
-    // Get Chat
-    let currentChat: Chat = this.chats.find(
-      (c) => c.id === this.selectedChatID,
-    )!;
-
-    // Add User Message
-    currentChat.messages.push(userMessage);
+    });
     currentChat.message_count += 1;
 
-    // Assistant Message Placeholder
-    let assistantMessage: Message | undefined;
-
     try {
       await sendMessage(
         this.selectedChatID === IN_PROGRESS_UUID ? '' : this.selectedChatID,
         { model: this.selectedModel, prompt: message },
         (chunk: MessageChunk) => {
-          // Handle Chat
-          if (chunk.chat) {
-            Object.assign(currentChat, {
-              ...chunk.chat,
-              messages: currentChat.messages,
-            });
-            this.selectedChatID = chunk.chat.id;
-            this.updateHash(chunk.chat.id);
-          }
-
-          // Handle User Message
-          if (chunk.user_message) {
-            Object.assign(userMessage, chunk.user_message);
-          }
-
-          // Handle Assistant Message
-          if (chunk.assistant_message) {
-            if (!assistantMessage) {
-              assistantMessage = chunk.assistant_message;
-              currentChat.messages.push(assistantMessage);
-            } else {
-              const index = currentChat.messages.findIndex(
-                (m) => m.id === assistantMessage!.id,
-              );
-              if (index !== -1) {
-                currentChat.messages[index] = {
-                  ...assistantMessage,
-                  ...chunk.assistant_message,
-                };
-                currentChat.messages = [...currentChat.messages];
-              }
-            }
-          }
+          if (chunk.chat) this.activeStreamChatID = chunk.chat.id;
+          this.applyMessageChunk(chunk);
         },
       );
     } catch (err) {
@@ -176,9 +139,56 @@ Alpine.data('chatManager', () => ({
       this.error = parseError(err);
     } finally {
       this.loading = false;
+      this.activeStreamChatID = null;
     }
   },
 
+  applyMessageChunk(chunk: MessageChunk) {
+    // Handle Chat
+    if (chunk.chat) {
+      let chat = this.chats.find((c) => c.id === chunk.chat!.id);
+      if (!chat) chat = this.chats.find((c) => c.id === IN_PROGRESS_UUID);
+      if (!chat) {
+        chat = { ...chunk.chat, messages: chunk.chat.messages || [] };
+        this.chats.unshift(chat);
+      } else {
+        Object.assign(chat, chunk.chat);
+        chat.messages = chunk.chat.messages?.length
+          ? chunk.chat.messages
+          : chat.messages;
+      }
+      this.selectedChatID = chunk.chat.id;
+      this.updateHash(chunk.chat.id);
+    }
+
+    const chatID = chunk.chat?.id || this.selectedChatID;
+    const currentChat = this.chats.find((c) => c.id === chatID);
+    if (!currentChat) return;
+
+    // Handle Messages
+    if (chunk.user_message) this.upsertMessage(currentChat, chunk.user_message);
+    if (chunk.assistant_message)
+      this.upsertMessage(currentChat, chunk.assistant_message);
+  },
+
+  upsertMessage(chat: Chat, message: Message) {
+    // Upsert Message
+    const existingIndex = chat.messages.findIndex(
+      (m) =>
+        m.id === message.id ||
+        (m.id === IN_PROGRESS_UUID && m.role === message.role),
+    );
+    if (existingIndex === -1) {
+      chat.messages.push(message);
+    } else {
+      chat.messages[existingIndex] = {
+        ...chat.messages[existingIndex],
+        ...message,
+      };
+    }
+    chat.messages = [...chat.messages];
+  },
+
   updateHash(chatID: string | null) {
     const newRoute = CHAT_ROUTE + (chatID ? '/' + chatID : '');
     window.history.pushState(null, '', newRoute);
@@ -202,13 +212,46 @@ Alpine.data('chatManager', () => ({
         (c) => c.id == this.selectedChatID,
       );
 
-      this.chats[chatIndex].messages = response.messages || [];
+      if (chatIndex === -1) return;
+      this.chats[chatIndex] = {
+        ...this.chats[chatIndex],
+        ...response,
+        messages: response.messages || [],
+      };
+      await this.reconnectChatStream(response);
     } catch (err) {
       console.error('Error loading chat messages:', err);
       this.error = 'Failed to load messages';
     }
   },
 
+  async reconnectChatStream(chat: Chat) {
+    const latestMessage = chat.messages[chat.messages.length - 1];
+    if (
+      !latestMessage ||
+      latestMessage.role !== 'assistant' ||
+      latestMessage.status !== 'streaming' ||
+      this.activeStreamChatID === chat.id
+    ) {
+      return;
+    }
+
+    // Reconnect Stream
+    this.loading = true;
+    this.activeStreamChatID = chat.id;
+    try {
+      await streamChatUpdates(chat.id, (chunk: MessageChunk) =>
+        this.applyMessageChunk(chunk),
+      );
+    } catch (err) {
+      console.error('Error reconnecting chat stream:', err);
+      this.error = parseError(err);
+    } finally {
+      this.loading = false;
+      this.activeStreamChatID = null;
+    }
+  },
+
   get models(): Model[] {
     if (!this.settings.text_generation_selector) return this._models;
     return applyFilter(this._models, this.settings.text_generation_selector);