Streaming an AI Agent Into a Chat UI: The Parts That Are Not the Model
Most write-ups about putting an AI agent in a chat product stop at the prompt. In practice, the model is the part you change least. The parts that break are message identity, tool-call rendering, thread context, and what the client does when a stream dies halfway through a sentence.
I work on messaging infrastructure at CometChat, where AI assistants and bots sit in the same conversation as human participants. That constraint is the interesting one: an agent reply is not a separate UI mode, it is a message in a thread that also has read receipts, reactions, threads, moderation, and history sync. Anything you invent for the agent has to survive contact with all of that.
Here is what I would insist on before wiring an agent into a chat surface.
1. A streaming reply is one message with many revisions
The naive implementation appends a message per chunk. It looks fine in a demo and falls apart the moment a second device is open, because now that device receives forty messages instead of one.
Model the reply as a single message with a stable identifier created at the moment the request is accepted, and a mutable body plus a state field:
-
pending- accepted, nothing generated yet -
streaming- partial content, safe to render -
complete- final content, safe to cache and index -
failed- render the partial content plus an explicit error affordance
The client that started the request updates optimistically from the token stream. Every other client learns about the message through the same real-time channel it uses for human messages, and reconciles by identifier. In CometChat terms this is the difference between sending many messages and updating one message repeatedly, and it decides whether your history stays readable a week later.
Practical detail: throttle the update you broadcast. Rendering every token locally is cheap; fanning out every token to all participants and persisting each revision is not. Local stream at token granularity, remote update on an interval or on sentence boundaries.
2. Tool calls are message parts, not log lines
An agent that searches a knowledge base, files a ticket, or looks up an order produces intermediate steps. If you flatten those into prose, users cannot tell what actually happened, and you cannot audit it afterwards.
Give the message structured parts:
{
""id"": ""msg_01H..."",
""state"": ""streaming"",
""parts"": [
{ ""type"": ""text"", ""text"": ""Checking your last invoice."" },
{ ""type"": ""tool_use"", ""name"": ""lookup_invoice"", ""args"": { ""id"": ""..."" }, ""status"": ""running"" },
{ ""type"": ""text"", ""text"": """" }
]
}
Two things fall out of this for free. The UI can collapse tool activity behind a disclosure instead of interleaving it with the answer. And moderation, analytics, and export can treat tool_use differently from text, which matters when the tool arguments contain customer data you do not want in a search index.
The metadata field on a chat message is usually the right place to carry this. Keep the human-readable summary in the message body so that clients you have not updated yet still show something sensible.
3. Thread context is a retrieval problem, not a window problem
""Send the last N messages"" is the default and it is wrong in both directions. In a support conversation, N=50 can be six months of small talk; in a fast group chat, N=50 is four minutes.
What has worked better as a starting policy:
- always include the last few turns verbatim, because coreference lives there (""no, the other one"")
- include a rolling summary of the conversation, regenerated on a cadence rather than per message
- retrieve by relevance from the rest of the thread and from your own docs, and label retrieved content as reference material, not as conversation
- carry participant and channel facts explicitly (locale, plan, entitlement) instead of hoping they appear in the transcript
And treat everything from the thread as untrusted input. A message in a group chat can contain instructions aimed at your agent. If your agent has tools that write - refunds, escalations, invites - then a user-authored message is an attacker-authored message. Authorise tool calls against the requesting user's own permissions, not the agent's.
4. Failure states need a design, not a spinner
Streams break: the socket drops, the model provider rate-limits you, generation is cut mid-token. Decide in advance what each one looks like.
| Failure | Behaviour |
|---|---|
| Socket drops mid-stream | Keep the partial message, mark it interrupted, resume or offer retry by message id |
| Provider error before any token | Move message to failed, show retry, do not leave an empty bubble |
| Timeout with no response | Read back the message state before retrying; a missing response does not prove nothing was generated |
| Retry after partial output | Replace the body of the same message id, never append a second reply |
That last row is the one teams get wrong. Retries must be idempotent at the message level, keyed on the request identifier the client generated. Otherwise every flaky connection produces duplicate assistant replies, and the transcript stops being a record of what happened.
5. The agent is a participant, so give it an identity
If the agent posts as a bot user with its own identifier, everything downstream works: mentions, per-conversation muting, permissions, rate limits, moderation rules, and analytics that can separate human from generated volume. If it posts as a magic side channel, you will re-implement each of those badly.
This is also how you keep the door open for more than one agent. A triage bot and a summariser bot in the same channel are just two participants with different tool scopes.
Why this matters more than prompt tuning
You can swap models in an afternoon. You cannot easily change the shape of a million stored messages. The decisions that age well are the boring ones: stable message identity, structured parts, explicit states, idempotent retries, and an agent that lives inside your existing permission model rather than beside it.
If you are adding an assistant to a product that already has real conversations in it, design the message first and the prompt second.
Top comments (0)