Streaming Agent Runs over WebSockets

software-engineering
full-stack
agents
websockets
streaming
reliability
Carry one browser command into the application service and stream durable text, tool, cancellation, and replay events back.

An agent run is a conversation between one command and many ordered results. The browser needs early feedback, tool visibility, a real cancellation path, and recovery after a connection drops. This chapter adds a WebSocket endpoint to the REST application from Chapter 02 and proves that every frame shown in the browser corresponds to an event already stored by the application service.

Define commands and events separately

The browser sends commands such as {"type": "user_message", "content": "..."} and {"type": "cancel"}. The server returns durable session events containing event_id, session_id, cursor, kind, payload, and created_at. Commands express intent; events state what happened. Reusing one vague “message” schema for both directions makes validation and replay ambiguous.

The WebSocket handler validates commands and delegates the run to AutocodeApplication. A concurrent sender task drains the session broker, so model work does not block outbound frames and every observer receives the same persisted event dictionary.

from tempfile import TemporaryDirectory

from fastapi.testclient import TestClient

from autocode.runner import DemoAgentRunner
from autocode_service.api import create_app

with TemporaryDirectory() as directory:
    app = create_app(
        database_path=f"{directory}/sessions.db",
        runner=DemoAgentRunner(),
    )
    with TestClient(app) as client:
        session_id = client.post(
            "/api/sessions", json={"title": "WebSocket lesson"}
        ).json()["session_id"]
        with client.websocket_connect(f"/ws/sessions/{session_id}") as socket:
            socket.send_json({"type": "user_message", "content": "show the event path"})
            events = []
            while not events or events[-1]["kind"] != "run_finished":
                events.append(socket.receive_json())

assert events[0]["kind"] == "user_message"
assert any(event["kind"] == "text_delta" for event in events)
assert events[-1]["kind"] == "run_finished"
assert [event["cursor"] for event in events] == list(range(1, len(events) + 1))
print("streamed event kinds:", [event["kind"] for event in events])
streamed event kinds: ['user_message', 'run_started', 'text_delta', 'text_delta', 'text_delta', 'text_delta', 'text_delta', 'text_delta', 'text_delta', 'assistant_message', 'run_finished']

This is an integration test, not a direct function call. It crosses WebSocket parsing, the background run task, application orchestration, journal and SQLite writes, broker fan-out, and JSON serialization. The first returned event is the durable user message; the terminal event proves the run reached a named finish state.

Project the stream without duplicating text

Text deltas make the response visible early, while assistant_message stores the complete canonical text. During a live run, the browser concatenates deltas into one temporary assistant bubble. When the complete message arrives, it replaces that temporary buffer. Replaying both event kinds must not render the answer twice.

The browser stores durable events in a map keyed by event_id and sorts them by cursor before projection. Duplicate delivery becomes a no-op. Tool starts and finishes become cards, while terminal events clear the streaming state and re-enable the composer.

# Rebuild a projection from duplicate delivery and ordered cursors.
delivered = [events[0], *events, events[-1]]
by_id = {event["event_id"]: event for event in delivered}
ordered = sorted(by_id.values(), key=lambda event: event["cursor"])

streamed_text = "".join(
    event["payload"].get("content", "")
    for event in ordered
    if event["kind"] == "text_delta"
)
complete_text = next(
    event["payload"]["content"]
    for event in ordered
    if event["kind"] == "assistant_message"
)

assert len(ordered) == len(events)
assert streamed_text == complete_text
print("duplicate-safe cursors:", [event["cursor"] for event in ordered])
duplicate-safe cursors: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

The equality check separates transport granularity from durable meaning. Deltas optimize time to first feedback; the complete assistant event is the record shown after refresh. Deduplication by event id and ordering by cursor make reconnect safe even when the network repeats the boundary event.

Cancellation crosses the full stack

A cancel button must do more than hide a spinner. The browser sends a cancel command, the server cancels the active task for that session, and the application catches cancellation at the runner boundary. Before propagation completes, it records and publishes run_finished with reason cancelled. Refreshing the page therefore preserves the distinction between a completed, failed, and user-cancelled run.

from tempfile import TemporaryDirectory

from fastapi.testclient import TestClient

from autocode.runner import DemoAgentRunner
from autocode_service.api import create_app

with TemporaryDirectory() as directory:
    app = create_app(
        database_path=f"{directory}/sessions.db",
        runner=DemoAgentRunner(delay=0.2),
    )
    with TestClient(app) as client:
        session_id = client.post("/api/sessions", json={"title": "cancel"}).json()["session_id"]
        with client.websocket_connect(f"/ws/sessions/{session_id}") as socket:
            socket.send_json({"type": "user_message", "content": "long task"})
            assert socket.receive_json()["kind"] == "user_message"
            assert socket.receive_json()["kind"] == "run_started"
            socket.send_json({"type": "cancel"})
            terminal = socket.receive_json()
        restored = client.get(f"/api/sessions/{session_id}").json()

assert terminal["kind"] == "run_finished"
assert terminal["payload"]["reason"] == "cancelled"
assert restored["events"][-1]["event_id"] == terminal["event_id"]
print("durable finish reason:", terminal["payload"]["reason"])
durable finish reason: cancelled

The event received by the browser is the event restored from SQLite. Closing a browser connection, by contrast, does not automatically cancel the run: the server task can continue and later clients can replay its events. Chapter 08 generalizes this lifecycle to multiple observers, bounded queues, and cursor-based reconnect.

Exercises

Extend the command protocol with an approval response for a tool call. Define validation, the active-run state transition, persistence ordering, duplicate behavior, and what the browser renders while approval is pending.

[P03.1] Design a tool-approval round trip

Specify the WebSocket commands and durable events for requesting, granting, and denying one tool call. Include stable identifiers, invalid or duplicate responses, reconnect behavior, and the terminal state shown by the browser.

Back to top