WebSocket & Real-Time Communication

Role

You are an elite real-time systems engineer. You build low-latency, reliable communication

systems that handle thousands of concurrent connections with proper lifecycle management.

Part 1: FastAPI WebSocket Server

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import Dict, Set
import json, asyncio

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        self.active: Dict[str, WebSocket] = {}       # user_id -> ws
        self.rooms: Dict[str, Set[str]] = {}          # room -> {user_ids}

    async def connect(self, ws: WebSocket, user_id: str):
        await ws.accept()
        self.active[user_id] = ws

    def disconnect(self, user_id: str):
        self.active.pop(user_id, None)
        for room in list(self.rooms):
            self.rooms[room].discard(user_id)
            if not self.rooms[room]:
                del self.rooms[room]

    def join_room(self, user_id: str, room: str):
        self.rooms.setdefault(room, set()).add(user_id)

    def leave_room(self, user_id: str, room: str):
        if room in self.rooms:
            self.rooms[room].discard(user_id)

    async def send_to_user(self, user_id: str, data: dict):
        if ws := self.active.get(user_id):
            await ws.send_json(data)

    async def broadcast_room(self, room: str, data: dict, exclude: str | None = None):
        for uid in self.rooms.get(room, set()):
            if uid != exclude and uid in self.active:
                await self.active[uid].send_json(data)

    async def broadcast_all(self, data: dict):
        for ws in self.active.values():
            await ws.send_json(data)

manager = ConnectionManager()

@app.websocket("/ws/{user_id}")
async def websocket_endpoint(ws: WebSocket, user_id: str):
    await manager.connect(ws, user_id)
    try:
        while True:
            raw = await ws.receive_text()
            msg = json.loads(raw)
            match msg.get("type"):
                case "join":
                    manager.join_room(user_id, msg["room"])
                case "leave":
                    manager.leave_room(user_id, msg["room"])
                case "message":
                    await manager.broadcast_room(
                        msg["room"],
                        {"type": "message", "from": user_id, "text": msg["text"]},
                        exclude=user_id,
                    )
                case "ping":
                    await ws.send_json({"type": "pong", "ts": msg.get("ts")})
    except WebSocketDisconnect:
        manager.disconnect(user_id)
        await manager.broadcast_all({"type": "user_left", "user": user_id})

Part 2: Heartbeat & Reconnection

import asyncio, time

class HeartbeatManager:
    """Server-side heartbeat to detect dead connections."""
    def __init__(self, manager: ConnectionManager, interval: int = 30, timeout: int = 10):
        self.manager = manager
        self.interval = interval
        self.timeout = timeout
        self.last_pong: Dict[str, float] = {}

    async def start(self):
        while True:
            await asyncio.sleep(self.interval)
            dead = []
            for user_id, ws in list(self.manager.active.items()):
                try:
                    await ws.send_json({"type": "ping", "ts": time.time()})
                    # If no pong received within timeout, mark dead
                    last = self.last_pong.get(user_id, 0)
                    if time.time() - last > self.interval + self.timeout:
                        dead.append(user_id)
                except Exception:
                    dead.append(user_id)
            for uid in dead:
                self.manager.disconnect(uid)

    def record_pong(self, user_id: str):
        self.last_pong[user_id] = time.time()

Client-side reconnection with exponential backoff:

class ReconnectingWebSocket {
  constructor(url, options = {}) {
    this.url = url;
    this.maxRetries = options.maxRetries || 10;
    this.baseDelay = options.baseDelay || 1000;
    this.maxDelay = options.maxDelay || 30000;
    this.retries = 0;
    this.handlers = { message: [], open: [], close: [] };
    this.connect();
  }

  connect() {
    this.ws = new WebSocket(this.url);
    this.ws.onopen = () => {
      this.retries = 0;
      this.startHeartbeat();
      this.handlers.open.forEach(h => h());
    };
    this.ws.onmessage = (e) => {
      const data = JSON.parse(e.data);
      if (data.type === 'ping') {
        this.ws.send(JSON.stringify({ type: 'pong', ts: data.ts }));
        return;
      }
      this.handlers.message.forEach(h => h(data));
    };
    this.ws.onclose = () => {
      this.stopHeartbeat();
      if (this.retries < this.maxRetries) {
        const delay = Math.min(this.baseDelay * 2 ** this.retries, this.maxDelay);
        setTimeout(() => this.connect(), delay);
        this.retries++;
      }
      this.handlers.close.forEach(h => h());
    };
  }

  startHeartbeat() {
    this._hb = setInterval(() => {
      if (this.ws.readyState === WebSocket.OPEN) {
        this.ws.send(JSON.stringify({ type: 'ping', ts: Date.now() }));
      }
    }, 25000);
  }

  stopHeartbeat() { clearInterval(this._hb); }
  on(event, handler) { this.handlers[event]?.push(handler); }
  send(data) { this.ws.send(JSON.stringify(data)); }
}

Part 3: Server-Sent Events (SSE)

from fastapi import Request
from fastapi.responses import StreamingResponse
import asyncio, json

async def event_generator(request: Request, user_id: str):
    queue = asyncio.Queue()
    # Register this queue for the user
    sse_clients[user_id] = queue
    try:
        while True:
            if await request.is_disconnected():
                break
            try:
                data = await asyncio.wait_for(queue.get(), timeout=15)
                yield f"event: {data['event']}\ndata: {json.dumps(data['payload'])}\n\n"
            except asyncio.TimeoutError:
                yield f": keepalive\n\n"  # comment line keeps connection alive
    finally:
        sse_clients.pop(user_id, None)

sse_clients: Dict[str, asyncio.Queue] = {}

@app.get("/events/{user_id}")
async def sse_endpoint(request: Request, user_id: str):
    return StreamingResponse(
        event_generator(request, user_id),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )

# Push an event to a user
async def push_event(user_id: str, event: str, payload: dict):
    if q := sse_clients.get(user_id):
        await q.put({"event": event, "payload": payload})

Part 4: Pub/Sub with Redis

import redis.asyncio as redis
import json

class PubSub:
    def __init__(self, redis_url: str = "redis://localhost:6379"):
        self.redis = redis.from_url(redis_url)

    async def publish(self, channel: str, data: dict):
        await self.redis.publish(channel, json.dumps(data))

    async def subscribe(self, channel: str):
        pubsub = self.redis.pubsub()
        await pubsub.subscribe(channel)
        async for message in pubsub.listen():
            if message["type"] == "message":
                yield json.loads(message["data"])

# Bridge Redis pub/sub to WebSocket rooms
async def redis_to_ws_bridge(pubsub: PubSub, room: str):
    async for msg in pubsub.subscribe(f"room:{room}"):
        await manager.broadcast_room(room, msg)

Part 5: Long Polling Fallback

from fastapi import Query
import asyncio, time

pending_messages: Dict[str, asyncio.Queue] = {}

@app.get("/poll/{user_id}")
async def long_poll(user_id: str, timeout: int = Query(default=30, le=60)):
    queue = pending_messages.setdefault(user_id, asyncio.Queue())
    try:
        data = await asyncio.wait_for(queue.get(), timeout=timeout)
        return {"status": "message", "data": data}
    except asyncio.TimeoutError:
        return {"status": "timeout"}

Part 6: Binary Protocol for High Throughput

import struct

# Simple binary frame: [type:1B][length:4B][payload:NB]
MSG_TYPES = {"text": 0x01, "binary": 0x02, "control": 0x03}

def encode_frame(msg_type: str, payload: bytes) -> bytes:
    type_byte = MSG_TYPES[msg_type]
    return struct.pack("!BI", type_byte, len(payload)) + payload

def decode_frame(data: bytes) -> tuple[str, bytes]:
    type_byte, length = struct.unpack("!BI", data[:5])
    type_name = {v: k for k, v in MSG_TYPES.items()}[type_byte]
    return type_name, data[5:5 + length]

@app.websocket("/ws/binary/{user_id}")
async def binary_ws(ws: WebSocket, user_id: str):
    await ws.accept()
    try:
        while True:
            raw = await ws.receive_bytes()
            msg_type, payload = decode_frame(raw)
            if msg_type == "control":
                # Handle control messages (ping, subscribe, etc.)
                pass
            else:
                # Echo back or route
                await ws.send_bytes(encode_frame("text", payload))
    except WebSocketDisconnect:
        pass

Part 7: Scaling with Multiple Workers

# For multi-process deployments, use Redis as message bus
# Each worker subscribes to Redis and forwards to local WebSocket clients

async def worker_listener(worker_id: str, pubsub: PubSub):
    """Each uvicorn worker runs this to relay messages."""
    async for msg in pubsub.subscribe("ws:broadcast"):
        room = msg.get("room")
        if room:
            await manager.broadcast_room(room, msg["data"])
        else:
            await manager.broadcast_all(msg["data"])

# On startup
@app.on_event("startup")
async def startup():
    pubsub = PubSub()
    asyncio.create_task(worker_listener("w1", pubsub))
    heartbeat = HeartbeatManager(manager)
    asyncio.create_task(heartbeat.start())

AXE MCP Server Integration

Every skill in the AXE Skills Hub runs with access to the AXE MCP Server — giving it the full fleet intelligence toolkit automatically. No setup required; tools are available in any AXE-powered session.

Core Tools Available

Quick Start

# In any AXE session, tools are pre-loaded. Example chaining:

# 1. Search for context
results = qdrant_search("user query here", collection="axe_persistent_memory")

# 2. Fetch live data if needed
content = web_fetch("https://docs.example.com/api")

# 3. Write result to memory for next session
write_memory("shared/last_result.md", output)

# 4. Log quality output to Edge training pipeline
hydra_add(prompt=user_query, response=output, score=0.9, source="skill-name")

Edge Training Integration

High-quality skill outputs are automatically eligible for Edge model training via hydra_add. When a response scores ≥0.85 in evals, pipe it to the Hydra pipeline to compound Edge's knowledge. This is how skills make Edge smarter over time.

# After generating a high-quality response:
hydra_add(
    prompt=user_input,
    response=final_output,
    score=0.9,          # eval score
    source="skill-name" # tracks provenance
)