WebSocket streaming for multi-agent LLM pipelines

Polling an agent graph for status is a waste of round trips. websocket streaming multi-agent llm pipelines let you push token deltas, intermediate tool calls, and handoff events over a single duplex connection, which matches the concurrent nature of agent execution far better than half-open SSE channels. The moment you have more than one agent reacting to user input in real time, unidirectional streams break down.

Why WebSockets beat SSE for agent graphs

SSE is unidirectional. It works when one server pushes to one client and the client never talks back. Multi-agent flows invert that assumption constantly: the client sends interrupts, retries, new context, or explicit agent-selection mid-run. WebSockets give you full duplex without spawning extra HTTP requests.

The tradeoff is real. SSE rides on HTTP semantics—timeouts, caching, and reconnection are partially handled by the browser. With WebSockets you own the connection state: heartbeat, reconnect, frame sizing, and graceful shutdown are all your code. For interactive agent UX the duplex win is worth it, but don’t adopt it for batch jobs that outlive the browser tab.

Connection lifecycle and auth

Open the socket after a normal HTTPS handshake. Pass a signed JWT in the query string or as the first application message. Never put secrets in the URL fragment—they leak to logs.

# server: websockets + jwt validation
import websockets, jwt

async def gateway(websocket, path):
    token = dict(websockets.urlparse(path).query).get("token")
    try:
        jwt.decode(token, KEY, algorithms=["EdDSA"])
    except jwt.InvalidTokenError:
        await websocket.close(4401, "unauthorized")
        return
    await agent_loop(websocket)

The client connects with new WebSocket(\wss://host/agent?token=${tok}`). Send a hello frame with a session id before any agent invocation. Use subprotocols (Sec-WebSocket-Protocol) to version your API; agent-v1` is clearer than documenting magic fields.

Message protocol design

Define a strict JSON envelope. Raw text frames force you to reconstruct state from ambiguous strings. You need explicit types for tokens, tool calls, errors, and control signals.

{
  "type": "token",
  "agent": "planner",
  "run_id": "r1",
  "delta": "Thinking about"
}
{
  "type": "tool_call",
  "agent": "retriever",
  "run_id": "r1",
  "name": "vector_search",
  "args": {"q": "postgres replication"}
}
{
  "type": "control",
  "action": "cancel",
  "run_id": "r1"
}

Keep deltas small. Batch if a frame exceeds 4 KB to reduce syscall overhead. Version the envelope ("v": 1) so you can evolve without breaking old clients. If you ever need binary (e.g., audio), use a separate type: "binary" and send bytes frames, but keep control planes JSON.

Orchestrating multiple agents over one socket

The core of websocket streaming multi-agent llm pipelines is a run_id-scoped context. Spawn agents as asyncio tasks under a single broker. The socket coroutine reads client messages and dispatches to the right run.

async def agent_loop(ws):
    runs = {}
    async for msg in ws:
        data = json.loads(msg)
        if data["type"] == "start":
            runs[data["run_id"]] = asyncio.create_task(
                run_agent(ws, data)
            )
        elif data["type"] == "control" and data["action"] == "cancel":
            task = runs.pop(data["run_id"], None)
            if task:
                task.cancel()

Each agent streams via await ws.send(json.dumps(...)). Because asyncio is single-threaded, sends are naturally serialized—no extra lock. For a supervisor pattern, one agent can emit a spawn control frame that the loop turns into another task.

If you front the model calls with a gateway such as n4n.ai, it honors client routing directives and automatically fails over when a provider is rate-limited, so your socket layer only forwards the resulting stream instead of branching per provider.

Handling backpressure and cancellation

Client-side buffering

Browsers buffer WebSocket data at the network layer, but your UI renderer does not. Use a queue and drain at animation-frame rate to avoid jank.

const queue: string[] = [];
ws.onmessage = (e) => queue.push((JSON.parse(e.data) as any).delta);
function flush() {
  const text = queue.splice(0, queue.length).join("");
  document.getElementById("out")!.innerText += text;
  requestAnimationFrame(flush);
}
requestAnimationFrame(flush);

Server-side agent cancellation

When the client sends cancel, propagate to the running LLM call immediately. Most SDKs accept asyncio.CancelledError or an event flag. If you are mid-stream from a provider, close the upstream connection to stop token billing.

Pitfall: ignoring cancellation leaves orphaned agents consuming tokens and holding sockets. Always task.cancel() and await task inside a try/except asyncio.CancelledError. Monitor ws.transport.get_write_buffer_size(); if it climbs past a few hundred KB, the client is slow and you should pause agent loops or drop the connection.

Fallback and provider routing

Multi-agent systems call different models for different roles. A planner might use a reasoning model; a summarizer uses a cheap one. Your socket server should accept a model field per agent start and validate it against an allowlist.

async def call_model(ws, run_id, agent, prompt, model):
    try:
        async for chunk in client.chat.completions.create(
            model=model, messages=prompt, stream=True
        ):
            await ws.send(json.dumps({"type":"token","agent":agent,
                "run_id":run_id,"delta":chunk.choices[0].delta.content}))
    except RateLimitError:
        await call_model(ws, run_id, agent, prompt, "fallback-model")

A unified OpenAI-compatible endpoint removes per-provider branching. You still need to map agent roles to model names, but the transport code stays identical.

Scaling beyond a single process

The loop above is single-process. For production, separate socket termination from agent execution using a pub/sub broker (Redis or NATS). The WebSocket worker subscribes to run_id channels and forwards frames; agent workers publish to those channels. This lets you scale agents horizontally without sticky connections.

# pseudo: redis pubsub forwarder
async def ws_forwarder(ws, redis, run_id):
    channel = await redis.subscribe(run_id)
    async for msg in channel:
        await ws.send(msg)

Heartbeat: send {type:"ping"} every 20s. Proxies kill idle sockets at 60s. Close cleanly on pong timeout.

Common pitfalls

• No heartbeat. Silent dead sockets waste resources. Implement app-level ping/pong.
• Massive frames. A 2 MB tool result as one frame blocks the event loop and blows client buffers.
• Ignoring backpressure. If ws.send awaits too long, agent tasks pile up. Cap concurrency per socket.
• Mixed run IDs. Always scope state per run_id; global variables cause cross-talk between sessions.
• Trusting client model strings. Validate against an allowlist; otherwise users invoke expensive endpoints.
• No replay log. If the socket drops, you lose in-flight agent state. Persist run events to a store so a reconnect can resume.

Minimal reference implementation

Server excerpt:

import asyncio, json, websockets

async def run_agent(ws, data):
    for w in ["Plan:", " step1", " step2"]:
        await ws.send(json.dumps({"type":"token","agent":data["agent"],
            "run_id":data["run_id"],"delta":w}))
        await asyncio.sleep(0.1)

async def agent_loop(ws):
    runs = {}
    async for msg in ws:
        d = json.loads(msg)
        if d["type"] == "start":
            runs[d["run_id"]] = asyncio.create_task(run_agent(ws, d))
        elif d["type"] == "control":
            task = runs.pop(d["run_id"], None)
            if task: task.cancel()

async def main():
    async with websockets.serve(agent_loop, "0.0.0.0", 8765):
        await asyncio.Future()

asyncio.run(main())

Client:

const ws = new WebSocket("wss://localhost:8765");
ws.onopen = () => ws.send(JSON.stringify({
  type: "start", agent: "planner", run_id: "r1"}));
ws.onmessage = (e) => console.log(JSON.parse(e.data));

That skeleton lacks auth, heartbeats, and routing, but shows the dispatch shape.

When to avoid WebSockets

If your agents run longer than the client session—overnight batch orchestration, deferred human approval—use a task queue with SSE or polling. WebSockets are not durable; a dropped connection mid-run means rebuilding state from a persisted log. For short-lived interactive agent graphs, websocket streaming multi-agent llm pipelines cut latency and simplify interrupt handling. Choose the protocol that matches your session lifetime, not the trend.