Update app.py

app.py

@@ -1,7 +1,7 @@
 """
 MindStudio WebSocket Streaming Client — Flask API Endpoint
-FIXED: 'ClientConnection' object has no attribute 'closed'
-       Compatible with websockets >= 11.x and older versions
+FIXED: Per-thread continuation tokens for proper conversation memory
+FIXED: WebSocket compatibility
 """
 
 import asyncio
@@ -26,7 +26,6 @@ INITIAL_CONTINUATION_TOKEN = os.getenv(
     "155b2f71-9738-4260-8849-46e3eb0dd384::2b71a3fd-f4f0-4439-8163-60cd53d28218"
 )
 TARGET_MODEL = "claude-4-7-opus"
-
 WS_URL = "wss://api-socket.mindstudio.ai/"
 
 # ─────────────────────────────────────────────
@@ -38,7 +37,7 @@ _clients_lock = threading.Lock()
 
 
 # ─────────────────────────────────────────────
-# HELPER
+# HELPERS
 # ─────────────────────────────────────────────
 
 def safe_get_model(debug_info: dict):
@@ -49,24 +48,16 @@ def safe_get_model(debug_info: dict):
 
 
 def check_ws_open(ws) -> bool:
-    """
-    Universal check for WebSocket open state.
-    Handles both old and new websockets library versions.
-    """
     if ws is None:
         return False
     try:
-        # websockets >= 11.x — ClientConnection
         if hasattr(ws, 'close_code'):
             return ws.close_code is None
-        # websockets < 11.x — WebSocketClientProtocol
         if hasattr(ws, 'closed'):
             return not ws.closed
-        # websockets >= 13.x — may use 'state'
         if hasattr(ws, 'state'):
             import websockets.connection
             return ws.state == websockets.connection.State.OPEN
-        # fallback
         return True
     except Exception:
         return False
@@ -80,7 +71,11 @@ class MindStudioClient:
     def __init__(self, access_token: str, model: str = TARGET_MODEL):
         self.access_token = access_token
         self.model = model
-        self.continuation_token = INITIAL_CONTINUATION_TOKEN
+
+        # ✅ FIX: Per-thread continuation tokens instead of one global token
+        self._continuation_tokens: dict = {}
+        self._tokens_lock = threading.Lock()
+
         self.ws = None
         self.loop = asyncio.new_event_loop()
 
@@ -111,7 +106,6 @@ class MindStudioClient:
         print(f"[WS] Connected! Type: {type(self.ws).__name__}")
         print(f"[WS] Subprotocol: {self.ws.subprotocol}")
 
-        # Log which attributes are available for debugging
         has_closed    = hasattr(self.ws, 'closed')
         has_closecode = hasattr(self.ws, 'close_code')
         has_state     = hasattr(self.ws, 'state')
@@ -120,6 +114,17 @@ class MindStudioClient:
     def is_connected(self) -> bool:
         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")
+
     def send_message_stream(self, message: str, thread_id: str):
         import queue
 
@@ -129,7 +134,7 @@ class MindStudioClient:
         async def _run():
             try:
                 listen_task = asyncio.create_task(
-                    self._listen_for_response(chunk_queue, done_event)
+                    self._listen_for_response(chunk_queue, done_event, thread_id)
                 )
                 await asyncio.sleep(0.1)
                 await self._post_message(message, thread_id)
@@ -152,9 +157,14 @@ 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]}...")
+
         payload = {
             "threadId": thread_id,
-            "continuationToken": self.continuation_token,
+            "continuationToken": continuation_token,
             "payload": {
                 "message": message,
                 "imageUrl": "",
@@ -176,7 +186,8 @@ class MindStudioClient:
         if resp.status_code not in (200, 204):
             print(f"[POST] Body: {resp.text}")
 
-    async def _listen_for_response(self, chunk_queue, done_event):
+    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
@@ -205,8 +216,8 @@ class MindStudioClient:
                 data = event.get("appThreadContinuationActionUpdated", {})
                 new_token = data.get("continuationAction", {}).get("token")
                 if new_token:
-                    self.continuation_token = new_token
-                    print(f"[WS] Continuation token updated")
+                    # ✅ FIX: Save token per thread, not globally
+                    self.set_continuation_token(thread_id, new_token)
 
                 underlying = (
                     data.get("continuationAction", {})