Update app.py

app.py

@@ -1061,16 +1061,22 @@ async def tiers():
 async def websocket_chat(ws: WebSocket):
     ip = ws.client.host
 
+    # Accept connection
     await ws.accept()
 
-    # rate limit auth attempts
+    # Basic connection rate limiting
     if not check_ws_auth_rate_limit(ip):
         await ws.close(code=4408)
         return
 
     try:
+        # Expect an auth message first
         auth_msg = await ws.receive_text()
-        auth_data = json.loads(auth_msg)
+        try:
+            auth_data = json.loads(auth_msg)
+        except Exception:
+            await ws.close(code=4400)
+            return
 
         provided_key = auth_data.get("key")
 
@@ -1078,45 +1084,97 @@ async def websocket_chat(ws: WebSocket):
             await ws.close(code=4403)
             return
 
-        # authenticated
-        await ws.send_json({"type": "auth", "status": "ok"})
-
-        while True:
-            msg = await ws.receive_text()
-            data = json.loads(msg)
-
-            body = data.get("body")
-            headers = data.get("headers", {})
+        # Auth success
+        await ws.send_json({
+            "type": "auth",
+            "status": "ok"
+        })
+
+        # Internal endpoint (avoids Hugging Face proxy redirect)
+        internal_url = "http://127.0.0.1:7860/gen/chat/completions"
+
+        # Persistent HTTP client for streaming
+        async with httpx.AsyncClient(
+            timeout=None,
+            follow_redirects=False
+        ) as client:
+
+            while True:
+                try:
+                    msg = await ws.receive_text()
+                except WebSocketDisconnect:
+                    break
+
+                try:
+                    data = json.loads(msg)
+                except Exception:
+                    await ws.send_json({"error": "Invalid JSON"})
+                    continue
+
+                body = data.get("body")
+                headers = data.get("headers") or {}
+
+                if not body:
+                    await ws.send_json({"error": "Missing body"})
+                    continue
+
+                try:
+                    async with client.stream(
+                        "POST",
+                        internal_url,
+                        json=body,
+                        headers=headers
+                    ) as response:
 
-            if not body:
-                await ws.send_json({"error": "Missing body"})
-                continue
+                        # Handle upstream errors
+                        if response.status_code >= 400:
+                            error_text = ""
+                            try:
+                                error_text = (await response.aread()).decode(
+                                    "utf-8", errors="replace"
+                                )[:500]
+                            except Exception:
+                                pass
 
-            url = str(ws.url).replace("ws://", "http://").replace("wss://", "https://")
-            url = url.split("/ws/chat")[0] + "/gen/chat/completions"
+                            await ws.send_json({
+                                "error": "Upstream request failed",
+                                "status": response.status_code,
+                                "detail": error_text
+                            })
+                            continue
 
-            async with httpx.AsyncClient(timeout=None) as client:
-                async with client.stream(
-                    "POST",
-                    url,
-                    json=body,
-                    headers=headers,
-                ) as r:
+                        # Stream tokens/lines to the websocket
+                        async for line in response.aiter_lines():
+                            if not line:
+                                continue
 
-                    async for line in r.aiter_lines():
-                        if not line:
-                            continue
+                            try:
+                                await ws.send_text(line)
+                            except RuntimeError:
+                                break
 
-                        await ws.send_text(line)
+                except Exception as stream_error:
+                    await ws.send_json({
+                        "error": "Streaming error",
+                        "detail": str(stream_error)
+                    })
 
     except WebSocketDisconnect:
         return
+
     except Exception as e:
         try:
-            await ws.send_json({"error": str(e)})
-        except:
+            await ws.send_json({
+                "error": "Server error",
+                "detail": str(e)
+            })
+        except Exception:
+            pass
+
+        try:
+            await ws.close()
+        except Exception:
             pass
-        await ws.close()
 
 @app.get("/portal")
 @app.post("/portal")