Presave state when refresh

app.py

@@ -48,6 +48,7 @@ class ChatMessageResponse(BaseModel):
     content: str
     timestamp: float
     created_at: str
+    sources: Optional[List[Dict[str, Any]]] = None
 
 class ChatHistoryResponse(BaseModel):
     messages: List[ChatMessageResponse]
@@ -298,19 +299,31 @@ async def save_chat_message(
     project_id: str = Form(...),
     role: str = Form(...),
     content: str = Form(...),
-    timestamp: Optional[float] = Form(None)
+    timestamp: Optional[float] = Form(None),
+    sources: Optional[str] = Form(None)
 ):
     """Save a chat message to the session"""
     if role not in ["user", "assistant"]:
         raise HTTPException(400, detail="Invalid role")
     
+    # Parse optional sources JSON
+    parsed_sources: Optional[List[Dict[str, Any]]] = None
+    if sources:
+        try:
+            parsed = json.loads(sources)
+            if isinstance(parsed, list):
+                parsed_sources = parsed
+        except Exception:
+            parsed_sources = None
+
     message = {
         "user_id": user_id,
         "project_id": project_id,
         "role": role,
         "content": content,
         "timestamp": timestamp or time.time(),
-        "created_at": datetime.now(timezone.utc)
+        "created_at": datetime.now(timezone.utc),
+        **({"sources": parsed_sources} if parsed_sources is not None else {})
     }
     
     rag.db["chat_sessions"].insert_one(message)
@@ -332,7 +345,8 @@ async def get_chat_history(user_id: str, project_id: str, limit: int = 100):
             role=message["role"],
             content=message["content"],
             timestamp=message["timestamp"],
-            created_at=message["created_at"].isoformat() if isinstance(message["created_at"], datetime) else str(message["created_at"])
+            created_at=message["created_at"].isoformat() if isinstance(message["created_at"], datetime) else str(message["created_at"]),
+            sources=message.get("sources")
         ))
     
     return ChatHistoryResponse(messages=messages)

static/projects.js

@@ -279,6 +279,11 @@
         messages.forEach(msg => {
           if (typeof window.appendMessage === 'function') {
             window.appendMessage(msg.role, msg.content);
+            if (msg.role === 'assistant' && Array.isArray(msg.sources) && msg.sources.length) {
+              if (typeof window.appendSources === 'function') {
+                window.appendSources(msg.sources);
+              }
+            }
           } else {
             // Fallback if not yet defined
             const messagesContainer = document.getElementById('messages');
@@ -286,6 +291,13 @@
             messageDiv.className = `msg ${msg.role}`;
             messageDiv.textContent = msg.content;
             messagesContainer.appendChild(messageDiv);
+            // Minimal fallback for sources
+            if (msg.role === 'assistant' && Array.isArray(msg.sources) && msg.sources.length) {
+              const srcDiv = document.createElement('div');
+              srcDiv.className = 'sources';
+              srcDiv.textContent = 'Sources: ' + msg.sources.map(s => s.filename).join(', ');
+              messagesContainer.appendChild(srcDiv);
+            }
           }
         });
         

static/script.js

@@ -532,8 +532,16 @@
         if (response.ok) {
           thinkingMsg.remove();
           appendMessage('assistant', data.answer || 'No answer received');
-          await saveChatMessage(user.user_id, currentProject.project_id, 'assistant', data.answer || 'No answer received');
-          if (data.sources && data.sources.length > 0) appendSources(data.sources);
+          if (data.sources && data.sources.length > 0) {
+            appendSources(data.sources);
+          }
+          await saveChatMessage(
+            user.user_id,
+            currentProject.project_id,
+            'assistant',
+            data.answer || 'No answer received',
+            (data.sources && data.sources.length > 0) ? data.sources : null
+          );
         } else {
           throw new Error(data.detail || 'Failed to get answer');
         }
@@ -641,7 +649,7 @@
     }
   });
 
-  async function saveChatMessage(userId, projectId, role, content) {
+  async function saveChatMessage(userId, projectId, role, content, sources = null) {
     try {
       const formData = new FormData();
       formData.append('user_id', userId);
@@ -649,6 +657,9 @@
       formData.append('role', role);
       formData.append('content', content);
       formData.append('timestamp', Date.now() / 1000);
+      if (sources) {
+        try { formData.append('sources', JSON.stringify(sources)); } catch {}
+      }
       
       await fetch('/chat/save', { method: 'POST', body: formData });
     } catch (error) {
@@ -692,6 +703,10 @@
     return messageDiv;
   }
 
+  // Expose markdown-aware appenders for use after refresh (projects.js)
+  window.appendMessage = appendMessage;
+  window.appendSources = appendSources;
+
   function addCopyButtonsToCodeBlocks(messageDiv) {
     const codeBlocks = messageDiv.querySelectorAll('pre code');
     codeBlocks.forEach((codeBlock, index) => {