Update app.py

app.py

@@ -1,6 +1,7 @@
 """
 MindStudio WebSocket Streaming Client — Flask API Endpoint
-FIXED: Per-thread continuation tokens for proper conversation memory
+FIXED: Race condition — AI was answering previous message (buffer drain fix)
+FIXED: Per-thread continuation tokens
 FIXED: WebSocket compatibility
 """
 
@@ -72,7 +73,7 @@ class MindStudioClient:
         self.access_token = access_token
         self.model = model
 
-        # ✅ FIX: Per-thread continuation tokens instead of one global token
+        # Per-thread continuation tokens
         self._continuation_tokens: dict = {}
         self._tokens_lock = threading.Lock()
 
@@ -115,16 +116,35 @@ class MindStudioClient:
         return check_ws_open(self.ws)
 
     def get_continuation_token(self, thread_id: str) -> str:
-        """✅ Get the continuation token for a specific thread."""
         with self._tokens_lock:
             return self._continuation_tokens.get(thread_id, INITIAL_CONTINUATION_TOKEN)
 
     def set_continuation_token(self, thread_id: str, token: str):
-        """✅ Save the continuation token for a specific thread."""
         with self._tokens_lock:
             self._continuation_tokens[thread_id] = token
             print(f"[TOKEN] Thread {thread_id[:8]}... token updated")
 
+    # ─────────────────────────────────────────
+    # ✅ FIX: Drain leftover WS buffer
+    # ─────────────────────────────────────────
+    async def _drain_ws_buffer(self):
+        """Flush any leftover WebSocket messages from a previous response."""
+        flushed = 0
+        try:
+            while True:
+                raw = await asyncio.wait_for(self.ws.recv(), timeout=0.05)
+                flushed += 1
+                print(f"[DRAIN] Flushed leftover msg #{flushed}: {str(raw)[:80]}...")
+        except asyncio.TimeoutError:
+            pass
+        except Exception:
+            pass
+        if flushed:
+            print(f"[DRAIN] Total flushed: {flushed} leftover message(s)")
+
+    # ─────────────────────────────────────────
+    # ✅ FIX: Post FIRST, then listen
+    # ─────────────────────────────────────────
     def send_message_stream(self, message: str, thread_id: str):
         import queue
 
@@ -133,12 +153,18 @@ class MindStudioClient:
 
         async def _run():
             try:
-                listen_task = asyncio.create_task(
-                    self._listen_for_response(chunk_queue, done_event, thread_id)
-                )
-                await asyncio.sleep(0.1)
+                # ✅ STEP 1: Drain any leftover buffer from previous message
+                await self._drain_ws_buffer()
+
+                # ✅ STEP 2: Post the new message
                 await self._post_message(message, thread_id)
-                await listen_task
+
+                # ✅ STEP 3: Small wait for server to process
+                await asyncio.sleep(0.15)
+
+                # ✅ STEP 4: Listen for THIS message's response only
+                await self._listen_for_response(chunk_queue, done_event, thread_id)
+
             except Exception as e:
                 print(f"[ERROR] send_message_stream: {e}")
                 chunk_queue.put(f"[ERROR] {e}")
@@ -158,7 +184,6 @@ class MindStudioClient:
     async def _post_message(self, message: str, thread_id: str):
         post_url = f"https://v1.mindstudio-api.com/v1/apps/load/{APP_ID}/threads/post"
 
-        # ✅ FIX: Use per-thread continuation token
         continuation_token = self.get_continuation_token(thread_id)
         print(f"[POST] Using token: {continuation_token[:30]}... for thread: {thread_id[:8]}...")
 
@@ -187,7 +212,6 @@ class MindStudioClient:
             print(f"[POST] Body: {resp.text}")
 
     async def _listen_for_response(self, chunk_queue, done_event, thread_id: str):
-        """✅ FIX: Accept thread_id to update its specific continuation token."""
         full_response = ""
         previous_content_length = 0
         system_post_id = None
@@ -216,7 +240,6 @@ class MindStudioClient:
                 data = event.get("appThreadContinuationActionUpdated", {})
                 new_token = data.get("continuationAction", {}).get("token")
                 if new_token:
-                    # ✅ FIX: Save token per thread, not globally
                     self.set_continuation_token(thread_id, new_token)
 
                 underlying = (