fix(streaming): keep SSE stream alive when the agent is slow

[kiro-agent]

This pull request was created by @kiro-agent on behalf of @galuszkm 👻

Comment with /kiro fix to address specific feedback or /kiro all to address everything.
_{Learn about Kiro Web}

---

Description

When an agent is slow to produce the next token (e.g. a Bedrock model that freezes for ~30s mid-turn), the SSE stream to the chat UI is dropped. The last thing the client receives is a single : ping heartbeat, after which the connection ends and the response the model eventually produces is never delivered.

Root cause (chat backend, agents/streaming.py::stream_turn). The heartbeat was implemented as:

event = await asyncio.wait_for(aiter.__anext__(), timeout=_HEARTBEAT_INTERVAL)

On timeout, asyncio.wait_for cancels the in-flight pull. That CancelledError propagates into the upstream client generator (AgentCoreClient.invoke / AsyncLocalClient.invoke), runs its finally block — which closes the botocore StreamingBody / httpx stream — and terminates the generator. The next loop iteration calls aiter.__anext__() on the now-finished generator, gets StopAsyncIteration, and breaks. Net effect: exactly one : ping is emitted and then the whole stream closes cleanly, so any event produced once the model "unfreezes" is lost. The strands-compose-agentcore clients and ASGI app are not at fault — the bug is entirely in this repo's SSE relay.

Fix. Pull the next upstream event with a persistent task and wait on it with asyncio.wait, which returns on timeout without cancelling the task. Heartbeats are emitted while the same in-flight read stays alive, so a slow/frozen model no longer drops the stream. The pending task is cancelled in a finally before the generator is closed, otherwise a mid-stream client disconnect during a heartbeat would raise RuntimeError: async generator is already running.

Before: ['tokenA', ':ping', 'STOP'] (stream dropped after one ping).
After: ['tokenA', ':ping', ':ping', ':ping', 'tokenB', ':ping', 'SESSION_END'] (heartbeats sent, upstream survives, real event delivered).

Changes:

• Replace the wait_for heartbeat loop in stream_turn with a persistent pull task + asyncio.wait, and cancel the pending task before aclose().
• Add a delays option to the test FakeAgentClient so it can simulate a stalled model.
• Add a regression test that stalls the client past the heartbeat interval and asserts a heartbeat is sent, the post-stall SESSION_END is delivered, and the assistant turn is persisted.

Related Issues

Type of Change

• Bug fix

Testing

How have you tested the change?

• I ran lint + type check (uvx ruff format, uvx ruff check, uvx ty check) — just is not installed in this environment, so the underlying commands were run directly. All pass.
• I ran uv run pytest -m "not postgres" (the fast test tier) — 255 passed, 7 deselected.
• I added or updated tests that prove my fix is effective (test_streaming_slow_agent_sends_heartbeat_without_dropping_stream).
• I verified existing examples in examples/ still work (not exercised for this backend-only change).

Checklist

• I have read the CONTRIBUTING document
• I have updated the documentation accordingly (no docs change needed)
• My changes generate no new warnings
• Any dependent changes have been merged and published

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.