A WebSocket Copilot Streaming Tokens and Tool Calls over One Socket

Stream LLM Tokens and Tool Calls over One WebSocket

Introduction

I built my first AI bot with standard HTTP streams. It worked well for plain text.

Then I added tool calls. I added user control steps.

Everything broke fast.

HTTP streams send data in one direction. The server streams text down to the browser.

If a user clicked stop, I sent a new HTTP call. If a tool needed approval, the app made another HTTP call.

Managing multiple HTTP routes for one chat felt messy. It created race bugs on my backend.

I switched to a single WebSocket link. A WebSocket gives you a full two-way link.

You can stream text tokens down. You can send tool calls down.

You can receive user approvals up. You can handle stop signals up.

Everything runs over one socket.

In my experience, building a socket bot requires clear rules. You need clean event types. You need strict frame limits. You need safe reconnect logic.

Here is what I learned from building WebSocket bots in production.

Why Choose WebSockets for AI Bots?

Before looking at wire events, let us compare stream types.

The Limits of One-Way HTTP Streams

Server-Sent Events work well for simple chat apps. The server streams text tokens to the browser.

The stream ends when the answer finishes.

However, real AI bots need two-way communication.

Users hit stop to halt long answers. Tool steps ask users for approval before running.

The app sends live cursor moves to the model.

With HTTP streams, every user action needs a new HTTP POST call. This setup causes race conditions.

It leads to state bugs on your server.

The Power of Two-Way Socket Links

A WebSocket holds one open link. Both sides can send data at any time.

Your backend streams text tokens down to the browser. It streams tool calls down as well.

At the same time, the browser sends user edits up. It sends stop signals up.

One socket keeps your state in sync. It simplifies your code.

4 Protocol Rules for One-Socket Bots

1. Clean Event Envelopes

Do not send raw text down the socket. Wrap every message in a clean JSON event frame.

Your backend should send four main event types:

• token_delta: A small text chunk from the AI model.
• tool_start: A signal that a tool call is starting.
• tool_args_delta: A streaming chunk of tool argument JSON.
• turn_end: A marker showing the model turn is done.

The client app should send three event types up:

• user_message: Fresh text entered by the user.
• tool_result: The output from a client tool run.
• user_interrupt: A stop signal to halt model output.

Using clear event names makes frontend code simple and clean.

2. Strict Frame Limits and Partial JSON

A model answer can include text and tool calls. You must mark where each section starts and ends.

When an LLM calls a tool, it streams argument strings in chunks. Your frontend must parse partial JSON strings on the fly.

I use three rules for tool boundaries:

• Send a tool_start event with a unique tool ID.
• Stream argument chunks using tool_args_delta events with that ID.
• Send a tool_end event when the argument payload ends.

This boundary rule stops your UI from breaking while streaming JSON.

3. Fast User Interrupts and Stop Signals

When a user clicks stop, your server must halt model generation right away.

Running extra model steps wastes server money. It wastes GPU tokens.

With a WebSocket, the client sends a user_interrupt frame over the open link.

Your backend socket listener receives the frame. It cancels the model API call right away.

It clears the active tool queue. It sends a turn_cancelled frame back to the browser.

This two-way setup halts running calls in under 50 milliseconds.

4. Reconnects and Stream Sync

Mobile Wi-Fi drops often. A good socket engine must handle dropped links safely.

It must not lose chat state.

To make socket streams safe against drops:

• Assign a sequence number to every outgoing event frame.
• Store recent frames in a fast Redis cache on your server.
• When a browser reconnects, it sends its last sequence ID.
• Your server replays missing frames from Redis to update the UI.

This stream sync rule keeps the chat clean during network drops.

Real-World Production Challenges

Race Conditions Between Tokens and Tools

An LLM might stream text tokens first. Then it calls a tool in the middle of a sentence.

If your frontend handles frames out of order, UI elements will overlap. Text widgets and tool boxes will render wrong.

Always enforce ordered event queues on the client side. Process each event frame in strict order before updating the UI state.

Managing Socket Delay

An LLM can stream tokens very fast. A mobile client might render text slowly.

When this happens, socket memory buffers fill up fast.

Monitor socket buffer sizes on your server. If a client falls behind, pause token streaming briefly.

Resume streaming once the client buffer clears.

Clearing Idle Sockets

Open WebSocket links use server memory. Sockets that sit open for hours waste resources.

Set up a ping-pong heartbeat on your server:

• Send a ping frame every 30 seconds.
• Close the socket if the browser fails to send a pong reply.
• Clean up related memory buffers right away.

Heartbeats keep your backend server fast and lean.

WebSocket Bot vs. Standard HTTP Stream

Two-way sockets give you total control over live chat flows.

Actionable Takeaways

• Use JSON event envelopes: Wrap all data in typed frames like token_delta and tool_start.
• Track sequence IDs: Number every frame so client reconnects can replay missing data.
• Handle interrupts upstream: Cancel model calls instantly when receiving a stop frame.
• Enforce UI event queues: Process socket frames in strict order on the frontend.
• Set up heartbeat pings: Close idle sockets after 30 seconds to save server memory.

Final Thoughts

Building a WebSocket bot takes more work than setting up an HTTP stream. You must manage socket servers, event rules, and frame boundaries.

However, the results are worth the effort. A single socket gives you instant stops, smooth tool steps, and a clean user experience.

If you build live AI assistants today, a full-duplex socket setup is the most reliable choice.

Frequently Asked Questions (FAQs)

Why is a WebSocket better than SSE for AI bots?

WebSockets allow two-way messaging. You can stream tokens down while receiving user stop commands and tool approvals up the same link.

How do you handle partial JSON streaming for tool calls?

Send a tool_start event with a tool ID. Stream argument chunks with matching IDs. Mark the end with a tool_end event.

How do you handle lost Wi-Fi connections on WebSockets?

Number every socket event frame. Store frames in Redis. Replay missed frames when the client reconnects with its last frame ID.

Do WebSockets use more server resources than HTTP?

Yes. WebSockets hold persistent TCP links open. You should use lightweight socket servers and heartbeat timers to clear idle links.

External References

• MDN Web Docs: The WebSocket API
• RFC 6455: The WebSocket Protocol