Building a WebSocket Chat Server with Kiro

Use Kiro to generate a FastAPI WebSocket server that streams Strands Agents replies to every connected user in real time.

Ribhav

Published Oct 20, 2026

Last modified Oct 20, 2026

A chat assistant feels responsive when its answer appears while the model is still writing it. Getting there takes more than one API call. You need a connection that stays open, a way to keep each user's conversation separate, and asynchronous code that sends text as it arrives without blocking everyone else.

In this post, we use Kiro to build that server in Python with FastAPI and WebSockets, then connect it to an AI agent built with the Strands Agents SDK and Amazon Bedrock. We walk through the code Kiro generated for the WebSocket endpoint, message handling and connection management, and explain the async patterns it used to stream responses.

The post is written for Python developers, but you don't need to know Python well to follow it. Where a Python feature matters, we explain it.

What we will build

1
2
3
4
5
flowchart LR
    CLI["Python CLI client"] -- "WebSocket JSON frames" --> API["FastAPI /ws/chat"]
    API --> CM["ConnectionManager"]
    CM --> AG["One Strands Agent per connection"]
    AG -- "stream_async" --> BR["Amazon Bedrock"]

Description: a command-line client opens a WebSocket to a FastAPI endpoint. The endpoint registers the socket with a connection manager, which pairs it with its own Strands Agent. The agent streams text from Amazon Bedrock, and the endpoint sends each piece back only to the socket that asked.

The client and server exchange small JSON messages, called frames. The client sends:

1
{"type": "message", "text": "Explain streaming briefly"}

The server answers with one or more token frames, then a done frame:

1
2
3
{"type": "token", "text": "Streaming "}
{"type": "token", "text": "delivers partial results early."}
{"type": "done"}

If something goes wrong, the server sends {"type": "error", "message": "..."} and keeps the connection open so the user can try again.

The scope is deliberately small. The server runs on one computer, keeps chat history only while a user is connected, and has no sign-in. That keeps the focus on the parts that are hard to get right: isolation between users, asynchronous Python and streaming.

How we worked with Kiro

We didn't ask Kiro for the whole server in one prompt. We used Kiro with the open-source AI-DLC workflow , which turns a request into a short sequence of steps and stops for approval at each one:

Clarify. Kiro asked who the post was for, what "done" meant, and what to leave out. We agreed on six pass/fail checks, including "the reply appears piece by piece" and "one person disconnecting doesn't break the others."
Plan. Kiro proposed a build order of thin layers: the JSON protocol first, then connections and agents, then the WebSocket endpoint, then the client. Each layer came with its own tests.
Generate and test. Kiro wrote each layer, then wrote and ran its tests before moving on.
Review. A separate review pass probed the running server with bad input and disconnects, and sent the code back with findings. We cover what it caught later in this post.

The "What we asked Kiro" boxes below are condensed from that plan, so you can reuse them as prompts in your own project. The code snippets are trimmed to show the shape of what Kiro generated. The full files are in the repository linked at the end.

Setup

You need:

Python 3.10 or later
AWS credentials available through the standard AWS credential chain (a profile, environment variables or IAM Identity Center)
Access to an Amazon Bedrock model supported by Strands

No credentials go in the code. Create a virtual environment and install the project with its test tools:

1
2
python3 -m venv .venv
.venv/bin/pip install -e ".[dev]"

The project depends on four packages: fastapi, uvicorn[standard] to serve it, websockets for the client, and strands-agents.

Step 1: Scaffold the WebSocket endpoint

FastAPI handles HTTP APIs and WebSockets. The decorator @app.websocket("/ws/chat") connects a Python function to incoming WebSocket connections on that path.

What we asked Kiro

Create a FastAPI WebSocket endpoint at /ws/chat. Accept each socket, register one isolated connection, validate every incoming text frame, and keep the receive loop open after validation errors. Bind the server to 127.0.0.1 by default.

Simplified from chat_server/app.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
@app.websocket("/ws/chat")
async def chat(websocket: WebSocket) -> None:
    await websocket.accept()
    try:
        agent = agent_builder(settings)              # one agent per connection
    except Exception as error:
        await websocket.send_text(error_frame(_readable_setup_error(error)))
        await websocket.close(code=1011)             # server error; never registered
        return

    connection = manager.register(websocket, agent)
    try:
        while True:
            event = await websocket.receive()        # text, binary or disconnect
            if event["type"] == "websocket.disconnect":
                break
            if event.get("bytes") is not None:
                await websocket.send_text(error_frame("Send JSON text."))
                continue
            try:
                message = parse_client_frame(event["text"], settings.max_message_chars)
            except ProtocolError as error:
                await websocket.send_text(error_frame(str(error)))
                continue
            if manager.is_busy(connection):
                await websocket.send_text(error_frame("A reply is already in progress."))
                continue
            manager.start_reply(connection, _stream_reply(connection, message.text))
    finally:
        await manager.disconnect(connection)         # cancels any reply in progress

Here is what the handler does, top to bottom:

It accepts the socket, then builds the agent. If agent setup fails, for example because there are no AWS credentials, the user gets a readable error frame and the socket closes with code 1011 (server error). The failed connection is never registered.
It reads raw events with websocket.receive(). That lets it tell text, binary and disconnect events apart. Binary frames get an error, and the loop continues.
It validates each message with parse_client_frame (Step 2) and turns any problem into an error frame.
It starts the reply in the background with manager.start_reply, so the loop can keep listening while the agent writes. If a reply is already running on this socket, it refuses the new message.
It always cleans up. The finally block runs on every exit path, including a disconnect, and cancels any reply still in progress.

The word async before def marks a function that can pause. Each await is a point where it pauses and lets the server work on other connections, which is how one Python process serves many users at once. We come back to this in "How Kiro handles async Python."

Step 2: Handle and validate messages

Kiro kept message validation out of the WebSocket code, in its own module. That makes it easy to test without a network.

What we asked Kiro

Add a protocol module for message, token, done and error JSON frames. Reject malformed JSON, non-object values, unknown types, missing or non-string text, blank text, and text over 4,000 characters with readable errors.

Simplified from chat_server/protocol.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def parse_client_frame(raw: str, max_chars: int) -> ClientMessage:
    try:
        frame = json.loads(raw)
    except json.JSONDecodeError as error:
        raise ProtocolError("Message must be valid JSON.") from error

    if not isinstance(frame, dict) or frame.get("type") != "message":
        raise ProtocolError("Message must be a JSON object of type 'message'.")
    text = frame.get("text")
    if not isinstance(text, str) or not text.strip():
        raise ProtocolError("Message text must be a non-blank string.")
    if len(text) > max_chars:
        raise ProtocolError(f"Message text must be at most {max_chars} characters.")
    return ClientMessage(text=text)

Every bad input raises one exception type, ProtocolError, with a message a person can read. The endpoint catches it, sends the message back as an error frame and waits for the next one. The 4,000-character limit is configurable with the CHAT_MAX_MESSAGE_CHARS environment variable, and the tests check both sides of the boundary: 4,000 characters pass and 4,001 fail.

Step 3: Manage connections

Each connection owns four things: an ID, its socket, its own agent and at most one running reply. Nothing is shared or broadcast, which is what keeps one user's answer out of another user's chat.

What we asked Kiro

Generate a connection manager that assigns UUIDs, stores one socket and one agent per connection, starts at most one reply task per connection, and cancels and awaits that task when the socket disconnects.

Simplified from chat_server/connection_manager.py

1
2
3
4
5
6
@dataclass(slots=True)
class Connection:
    websocket: WebSocket
    agent: Agent                                   # private conversation history
    id: str = field(default_factory=lambda: str(uuid4()))
    reply_task: asyncio.Task[None] | None = None   # at most one reply at a time

A dataclass generates the constructor and other boilerplate from the field list. The reply runs as an asyncio.Task, which lets a coroutine make progress in the background while the endpoint keeps receiving frames:

1
2
3
4
5
6
7
def start_reply(self, connection, reply) -> asyncio.Task[None]:
    if self.is_busy(connection):
        reply.close()                              # avoid a "never awaited" warning
        raise RuntimeError("A reply is already in progress.")
    task = asyncio.create_task(reply)              # runs in the background
    connection.reply_task = task                   # cleared again when it finishes
    return task

When the socket closes, disconnect stops that connection's reply:

1
2
3
4
5
6
7
async def disconnect(self, connection) -> None:
    task = connection.reply_task
    if task and not task.done():
        task.cancel()                              # ask the reply to stop...
        with suppress(asyncio.CancelledError):
            await task                             # ...and wait until it has
    self._connections.pop(connection.id, None)

Two details are worth noticing:

reply.close() on the refusal path. If the manager refuses a second reply, it closes the coroutine it was given. Otherwise Python would warn that a coroutine was never awaited.
Cancel, then await. task.cancel() only asks the task to stop. Awaiting it afterwards waits until the task has actually finished cleaning up, so the connection is removed only after its reply has stopped.

Step 4: Connect a Strands agent and stream its reply

Strands Agents is an open-source SDK for building AI agents. An Agent holds the conversation history and calls the model, and its stream_async method returns events as the model produces them. Here we use the Strands default model provider, Amazon Bedrock.

What we asked Kiro

Check the installed Strands Agents SDK source. Create a fresh Agent for every connection, use BedrockModel only when a model ID or Region override is set, turn off console printing, and stream event["data"] text to the owning socket.

Asking Kiro to read the installed SDK, rather than rely on what it remembered, mattered. It confirmed the exact names it then used: callback_handler=None turns off the SDK's default printing to the console, stream_async is an async iterator, and text events carry their text under data.

Simplified from chat_server/agent_factory.py

1
2
3
4
5
6
def build_agent(settings: Settings) -> Agent:
    # callback_handler=None stops Strands printing replies to the server console.
    if settings.bedrock_model_id or settings.aws_region:
        model = BedrockModel(**overrides)          # only the values that are set
        return Agent(model=model, callback_handler=None)
    return Agent(callback_handler=None)            # Strands default: Amazon Bedrock

A new agent per connection means each user has their own conversation history, which ends when they disconnect. The optional BEDROCK_MODEL_ID and AWS_REGION environment variables change the model or Region without touching the code.

Simplified from chat_server/app.py

This is the function that does the streaming:

1
2
3
4
5
6
7
8
9
10
11
async def _stream_reply(connection: Connection, prompt: str) -> None:
    try:
        async for event in connection.agent.stream_async(prompt):
            if text := event.get("data"):          # text pieces only
                await connection.websocket.send_text(token_frame(text))
        await connection.websocket.send_text(done_frame())
    except asyncio.CancelledError:
        raise                                      # let a disconnect stop the task
    except Exception as error:
        logger.exception("Agent reply failed")
        await connection.websocket.send_text(error_frame(_readable_model_error(error)))

async for waits for the next event from the agent and handles it as soon as it arrives. Each piece of text becomes a token frame for this connection's socket, and a done frame follows the last one. Nothing waits for the full answer.

The error handling is ordered on purpose. CancelledError is caught first and re-raised, so a disconnect still stops the task. Any other failure, such as Bedrock throttling or a missing model permission, is logged by type and sent to the user as a readable error frame. The connection stays open for the next message.

Step 5: A small command-line client

To try the server, Kiro also generated a terminal client in client/chat_client.py. Two parts of it show the async ideas on the client side:

1
2
3
4
text = await asyncio.to_thread(input, "you> ")    # blocking input, off the event loop
...
if frame["type"] == "token":
    print(frame["text"], end="", flush=True)      # show each piece immediately

Python's built-in input() blocks while it waits for the user to type. Running it with asyncio.to_thread moves that wait to a worker thread, so the event loop stays free to handle the socket. On the output side, each token frame is printed with no newline and flush=True, so text appears on screen the moment it arrives instead of sitting in a buffer until the line ends.

How Kiro handles async Python and streaming

Looking across the generated code, Kiro used a small set of async patterns consistently:

Pattern	Where	Why it matters
`async def` and `await`	Every handler and send	Each connection pauses at a well-defined point while waiting, so one process serves many users without a thread per socket.
`asyncio.create_task`	`start_reply`	The reply runs in the background, so the endpoint can still notice a disconnect or refuse a second message mid-reply.
`async for` over `stream_async`	`_stream_reply`	Each model event is forwarded as soon as it arrives. This is what makes the reply appear piece by piece.
Cancel, then await	`disconnect`	A disconnect stops only that user's reply, and cleanup finishes before the connection is forgotten.
Re-raise `CancelledError`	`_stream_reply`	A broad `except Exception` doesn't accidentally swallow cancellation.
`asyncio.to_thread`	Client input	A blocking call doesn't freeze the event loop.

The review step mattered as much as the generation step. Kiro's review pass ran probes against the live server, and its first verdict was "not ready." The findings it raised, and that Kiro then fixed, are the kind of bugs that slip through in async WebSocket code:

Agent setup failures broke the handshake. The first version built the agent before accepting the socket and outside any error handler. A missing credential produced an HTTP 500 and a stack trace instead of a readable message. The fix is the accept-then-build order you saw in Step 1.
Binary frames crashed the connection. The first version used websocket.receive_text(), which raised a KeyError when a client sent binary data. Switching to websocket.receive() and checking the event type fixed it, and a new test sends binary data and then completes a normal request on the same socket.
A test passed for the wrong reason. The first client test only checked the final printed text, so it would still pass if the client buffered the whole reply. Kiro replaced it with a test that confirms the first piece is on screen before done arrives.

Test streaming without calling the model

The tests run with no AWS credentials and no network. To do that, Kiro generated a stand-in model in tests/conftest.py that subclasses the Strands Model class and emits the same stream events Bedrock does, with an optional delay between pieces:

1
2
3
4
5
6
7
class ScriptedModel(Model):
    async def stream(self, messages, **kwargs):
        yield {"messageStart": {"role": "assistant"}}
        for piece in self.pieces:
            await asyncio.sleep(self.delay)        # makes timing observable
            yield {"contentBlockDelta": {"delta": {"text": piece}}}
        yield {"messageStop": {"stopReason": "end_turn"}}

Because the real Agent runs on top of this model, the tests exercise the actual Strands streaming path, not a mock of it. The delay makes timing observable, and a failure option (not shown) lets the tests check that a model error becomes an error frame while the connection survives. One test connects five clients at once, disconnects one mid-reply, and checks that the other four each receive only their own complete answer.

The suite has 26 tests. They passed on Python 3.15.0 and on Python 3.10.18 with every dependency at its minimum declared version.

Run it

Start the server:

1
.venv/bin/python -m chat_server

In a second terminal, start the client:

1
.venv/bin/python -m client.chat_client

Type a message and press Enter. Type /quit or press Ctrl+C to exit. The server listens on ws://127.0.0.1:8000/ws/chat by default.

To run the offline tests with your AWS credentials hidden:

1
env -u AWS_PROFILE -u AWS_ACCESS_KEY_ID -u AWS_SECRET_ACCESS_KEY -u AWS_SESSION_TOKEN .venv/bin/python -m pytest -v

Sample session

Measured numbers

scripts/benchmark.py --mode load connects 10 clients at once. Each sends 5 messages to the stand-in model, which returns 20 pieces 20 ms apart. On macOS (arm64) with Python 3.15.0, the 50 samples measured:

Measurement	Samples	Median	95th percentile
Time to first token	50	22.41 ms	25.69 ms
Time to `done`	50	439.12 ms	450.30 ms

These numbers cover the WebSocket, FastAPI, Strands and scheduling path on one machine. They don't include Bedrock latency. The 20 pieces take 400 ms by design, so the remaining time is the server's own overhead.

Known limitations

This is a local proof of concept. Before using it beyond your own machine:

There is no sign-in. Keep the default 127.0.0.1 bind until you add and test authentication on the WebSocket handshake.
The client prints model text as-is. A terminal can interpret control sequences hidden in text, so strip control characters from model output before printing it.
The client reads input before it reads the socket. If the server closes the connection right away, for example because the agent couldn't be created, the client only shows the error after the user types something.

Taking it further

Sign-in: authenticate the WebSocket handshake with your identity provider, and test that anonymous connections are refused before you change the bind address.
Saved history: give each conversation a durable ID and store messages in an encrypted database, loading only history the signed-in user may read.
Cloud deployment: run the app on managed compute that supports long-lived WebSockets, behind TLS, with health checks, logs and scaling limits.
A web client: replace the terminal with a browser client that renders tokens as they arrive. The JSON protocol stays the same.
Tools: Strands agents can call tools. Add one, and stream a status frame while it runs.

Conclusion

Kiro turned a real-time agent server into four layers we could test one at a time: protocol, connections, agent streaming and client. The generated code uses a handful of async patterns consistently, and Kiro's review pass caught the WebSocket edge cases that are easiest to miss. The result streams each user's reply as it is written, keeps conversations separate, and has a test suite that proves it without calling the model.

The opinions expressed in this post are those of the author and do not necessarily reflect the views of Amazon Web Services.

# kiro # strands-agents # amazon-bedrock # python # websockets

Any opinions in this article are those of the individual author and may not reflect the opinions of AWS.

Enjoyed reading this content? Let the author know!

Your likes, comments, shares, and saves help creators reach more builders.

Loading recommendations

Loading article

Builder Center

Building a WebSocket Chat Server with Kiro

Ribhav

What we will build

How we worked with Kiro

Setup

Step 1: Scaffold the WebSocket endpoint

Step 2: Handle and validate messages

Step 3: Manage connections

Step 4: Connect a Strands agent and stream its reply

Step 5: A small command-line client

How Kiro handles async Python and streaming

Test streaming without calling the model

Run it

Sample session

Measured numbers

Known limitations

Taking it further

Conclusion