Update app.js

app.js

@@ -53,12 +53,18 @@ async function initDB() {
         console.log('✅ Connected to Supabase');
         if (dbChats) {
             dbChats.forEach(c => {
-                memoryChats[c.id] = { ...c, isGenerating: false };
+                // Ensure legacy chats don't crash if they lack the new columns
+                memoryChats[c.id] = { 
+                    ...c, 
+                    inputTokens: c.inputTokens || 0,
+                    outputTokens: c.outputTokens || 0,
+                    isGenerating: false 
+                };
             });
             console.log(`✅ Hydrated ${dbChats.length} chats from DB.`);
         }
 
-        // Sync to DB every 15 seconds (occasional sync)
+        // Sync to DB every 15 seconds
         setInterval(async () => {
             if (dirtyChats.size === 0) return;
             const toSync = Array.from(dirtyChats);
@@ -71,6 +77,8 @@ async function initDB() {
                     id: chat.id,
                     title: chat.title,
                     totalTokens: chat.totalTokens,
+                    inputTokens: chat.inputTokens,
+                    outputTokens: chat.outputTokens,
                     messages: chat.messages,
                     updatedAt: chat.updatedAt
                 };
@@ -97,7 +105,12 @@ initDB();
 // Get all chats (Sidebar)
 app.get('/api/chats', (req, res) => {
     const chatsList = Object.values(memoryChats).map(c => ({
-        id: c.id, title: c.title, totalTokens: c.totalTokens, updatedAt: c.updatedAt
+        id: c.id, 
+        title: c.title, 
+        totalTokens: c.totalTokens, 
+        inputTokens: c.inputTokens,
+        outputTokens: c.outputTokens,
+        updatedAt: c.updatedAt
     })).sort((a, b) => new Date(b.updatedAt) - new Date(a.updatedAt));
     res.json(chatsList);
 });
@@ -116,6 +129,8 @@ app.post('/api/chats', (req, res) => {
         id: newId,
         title: "New Chat",
         totalTokens: 0,
+        inputTokens: 0,
+        outputTokens: 0,
         messages:[],
         isGenerating: false,
         updatedAt: new Date().toISOString()
@@ -124,6 +139,20 @@ app.post('/api/chats', (req, res) => {
     res.json(memoryChats[newId]);
 });
 
+// Update chat title manually
+app.put('/api/chats/:id/title', (req, res) => {
+    const { id } = req.params;
+    const { title } = req.body;
+    
+    if (!memoryChats[id]) return res.status(404).json({ error: "Chat not found" });
+    if (!title || typeof title !== 'string') return res.status(400).json({ error: "Invalid title" });
+    
+    memoryChats[id].title = title.trim();
+    dirtyChats.add(id);
+    
+    res.json({ success: true, title: memoryChats[id].title });
+});
+
 // Delete chat permanently
 app.delete('/api/chats/:id', async (req, res) => {
     const { id } = req.params;
@@ -141,12 +170,12 @@ app.post('/api/chats/:id/stream', async (req, res) => {
 
     if (!memoryChats[id]) return res.status(404).send("Chat not found");
 
-    // Title generation on first message
-    if (memoryChats[id].messages.length === 0) {
+    // Title generation ONLY if it hasn't been manually renamed from defaults
+    if (memoryChats[id].messages.length === 0 && memoryChats[id].title === "New Chat") {
         memoryChats[id].title = prompt.substring(0, 30) + (prompt.length > 30 ? '...' : '');
     }
 
-    // 1. SAVE TO MEMORY (Images excluded from storage to save DB space)
+    // 1. SAVE TO MEMORY
     memoryChats[id].messages.push({ role: "user", content: prompt });
     
     const aiMessage = { role: "assistant", content: "", reasoning: "" };
@@ -160,7 +189,9 @@ app.post('/api/chats/:id/stream', async (req, res) => {
     res.setHeader('X-Accel-Buffering', 'no'); 
     res.flushHeaders();
 
-    let totalTokenCount = 0;
+    let streamInputTokens = 0;
+    let streamOutputTokens = 0;
+    let streamTotalTokens = 0;
 
     try {
         const bedrockModelId = getBedrockModelId(model);
@@ -175,7 +206,7 @@ app.post('/api/chats/:id/stream', async (req, res) => {
             contentBlock = [...imageBlocks, ...contentBlock];
         }
 
-        // Map history for Bedrock (excluding the newly pushed user message which we manually constructed)
+        // Map history for Bedrock
         const historicalMessages = memoryChats[id].messages.slice(0, -2).map(m => ({
             role: m.role, content:[{ text: m.content }]
         }));
@@ -207,16 +238,25 @@ app.post('/api/chats/:id/stream', async (req, res) => {
                 }
             }
             if (chunk.metadata && chunk.metadata.usage) {
-                totalTokenCount = (chunk.metadata.usage.inputTokens || 0) + (chunk.metadata.usage.outputTokens || 0);
+                streamInputTokens = chunk.metadata.usage.inputTokens || 0;
+                streamOutputTokens = chunk.metadata.usage.outputTokens || 0;
+                streamTotalTokens = streamInputTokens + streamOutputTokens;
             }
         }
 
-        // Clean up & finalize memory
-        memoryChats[id].totalTokens += totalTokenCount;
+        // Clean up & finalize memory with granular token usage
+        memoryChats[id].inputTokens += streamInputTokens;
+        memoryChats[id].outputTokens += streamOutputTokens;
+        memoryChats[id].totalTokens += streamTotalTokens;
         memoryChats[id].isGenerating = false;
         dirtyChats.add(id);
         
-        res.write(`__USAGE__${JSON.stringify({ totalTokenCount })}`);
+        // Send granular usage to frontend
+        res.write(`__USAGE__${JSON.stringify({ 
+            inputTokens: streamInputTokens, 
+            outputTokens: streamOutputTokens, 
+            totalTokens: streamTotalTokens 
+        })}`);
         res.end();
 
     } catch (err) {