Real-Time AI Features: WebSockets, Streaming, and Partial Results
Switching an AI writing assistant from request-response to streamed token output was an easy win for perceived latency — users saw text appearing within a few hundred milliseconds instead of staring at a spinner for four or five seconds waiting for the full response. That part worked as expected. What we underestimated was how much downstream logic assumed a complete response and broke once output started arriving incrementally.
The first casualty was our content moderation check, which scanned the full response for policy violations before displaying it. With streaming, there was no 'full response' to scan before display — by the time moderation ran, the problematic content was already on screen. We moved to a sliding-window moderation check that scans completed sentences as they stream in, with the ability to retroactively redact a sentence if a check flags it a beat after it rendered, which is a real trade-off but a rare-enough case to accept.
We used WebSockets rather than server-sent events specifically because the feature also needed bidirectional communication — users could interrupt a generating response mid-stream to redirect it, which meant the client needed to send a cancellation signal back over the same connection rather than opening a separate request. That interruption path had its own edge case: canceling a stream needed to actually stop token generation server-side and release the model connection, not just stop forwarding tokens to a client that had lost interest, or we'd silently waste capacity on abandoned requests.
Partial-result handling on the frontend needed more care than we expected too. If the connection dropped mid-stream, the UI needed to distinguish between 'still generating, reconnecting' and 'generation failed, here's what we got.' We added a heartbeat and an explicit stream-completion marker rather than inferring completion from the connection simply closing, since a dropped connection and a finished response looked identical without it.
The result was worth the added complexity — time-to-first-token dropped from roughly 4 seconds to under 400 milliseconds, and user-reported 'is this working' complaints during long generations dropped close to zero. But streaming isn't a frontend-only change; it touches moderation, cancellation, error handling, and reconnection logic that all assumed atomic responses, and each of those needed its own redesign.
