Review Fix: Improve src/infrastructure/chat_store.py

src/dashboard/routes/chat_api.py

@@ -101,7 +101,7 @@ async def _process_chat(user_msg: str) -> dict | JSONResponse:
     try:
         response_text = await agent_chat(
             _build_context_prefix() + user_msg,
-            session_id="mobile",
+            session_id=body.get("session_id", "mobile"),
         )
         message_log.append(role="user", content=user_msg, timestamp=timestamp, source="api")
         message_log.append(role="agent", content=response_text, timestamp=timestamp, source="api")
@@ -165,6 +165,11 @@ async def api_upload(file: UploadFile = File(...)):
     if not str(resolved).startswith(str(upload_root)):
         raise HTTPException(status_code=400, detail="Invalid file name")
 
+    # Validate MIME type
+    allowed_types = ["image/png", "image/jpeg", "image/gif", "application/pdf", "text/plain"]
+    if file.content_type not in allowed_types:
+        raise HTTPException(status_code=400, detail=f"File type {file.content_type} not allowed")
+
     contents = await file.read()
     if len(contents) > _MAX_UPLOAD_SIZE:
         raise HTTPException(status_code=413, detail="File too large (max 50 MB)")

src/infrastructure/chat_store.py

@@ -60,7 +60,12 @@ class MessageLog:
         self._conn: sqlite3.Connection | None = None
 
     # Lazy connection — opened on first use, not at import time.
-    def _ensure_conn(self) -> sqlite3.Connection:
+    @contextmanager
+    def _get_conn(self) -> Generator[sqlite3.Connection, None, None]:
+        path = self._db_path or DB_PATH
+        with closing(sqlite3.connect(str(path), check_same_thread=False)) as conn:
+            conn.row_factory = sqlite3.Row
+            yield conn
         if self._conn is None:
             # Open a persistent connection for the class instance
             path = self._db_path or DB_PATH

src/infrastructure/ws_manager/handler.py

@@ -79,7 +79,17 @@ class WebSocketManager:
         message = ws_event.to_json()
         disconnected = []
 
-        for ws in self._connections:
+        import asyncio
+        tasks = [ws.send_text(message) for ws in self._connections]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        
+        disconnected = []
+        for ws, result in zip(self._connections, results):
+            if isinstance(result, Exception):
+                logger.warning(f"WebSocket send error: {result}")
+                disconnected.append(ws)
+
+        # Skip the old loop
             try:
                 await ws.send_text(message)
             except ConnectionError: