# WebSocket vs Heartbeat Polling for Real-Time Notifications in AI-Trader: Architecture and Implementation

> Explore WebSocket vs heartbeat polling for real-time AI-Trader notifications. Discover how AI-Trader uses WebSockets for instant updates and polling for agent health checks.

- Repository: [✨Data Intelligence Lab@HKU✨/AI-Trader](https://github.com/HKUDS/AI-Trader)
- Tags: architecture
- Published: 2026-05-09

---

**AI-Trader employs WebSocket connections for instantaneous server-to-client push notifications while utilizing HTTP heartbeat polling to monitor autonomous agent liveliness, enabling both real-time UI updates and scalable stateless health checks.**

The HKUDS/AI-Trader repository maintains bidirectional communication between the trading backend and its clients through two distinct mechanisms. Understanding when the system uses **WebSocket** persistent connections versus **heartbeat polling** helps developers optimize agent performance and ensure reliable market data delivery.

## Communication Patterns and Architectural Design

AI-Trader deliberately separates concerns between immediate data streaming and presence detection. Each pattern addresses specific scalability and latency requirements inherent in algorithmic trading environments.

### WebSocket Push Architecture

The **WebSocket** implementation provides a persistent, full-duplex communication channel defined in [`service/server/routes_shared.py`](https://github.com/HKUDS/AI-Trader/blob/main/service/server/routes_shared.py). This approach allows the server to emit market intelligence streams and challenge notifications instantly without waiting for client requests.

Key implementation details include:

- **Connection State Management**: The `WsState` dataclass maintains a `ws_connections` dictionary mapping client IDs to active `WebSocket` objects
- **Event Broadcasting**: Server-side helpers in [`service/server/services.py`](https://github.com/HKUDS/AI-Trader/blob/main/service/server/services.py) iterate over the connection map to broadcast JSON events to specific clients or all connected UIs
- **Latency Characteristics**: Near-zero latency delivery ensures time-critical trading signals reach users immediately upon generation

### Heartbeat Polling Mechanism

The **heartbeat polling** system operates through stateless HTTP POST requests, implemented in [`service/server/routes_agent.py`](https://github.com/HKUDS/AI-Trader/blob/main/service/server/routes_agent.py). Autonomous trading agents call the `/api/claw/agents/heartbeat` endpoint at regular intervals to report their operational status.

Key characteristics include:

- **Stateless Scalability**: Each heartbeat request is independent, allowing load balancing across multiple server instances without connection state synchronization
- **Database Persistence**: The handler updates an agent's `last_seen` timestamp in the database, enabling the backend to mark agents as "online" or "offline" based on recency
- **Polling Latency**: Notification delay equals the polling interval (typically every 10-15 seconds), acceptable for health monitoring but insufficient for market data streaming

## Implementation Details in AI-Trader

### WebSocket Endpoint Configuration

In [`service/server/routes_shared.py`](https://github.com/HKUDS/AI-Trader/blob/main/service/server/routes_shared.py), the application registers WebSocket endpoints using FastAPI's native support. The connection manager stores active sockets in a module-level state container:

```python

# service/server/routes_shared.py

from fastapi import HTTPException, WebSocket
from dataclasses import dataclass, field

@dataclass
class WsState:
    # Keep a mapping of client IDs → active WebSocket connections

    ws_connections: dict[int, WebSocket] = field(default_factory=dict)

ws_state = WsState()

async def websocket_endpoint(websocket: WebSocket, client_id: int):
    await websocket.accept()
    ws_state.ws_connections[client_id] = websocket
    try:
        while True:
            # Keep the connection alive; server may ignore incoming messages

            await websocket.receive_text()
    except Exception:
        # On disconnect, clean up

        ws_state.ws_connections.pop(client_id, None)

```

The server pushes notifications through these connections using helper functions in [`service/server/services.py`](https://github.com/HKUDS/AI-Trader/blob/main/service/server/services.py) that iterate over `ws_state.ws_connections.values()` to deliver JSON payloads.

### Heartbeat Handler Implementation

The agent heartbeat endpoint in [`service/server/routes_agent.py`](https://github.com/HKUDS/AI-Trader/blob/main/service/server/routes_agent.py) provides a lightweight health-check interface:

```python

# service/server/routes_agent.py

@app.post('/api/claw/agents/heartbeat')
async def agent_heartbeat(authorization: str = Header(None)):
    """
    Agents call this endpoint every N seconds.
    The server records the timestamp so the agent can be shown as online.
    """
    # Verify the agent token, then update its last_seen field in the DB

    await db.update_agent_last_seen(agent_id, datetime.utcnow())
    return {"status": "ok"}

```

This design allows agents running in containerized or restricted environments to maintain presence without maintaining persistent sockets or handling complex reconnection logic.

## Scalability and Performance Considerations

**WebSocket** connections require the server to maintain an open socket pool for each connected client, consuming memory and file descriptors. The system must handle disconnections, reconnections, and connection state synchronization across potential horizontal scaling events.

**Heartbeat polling** scales horizontally without state sharing between server instances, as each request is self-contained. However, the approach introduces inherent latency equal to the polling interval, making it unsuitable for real-time market data but optimal for simple liveness probes.

## Summary

- **WebSocket** in [`service/server/routes_shared.py`](https://github.com/HKUDS/AI-Trader/blob/main/service/server/routes_shared.py) enables instantaneous push notifications to UI clients through persistent connections managed via the `ws_connections` dictionary
- **Heartbeat polling** via `POST /api/claw/agents/heartbeat` in [`service/server/routes_agent.py`](https://github.com/HKUDS/AI-Trader/blob/main/service/server/routes_agent.py) provides scalable, stateless health monitoring for autonomous trading agents
- WebSocket delivers near-zero latency for time-critical market data, while heartbeat polling trades immediacy for horizontal scalability and implementation simplicity
- The dual-mode architecture separates concerns: WebSocket handles real-time frontend updates, while heartbeat ensures reliable backend agent bookkeeping

## Frequently Asked Questions

### Can autonomous trading agents use WebSocket instead of heartbeat polling?

Agents technically could establish WebSocket connections, but the heartbeat design specifically accommodates containerized and ephemeral execution environments where maintaining persistent connections adds complexity. The stateless HTTP approach in [`routes_agent.py`](https://github.com/HKUDS/AI-Trader/blob/main/routes_agent.py) allows agents to function behind strict firewalls or in serverless contexts without managing socket lifecycles.

### What happens when a WebSocket connection drops unexpectedly?

When a disconnect occurs, the exception handler in `websocket_endpoint` removes the client ID from `ws_state.ws_connections`, immediately freeing resources. The frontend must implement reconnection logic to establish a new socket, whereas the server treats the absence as a temporary disconnection unless paired with heartbeat data indicating the agent itself is offline.

### How does the server handle concurrent WebSocket connections under load?

The current implementation stores connections in an in-memory dictionary within the `WsState` dataclass. For high availability deployments, this requires sticky session configuration or shared state infrastructure (such as Redis) to ensure notifications reach the correct server instance holding the client's socket.

### What is the recommended heartbeat interval for production trading agents?

The raw source indicates agents typically poll every 10 to 15 seconds. This interval balances timely failure detection with minimal API load, though specific deployments may adjust this based on network reliability and the criticality of rapid offline detection versus request volume reduction.