> ## Documentation Index
> Fetch the complete documentation index at: https://docs.zeus.agentspro.cn/llms.txt
> Use this file to discover all available pages before exploring further.

# WebSocket Gateway

> WebSocket Gateway — Real-time connections and tool calls for browser extensions, desktop apps, and web clients

Zeus AI Backend manages all node connections through the WebSocket Gateway. Three types of clients connect via different WebSocket endpoints, using the JSON-RPC 2.0 protocol for tool call communication.

***

## WebSocket Endpoints

| Endpoint           | Client            | Connection URL                                                                     |
| ------------------ | ----------------- | ---------------------------------------------------------------------------------- |
| `WS /ws/extension` | Browser Extension | `wss://zeus-api.agentspro.cn/ws/extension?client_id={client_id}&node_id={node_id}` |
| `WS /ws/desktop`   | Desktop App       | `wss://zeus-api.agentspro.cn/ws/desktop?client_id={client_id}&node_id={node_id}`   |
| `WS /ws/web`       | Web Client        | `wss://zeus-api.agentspro.cn/ws/web?user_id={user_id}`                             |

***

## Connection Parameters

### Browser Extension / Desktop App

<ParamField query="client_id" type="string" required>
  Client ID. Format: `user_{user_id}` (browser extension) or `desktop_{user_id}` (desktop app). The server automatically parses and extracts the `user_id`.
</ParamField>

<ParamField query="node_id" type="string">
  Node ID that uniquely identifies the device. If not provided, the server auto-generates one (`ext_{user_id}` or `desktop_{user_id}`).
</ParamField>

### Web Client

<ParamField query="user_id" type="string" required>
  User ID, used for routing messages to the correct user.
</ParamField>

***

## Connection Lifecycle

```mermaid theme={null}
sequenceDiagram
    participant C as Client
    participant S as Server

    C->>S: WebSocket Connect
    S-->>C: Accept connection

    C->>S: register message (node info)
    S-->>C: registered response

    loop Keepalive
        C->>S: heartbeat
        S-->>C: heartbeat_ack
    end

    Note over C,S: Agent initiates tool call
    S->>C: JSON-RPC tools/call
    C-->>S: JSON-RPC response

    C->>S: disconnect (unregister node)
```

***

## Message Types

### 1. Node Registration (register)

After connecting, browser extensions and desktop apps should send a `register` message to register node information:

**Client → Server:**

```json theme={null}
{
  "type": "register",
  "node_id": "node_abc123",
  "node_name": "Chrome Extension",
  "node_type": "extension",
  "os": "macOS",
  "os_version": "14.0",
  "app_version": "1.2.0",
  "capabilities": ["browser_control", "screenshot"],
  "available_tools": ["click", "type", "screenshot", "navigate"],
  "max_concurrent_tasks": 3
}
```

<ResponseField name="node_id" type="string" required>
  Unique node ID
</ResponseField>

<ResponseField name="node_name" type="string">
  Display name for the node (defaults to "Unknown Node")
</ResponseField>

<ResponseField name="node_type" type="string" required>
  Node type: `extension` or `desktop`
</ResponseField>

<ResponseField name="os" type="string">
  Operating system name
</ResponseField>

<ResponseField name="os_version" type="string">
  Operating system version
</ResponseField>

<ResponseField name="app_version" type="string">
  Client application version
</ResponseField>

<ResponseField name="capabilities" type="string[]">
  List of node capabilities, e.g. `["browser_control", "screenshot", "file_system"]`
</ResponseField>

<ResponseField name="available_tools" type="string[]">
  List of available tools, e.g. `["click", "type", "screenshot", "navigate"]`
</ResponseField>

<ResponseField name="max_concurrent_tasks" type="integer">
  Maximum concurrent tasks (defaults to 3)
</ResponseField>

**Server → Client:**

```json theme={null}
{
  "type": "registered",
  "node_id": "node_abc123",
  "success": true
}
```

***

### 2. Heartbeat

Nodes send periodic heartbeats to maintain connection status. The server updates the node state in NodeManager.

**Client → Server:**

```json theme={null}
{
  "type": "heartbeat",
  "status": "online",
  "current_tasks": 0
}
```

<ResponseField name="status" type="string">
  Node status: `online`, `busy`, or `offline` (defaults to `online`)
</ResponseField>

<ResponseField name="current_tasks" type="integer">
  Number of tasks currently being executed (defaults to 0)
</ResponseField>

**Server → Client:**

```json theme={null}
{
  "type": "heartbeat_ack"
}
```

***

### 3. Ping/Pong Keepalive

All three client types (Extension, Desktop, Web) support ping/pong keepalive:

**Client → Server:**

```json theme={null}
{
  "type": "ping"
}
```

**Server → Client:**

```json theme={null}
{
  "type": "pong"
}
```

<Note>
  For Extension and Desktop clients, `ping` messages also trigger a NodeManager heartbeat update.
</Note>

***

### 4. Tool Calls (JSON-RPC 2.0)

The server sends JSON-RPC 2.0 requests to nodes via WebSocket to invoke tools. This follows the MCP (Model Context Protocol) format.

**Server → Client (Request):**

```json theme={null}
{
  "jsonrpc": "2.0",
  "id": "550e8400-e29b-41d4-a716-446655440000",
  "method": "tools/call",
  "params": {
    "name": "screenshot",
    "arguments": { "full_page": true },
    "session_id": "sess_abc123"
  }
}
```

<ResponseField name="id" type="string" required>
  Request ID (UUID), used for matching responses
</ResponseField>

<ResponseField name="method" type="string" required>
  Always `tools/call`
</ResponseField>

<ResponseField name="params.name" type="string" required>
  Name of the tool to call
</ResponseField>

<ResponseField name="params.arguments" type="object" required>
  Tool arguments
</ResponseField>

<ResponseField name="params.session_id" type="string">
  Associated session ID (optional)
</ResponseField>

**Client → Server (Success Response):**

```json theme={null}
{
  "jsonrpc": "2.0",
  "id": "550e8400-e29b-41d4-a716-446655440000",
  "result": {
    "content": [
      {
        "type": "text",
        "text": "Screenshot taken successfully"
      },
      {
        "type": "image",
        "data": "iVBORw0KGgo...",
        "mimeType": "image/png"
      }
    ],
    "isError": false,
    "_screenshot": "base64_screenshot_data..."
  }
}
```

**Client → Server (Error Response):**

```json theme={null}
{
  "jsonrpc": "2.0",
  "id": "550e8400-e29b-41d4-a716-446655440000",
  "error": {
    "code": -32603,
    "message": "Tool execution failed: element not found"
  }
}
```

#### MCP Result Content Types

Each item in the `result.content` array can be:

| type    | Fields             | Description               |
| ------- | ------------------ | ------------------------- |
| `text`  | `text`             | Text content              |
| `image` | `data`, `mimeType` | Base64-encoded image data |

<Note>
  When parsing responses, the server extracts `text`, `image`, and `_screenshot` fields, and maps `isError` to a `success` status. Default timeout is 60 seconds.
</Note>

***

### 5. Legacy MCP Response (mcp\_response)

<Warning>
  This format is deprecated and only kept for backward compatibility. New clients should use the JSON-RPC 2.0 format.
</Warning>

```json theme={null}
{
  "type": "mcp_response",
  "request_id": "req_001",
  "result": { "screenshot": "base64..." },
  "error": null
}
```

***

### 6. Status Update

Clients can send status update messages:

```json theme={null}
{
  "type": "status",
  "status": "busy"
}
```

***

### 7. Workflow Execution

#### Web Client Requests Workflow List

**Web Client → Server:**

```json theme={null}
{
  "type": "get_workflows"
}
```

The server forwards the request to the browser extension. The `workflows_list` response from the extension is automatically relayed to the web client. If the extension is not connected:

```json theme={null}
{
  "type": "workflows_list",
  "success": false,
  "error": "Extension not connected"
}
```

#### Web Client Executes Workflow

**Web Client → Server:**

```json theme={null}
{
  "type": "execute_workflow",
  "workflow_id": "wf_abc123",
  "variables": {
    "url": "https://example.com",
    "query": "search term"
  }
}
```

**Server → Web Client (Acknowledge):**

```json theme={null}
{
  "type": "workflow_started",
  "workflow_id": "wf_abc123"
}
```

**Server → Web Client (Completion):**

When the extension sends a `task_complete` message, the server notifies the web client:

```json theme={null}
{
  "type": "workflow_complete",
  "task_id": "task_456",
  "success": true,
  "result": { "data": "..." },
  "output": "Workflow completed successfully"
}
```

#### Extension Workflow Completion

**Extension → Server:**

```json theme={null}
{
  "type": "task_complete",
  "task_id": "task_456",
  "result": { "data": "..." },
  "output": "Workflow completed"
}
```

<Note>
  Default timeout for workflow execution is 300 seconds (5 minutes).
</Note>

***

## Error Handling

### Connection Rejection

The server closes the connection when required parameters are missing:

| Code   | Reason                                                     |
| ------ | ---------------------------------------------------------- |
| `4001` | Missing `client_id` (Extension/Desktop) or `user_id` (Web) |

### Disconnection Handling

When a WebSocket connection drops, the server automatically:

1. Removes the connection from ConnectionManager
2. Unregisters the node from NodeManager (Extension/Desktop)
3. Cleans up all pending requests

***

## ConnectionManager API

`ConnectionManager` is the core singleton of the WebSocket Gateway, managing the lifecycle of all connections.

### Connection Management

| Method                                                       | Description                                                              |
| ------------------------------------------------------------ | ------------------------------------------------------------------------ |
| `connect_extension(user_id, node_id, ws, register_request?)` | Register browser extension connection, optionally with registration info |
| `connect_desktop(user_id, node_id, ws, register_request?)`   | Register desktop app connection, optionally with registration info       |
| `connect_web(user_id, ws)`                                   | Register web client connection                                           |
| `disconnect_extension(user_id, node_id)`                     | Disconnect extension and unregister node                                 |
| `disconnect_desktop(user_id, node_id)`                       | Disconnect desktop and unregister node                                   |
| `disconnect_web(user_id, ws)`                                | Disconnect web client                                                    |

### Connection Queries

| Method                                       | Description                                                                                      |
| -------------------------------------------- | ------------------------------------------------------------------------------------------------ |
| `is_extension_connected(user_id, node_id?)`  | Check extension connection status. Without `node_id`, checks if user has any extension connected |
| `is_desktop_connected(user_id, node_id?)`    | Check desktop app connection status                                                              |
| `get_extension_node_ids(user_id)`            | Get list of all Extension node IDs for a user                                                    |
| `get_desktop_node_ids(user_id)`              | Get list of all Desktop node IDs for a user                                                      |
| `get_websocket(user_id, node_id, node_type)` | Get the WebSocket connection object for a specific node                                          |
| `get_stats()`                                | Get connection statistics                                                                        |

### Message Sending

| Method                                          | Description                                                                                                  |
| ----------------------------------------------- | ------------------------------------------------------------------------------------------------------------ |
| `send_to_extension(user_id, message, node_id?)` | Send message to browser extension. Specify `node_id` for a specific node, otherwise sends to first available |
| `send_to_desktop(user_id, message, node_id?)`   | Send message to desktop app. Same logic as above                                                             |
| `send_to_web(user_id, message)`                 | Send message to **all** web client connections for the user                                                  |

### Tool Calls

| Method                                                                                  | Description                                                     |
| --------------------------------------------------------------------------------------- | --------------------------------------------------------------- |
| `call_extension_tool(user_id, tool_name, tool_args, session_id?, node_id?, timeout=60)` | Call browser extension tool and wait for response               |
| `call_desktop_tool(user_id, tool_name, tool_args, session_id?, node_id?, timeout=60)`   | Call desktop app tool and wait for response                     |
| `call_tool(user_id, node_id, node_type, tool_name, tool_args, session_id?, timeout=60)` | Unified tool call method, automatically routes by `node_type`   |
| `execute_workflow(user_id, workflow_id, variables?, timeout=300)`                       | Execute a workflow on the browser extension                     |
| `resolve_request(request_id, result)`                                                   | Resolve a pending request (called by the message handling loop) |

### Connection Stats (get\_stats)

```json theme={null}
{
  "extension_connections": 3,
  "desktop_connections": 1,
  "web_connections": 5,
  "pending_requests": 2
}
```

***

## Data Structures

### Multi-Node Connection Model

ConnectionManager uses nested dictionaries to manage connections, supporting multiple nodes per user:

```
extension_connections: { user_id -> { node_id -> WebSocket } }
desktop_connections:   { user_id -> { node_id -> WebSocket } }
web_connections:       { user_id -> Set[WebSocket] }
pending_requests:      { request_id -> asyncio.Future }
```

* **Extension / Desktop**: Each user can connect multiple nodes (multiple devices), distinguished by `node_id`
* **Web**: Each user can have multiple web client connections (multiple tabs), stored in a Set

### NodeType Enum

| Value       | Description              |
| ----------- | ------------------------ |
| `extension` | Browser extension node   |
| `desktop`   | Desktop application node |

### NodeStatus Enum

| Value     | Description                         |
| --------- | ----------------------------------- |
| `online`  | Node is online and can accept tasks |
| `busy`    | Node is busy executing tasks        |
| `offline` | Node is offline                     |

***

## Connection Examples

### JavaScript (Browser Extension)

```javascript theme={null}
const ws = new WebSocket(
  'wss://zeus-api.agentspro.cn/ws/extension?client_id=user_abc123&node_id=ext_001'
);

ws.onopen = () => {
  // Register node info
  ws.send(JSON.stringify({
    type: 'register',
    node_id: 'ext_001',
    node_name: 'Chrome Extension',
    node_type: 'extension',
    os: navigator.platform,
    app_version: '1.0.0',
    capabilities: ['browser_control', 'screenshot'],
    available_tools: ['click', 'type', 'screenshot', 'navigate']
  }));

  // Start heartbeat
  setInterval(() => {
    ws.send(JSON.stringify({ type: 'heartbeat', status: 'online' }));
  }, 30000);
};

ws.onmessage = (event) => {
  const message = JSON.parse(event.data);

  // Handle JSON-RPC tool calls
  if (message.jsonrpc === '2.0' && message.method === 'tools/call') {
    const { id, params } = message;
    
    // Execute tool...
    executeToolCall(params.name, params.arguments).then(result => {
      ws.send(JSON.stringify({
        jsonrpc: '2.0',
        id: id,
        result: {
          content: [{ type: 'text', text: 'Tool executed successfully' }],
          isError: false
        }
      }));
    }).catch(error => {
      ws.send(JSON.stringify({
        jsonrpc: '2.0',
        id: id,
        error: { code: -32603, message: error.message }
      }));
    });
  }
};
```

### Python (Desktop App)

```python theme={null}
import asyncio
import websockets
import json

async def connect_desktop():
    uri = "wss://zeus-api.agentspro.cn/ws/desktop?client_id=desktop_abc123&node_id=desk_001"
    
    async with websockets.connect(uri) as ws:
        # Register node
        await ws.send(json.dumps({
            "type": "register",
            "node_id": "desk_001",
            "node_name": "My Desktop",
            "node_type": "desktop",
            "os": "macOS",
            "os_version": "14.0",
            "app_version": "1.0.0",
            "capabilities": ["file_system", "screen_capture"],
            "available_tools": ["read_file", "write_file", "screenshot"]
        }))
        
        async for message in ws:
            data = json.loads(message)
            
            if data.get("type") == "heartbeat_ack":
                continue
            
            # Handle JSON-RPC tool calls
            if data.get("jsonrpc") == "2.0" and "method" in data:
                request_id = data["id"]
                tool_name = data["params"]["name"]
                tool_args = data["params"]["arguments"]
                
                try:
                    result = await execute_tool(tool_name, tool_args)
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "result": {
                            "content": [{"type": "text", "text": str(result)}],
                            "isError": False
                        }
                    }))
                except Exception as e:
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "error": {"code": -32603, "message": str(e)}
                    }))
```