fix(chat): resolve imagine ws auth mismatch and improve connection diagnostics

Author: TQZHR

app/api/v1/admin.py

@@ -103,7 +103,15 @@ async def _verify_ws_api_key(websocket: WebSocket) -> bool:
     token = str(websocket.query_params.get("api_key") or "").strip()
     if not token:
         return False
-    return (api_key and token == api_key) or token in legacy_keys
+    if (api_key and token == api_key) or token in legacy_keys:
+        return True
+    try:
+        await api_key_manager.init()
+        if api_key_manager.validate_key(token):
+            return True
+    except Exception as e:
+        logger.warning(f"Imagine ws api_key validation fallback failed: {e}")
+    return False
 
 
 async def _collect_imagine_batch(token: str, prompt: str, aspect_ratio: str) -> list[str]:

app/static/chat/chat.html

@@ -199,7 +199,7 @@
   </main>
 
   <script src="/static/common/toast.js"></script>
-  <script src="/static/chat/chat.js?v=3"></script>
+  <script src="/static/chat/chat.js?v=4"></script>
 </body>
 
 </html>

app/static/chat/chat.js

@@ -498,7 +498,10 @@ function openImageContinuousSocket(socketIndex, runToken, prompt, aspectRatio, a
       socketState.active = status === 'running';
       updateImageContinuousStats();
       if (imageContinuousRunning) {
-        if (status === 'running') setImageStatusText('Running');
+        if (status === 'running') {
+          clearImageContinuousError();
+          setImageStatusText('Running');
+        }
         if (status === 'stopped') setImageStatusText('Stopped');
       }
       updateImageContinuousButtons();
@@ -507,6 +510,7 @@ function openImageContinuousSocket(socketIndex, runToken, prompt, aspectRatio, a
 
     if (msgType === 'image') {
       socketState.active = true;
+      clearImageContinuousError();
       appendWaterfallImage(data, socketIndex);
       if (imageContinuousRunning) setImageStatusText('Running');
       updateImageContinuousButtons();
@@ -545,7 +549,14 @@ function openImageContinuousSocket(socketIndex, runToken, prompt, aspectRatio, a
         setImageContinuousError('WebSocket auth rejected. Check API key.');
         setImageStatusText('Auth failed');
       } else if (socketState.hadError && stillActive <= 0 && stillOpen <= 0) {
-        setImageContinuousError(socketState.lastError || `WS${socketIndex + 1} connection error`);
+        const closeCode = Number(event?.code || 0);
+        const closeReason = String(event?.reason || '').trim();
+        if (closeCode > 0) {
+          const suffix = closeReason ? `: ${closeReason}` : '';
+          setImageContinuousError(`WebSocket closed (${closeCode})${suffix}`);
+        } else {
+          setImageContinuousError(socketState.lastError || `WS${socketIndex + 1} connection error`);
+        }
         setImageStatusText('Disconnected');
       }
 

app/static/chat/chat_admin.html

@@ -195,7 +195,7 @@
   <script src="/static/common/admin-auth.js"></script>
   <script src="/static/common/header.js?v=3"></script>
   <script src="/static/common/footer.js"></script>
-  <script src="/static/chat/chat.js?v=3"></script>
+  <script src="/static/chat/chat.js?v=4"></script>
 </body>
 
 </html>

tests/test_imagine_ws_api.py

@@ -49,6 +49,25 @@ def test_imagine_ws_ping_pong(monkeypatch: pytest.MonkeyPatch):
     assert msg == {"type": "pong"}
 
 
+def test_imagine_ws_accepts_managed_api_key(monkeypatch: pytest.MonkeyPatch):
+    client = _build_client(monkeypatch, api_key="global-key")
+
+    async def _fake_init():
+        return None
+
+    monkeypatch.setattr(admin_api.api_key_manager, "init", _fake_init)
+    monkeypatch.setattr(
+        admin_api.api_key_manager,
+        "validate_key",
+        lambda token: {"key": token, "is_active": True} if token == "managed-key" else None,
+    )
+
+    with client.websocket_connect("/api/v1/admin/imagine/ws?api_key=managed-key") as ws:
+        ws.send_json({"type": "ping"})
+        msg = ws.receive_json()
+    assert msg == {"type": "pong"}
+
+
 def test_imagine_ws_empty_prompt_error(monkeypatch: pytest.MonkeyPatch):
     client = _build_client(monkeypatch, api_key="valid-key")
     with client.websocket_connect("/api/v1/admin/imagine/ws?api_key=valid-key") as ws: