Update app.py

app.py

@@ -1,17 +1,7 @@
 """
 MindStudio WebSocket Streaming Client — Flask API Endpoint
-Exposes:
-  GET  /connect         → pre-connect WebSocket
-  POST /chat            → send message, stream response
-  
-Headers (on /chat):
-  Authorization: Bearer <access_token>
-
-Body (on /chat):
-  {
-    "message": "your message here",
-    "threadId": "your-thread-id"
-  }
+FIXED: 'ClientConnection' object has no attribute 'closed'
+       Compatible with websockets >= 11.x and older versions
 """
 
 import asyncio
@@ -40,17 +30,12 @@ TARGET_MODEL = "claude-4-7-opus"
 WS_URL = "wss://api-socket.mindstudio.ai/"
 
 # ─────────────────────────────────────────────
-# GLOBAL STATE — one WS client per access token
+# GLOBAL STATE
 # ─────────────────────────────────────────────
 
-# { access_token: MindStudioClient }
 _clients: dict = {}
 _clients_lock = threading.Lock()
 
-# Each client runs its own asyncio event loop in a dedicated thread
-# { access_token: asyncio.AbstractEventLoop }
-_loops: dict = {}
-
 
 # ─────────────────────────────────────────────
 # HELPER
@@ -63,6 +48,30 @@ def safe_get_model(debug_info: dict):
     return None
 
 
+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
+
+
 # ─────────────────────────────────────────────
 # MINDSTUDIO CLIENT
 # ─────────────────────────────────────────────
@@ -74,9 +83,7 @@ class MindStudioClient:
         self.continuation_token = INITIAL_CONTINUATION_TOKEN
         self.ws = None
         self.loop = asyncio.new_event_loop()
-        self._ws_ready = threading.Event()
 
-        # Start the event loop in a background thread
         self._thread = threading.Thread(target=self._run_loop, daemon=True)
         self._thread.start()
 
@@ -85,12 +92,11 @@ class MindStudioClient:
         self.loop.run_forever()
 
     def connect_sync(self):
-        """Connect WebSocket synchronously (blocks until connected)."""
         future = asyncio.run_coroutine_threadsafe(self._connect(), self.loop)
-        future.result(timeout=15)
+        future.result(timeout=20)
 
     async def _connect(self):
-        print(f"[WS] Connecting...")
+        print(f"[WS] Connecting to {WS_URL} ...")
         extra_headers = {
             "Origin": "https://app.mindstudio.ai",
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
@@ -102,21 +108,23 @@ class MindStudioClient:
             ping_interval=None,
             ping_timeout=None,
         )
-        print(f"[WS] Connected! Subprotocol: {self.ws.subprotocol}")
-        self._ws_ready.set()
+        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')
+        print(f"[WS] Attrs → closed={has_closed} close_code={has_closecode} state={has_state}")
 
     def is_connected(self) -> bool:
-        return self.ws is not None and not self.ws.closed
+        return check_ws_open(self.ws)
 
     def send_message_stream(self, message: str, thread_id: str):
-        """
-        Generator that yields SSE-formatted chunks of the AI response.
-        Runs async logic in the client's dedicated event loop.
-        """
         import queue
 
         chunk_queue = queue.Queue()
-        done_event = threading.Event()
+        done_event  = threading.Event()
 
         async def _run():
             try:
@@ -127,12 +135,12 @@ class MindStudioClient:
                 await self._post_message(message, thread_id)
                 await listen_task
             except Exception as e:
+                print(f"[ERROR] send_message_stream: {e}")
                 chunk_queue.put(f"[ERROR] {e}")
                 done_event.set()
 
         asyncio.run_coroutine_threadsafe(_run(), self.loop)
 
-        # Yield chunks as they arrive
         while not done_event.is_set() or not chunk_queue.empty():
             try:
                 chunk = chunk_queue.get(timeout=0.5)
@@ -165,6 +173,8 @@ class MindStudioClient:
         async with httpx.AsyncClient() as client:
             resp = await client.post(post_url, json=payload, headers=headers)
         print(f"[POST] Status: {resp.status_code}")
+        if resp.status_code not in (200, 204):
+            print(f"[POST] Body: {resp.text}")
 
     async def _listen_for_response(self, chunk_queue, done_event):
         full_response = ""
@@ -180,6 +190,9 @@ class MindStudioClient:
             except websockets.exceptions.ConnectionClosed as e:
                 print(f"[WS] Connection closed: {e}")
                 break
+            except Exception as e:
+                print(f"[WS] Recv error: {e}")
+                break
 
             try:
                 event = json.loads(raw_message)
@@ -188,12 +201,12 @@ class MindStudioClient:
 
             event_type = event.get("type", "")
 
-            # ── Update continuation token ──
             if event_type == "Apps/ThreadContinuationActionUpdated":
                 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")
 
                 underlying = (
                     data.get("continuationAction", {})
@@ -204,7 +217,6 @@ class MindStudioClient:
                     if not detected_model:
                         detected_model = underlying["model"]
 
-            # ── Capture system post ID ──
             elif event_type == "Apps/ThreadPostsCreated":
                 posts = event.get("appThreadPostsCreated", {}).get("posts", [])
                 for post in posts:
@@ -230,12 +242,11 @@ class MindStudioClient:
                             previous_content_length = len(content)
                             full_response = content
 
-            # ── Stream tokens ──
             elif event_type == "Apps/ThreadPostUpdated":
                 post_data = event.get("appThreadPostUpdated", {})
-                new_post = post_data.get("newPost", {})
-                post_id = post_data.get("postId")
-                chat_msg = new_post.get("chatMessage", {})
+                new_post  = post_data.get("newPost", {})
+                post_id   = post_data.get("postId")
+                chat_msg  = new_post.get("chatMessage", {})
 
                 if chat_msg.get("source") != "system":
                     continue
@@ -249,7 +260,7 @@ class MindStudioClient:
                         detected_model = model
 
                 current_content = chat_msg.get("content", "")
-                is_in_progress = chat_msg.get("isInProgress", True)
+                is_in_progress  = chat_msg.get("isInProgress", True)
 
                 if len(current_content) > previous_content_length:
                     delta = current_content[previous_content_length:]
@@ -264,21 +275,22 @@ class MindStudioClient:
 
         done_event.set()
 
-    async def close(self):
+    async def _close(self):
         if self.ws:
             await self.ws.close()
 
 
 # ─────────────────────────────────────────────
-# HELPER: get or create client for a token
+# HELPER: get or create client
 # ─────────────────────────────────────────────
 
 def get_or_create_client(access_token: str) -> MindStudioClient:
     with _clients_lock:
         client = _clients.get(access_token)
         if client and client.is_connected():
+            print("[CLIENT] Reusing existing connected client")
             return client
-        # Create new client and connect
+        print("[CLIENT] Creating new client and connecting...")
         client = MindStudioClient(access_token=access_token)
         client.connect_sync()
         _clients[access_token] = client
@@ -289,12 +301,17 @@ def get_or_create_client(access_token: str) -> MindStudioClient:
 # ROUTES
 # ─────────────────────────────────────────────
 
+@app.route("/health", methods=["GET"])
+def health():
+    return jsonify({
+        "status": "ok",
+        "model": TARGET_MODEL,
+        "websockets_version": websockets.__version__,
+    }), 200
+
+
 @app.route("/connect", methods=["GET"])
 def connect():
-    """
-    Pre-warm the WebSocket connection.
-    Pass Authorization: Bearer <token> header.
-    """
     auth = request.headers.get("Authorization", "")
     if not auth.startswith("Bearer "):
         return jsonify({"error": "Missing or invalid Authorization header"}), 401
@@ -302,26 +319,20 @@ def connect():
     access_token = auth.removeprefix("Bearer ").strip()
 
     try:
-        get_or_create_client(access_token)
-        return jsonify({"status": "connected", "model": TARGET_MODEL}), 200
+        client = get_or_create_client(access_token)
+        ws_type = type(client.ws).__name__ if client.ws else "None"
+        return jsonify({
+            "status": "connected",
+            "model": TARGET_MODEL,
+            "ws_object_type": ws_type,
+        }), 200
     except Exception as e:
+        print(f"[CONNECT ERROR] {type(e).__name__}: {e}")
         return jsonify({"error": str(e)}), 500
 
 
 @app.route("/chat", methods=["POST"])
 def chat():
-    """
-    Send a message and stream the AI response as SSE.
-
-    Headers:
-      Authorization: Bearer <access_token>
-
-    Body JSON:
-      {
-        "message": "Hello!",
-        "threadId": "your-thread-id"
-      }
-    """
     auth = request.headers.get("Authorization", "")
     if not auth.startswith("Bearer "):
         return jsonify({"error": "Missing or invalid Authorization header"}), 401
@@ -332,7 +343,7 @@ def chat():
     if not data:
         return jsonify({"error": "Invalid JSON body"}), 400
 
-    message = data.get("message", "").strip()
+    message   = data.get("message", "").strip()
     thread_id = data.get("threadId", "").strip()
 
     if not message:
@@ -361,19 +372,13 @@ def chat():
     )
 
 
-@app.route("/health", methods=["GET"])
-def health():
-    return jsonify({"status": "ok", "model": TARGET_MODEL}), 200
-
-
 # ─────────────────────────────────────────────
 # ENTRY POINT
 # ─────────────────────────────────────────────
 
 if __name__ == "__main__":
-    print("=" * 60)
-    print(f"  MindStudio Flask Streaming API")
-    print(f"  Model: {TARGET_MODEL}")
-    print(f"  Visit GET /connect first to pre-warm WebSocket")
-    print("=" * 60)
+    print("="*60)
+    print(f"  MindStudio Flask API  |  Model: {TARGET_MODEL}")
+    print(f"  websockets version: {websockets.__version__}")
+    print("="*60)
     app.run(host="0.0.0.0", port=7860, debug=False, threaded=True)