WebSocket Architecture for AI Chat

Web & API Engineering intermediate Building with AI 13 min read

Overview

WebSockets provide full-duplex, bidirectional communication between clients and servers over a single persistent TCP connection. While Server-Sent Events (SSE) dominate simple LLM token streaming, WebSockets have become increasingly important as AI applications evolve beyond basic chat into complex interactive systems that require real-time bidirectional communication — agent workflows with human-in-the-loop approval, collaborative editing, voice AI, and multi-device synchronization.

The WebSocket protocol (RFC 6455) starts as an HTTP request that “upgrades” to a persistent binary or text connection. Once established, either side can send messages at any time without the overhead of HTTP headers per message. This makes WebSockets ideal for scenarios where the client needs to send signals back to the server during an active AI session — canceling generation mid-stream, approving tool calls, providing additional context, or steering agent behavior in real time.

For AI engineers, understanding WebSocket architecture is critical because it underpins the real-time infrastructure of production AI applications. OpenAI’s Realtime API for voice uses WebSockets natively. Any AI application that goes beyond “send prompt, receive response” likely needs WebSocket capabilities for its interactive features, even if the token streaming itself uses SSE.

How It Works

The WebSocket Handshake

A WebSocket connection begins with an HTTP upgrade request. The client sends a standard HTTP request with special headers indicating it wants to upgrade to the WebSocket protocol:

GET /ws/chat HTTP/1.1
Host: api.example.com
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Sec-WebSocket-Version: 13
Authorization: Bearer <token>

The server responds with a 101 Switching Protocols status, and from that point forward, the connection operates as a persistent bidirectional channel. Both sides can send frames (text or binary) at any time without additional HTTP overhead.

Key Concept: Protocol Upgrade

WebSocket connections start as HTTP but “upgrade” to a different protocol. This means they work with existing HTTP infrastructure for the initial handshake, but load balancers and proxies must be configured to support WebSocket connections (long-lived, stateful) rather than treating them as standard HTTP requests.

Message Routing with Type-Based Dispatch

Production AI chat systems use type-based message routing, where every WebSocket message includes a type field that determines how the server processes it. This pattern enables multiplexing multiple concerns over a single connection:

// Client → Server messages
{ "type": "chat.message", "payload": { "content": "Explain quantum computing" } }
{ "type": "chat.cancel", "payload": { "message_id": "msg_123" } }
{ "type": "tool.approve", "payload": { "tool_call_id": "tc_456", "approved": true } }
{ "type": "session.config", "payload": { "model": "gpt-4", "temperature": 0.7 } }

// Server → Client messages
{ "type": "chat.token", "payload": { "message_id": "msg_123", "token": "Quantum" } }
{ "type": "chat.complete", "payload": { "message_id": "msg_123" } }
{ "type": "tool.request", "payload": { "tool_call_id": "tc_456", "tool": "web_search", "args": {...} } }
{ "type": "error", "payload": { "code": "rate_limited", "message": "..." } }

On the server, a dispatcher routes messages to the appropriate handler based on the type field:

ws.on('message', (raw) => {
  const msg = JSON.parse(raw);
  switch (msg.type) {
    case 'chat.message':
      handleChatMessage(ws, msg.payload);
      break;
    case 'chat.cancel':
      handleCancelGeneration(ws, msg.payload);
      break;
    case 'tool.approve':
      handleToolApproval(ws, msg.payload);
      break;
    default:
      ws.send(JSON.stringify({
        type: 'error',
        payload: { code: 'unknown_type', message: `Unknown message type: ${msg.type}` }
      }));
  }
});

Connection Lifecycle Management

Unlike HTTP where each request is stateless, WebSocket servers must actively manage connection state throughout the connection’s lifetime. This includes:

class ConnectionManager {
  constructor() {
    this.connections = new Map(); // userId → Set of WebSocket connections
  }

  addConnection(userId, ws) {
    if (!this.connections.has(userId)) {
      this.connections.set(userId, new Set());
    }
    this.connections.get(userId).add(ws);

    // Set up heartbeat
    ws.isAlive = true;
    ws.on('pong', () => { ws.isAlive = true; });
  }

  removeConnection(userId, ws) {
    const userConns = this.connections.get(userId);
    if (userConns) {
      userConns.delete(ws);
      if (userConns.size === 0) this.connections.delete(userId);
    }
  }

  // Run every 30 seconds
  heartbeat() {
    for (const [userId, conns] of this.connections) {
      for (const ws of conns) {
        if (!ws.isAlive) {
          this.removeConnection(userId, ws);
          ws.terminate();
          continue;
        }
        ws.isAlive = false;
        ws.ping();
      }
    }
  }
}

Reconnection Strategy

Unlike SSE which has built-in reconnection, WebSocket reconnection must be implemented manually on the client side. The standard approach uses exponential backoff with jitter to prevent thundering herd problems when many clients reconnect simultaneously:

class ReconnectingWebSocket {
  constructor(url) {
    this.url = url;
    this.attempt = 0;
    this.maxDelay = 30000; // 30 seconds max
    this.lastEventId = null;
    this.connect();
  }

  connect() {
    this.ws = new WebSocket(this.url);

    this.ws.onopen = () => {
      this.attempt = 0;
      // Resume from last known position
      if (this.lastEventId) {
        this.ws.send(JSON.stringify({
          type: 'session.resume',
          payload: { lastEventId: this.lastEventId }
        }));
      }
    };

    this.ws.onclose = () => {
      const delay = Math.min(1000 * Math.pow(2, this.attempt), this.maxDelay);
      const jitter = delay * 0.5 * Math.random();
      setTimeout(() => this.connect(), delay + jitter);
      this.attempt++;
    };
  }
}
Key Concept: Exponential Backoff with Jitter

When a server goes down, thousands of clients may try to reconnect simultaneously, creating a “thundering herd.” Exponential backoff (doubling the wait time each attempt) combined with random jitter (adding randomness to the delay) spreads reconnection attempts over time, preventing the server from being overwhelmed on recovery.

Scaling Across Multiple Server Instances

The biggest architectural challenge with WebSockets is horizontal scaling. Since WebSocket connections are stateful and pinned to a specific server instance, broadcasting messages to all users or routing messages between users connected to different servers requires a pub/sub layer:

// Using Redis Pub/Sub for cross-instance communication
const Redis = require('ioredis');
const pub = new Redis();
const sub = new Redis();

// When a message needs to reach all instances
async function broadcast(channel, message) {
  await pub.publish(channel, JSON.stringify(message));
}

// Each server instance subscribes to relevant channels
sub.subscribe('chat:global', 'chat:room:*');
sub.on('message', (channel, raw) => {
  const message = JSON.parse(raw);
  // Deliver to locally connected clients
  connectionManager.deliverToRoom(channel, message);
});

Each WebSocket connection consumes approximately 70 KiB of memory. At 1,000 connections, that’s ~68 MiB — manageable. At 100,000 connections, it climbs to ~6.7 GiB. Production systems typically use sticky sessions (routing a user to the same server instance) combined with Redis Pub/Sub or similar systems for cross-instance communication.

Practical Applications

Interactive Agent Workflows

Modern AI agents often need human approval before executing actions. WebSockets enable real-time tool approval workflows: the agent requests permission to call a tool, the server pushes a tool.request event to the client, the user approves or denies via a tool.approve message, and the agent continues — all without polling or page refreshes.

Real-Time Voice AI

OpenAI’s Realtime API uses WebSockets for bidirectional audio streaming. The client streams microphone audio to the server while simultaneously receiving generated audio back. This requires the full-duplex capability of WebSockets — SSE cannot handle simultaneous bidirectional audio streams.

Collaborative AI Editing

Applications where multiple users interact with the same AI session — pair programming with an AI assistant, collaborative document editing with AI suggestions — require WebSockets for real-time synchronization of state between all participants and the AI backend.

Hybrid Architecture: SSE + WebSocket

Many production AI applications use a hybrid approach: SSE for the high-throughput, unidirectional token stream (the “data plane”) and WebSockets for low-latency bidirectional control signals (the “control plane”). This separation leverages the simplicity and HTTP-compatibility of SSE for the bulk of the traffic while using WebSockets only where bidirectionality is required. Anthropic and OpenAI’s chat UIs both employ variations of this pattern.

Key Takeaways

Further Reading