FastAPI WebSocket Production: Real-Time AI Agent Communication

Build production WebSocket servers in FastAPI for agent-to-agent communication, connection pooling, backpressure handling, and scaling to thousands of concurrent connections.

0

REST APIs work great for request-response patterns, but they require repeated HTTP handshakes and polling when you need to stream AI reasoning, push live updates to multiple agents, or maintain persistent bidirectional connections at scale. WebSockets change that equation entirely. They give you a persistent, low-latency channel where both client and server can send data whenever needed, without the overhead of repeated HTTP handshakes.

If you’re building multi-agent systems, live dashboards for AI reasoning, or infrastructure where agents need to coordinate in real time, FastAPI’s WebSocket support paired with async/await makes this surprisingly elegant. The challenge isn’t the basics, it’s production: handling thousands of concurrent connections, managing backpressure when agents stream faster than consumers can process, pooling database connections efficiently, and gracefully degrading when things get tight.

Let’s walk through a real production pattern, from connection lifecycle to scaling strategies.

Why WebSockets for AI Agent Communication

Consider a typical multi-agent setup: Agent A processes a task and discovers it needs Agent B’s output. With REST, Agent A makes a blocking HTTP call to Agent B, waits for the response, and continues. If Agent B is slow or the network hiccups, Agent A blocks. Now scale that to dozens of agents coordinating in real time, and you’ve got latency accumulation and resource contention.

WebSockets flip this. Agent A opens a connection to Agent B, subscribes to events, and Agent B pushes updates as they arrive. No blocking, no polling, no wasted requests. Both sides can send at any time. For streaming AI responses (tokens arriving one at a time), this is essential. You can’t afford the latency of a new HTTP request per token.

The benefits compound:

  • Bidirectional: agents push and receive without waiting for a request
  • Low latency: persistent connection, no handshake overhead
  • Efficient streaming: tokens, events, or state changes flow as they happen
  • Event-driven: naturally suits pub-sub and event-driven architectures

Basic WebSocket Server in FastAPI

FastAPI’s WebSocket support is built on Starlette and integrates seamlessly with async/await. Here’s the simplest working example:

from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Echo: {data}")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        await websocket.close()

This accepts a connection, reads messages in a loop, echoes them back, and cleans up on disconnect. It’s a starting point, but production needs more: connection tracking, broadcast capabilities, backpressure handling, and database integration.

Connection Lifecycle and Management

In production, you need to track active connections, route messages to the right agents, and handle disconnects cleanly. A connection manager class gives you this:

from typing import Dict, List, Set
from fastapi import WebSocket
import asyncio

class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, List[WebSocket]] = {}
        self.connection_lock = asyncio.Lock()
    
    async def connect(self, agent_id: str, websocket: WebSocket):
        await websocket.accept()
        async with self.connection_lock:
            if agent_id not in self.active_connections:
                self.active_connections[agent_id] = []
            self.active_connections[agent_id].append(websocket)
    
    async def disconnect(self, agent_id: str, websocket: WebSocket):
        async with self.connection_lock:
            if agent_id in self.active_connections:
                self.active_connections[agent_id].remove(websocket)
                if not self.active_connections[agent_id]:
                    del self.active_connections[agent_id]
    
    async def broadcast_to_agent(self, agent_id: str, message: str):
        async with self.connection_lock:
            connections = self.active_connections.get(agent_id, [])
        
        # Send outside the lock to avoid blocking other operations
        for connection in connections:
            try:
                await connection.send_text(message)
            except Exception as e:
                print(f"Failed to send message: {e}")
    
    async def broadcast_to_all(self, message: str):
        async with self.connection_lock:
            all_connections = []
            for connections in self.active_connections.values():
                all_connections.extend(connections)
        
        for connection in all_connections:
            try:
                await connection.send_text(message)
            except Exception as e:
                print(f"Failed to send message: {e}")

manager = ConnectionManager()

@app.websocket("/ws/{agent_id}")
async def websocket_endpoint(websocket: WebSocket, agent_id: str):
    await manager.connect(agent_id, websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast_to_agent(agent_id, f"Message from {agent_id}: {data}")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        await manager.disconnect(agent_id, websocket)

This manager tracks connections by agent ID, handles broadcast operations safely with a lock, and ensures cleanup on disconnect. The lock is held only during dictionary operations, not during send, so one slow client doesn’t block others.

Backpressure Handling for High-Volume Streaming

When agents stream AI responses (thousands of tokens per second), a naive approach sends each token immediately and hopes the client keeps up. If the client falls behind, the send buffer grows, memory usage increases, and the server can become resource-constrained.

Backpressure handling means: if the client can’t keep up, the server slows down or drops messages gracefully. Here’s a practical pattern:

import asyncio
from collections import deque
from enum import Enum

class BackpressureStrategy(Enum):
    QUEUE = "queue"  # Buffer up to a limit, then drop
    DROP = "drop"    # Drop newest messages
    BLOCK = "block"  # Wait for client to catch up

class BackpressureManager:
    def __init__(self, max_queue_size: int = 1000, strategy: BackpressureStrategy = BackpressureStrategy.QUEUE):
        self.max_queue_size = max_queue_size
        self.strategy = strategy
        self.queue: deque = deque()
        self.send_event = asyncio.Event()
    
    async def add_message(self, message: str):
        if len(self.queue) >= self.max_queue_size:
            if self.strategy == BackpressureStrategy.DROP:
                # Silently drop oldest message
                self.queue.popleft()
            elif self.strategy == BackpressureStrategy.BLOCK:
                # Wait until queue drains
                while len(self.queue) >= self.max_queue_size:
                    await asyncio.sleep(0.01)
        
        self.queue.append(message)
        self.send_event.set()
    
    async def get_message(self):
        while not self.queue:
            self.send_event.clear()
            await self.send_event.wait()
        
        return self.queue.popleft()

@app.websocket("/ws/stream/{agent_id}")
async def stream_endpoint(websocket: WebSocket, agent_id: str):
    await websocket.accept()
    backpressure = BackpressureManager(max_queue_size=500, strategy=BackpressureStrategy.QUEUE)
    
    try:
        while True:
            message = await backpressure.get_message()
            await websocket.send_text(message)
    except Exception as e:
        print(f"Stream error: {e}")
    finally:
        await websocket.close()

This pattern decouples message production from sending. Messages are queued, and a separate loop sends them at the client’s pace. If the queue fills, you choose: drop old messages (useful for live dashboards where recent data matters most), drop new messages (preserve recent state), or block the producer (apply backpressure upstream).

Connection Pooling with asyncpg

Most agent systems need database integration: storing agent state, retrieving task context, logging interactions. With WebSockets handling thousands of concurrent connections, you need connection pooling to avoid exhausting database resources.

asyncpg is the best choice for async PostgreSQL in Python. Here’s how to integrate it with your WebSocket server:

import asyncpg
from contextlib import asynccontextmanager

class DatabasePool:
    def __init__(self, dsn: str, min_size: int = 10, max_size: int = 20):
        self.dsn = dsn
        self.min_size = min_size
        self.max_size = max_size
        self.pool = None
    
    async def initialize(self):
        self.pool = await asyncpg.create_pool(
            self.dsn,
            min_size=self.min_size,
            max_size=self.max_size,
            command_timeout=10,
        )
    
    async def close(self):
        if self.pool:
            await self.pool.close()
    
    @asynccontextmanager
    async def acquire(self):
        async with self.pool.acquire() as connection:
            yield connection

db_pool = DatabasePool("postgresql://user:password@localhost/agents")

@app.on_event("startup")
async def startup():
    await db_pool.initialize()

@app.on_event("shutdown")
async def shutdown():
    await db_pool.close()

@app.websocket("/ws/{agent_id}")
async def websocket_endpoint(websocket: WebSocket, agent_id: str):
    await websocket.accept()
    try:
        # Retrieve agent state from database
        async with db_pool.acquire() as conn:
            state = await conn.fetchrow(
                "SELECT state FROM agents WHERE id = $1",
                agent_id
            )
        
        if state:
            await websocket.send_json({"type": "state", "data": dict(state)})
        
        while True:
            data = await websocket.receive_text()
            # Update database
            async with db_pool.acquire() as conn:
                await conn.execute(
                    "UPDATE agents SET state = $1 WHERE id = $2",
                    data,
                    agent_id
                )
            
            await websocket.send_text(f"State updated")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        await websocket.close()

The pool manages connections efficiently: reuses them across requests, limits the total number to avoid overwhelming the database, and times out idle connections. Each WebSocket handler acquires a connection only when needed, then releases it back to the pool.

Graceful Degradation and Circuit Breaking

When the database is slow or the system is under load, you don’t want every WebSocket handler to block waiting for a connection. Implement a circuit breaker to detect failures and degrade gracefully:

from datetime import datetime, timedelta
from enum import Enum

class CircuitState(Enum):
    CLOSED = "closed"      # Normal operation
    OPEN = "open"         # Failing, reject requests
    HALF_OPEN = "half_open"  # Testing if recovered

class CircuitBreaker:
    def __init__(self, failure_threshold: int = 5, timeout_seconds: int = 60):
        self.failure_threshold = failure_threshold
        self.timeout_seconds = timeout_seconds
        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.last_failure_time = None
    
    async def call(self, coro):
        if self.state == CircuitState.OPEN:
            if datetime.now() - self.last_failure_time > timedelta(seconds=self.timeout_seconds):
                self.state = CircuitState.HALF_OPEN
                self.failure_count = 0
            else:
                raise Exception("Circuit breaker is open")
        
        try:
            result = await coro
            if self.state == CircuitState.HALF_OPEN:
                self.state = CircuitState.CLOSED
                self.failure_count = 0
            return result
        except Exception as e:
            self.failure_count += 1
            self.last_failure_time = datetime.now()
            if self.failure_count >= self.failure_threshold:
                self.state = CircuitState.OPEN
            raise

db_breaker = CircuitBreaker(failure_threshold=5, timeout_seconds=30)

@app.websocket("/ws/{agent_id}")
async def websocket_endpoint(websocket: WebSocket, agent_id: str):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            try:
                async with db_pool.acquire() as conn:
                    await db_breaker.call(
                        conn.execute(
                            "UPDATE agents SET state = $1 WHERE id = $2",
                            data,
                            agent_id
                        )
                    )
                await websocket.send_text("OK")
            except Exception as e:
                # Database unavailable, send cached response
                await websocket.send_text(f"Cached: {data}")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        await websocket.close()

When the database fails repeatedly, the circuit opens and rejects new requests immediately instead of waiting for timeouts. After a cooldown, it tries again. This prevents cascading failures and lets the system recover.

Deployment Patterns: Scaling to Production

WebSocket connections are stateful, so you can’t just spin up multiple server instances and use a simple load balancer. Each connection is tied to a specific server process. You have two main patterns:

Sticky Sessions with Load Balancer

Use a load balancer that routes all requests from the same client to the same server instance. This is simple but limits flexibility. Azure Load Balancer and AWS ALB both support session affinity.

Redis-Backed Message Bus

Better approach: use Redis as a message broker. Each server instance subscribes to a channel, and messages are routed through Redis. This decouples connection state from server instances:

import redis.asyncio as redis
import json

class RedisMessageBus:
    def __init__(self, redis_url: str = "redis://localhost"):
        self.redis_url = redis_url
        self.redis = None
    
    async def initialize(self):
        self.redis = await redis.from_url(self.redis_url)
    
    async def publish(self, channel: str, message: dict):
        await self.redis.publish(channel, json.dumps(message))
    
    async def subscribe(self, channel: str):
        pubsub = self.redis.pubsub()
        await pubsub.subscribe(channel)
        return pubsub

message_bus = RedisMessageBus()

@app.on_event("startup")
async def startup():
    await db_pool.initialize()
    await message_bus.initialize()

@app.websocket("/ws/{agent_id}")
async def websocket_endpoint(websocket: WebSocket, agent_id: str):
    await websocket.accept()
    pubsub = await message_bus.subscribe(f"agent:{agent_id}")
    
    async def receive_from_client():
        while True:
            data = await websocket.receive_text()
            await message_bus.publish(f"agent:{agent_id}", {"from": agent_id, "data": data})
    
    async def send_to_client():
        async for message in pubsub.listen():
            if message["type"] == "message":
                await websocket.send_text(message["data"])
    
    try:
        await asyncio.gather(
            receive_from_client(),
            send_to_client()
        )
    except Exception as e:
        print(f"Error: {e}")
    finally:
        await pubsub.close()
        await websocket.close()

Now any server instance can handle any connection. Messages flow through Redis, so agents on different servers can communicate seamlessly. This scales horizontally: add more instances, they all subscribe to the same channels, and load distributes naturally.

Deployment on Azure Container Instances and Kubernetes

For Azure Container Instances, deploy your FastAPI app in a container with appropriate resource limits:

# azure-container-instance.yaml
apiVersion: 2019-12-01
name: fastapi-agents
properties:
  containers:
  - name: fastapi-app
    properties:
      image: myregistry.azurecr.io/fastapi-agents:latest
      ports:
      - port: 8000
      resources:
        requests:
          cpu: 1.0
          memoryInGb: 1.5
      environmentVariables:
      - name: DATABASE_URL
        secureValue: postgresql://...
      - name: REDIS_URL
        secureValue: redis://...
  osType: Linux
  ipAddress:
    type: Public
    ports:
    - protocol: TCP
      port: 8000

For Kubernetes, use a Deployment with multiple replicas, a Service to route traffic, and an Ingress for external access. Include a liveness probe to detect dead connections:

# k8s-deployment.yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: fastapi-agents
spec:
  replicas: 3
  selector:
    matchLabels:
      app: fastapi-agents
  template:
    metadata:
      labels:
        app: fastapi-agents
    spec:
      containers:
      - name: fastapi-app
        image: myregistry.azurecr.io/fastapi-agents:latest
        ports:
        - containerPort: 8000
        livenessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 10
          periodSeconds: 10
        resources:
          requests:
            cpu: 500m
            memory: 512Mi
          limits:
            cpu: 1000m
            memory: 1Gi
        env:
        - name: DATABASE_URL
          valueFrom:
            secretKeyRef:
              name: agent-secrets
              key: database-url
        - name: REDIS_URL
          valueFrom:
            secretKeyRef:
              name: agent-secrets
              key: redis-url
---
apiVersion: v1
kind: Service
metadata:
  name: fastapi-agents-service
spec:
  selector:
    app: fastapi-agents
  ports:
  - protocol: TCP
    port: 80
    targetPort: 8000
  type: LoadBalancer

Include a health check endpoint that doesn’t require database access, so Kubernetes can detect pod failures independently:

@app.get("/health")
async def health():
    return {"status": "ok"}

Monitoring and Observability

With thousands of concurrent WebSocket connections, you need visibility into what’s happening. Log connection lifecycle events, track active connections, and monitor backpressure:

import logging
import time
from prometheus_client import Counter, Gauge, Histogram

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

ws_connections = Gauge('websocket_connections_total', 'Active WebSocket connections')
ws_messages_sent = Counter('websocket_messages_sent_total', 'Total messages sent', ['agent_id'])
ws_message_latency = Histogram('websocket_message_latency_seconds', 'Message send latency')

@app.websocket("/ws/{agent_id}")
async def websocket_endpoint(websocket: WebSocket, agent_id: str):
    await websocket.accept()
    ws_connections.inc()
    logger.info(f"Agent {agent_id} connected")
    
    try:
        while True:
            data = await websocket.receive_text()
            start = time.time()
            await websocket.send_text(f"Echo: {data}")
            latency = time.time() - start
            ws_message_latency.observe(latency)
            ws_messages_sent.labels(agent_id=agent_id).inc()
    except Exception as e:
        logger.error(f"Error with agent {agent_id}: {e}")
    finally:
        ws_connections.dec()
        logger.info(f"Agent {agent_id} disconnected")
        await websocket.close()

Export these metrics to Prometheus, and visualize them in Grafana. Set alerts for high connection count, message latency spikes, or circuit breaker opens.

Production Considerations

A few key points to keep in mind when deploying WebSocket servers:

Lock timing: Don’t hold async locks while sending messages. Send outside the lock, or you’ll block other operations. We showed this in the ConnectionManager example.

Queue sizing: Always set a max queue size for backpressure. Without it, memory grows until the server becomes resource-constrained under load.

Connection timeouts: WebSocket connections can hang if the client disconnects abruptly. Set read timeouts on your WebSocket handler to detect dead connections.

Database connection exhaustion: Each WebSocket handler that queries the database needs a connection. With 10,000 concurrent connections and a pool of 20, you’ll quickly reach the limit. Use connection pooling and keep queries fast.

Stateless scaling without a message bus: If you deploy multiple server instances without Redis or another message bus, agents on different servers can’t communicate. Use Redis or a similar system to broker messages across instances.

Conclusion

FastAPI’s WebSocket support, combined with Python’s async/await, makes building real-time agent communication systems straightforward. The production challenges, connection pooling, backpressure, and scaling, are solvable with patterns we’ve covered: connection managers for lifecycle handling, backpressure strategies for high-volume streaming, asyncpg connection pooling for database integration, circuit breakers for graceful degradation, and Redis message buses for horizontal scaling.

Start with a basic connection manager and backpressure handling. Add database integration and circuit breakers as you grow. Deploy with sticky sessions or Redis-backed routing, depending on your infrastructure. Monitor connection counts, message latency, and error rates. With these patterns in place, you can build WebSocket servers that handle thousands of concurrent agent connections reliably.

How do WebSockets differ from REST APIs for agent communication?

REST APIs use a request-response model where the client must initiate every interaction, causing latency and overhead when agents need to push updates or stream data. WebSockets maintain a persistent, bidirectional connection where both client and server can send data at any time, making them ideal for real-time agent coordination and token streaming.

What’s the best backpressure strategy for streaming AI responses?

For live dashboards and non-critical streams, drop old messages when the queue fills, so new data always flows. For critical agent-to-agent communication, use the block strategy to apply backpressure upstream and slow the producer. The queue strategy with a size limit offers a middle ground, buffering a reasonable amount before deciding.

Can I run a WebSocket server on a single instance, or do I need horizontal scaling?

A single well-tuned FastAPI instance can handle thousands of concurrent WebSocket connections. Scale horizontally when you need fault tolerance, load distribution, or want to exceed a single machine’s capacity. Use Redis as a message bus to coordinate across multiple instances without sticky sessions.

How do I monitor WebSocket connections in production?

Expose Prometheus metrics for active connections, messages sent, and message latency. Log connection lifecycle events (connect, disconnect, errors). Set up Grafana dashboards and alerts for high connection count, latency spikes, or circuit breaker opens. Include a health check endpoint that doesn’t require database access so Kubernetes can detect pod failures.

What’s the relationship between connection pooling and WebSocket scaling?

Each WebSocket connection that queries the database needs a database connection from the pool. With thousands of WebSocket clients, a small pool quickly reaches its limit. Use asyncpg with appropriate pool sizing (min 10, max 20 is typical), keep queries fast, and acquire connections only when needed. Without pooling, you’ll run out of database connections and start dropping clients.

Leave a Reply

Your email address will not be published. Required fields are marked *