Repository navigation
fix: resume ChatUI streams after refresh - #9259
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the ability to resume active webchat runs as server-sent event (SSE) streams. It refactors the backend queue management and introduces a ChatRunState dataclass to decouple chat runs from their subscribers, allowing clients to reconnect and resume active streams. On the frontend, it adds support for restoring and resuming active runs. The review feedback suggests several robustness improvements, including defensive checks for streaming error payloads, stringifying message IDs before comparison to avoid type mismatches, and optimizing the retry loop in the frontend stream resume logic when an abort signal is triggered.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
dosubot
Bot
added
size:XXL
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- In ChatService._consume_chat_run, the pending_refs variable is never updated but is used in the should_save logic and as a fallback for extract_web_search_refs, which means refs are effectively ignored; consider wiring through any reference data from the stream and reusing the previous flush_pending_bot_message semantics to avoid losing ref information.
- The logic for constructing active run snapshots (get_active_chat_runs and the snapshot in _subscribe_chat_run) is duplicated and could drift over time; extracting a shared helper for building these snapshot payloads would reduce duplication and make evolution of the resume semantics safer.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In ChatService._consume_chat_run, the pending_refs variable is never updated but is used in the should_save logic and as a fallback for extract_web_search_refs, which means refs are effectively ignored; consider wiring through any reference data from the stream and reusing the previous flush_pending_bot_message semantics to avoid losing ref information.
- The logic for constructing active run snapshots (get_active_chat_runs and the snapshot in _subscribe_chat_run) is duplicated and could drift over time; extracting a shared helper for building these snapshot payloads would reduce duplication and make evolution of the resume semantics safer.
## Individual Comments
### Comment 1
<location path="astrbot/dashboard/services/chat_service.py" line_range="887-896" />
<code_context>
+ async def _consume_chat_run(self, run: ChatRunState) -> None:
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Pending refs are never updated from streamed data and only overwritten by extraction errors.
In `_consume_chat_run`, `pending_refs` is created and passed to `flush_pending_bot_message` but is only ever set when `extract_web_search_refs` raises, so it stays `{}` and the fallback has no effect. If you intended to preserve refs from other sources (e.g., tool outputs or intermediate messages), you should update `pending_refs` whenever new refs are found; otherwise, consider removing it to avoid misleading, unused state.
Suggested implementation:
```python
pending_accumulator = BotMessageAccumulator()
display_accumulator = BotMessageAccumulator()
pending_agent_stats = {}
```
To fully remove the misleading `pending_refs` state, you should also:
1. Search for all uses of `pending_refs` in `astrbot/dashboard/services/chat_service.py` (e.g., parameters to `flush_pending_bot_message` or tracking structures inside `_consume_chat_run`) and delete the variable and parameter.
2. Update the signature of `flush_pending_bot_message` (if it accepts `pending_refs`) to drop that parameter, and remove any logic that assumes `pending_refs` is passed in.
3. If there is error-handling code that sets `pending_refs` based on `extract_web_search_refs` failures, remove that logic or refactor it so that errors are handled without relying on `pending_refs`.
These cleanups will ensure there is no unused or misleading state related to pending refs and will keep the implementation consistent with the current behavior (no ref preservation across streamed messages).
</issue_to_address>
### Comment 2
<location path="astrbot/core/platform/sources/webchat/webchat_queue_mgr.py" line_range="48" />
<code_context>
self._conversation_back_requests[conversation_id].add(request_id)
return self.back_queues[request_id]
+ async def put_back_queue(self, request_id: str, data: object) -> bool:
+ """Write a response while the request queue remains active.
+
</code_context>
<issue_to_address>
**issue (complexity):** Consider extracting the close-aware put logic into a helper or small abstraction so `put_back_queue` expresses only the high-level intent instead of detailed task coordination.
The `put_back_queue` coordination logic is doing the right thing semantically, but the per‑call two‑task race and manual cancellation makes the queue manager harder to read and maintain. You can keep the current behavior while encapsulating the complexity in a small helper/wrapper, so the queue manager only expresses the high‑level intent.
### 1. Extract the close‑aware put behavior into a helper
Move the low‑level concurrency into a private method. This keeps `put_back_queue` simple and avoids repeating the task/cancel pattern if you ever need it elsewhere:
```python
class WebChatQueueMgr:
...
async def _put_while_open(
self,
queue: asyncio.Queue,
close_event: asyncio.Event,
data: object,
) -> bool:
try:
queue.put_nowait(data)
return True
except asyncio.QueueFull:
pass
put_task = asyncio.create_task(queue.put(data))
close_task = asyncio.create_task(close_event.wait())
try:
done, _ = await asyncio.wait(
{put_task, close_task},
return_when=asyncio.FIRST_COMPLETED,
)
if close_task in done:
return False
await put_task
return True
finally:
for task in (put_task, close_task):
if not task.done():
task.cancel()
await asyncio.gather(put_task, close_task, return_exceptions=True)
async def put_back_queue(self, request_id: str, data: object) -> bool:
queue = self.back_queues.get(request_id)
close_event = self._back_queue_close_events.get(request_id)
if queue is None or close_event is None or close_event.is_set():
return False
return await self._put_while_open(queue, close_event, data)
```
This preserves all current semantics (including unblocking immediately on close) but makes `put_back_queue` much easier to understand: it now reads as “put while open” instead of exposing task management details.
### 2. Optionally encapsulate queues + close events into a small abstraction
If you find yourself needing similar behavior for front queues later, consider a focused wrapper. The manager then only stores and uses that abstraction, reducing asymmetry between front/back queues:
```python
class ClosableQueue:
def __init__(self, maxsize: int):
self.queue = asyncio.Queue(maxsize=maxsize)
self.close_event = asyncio.Event()
async def put(self, item: object) -> bool:
if self.close_event.is_set():
return False
# reuse _put_while_open logic here or inline it
...
def close(self) -> None:
self.close_event.set()
```
Then the manager would hold `dict[str, ClosableQueue]` and `put_back_queue` becomes a direct call to `ClosableQueue.put`, avoiding per‑call manual coordination code in the manager itself.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| async def _consume_chat_run(self, run: ChatRunState) -> None: | ||
| """Drain runner output, persist it, and fan it out to subscribers. | ||
|
|
||
| Args: | ||
| run: Chat run owning the producer queue and durable state. | ||
| """ | ||
| pending_accumulator = BotMessageAccumulator() | ||
| display_accumulator = BotMessageAccumulator() | ||
| pending_agent_stats = {} | ||
| pending_refs = {} |
There was a problem hiding this comment.
suggestion (bug_risk): Pending refs are never updated from streamed data and only overwritten by extraction errors.
In _consume_chat_run, pending_refs is created and passed to flush_pending_bot_message but is only ever set when extract_web_search_refs raises, so it stays {} and the fallback has no effect. If you intended to preserve refs from other sources (e.g., tool outputs or intermediate messages), you should update pending_refs whenever new refs are found; otherwise, consider removing it to avoid misleading, unused state.
Suggested implementation:
pending_accumulator = BotMessageAccumulator()
display_accumulator = BotMessageAccumulator()
pending_agent_stats = {}To fully remove the misleading pending_refs state, you should also:
- Search for all uses of
pending_refsinastrbot/dashboard/services/chat_service.py(e.g., parameters toflush_pending_bot_messageor tracking structures inside_consume_chat_run) and delete the variable and parameter. - Update the signature of
flush_pending_bot_message(if it acceptspending_refs) to drop that parameter, and remove any logic that assumespending_refsis passed in. - If there is error-handling code that sets
pending_refsbased onextract_web_search_refsfailures, remove that logic or refactor it so that errors are handled without relying onpending_refs.
These cleanups will ensure there is no unused or misleading state related to pending refs and will keep the implementation consistent with the current behavior (no ref preservation across streamed messages).
| self._conversation_back_requests[conversation_id].add(request_id) | ||
| return self.back_queues[request_id] | ||
|
|
||
| async def put_back_queue(self, request_id: str, data: object) -> bool: |
There was a problem hiding this comment.
issue (complexity): Consider extracting the close-aware put logic into a helper or small abstraction so put_back_queue expresses only the high-level intent instead of detailed task coordination.
The put_back_queue coordination logic is doing the right thing semantically, but the per‑call two‑task race and manual cancellation makes the queue manager harder to read and maintain. You can keep the current behavior while encapsulating the complexity in a small helper/wrapper, so the queue manager only expresses the high‑level intent.
1. Extract the close‑aware put behavior into a helper
Move the low‑level concurrency into a private method. This keeps put_back_queue simple and avoids repeating the task/cancel pattern if you ever need it elsewhere:
class WebChatQueueMgr:
...
async def _put_while_open(
self,
queue: asyncio.Queue,
close_event: asyncio.Event,
data: object,
) -> bool:
try:
queue.put_nowait(data)
return True
except asyncio.QueueFull:
pass
put_task = asyncio.create_task(queue.put(data))
close_task = asyncio.create_task(close_event.wait())
try:
done, _ = await asyncio.wait(
{put_task, close_task},
return_when=asyncio.FIRST_COMPLETED,
)
if close_task in done:
return False
await put_task
return True
finally:
for task in (put_task, close_task):
if not task.done():
task.cancel()
await asyncio.gather(put_task, close_task, return_exceptions=True)
async def put_back_queue(self, request_id: str, data: object) -> bool:
queue = self.back_queues.get(request_id)
close_event = self._back_queue_close_events.get(request_id)
if queue is None or close_event is None or close_event.is_set():
return False
return await self._put_while_open(queue, close_event, data)This preserves all current semantics (including unblocking immediately on close) but makes put_back_queue much easier to understand: it now reads as “put while open” instead of exposing task management details.
2. Optionally encapsulate queues + close events into a small abstraction
If you find yourself needing similar behavior for front queues later, consider a focused wrapper. The manager then only stores and uses that abstraction, reducing asymmetry between front/back queues:
class ClosableQueue:
def __init__(self, maxsize: int):
self.queue = asyncio.Queue(maxsize=maxsize)
self.close_event = asyncio.Event()
async def put(self, item: object) -> bool:
if self.close_event.is_set():
return False
# reuse _put_while_open logic here or inline it
...
def close(self) -> None:
self.close_event.set()Then the manager would hold dict[str, ClosableQueue] and put_back_queue becomes a direct call to ClosableQueue.put, avoiding per‑call manual coordination code in the manager itself.
* fix: resume ChatUI streams after refresh * fix: hide chat run service errors * fix: harden resumed chat streams
Summary
Root cause
The Chat service consumed runner output inside the request-owned SSE generator. Refreshing ChatUI cancelled that generator, removed the response queue, and left later runner output without a consumer or persisted bot response. The frontend also treated any resumed-stream EOF as completion and selected the newest active run, which could be an empty queued follow-up rather than the run producing tokens.
Compatibility
This is not a breaking change. The existing
POST /chat/sendrequest fields and legacy SSE event shapes remain unchanged. Resume metadata andGET /chat/runs/{run_id}/streamare additive.Validation
uv run ruff format --check .uv run ruff check .uv run pytest -q tests/test_chat_route.py tests/test_webchat_queue_lifecycle.py tests/test_astr_agent_run_util.py tests/test_platform_image_format_preservation.py tests/unit/test_webchat_message_parts.py(16 passed)cd dashboard && pnpm typecheckcd dashboard && pnpm exec prettier --check src/composables/useMessages.tsFixes #9249.
This takes a different approach from #9256: disconnecting the browser no longer stops the active agent run; ChatUI can reattach and continue receiving its output.
Summary by Sourcery
Decouple WebChat chat run streaming from the original SSE response so runs continue after browser disconnects and can be resumed via a dedicated API, while keeping legacy stream behavior intact.
New Features:
Bug Fixes:
Enhancements:
Documentation:
Tests: