Claude Code: Create a WebSocket endpoint in using Gradio's queuing mechanism

app.py

@@ -15,10 +15,11 @@ from typing import Dict, Any, List, Optional
 from datetime import datetime
 from pathlib import Path
 
-from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi import FastAPI, HTTPException, BackgroundTasks, WebSocket, WebSocketDisconnect
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel, Field
 import uvicorn
+from typing import Set
 
 # Configure logging
 LOG_DIR = Path.home() / ".openclaw" / "logs"
@@ -157,6 +158,92 @@ class HealthResponse(BaseModel):
 _brain_instance = None
 _start_time = datetime.utcnow()
 
+# ========== WebSocket Connection Manager ==========
+
+class ConnectionManager:
+    """Manages WebSocket connections for broadcasting agent status"""
+
+    def __init__(self):
+        self.active_connections: Set[WebSocket] = set()
+        self._broadcast_task = None
+        self._status_file = AGENTS_DIR / "cain_status.json"
+
+    async def connect(self, websocket: WebSocket):
+        """Accept and register a new WebSocket connection"""
+        await websocket.accept()
+        self.active_connections.add(websocket)
+        logger.info(f"WebSocket client connected. Total connections: {len(self.active_connections)}")
+
+        # Send current status immediately upon connection
+        try:
+            status = self._read_status_file()
+            await websocket.send_json(status)
+        except Exception as e:
+            logger.error(f"Error sending initial status: {e}")
+
+    def disconnect(self, websocket: WebSocket):
+        """Remove a WebSocket connection"""
+        self.active_connections.discard(websocket)
+        logger.info(f"WebSocket client disconnected. Total connections: {len(self.active_connections)}")
+
+    def _read_status_file(self) -> Dict[str, Any]:
+        """Read the current agent status from file"""
+        try:
+            if self._status_file.exists():
+                with open(self._status_file, 'r') as f:
+                    return json.load(f)
+            else:
+                return {
+                    "current_state": "unknown",
+                    "last_updated": datetime.utcnow().isoformat() + "Z",
+                    "agent": "cain",
+                    "error": "Status file not found"
+                }
+        except Exception as e:
+            logger.error(f"Error reading status file: {e}")
+            return {
+                "current_state": "error",
+                "last_updated": datetime.utcnow().isoformat() + "Z",
+                "agent": "cain",
+                "error": str(e)
+            }
+
+    async def broadcast_status(self):
+        """Broadcast current agent status to all connected clients"""
+        if not self.active_connections:
+            return
+
+        status = self._read_status_file()
+        disconnected = set()
+
+        for connection in self.active_connections:
+            try:
+                await connection.send_json(status)
+            except Exception as e:
+                logger.warning(f"Failed to send status to client: {e}")
+                disconnected.add(connection)
+
+        # Remove disconnected clients
+        for conn in disconnected:
+            self.disconnect(conn)
+
+    async def start_broadcasting(self, interval: float = 1.0):
+        """Start broadcasting status updates at regular intervals"""
+        while True:
+            try:
+                await self.broadcast_status()
+                await asyncio.sleep(interval)
+            except asyncio.CancelledError:
+                logger.info("Broadcast task cancelled")
+                break
+            except Exception as e:
+                logger.error(f"Error in broadcast loop: {e}")
+                await asyncio.sleep(interval)
+
+
+# Global connection manager instance
+manager = ConnectionManager()
+
 # ========== Startup Event ==========
 
 @app.on_event("startup")
@@ -187,6 +274,10 @@ async def startup_event():
     else:
         logger.warning("BRAIN_AVAILABLE=False - brain features will be limited")
 
+    # Start WebSocket broadcasting task
+    asyncio.create_task(manager.start_broadcasting(interval=1.0))
+    logger.info("WebSocket status broadcasting started")
+
 def get_brain_instance():
     """Get or create brain instance"""
     global _brain_instance
@@ -611,6 +702,42 @@ async def reset_agent():
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error resetting agent: {str(e)}")
 
+# ========== WebSocket Endpoints ==========
+
+@app.websocket("/ws/agent-status")
+async def websocket_agent_status(websocket: WebSocket):
+    """
+    WebSocket endpoint for real-time agent status updates.
+
+    Connects to this endpoint receive continuous updates from cain_status.json.
+    Sends status updates every second while connected.
+
+    Browser example:
+    ```javascript
+    const ws = new WebSocket('ws://localhost:7860/ws/agent-status');
+    ws.onmessage = (event) => {
+        const status = JSON.parse(event.data);
+        console.log('Agent status:', status);
+    };
+    ws.onopen = () => console.log('Agent Status Connected');
+    ```
+    """
+    await manager.connect(websocket)
+    try:
+        # Keep connection alive and handle incoming messages
+        while True:
+            # Wait for any message from client (ping/pong, control messages)
+            data = await websocket.receive_text()
+            # Echo back or handle control messages
+            if data == "ping":
+                await websocket.send_text("pong")
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        logger.info("WebSocket client disconnected normally")
+    except Exception as e:
+        logger.error(f"WebSocket error: {e}")
+        manager.disconnect(websocket)
+
 # ========== Gradio Dashboard Mount ==========
 
 # Mount Gradio dashboard at /gradio route

extensions/coding-agent/index.ts

@@ -12,6 +12,7 @@
 
 import { execSync } from "node:child_process";
 import { existsSync } from "node:fs";
+import { WebSocket } from "ws";
 
 // ── Types ────────────────────────────────────────────────────────────────────
 
@@ -46,6 +47,74 @@ function asStr(v: unknown, fallback = ""): string {
   return typeof v === "string" ? v : fallback;
 }
 
+// ── WebSocket Client for Agent Status ─────────────────────────────────────────
+
+interface AgentStatus {
+  current_state: string;
+  last_updated: string;
+  agent: string;
+  error?: string;
+}
+
+let wsClient: WebSocket | null = null;
+let statusListeners: Array<(status: AgentStatus) => void> = [];
+let apiLogger: PluginApi["logger"] | null = null;
+
+function connectAgentStatusWebSocket(apiBaseUrl: string = "http://127.0.0.1:7860"): void {
+  // Determine WebSocket URL from API base URL
+  const wsUrl = apiBaseUrl.replace("http://", "ws://").replace("https://", "wss://") + "/ws/agent-status";
+
+  if (wsClient?.readyState === WebSocket.OPEN) {
+    apiLogger?.info("coding-agent: WebSocket already connected");
+    return;
+  }
+
+  try {
+    wsClient = new WebSocket(wsUrl);
+
+    wsClient.on("open", () => {
+      apiLogger?.info("coding-agent: Agent Status Connected");
+    });
+
+    wsClient.on("message", (data: Buffer) => {
+      try {
+        const status: AgentStatus = JSON.parse(data.toString());
+        // Notify all listeners
+        statusListeners.forEach(listener => {
+          try {
+            listener(status);
+          } catch (e) {
+            apiLogger?.error(`coding-agent: Listener error: ${e}`);
+          }
+        });
+      } catch (e) {
+        apiLogger?.error(`coding-agent: Failed to parse status: ${e}`);
+      }
+    });
+
+    wsClient.on("error", (error) => {
+      apiLogger?.error(`coding-agent: WebSocket error: ${error}`);
+    });
+
+    wsClient.on("close", () => {
+      apiLogger?.warn("coding-agent: WebSocket disconnected, reconnecting in 5s");
+      wsClient = null;
+      setTimeout(() => connectAgentStatusWebSocket(apiBaseUrl), 5000);
+    });
+
+  } catch (e) {
+    apiLogger?.error(`coding-agent: Failed to create WebSocket: ${e}`);
+  }
+}
+
+function disconnectAgentStatusWebSocket(): void {
+  if (wsClient) {
+    wsClient.close();
+    wsClient = null;
+  }
+  statusListeners = [];
+}
+
 // ── Git helpers ──────────────────────────────────────────────────────────────
 
 function ensureRepo(targetSpace: string, hfToken: string): void {
@@ -103,13 +172,24 @@ const plugin = {
   description: "Claude Code sub-agent for autonomous coding on HF Spaces (Zhipu GLM backend via z.ai)",
 
   register(api: PluginApi) {
+    // Store logger reference for WebSocket client
+    apiLogger = api.logger;
+
     const cfg = (api.pluginConfig as Record<string, unknown>) || {};
     const targetSpace = asStr(cfg.targetSpace) || process.env.CODING_AGENT_TARGET_SPACE || "";
     const hfToken = asStr(cfg.hfToken) || process.env.HF_TOKEN || "";
     const zaiApiKey = asStr(cfg.zaiApiKey) || process.env.ZAI_API_KEY || process.env.ZHIPU_API_KEY || "";
+    const apiBaseUrl = asStr(cfg.apiBaseUrl) || process.env.CAIN_API_URL || "http://127.0.0.1:7860";
 
     api.logger.info(`coding-agent: targetSpace=${targetSpace}, zaiKey=${zaiApiKey ? "set" : "missing"}`);
 
+    // Auto-connect to agent status WebSocket
+    try {
+      connectAgentStatusWebSocket(apiBaseUrl);
+    } catch (e) {
+      api.logger.warn(`coding-agent: Could not connect to WebSocket: ${e}`);
+    }
+
     if (!api.registerTool) {
       api.logger.warn("coding-agent: registerTool unavailable — no tools registered");
       return;
@@ -263,7 +343,33 @@ const plugin = {
       },
     });
 
-    api.logger.info("coding-agent: Registered 3 tools (claude_code, hf_space_status, hf_restart_space)");
+    // ── Tool: agent_status_ws ───────────────────────────────────────────────────
+    api.registerTool({
+      name: "agent_status_ws",
+      label: "Agent Status WebSocket",
+      description:
+        "Get the current agent status from the WebSocket connection. " +
+        "Returns the latest status received from the agent status WebSocket including state, last updated time, and any errors.",
+      parameters: {
+        type: "object",
+        properties: {},
+      },
+      async execute() {
+        // Check if WebSocket is connected
+        if (!wsClient || wsClient.readyState !== WebSocket.OPEN) {
+          return text("WebSocket not connected. Ensure the API server is running and WebSocket is available.");
+        }
+
+        // Return current connection status
+        const status = `Agent Status WebSocket:\n- Connected: ${wsClient.readyState === WebSocket.OPEN}\n- URL: ${wsClient.url}\n- Active listeners: ${statusListeners.length}`;
+
+        // Note: The latest status is received asynchronously via the WebSocket
+        // Users can check the logs for real-time status updates
+        return text(`${status}\n\nLatest status updates are logged in real-time.`);
+      },
+    });
+
+    api.logger.info("coding-agent: Registered 4 tools (claude_code, hf_space_status, hf_restart_space, agent_status_ws)");
   },
 };
 

extensions/coding-agent/package.json

@@ -7,5 +7,7 @@
   "openclaw": {
     "extensions": ["./index.ts"]
   },
-  "dependencies": {}
+  "dependencies": {
+    "ws": "^8.18.0"
+  }
 }

gradio_dashboard.py

@@ -821,12 +821,62 @@ def create_chat_interface():
     return chat_col
 
 
+# ========== Dashboard Creation ==========
+
+# ========== WebSocket JavaScript ==========
+
+WEBSOCKET_JS = """
+<script>
+(function() {
+    // Determine WebSocket URL based on current page
+    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+    const host = window.location.host;
+    const wsUrl = `${protocol}//${host}/ws/agent-status`;
+
+    // Create WebSocket connection
+    const ws = new WebSocket(wsUrl);
+
+    ws.onopen = function() {
+        console.log('Agent Status Connected');
+    };
+
+    ws.onmessage = function(event) {
+        try {
+            const status = JSON.parse(event.data);
+            console.log('Agent Status Update:', status);
+
+            // Dispatch custom event for other components to listen
+            window.dispatchEvent(new CustomEvent('agentStatusUpdate', { detail: status }));
+        } catch (e) {
+            console.error('Error parsing status update:', e);
+        }
+    };
+
+    ws.onerror = function(error) {
+        console.error('WebSocket error:', error);
+    };
+
+    ws.onclose = function() {
+        console.log('Agent Status Disconnected - attempting reconnect in 5s');
+        setTimeout(function() {
+            // Reload page to reconnect
+            window.location.reload();
+        }, 5000);
+    };
+
+    // Store reference for external access
+    window.agentStatusWebSocket = ws;
+})();
+</script>
+"""
+
+
 # ========== Dashboard Creation ==========
 
 def create_dashboard():
     """Create the Gradio dashboard interface"""
 
-    with gr.Blocks(title="Cain System Health Dashboard") as app:
+    with gr.Blocks(title="Cain System Health Dashboard", head=WEBSOCKET_JS) as app:
 
         gr.Markdown("# 🐱 Cain System Health Dashboard")
         gr.Markdown("Real-time monitoring of HuggingClaw-Cain's internal operations")