Update public/index.html

public/index.html

@@ -124,403 +124,517 @@
         </div>
     </main>
 
-    <script>
-        // NEW FEATURE: Toggle high-cost rendering during stream to prevent memory crashes
-        let enableHighlighting = true;
-        
-        marked.setOptions({
-            highlight: function(code, lang) {
-                if (!enableHighlighting) return code; // Skip highlighting while streaming!
-                const language = hljs.getLanguage(lang) ? lang : 'plaintext';
-                try { return hljs.highlight(code, { language }).value; } catch(e) { return code; }
-            }
-        });
-
-        let currentChatId = null;
-        let attachedImages =[];
-        let pollingInterval = null;
-        let currentTokens = { total: 0, in: 0, out: 0 };
-
-        function toggleInputState(isGenerating) {
-            const sendBtn = document.getElementById('send-btn');
-            const stopBtn = document.getElementById('stop-btn');
-            if (isGenerating) {
-                sendBtn.classList.add('hidden');
-                stopBtn.classList.remove('hidden');
-            } else {
-                sendBtn.classList.remove('hidden');
-                stopBtn.classList.add('hidden');
-            }
-        }
+    <!-- ================================= -->
+<!-- CLIENT SCRIPT (REPLACE SCRIPT TAG) -->
+<!-- ================================= -->
+
+<script>
+
+marked.setOptions({
+  highlight(code, lang) {
+    const language =
+      hljs.getLanguage(lang)
+        ? lang
+        : 'plaintext';
+
+    return hljs.highlight(code, {
+      language
+    }).value;
+  }
+});
+
+let currentChatId = null;
+let attachedImages = [];
+
+const chatWindow =
+  document.getElementById('chat-window');
+
+const scrollBtn =
+  document.getElementById('scroll-bottom-btn');
+
+// =================================
+// SAFE HTML
+// =================================
+
+function escapeHtml(str = "") {
+  return str
+    .replaceAll("&", "&amp;")
+    .replaceAll("<", "&lt;")
+    .replaceAll(">", "&gt;")
+    .replaceAll('"', "&quot;")
+    .replaceAll("'", "&#39;");
+}
+
+// =================================
+// TOKEN DISPLAY
+// =================================
+
+function updateTokenDisplay(
+  total,
+  input,
+  output
+) {
+
+  document.getElementById('token-total')
+    .innerText =
+      (total || 0).toLocaleString();
+
+  document.getElementById('token-in')
+    .innerText =
+      (input || 0).toLocaleString();
+
+  document.getElementById('token-out')
+    .innerText =
+      (output || 0).toLocaleString();
+}
+
+// =================================
+// RENDER SINGLE MESSAGE
+// =================================
+
+function createMessageElement(
+  role,
+  content = "",
+  reasoning = ""
+) {
 
-        async function stopGeneration() {
-            if (!currentChatId) return;
-            toggleInputState(false);
-            try {
-                await fetch(`/api/chats/${currentChatId}/stop`, { method: 'POST' });
-            } catch (err) {
-                console.error("Failed to send stop signal", err);
-            }
-        }
+  const wrapper =
+    document.createElement('div');
 
-        const chatWindow = document.getElementById('chat-window');
-        const scrollBtn = document.getElementById('scroll-bottom-btn');
+  wrapper.className =
+    `flex ${
+      role === 'user'
+        ? 'justify-end'
+        : 'justify-start'
+    } w-full`;
 
-        chatWindow.addEventListener('scroll', () => {
-            const distanceToBottom = chatWindow.scrollHeight - chatWindow.scrollTop - chatWindow.clientHeight;
-            if (distanceToBottom > 100) scrollBtn.classList.remove('hidden');
-            else scrollBtn.classList.add('hidden');
-        });
+  const bubble =
+    document.createElement('div');
 
-        function scrollToBottom() {
-            chatWindow.scrollTo({ top: chatWindow.scrollHeight, behavior: 'smooth' });
-        }
+  bubble.className =
+    `max-w-[95%] lg:max-w-[85%]
+    rounded-2xl p-4 shadow-sm
+    ${
+      role === 'user'
+        ? 'bg-blue-600 text-white rounded-br-sm'
+        : 'bg-gray-850 border border-gray-800 text-gray-200 markdown-body rounded-bl-sm'
+    }`;
 
-        function enableTitleEdit() {
-            if (!currentChatId) return;
-            const currentTitle = document.getElementById('current-chat-title').innerText;
-            document.getElementById('title-display').classList.add('hidden');
-            document.getElementById('title-edit').classList.remove('hidden');
-            document.getElementById('title-edit').classList.add('flex');
-            const input = document.getElementById('title-input');
-            input.value = currentTitle;
-            input.focus();
-            input.select();
-        }
+  if (role === 'user') {
 
-        function cancelTitleEdit() {
-            document.getElementById('title-edit').classList.add('hidden');
-            document.getElementById('title-edit').classList.remove('flex');
-            document.getElementById('title-display').classList.remove('hidden');
-        }
+    bubble.innerHTML =
+      `<div class="whitespace-pre-wrap">${escapeHtml(content)}</div>`;
 
-        async function saveTitle() {
-            if (!currentChatId) return cancelTitleEdit();
-            const newTitle = document.getElementById('title-input').value.trim();
-            if (!newTitle) return cancelTitleEdit();
+  } else {
 
-            document.getElementById('current-chat-title').innerText = newTitle;
-            cancelTitleEdit();
+    let html = "";
 
-            await fetch(`/api/chats/${currentChatId}/title`, {
-                method: 'PUT',
-                headers: { 'Content-Type': 'application/json' },
-                body: JSON.stringify({ title: newTitle })
-            });
-            loadSidebar(); 
-        }
+    if (reasoning) {
+      html += `
+      <div class="reasoning-block">
+        <i>Thinking Process</i><br/>
+        ${DOMPurify.sanitize(
+          marked.parse(reasoning)
+        )}
+      </div>`;
+    }
 
-        document.getElementById('title-input').addEventListener('keydown', (e) => {
-            if (e.key === 'Enter') saveTitle();
-            if (e.key === 'Escape') cancelTitleEdit();
-        });
+    html += DOMPurify.sanitize(
+      marked.parse(content)
+    );
 
-        function toggleSidebar() {
-            const sidebar = document.getElementById('sidebar');
-            const overlay = document.getElementById('mobile-overlay');
-            sidebar.classList.toggle('-translate-x-full');
-            overlay.classList.toggle('hidden');
-        }
+    bubble.innerHTML = html;
+  }
 
-        const textarea = document.getElementById('message-input');
-        textarea.addEventListener('input', function() {
-            this.style.height = 'auto';
-            this.style.height = (this.scrollHeight < 128 ? this.scrollHeight : 128) + 'px';
-        });
-
-        function updateTokenDisplay(total, input, output) {
-            currentTokens = { total, in: input, out: output };
-            document.getElementById('token-total').innerText = (total || 0).toLocaleString();
-            document.getElementById('token-in').innerText = (input || 0).toLocaleString();
-            document.getElementById('token-out').innerText = (output || 0).toLocaleString();
-        }
+  wrapper.appendChild(bubble);
 
-        async function loadSidebar() {
-            const res = await fetch('/api/chats');
-            const chats = await res.json();
-            const list = document.getElementById('chat-list');
-            list.innerHTML = chats.map(c => `
-                <div onclick="selectChat('${c.id}')" class="p-3 rounded-xl cursor-pointer transition border border-transparent ${c.id === currentChatId ? 'bg-gray-800 border-gray-700' : 'hover:bg-gray-800/50'}">
-                    <div class="text-sm font-medium truncate text-gray-200">${c.title}</div>
-                    <div class="text-[10px] text-gray-500 mt-1 flex justify-between">
-                        <span>Total: ${c.totalTokens.toLocaleString()}</span>
-                        <span class="opacity-70">↑${(c.outputTokens || 0).toLocaleString()}</span>
-                    </div>
-                </div>
-            `).join('');
-        }
+  return wrapper;
+}
 
-        async function createNewChat() {
-            const res = await fetch('/api/chats', { method: 'POST' });
-            const chat = await res.json();
-            selectChat(chat.id);
-            loadSidebar();
-            if(window.innerWidth < 1024) toggleSidebar();
-        }
+// =================================
+// RENDER CHAT
+// =================================
 
-        async function deleteCurrentChat() {
-            if (!currentChatId) return;
-            if (confirm("Permanently delete this chat?")) {
-                await fetch(`/api/chats/${currentChatId}`, { method: 'DELETE' });
-                currentChatId = null;
-                document.getElementById('chat-window').innerHTML = '';
-                document.getElementById('current-chat-title').innerText = 'Select or create a chat';
-                document.getElementById('edit-title-btn').classList.add('hidden');
-                updateTokenDisplay(0, 0, 0);
-                toggleInputState(false);
-                loadSidebar();
-            }
-        }
+function renderMessages(messages) {
 
-        async function selectChat(id) {
-            currentChatId = id;
-            if (pollingInterval) clearInterval(pollingInterval);
-
-            const res = await fetch(`/api/chats/${id}`);
-            const chat = await res.json();
-            
-            document.getElementById('current-chat-title').innerText = chat.title;
-            document.getElementById('edit-title-btn').classList.remove('hidden'); 
-            
-            updateTokenDisplay(chat.totalTokens, chat.inputTokens, chat.outputTokens);
-            cancelTitleEdit(); 
-            renderMessages(chat.messages);
-            loadSidebar(); 
-
-            if(window.innerWidth < 1024 && !document.getElementById('sidebar').classList.contains('-translate-x-full')) {
-                toggleSidebar();
-            }
-
-            toggleInputState(chat.isGenerating);
-            if (chat.isGenerating) pollGeneratingChat(id);
-        }
+  chatWindow.innerHTML = '';
 
-        function renderMessages(messages) {
-            const container = document.getElementById('chat-window');
-            const isScrolledToBottom = container.scrollHeight - container.clientHeight <= container.scrollTop + 50;
-
-            container.innerHTML = messages.map(m => {
-                let html = '';
-                if (m.reasoning) {
-                    html += `<div class="reasoning-block"><i>Thinking Process</i><br/>${marked.parse(m.reasoning)}</div>`;
-                }
-                if (m.content) {
-                    html += DOMPurify.sanitize(marked.parse(m.content));
-                }
-                
-                return `
-                    <div class="flex ${m.role === 'user' ? 'justify-end' : 'justify-start'} w-full">
-                        <div class="max-w-[95%] lg:max-w-[85%] rounded-2xl p-4 shadow-sm ${m.role === 'user' ? 'bg-blue-600 text-white rounded-br-sm' : 'bg-gray-850 border border-gray-800 text-gray-200 markdown-body rounded-bl-sm'}">
-                            ${m.role === 'user' ? `<div class="whitespace-pre-wrap">${m.content}</div>` : html}
-                        </div>
-                    </div>
-                `;
-            }).join('');
-            
-            document.querySelectorAll('.markdown-body pre').forEach(pre => {
-                if (!pre.querySelector('.copy-btn')) {
-                    const btn = document.createElement('button');
-                    btn.className = 'copy-btn';
-                    btn.innerText = 'Copy';
-                    btn.onclick = () => {
-                        const codeText = pre.querySelector('code')?.innerText || pre.innerText.replace('Copy', '');
-                        navigator.clipboard.writeText(codeText.trim());
-                        btn.innerText = 'Copied!';
-                        setTimeout(() => btn.innerText = 'Copy', 2000);
-                    };
-                    pre.appendChild(btn);
-                }
-            });
-
-            // Prevent unwanted jump to bottom if chat is empty or user scrolled up
-            if (isScrolledToBottom && messages.length > 0) scrollToBottom();
+  messages.forEach(m => {
+
+    chatWindow.appendChild(
+      createMessageElement(
+        m.role,
+        m.content,
+        m.reasoning
+      )
+    );
+  });
+
+  addCopyButtons();
+
+  scrollToBottom();
+}
+
+// =================================
+// COPY BUTTONS
+// =================================
+
+function addCopyButtons() {
+
+  document
+    .querySelectorAll('.markdown-body pre')
+    .forEach(pre => {
+
+      if (
+        pre.querySelector('.copy-btn')
+      ) return;
+
+      const btn =
+        document.createElement('button');
+
+      btn.className = 'copy-btn';
+
+      btn.innerText = 'Copy';
+
+      btn.onclick = () => {
+
+        const codeText =
+          pre.querySelector('code')
+            ?.innerText || '';
+
+        navigator.clipboard.writeText(
+          codeText.trim()
+        );
+
+        btn.innerText = 'Copied';
+
+        setTimeout(() => {
+          btn.innerText = 'Copy';
+        }, 2000);
+      };
+
+      pre.appendChild(btn);
+    });
+}
+
+// =================================
+// SCROLL
+// =================================
+
+function scrollToBottom() {
+
+  chatWindow.scrollTo({
+    top: chatWindow.scrollHeight,
+    behavior: 'smooth'
+  });
+}
+
+// =================================
+// SIDEBAR
+// =================================
+
+async function loadSidebar() {
+
+  const res =
+    await fetch('/api/chats');
+
+  const chats =
+    await res.json();
+
+  document.getElementById('chat-list')
+    .innerHTML = chats.map(c => `
+
+      <div
+        onclick="selectChat('${c.id}')"
+        class="p-3 rounded-xl cursor-pointer
+        ${
+          c.id === currentChatId
+            ? 'bg-gray-800'
+            : 'hover:bg-gray-800/50'
+        }">
+
+        <div class="text-sm font-medium truncate text-gray-200">
+          ${escapeHtml(c.title)}
+        </div>
+
+      </div>
+
+    `).join('');
+}
+
+// =================================
+// SELECT CHAT
+// =================================
+
+async function selectChat(id) {
+
+  currentChatId = id;
+
+  const res =
+    await fetch(`/api/chats/${id}`);
+
+  const chat =
+    await res.json();
+
+  document.getElementById(
+    'current-chat-title'
+  ).innerText = chat.title;
+
+  renderMessages(chat.messages);
+
+  updateTokenDisplay(
+    chat.totalTokens,
+    chat.inputTokens,
+    chat.outputTokens
+  );
+
+  loadSidebar();
+}
+
+// =================================
+// CREATE CHAT
+// =================================
+
+async function createNewChat() {
+
+  const res =
+    await fetch('/api/chats', {
+      method: 'POST'
+    });
+
+  const chat =
+    await res.json();
+
+  await selectChat(chat.id);
+}
+
+// =================================
+// SEND MESSAGE
+// =================================
+
+async function sendMessage() {
+
+  const input =
+    document.getElementById(
+      'message-input'
+    );
+
+  const text =
+    input.value.trim();
+
+  if (!text || !currentChatId)
+    return;
+
+  // ADD USER MESSAGE
+  const userElement =
+    createMessageElement(
+      'user',
+      text
+    );
+
+  chatWindow.appendChild(userElement);
+
+  // ADD AI PLACEHOLDER
+  const aiWrapper =
+    document.createElement('div');
+
+  aiWrapper.className =
+    'flex justify-start w-full';
+
+  aiWrapper.innerHTML = `
+    <div
+      class="max-w-[95%] lg:max-w-[85%]
+      rounded-2xl rounded-bl-sm p-4
+      bg-gray-850 border border-gray-800
+      text-gray-200 markdown-body"
+      id="streaming-ai-msg">
+
+      <span class="animate-pulse text-gray-400">
+        Thinking...
+      </span>
+
+    </div>
+  `;
+
+  chatWindow.appendChild(aiWrapper);
+
+  scrollToBottom();
+
+  input.value = '';
+
+  let aiContent = "";
+  let aiReasoning = "";
+
+  try {
+
+    const response =
+      await fetch(
+        `/api/chats/${currentChatId}/stream`,
+        {
+          method: 'POST',
+
+          headers: {
+            'Content-Type':
+              'application/json'
+          },
+
+          body: JSON.stringify({
+            model:
+              document.getElementById(
+                'model-select'
+              ).value,
+
+            prompt: text,
+
+            images: attachedImages
+          })
         }
+      );
 
-        function pollGeneratingChat(id) {
-            if (pollingInterval) clearInterval(pollingInterval);
-            
-            pollingInterval = setInterval(async () => {
-                const res = await fetch(`/api/chats/${id}`);
-                const chat = await res.json();
-                
-                renderMessages(chat.messages);
-                updateTokenDisplay(chat.totalTokens, chat.inputTokens, chat.outputTokens);
-                
-                if (chat.title !== document.getElementById('current-chat-title').innerText && document.getElementById('title-display').classList.contains('hidden') === false) {
-                     document.getElementById('current-chat-title').innerText = chat.title;
-                     loadSidebar();
-                }
-
-                if (!chat.isGenerating) {
-                    clearInterval(pollingInterval);
-                    if (currentChatId === id) toggleInputState(false);
-                }
-            }, 1500); 
+    const reader =
+      response.body
+        .pipeThrough(
+          new TextDecoderStream()
+        )
+        .getReader();
+
+    let buffer = "";
+
+    while (true) {
+
+      const {
+        value,
+        done
+      } = await reader.read();
+
+      if (done) break;
+
+      buffer += value;
+
+      const events =
+        buffer.split("\n\n");
+
+      buffer = events.pop();
+
+      for (const event of events) {
+
+        const lines =
+          event.split("\n");
+
+        let type = "";
+        let data = "";
+
+        for (const line of lines) {
+
+          if (
+            line.startsWith("event:")
+          ) {
+            type =
+              line.replace(
+                "event:",
+                ""
+              ).trim();
+          }
+
+          if (
+            line.startsWith("data:")
+          ) {
+            data =
+              line.replace(
+                "data:",
+                ""
+              ).trim();
+          }
         }
 
-        function handleImageSelect(event) {
-            const files = event.target.files;
-            const container = document.getElementById('image-preview-container');
-            container.classList.remove('hidden');
-            
-            Array.from(files).forEach(file => {
-                const reader = new FileReader();
-                reader.onload = (e) => {
-                    attachedImages.push(e.target.result);
-                    container.innerHTML += `
-                        <div class="relative group">
-                            <img src="${e.target.result}" class="h-14 w-14 object-cover rounded-lg border border-gray-600 shadow-sm">
-                        </div>`;
-                };
-                reader.readAsDataURL(file);
-            });
+        if (!data) continue;
+
+        const parsed =
+          JSON.parse(data);
+
+        if (type === 'thinking') {
+
+          aiReasoning += parsed.text;
+
+        } else if (type === 'text') {
+
+          aiContent += parsed.text;
+
+        } else if (type === 'usage') {
+
+          updateTokenDisplay(
+            parsed.totalTokens,
+            parsed.inputTokens,
+            parsed.outputTokens
+          );
         }
 
-        async function sendMessage() {
-            const input = document.getElementById('message-input');
-            const text = input.value.trim();
-            if (!text || !currentChatId) return;
-
-            const payload = {
-                model: document.getElementById('model-select').value,
-                prompt: text,
-                images: attachedImages
-            };
-
-            const container = document.getElementById('chat-window');
-            container.innerHTML += `
-                <div class="flex justify-end w-full">
-                    <div class="max-w-[95%] lg:max-w-[85%] rounded-2xl rounded-br-sm p-4 bg-blue-600 text-white shadow-sm whitespace-pre-wrap">${text}</div>
-                </div>
-                <div class="flex justify-start w-full" id="temp-ai-wrapper">
-                    <div class="max-w-[95%] lg:max-w-[85%] rounded-2xl rounded-bl-sm p-4 bg-gray-850 border border-gray-800 text-gray-200 markdown-body shadow-sm" id="temp-ai-msg">
-                        <span class="flex items-center gap-2 text-gray-400">
-                            <svg class="animate-spin h-4 w-4" viewBox="0 0 24 24"><circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4" fill="none"></circle><path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path></svg>
-                            Thinking...
-                        </span>
-                    </div>
-                </div>
-            `;
-            scrollToBottom();
-            
-            input.value = '';
-            input.style.height = 'auto'; 
-            attachedImages =[];
-            document.getElementById('image-preview-container').innerHTML = '';
-            document.getElementById('image-preview-container').classList.add('hidden');
-            
-            toggleInputState(true);
-
-            try {
-                const response = await fetch(`/api/chats/${currentChatId}/stream`, {
-                    method: 'POST',
-                    headers: { 'Content-Type': 'application/json' },
-                    body: JSON.stringify(payload)
-                });
-
-                if (!response.ok) throw new Error("API stream rejected. Chat might be busy.");
-
-                const reader = response.body.getReader();
-                const decoder = new TextDecoder("utf-8");
-                let aiContent = "";
-                let aiReasoning = "";
-                
-                enableHighlighting = false; // SPEED UP: Disable expensive highlighting during stream
-                let lastRenderTime = Date.now();
+        let html = "";
 
-                while (true) {
-                    const { value, done } = await reader.read();
-                    if (done) break;
-                    
-                    const chunkText = decoder.decode(value, { stream: true });
-                    const chunks = chunkText.split(/(__THINK__|__USAGE__)/);
-                    let i = 0;
-                    
-                    while (i < chunks.length) {
-                        if (chunks[i] === '__THINK__') {
-                            aiReasoning += chunks[i+1] || '';
-                            i += 2;
-                        } else if (chunks[i] === '__USAGE__') {
-                            try {
-                                if(chunks[i+1]) {
-                                    const usageData = JSON.parse(chunks[i+1]);
-                                    updateTokenDisplay(
-                                        currentTokens.total + usageData.totalTokens,
-                                        currentTokens.in + usageData.inputTokens,
-                                        currentTokens.out + usageData.outputTokens
-                                    );
-                                }
-                            } catch(e) {}
-                            i += 2;
-                        } else {
-                            if (chunks[i]) aiContent += chunks[i];
-                            i++;
-                        }
-                    }
-
-                    // THROTTLE RENDER: Only update DOM max 5 times a second to prevent memory crashes
-                    if (Date.now() - lastRenderTime > 200) {
-                        let tempHtml = "";
-                        if (aiReasoning) tempHtml += `<div class="reasoning-block"><i>Thinking Process</i><br/>${marked.parse(aiReasoning)}</div>`;
-                        if (aiContent) tempHtml += DOMPurify.sanitize(marked.parse(aiContent));
-                        
-                        document.getElementById('temp-ai-msg').innerHTML = tempHtml || "<span class='animate-pulse text-gray-400'>Thinking...</span>";
-                        
-                        const distanceToBottom = container.scrollHeight - container.scrollTop - container.clientHeight;
-                        if (distanceToBottom < 100) scrollToBottom();
-                        
-                        lastRenderTime = Date.now();
-                    }
-                }
-
-                // FINAL RENDER: Stream is complete, turn highlighting back on and do a final pristine render
-                enableHighlighting = true;
-                
-                let finalHtml = "";
-                if (aiReasoning) finalHtml += `<div class="reasoning-block"><i>Thinking Process</i><br/>${marked.parse(aiReasoning)}</div>`;
-                if (aiContent) finalHtml += DOMPurify.sanitize(marked.parse(aiContent));
-                
-                const finalMsgElement = document.getElementById('temp-ai-msg');
-                finalMsgElement.innerHTML = finalHtml || "Done.";
-                finalMsgElement.removeAttribute('id'); // Remove ID so we don't accidentally overwrite it
-                
-                // Attach copy buttons explicitly since we aren't calling renderMessages()
-                finalMsgElement.querySelectorAll('pre').forEach(pre => {
-                    if (!pre.querySelector('.copy-btn')) {
-                        const btn = document.createElement('button');
-                        btn.className = 'copy-btn';
-                        btn.innerText = 'Copy';
-                        btn.onclick = () => {
-                            const codeText = pre.querySelector('code')?.innerText || pre.innerText.replace('Copy', '');
-                            navigator.clipboard.writeText(codeText.trim());
-                            btn.innerText = 'Copied!';
-                            setTimeout(() => btn.innerText = 'Copy', 2000);
-                        };
-                        pre.appendChild(btn);
-                    }
-                });
-
-                scrollToBottom();
-                loadSidebar(); // Sync sidebar updates but PREVENT wiping the main chat DOM 
-
-            } catch (error) {
-                console.error(error);
-                document.getElementById('temp-ai-msg').innerHTML = "<span class='text-red-400'>Connection lost or chat busy. Reload to sync memory.</span>";
-            } finally {
-                if (!pollingInterval) toggleInputState(false);
-            }
+        if (aiReasoning) {
+          html += `
+            <div class="reasoning-block">
+              <i>Thinking Process</i><br/>
+              ${DOMPurify.sanitize(
+                marked.parse(aiReasoning)
+              )}
+            </div>
+          `;
         }
 
-        document.getElementById('message-input').addEventListener('keydown', (e) => {
-            if (e.key === 'Enter' && !e.shiftKey) { 
-                e.preventDefault(); 
-                if (!document.getElementById('send-btn').classList.contains('hidden')) {
-                    sendMessage(); 
-                }
-            }
-        });
+        html += DOMPurify.sanitize(
+          marked.parse(aiContent)
+        );
 
-        loadSidebar();
-    </script>
+        document.getElementById(
+          'streaming-ai-msg'
+        ).innerHTML = html;
+
+        addCopyButtons();
+
+        scrollToBottom();
+      }
+    }
+
+  } catch (err) {
+
+    console.error(err);
+
+    document.getElementById(
+      'streaming-ai-msg'
+    ).innerHTML = `
+      <span class="text-red-400">
+        Connection lost.
+      </span>
+    `;
+  }
+}
+
+// =================================
+// ENTER SEND
+// =================================
+
+document.getElementById(
+  'message-input'
+).addEventListener('keydown', e => {
+
+  if (
+    e.key === 'Enter' &&
+    !e.shiftKey
+  ) {
+
+    e.preventDefault();
+
+    sendMessage();
+  }
+});
+
+// =================================
+// INIT
+// =================================
+
+loadSidebar();
+
+</script>
 </body>
 </html>