Update public/index.html

public/index.html

@@ -60,7 +60,7 @@
         <!-- Header -->
         <header class="h-14 min-h-[3.5rem] border-b border-gray-800 flex items-center px-4 lg:px-6 justify-between bg-gray-900/80 backdrop-blur-sm z-10">
             
-            <div class="flex items-center gap-3 overflow-hidden group w-full max-w-[50%]">
+            <div class="flex items-center gap-3 overflow-hidden group w-full max-w-[45%]">
                 <button onclick="toggleSidebar()" class="lg:hidden p-1 -ml-1 text-gray-400 hover:text-white shrink-0">
                     <svg class="w-6 h-6" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 6h16M4 12h16M4 18h16"></path></svg>
                 </button>
@@ -85,10 +85,13 @@
                 </div>
             </div>
 
-            <div class="flex items-center gap-2 lg:gap-4 shrink-0 pl-2">
-                <span class="hidden sm:inline-block text-xs text-blue-300 bg-blue-900/30 border border-blue-800/50 px-2 py-1 rounded-md font-mono">
-                    Tokens: <span id="token-counter">0</span>
-                </span>
+            <div class="flex items-center gap-3 lg:gap-4 shrink-0 pl-2">
+                <!-- Detailed Token Display -->
+                <div class="hidden sm:flex flex-col items-end text-[10px] text-blue-300 bg-blue-900/30 border border-blue-800/50 px-2 py-1 rounded-md font-mono leading-tight">
+                    <span class="font-semibold text-[11px]">Total: <span id="token-total">0</span></span>
+                    <span class="text-blue-400/70">In: <span id="token-in">0</span> | Out: <span id="token-out">0</span></span>
+                </div>
+
                 <select id="model-select" class="bg-gray-800 text-xs lg:text-sm rounded-md border border-gray-700 px-2 py-1.5 outline-none focus:border-blue-500 transition-colors cursor-pointer text-gray-200">
                     <option value="claude">Claude Sonnet 4.6</option>
                     <option value="haiku">Claude Haiku</option>
@@ -133,6 +136,9 @@
         let currentChatId = null;
         let attachedImages =[];
         let pollingInterval = null;
+        
+        // Local token state variables
+        let currentTokens = { total: 0, in: 0, out: 0 };
 
         // --- TITLE EDITING LOGIC ---
         function enableTitleEdit() {
@@ -167,16 +173,14 @@
                 headers: { 'Content-Type': 'application/json' },
                 body: JSON.stringify({ title: newTitle })
             });
-            loadSidebar(); // Refresh sidebar to show new name
+            loadSidebar(); 
         }
 
-        // Handle Enter and Escape keys for title input
         document.getElementById('title-input').addEventListener('keydown', (e) => {
             if (e.key === 'Enter') saveTitle();
             if (e.key === 'Escape') cancelTitleEdit();
         });
 
-
         // --- UI CORE LOGIC ---
         function toggleSidebar() {
             const sidebar = document.getElementById('sidebar');
@@ -191,6 +195,13 @@
             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();
+        }
+
         async function loadSidebar() {
             const res = await fetch('/api/chats');
             const chats = await res.json();
@@ -198,7 +209,10 @@
             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-xs text-gray-500 mt-1">${c.totalTokens.toLocaleString()} tokens</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('');
         }
@@ -219,195 +233,5 @@
                 document.getElementById('chat-window').innerHTML = '';
                 document.getElementById('current-chat-title').innerText = 'Select or create a chat';
                 document.getElementById('edit-title-btn').classList.add('hidden');
-                document.getElementById('token-counter').innerText = '0';
-                loadSidebar();
-            }
-        }
-
-        async function selectChat(id) {
-            currentChatId = id;
-            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'); // Show edit icon for active chat
-            document.getElementById('token-counter').innerText = chat.totalTokens.toLocaleString();
-            
-            cancelTitleEdit(); // Ensure we reset title UI state if switching
-            renderMessages(chat.messages);
-            loadSidebar(); 
-
-            if(window.innerWidth < 1024 && !document.getElementById('sidebar').classList.contains('-translate-x-full')) {
-                toggleSidebar();
-            }
-            if (chat.isGenerating) pollGeneratingChat(id);
-        }
-
-        function renderMessages(messages) {
-            const container = document.getElementById('chat-window');
-            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);
-                }
-            });
-            container.scrollTop = container.scrollHeight;
-        }
-
-        function pollGeneratingChat(id) {
-            pollingInterval = setInterval(async () => {
-                const res = await fetch(`/api/chats/${id}`);
-                const chat = await res.json();
-                renderMessages(chat.messages);
-                document.getElementById('token-counter').innerText = chat.totalTokens.toLocaleString();
-                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);
-            }, 1500);
-        }
-
-        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);
-            });
-        }
-
-        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>
-            `;
-            container.scrollTop = container.scrollHeight;
-            
-            input.value = '';
-            input.style.height = 'auto'; 
-            attachedImages =[];
-            document.getElementById('image-preview-container').innerHTML = '';
-            document.getElementById('image-preview-container').classList.add('hidden');
-            document.getElementById('send-btn').disabled = true;
-
-            try {
-                const response = await fetch(`/api/chats/${currentChatId}/stream`, {
-                    method: 'POST',
-                    headers: { 'Content-Type': 'application/json' },
-                    body: JSON.stringify(payload)
-                });
-
-                const reader = response.body.getReader();
-                const decoder = new TextDecoder("utf-8");
-                let aiContent = "";
-                let aiReasoning = "";
-
-                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__') {
-                            const usageData = JSON.parse(chunks[i+1]);
-                            document.getElementById('token-counter').innerText = 
-                                (parseInt(document.getElementById('token-counter').innerText.replace(/,/g, '')) + usageData.totalTokenCount).toLocaleString();
-                            i += 2;
-                        } else {
-                            aiContent += chunks[i];
-                            i++;
-                        }
-                    }
-
-                    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>";
-                    container.scrollTop = container.scrollHeight;
-                }
-
-                selectChat(currentChatId);
-
-            } catch (error) {
-                console.error(error);
-                document.getElementById('temp-ai-msg').innerHTML = "<span class='text-red-400'>Error connecting to API.</span>";
-            } finally {
-                document.getElementById('send-btn').disabled = false;
-            }
-        }
-
-        document.getElementById('message-input').addEventListener('keydown', (e) => {
-            if (e.key === 'Enter' && !e.shiftKey) { 
-                e.preventDefault(); 
-                sendMessage(); 
-            }
-        });
-
-        loadSidebar();
-    </script>
-</body>
-</html>
+                updateTokenDisplay(0, 0, 0);
+                loadSid