From 5005bca353849e22386c1321dbd705d3cbec02f1 Mon Sep 17 00:00:00 2001 From: axelray-dev <110029405+axelray-dev@users.noreply.github.com> Date: Sun, 28 Jun 2026 07:48:13 +0800 Subject: [PATCH] fix(webui): clear stuck streaming after reconnect and improve stop reliability After a gateway restart or websocket reconnect, the UI stays stuck in processing state because reconnecting clients only replay running status when a turn is active, never push idle when no turn is running. Fix _hydrate_after_subscribe to always push goal_status (running with started_at when turn is active, idle when no turn is running) so the frontend can reset its processing indicator on reconnect. Also fix cmd_stop reporting 'No active task to stop' when a task is actually processing by draining the pending injection queue in addition to cancelling active tasks. This prevents mid-turn injection deadlocks and gives accurate task counts. --- nanobot/channels/websocket.py | 6 +- nanobot/command/builtin.py | 9 +++ .../channels/test_websocket_reconnect_idle.py | 61 ++++++++++++++ tests/command/test_stop_pending_queue.py | 81 +++++++++++++++++++ 4 files changed, 156 insertions(+), 1 deletion(-) create mode 100644 tests/channels/test_websocket_reconnect_idle.py create mode 100644 tests/command/test_stop_pending_queue.py diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index 1fcf7aa1..2d52d0e0 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -346,7 +346,11 @@ class WebSocketChannel(BaseChannel): async def _hydrate_after_subscribe(self, chat_id: str) -> None: """Replay goal/run strip state after subscribe (same-process refresh).""" await self._maybe_push_active_goal_state(chat_id) - await self._maybe_push_turn_run_wall_clock(chat_id) + t0 = websocket_turn_wall_started_at(chat_id) + if t0 is not None: + await self.send_goal_status(chat_id, "running", started_at=t0) + else: + await self.send_goal_status(chat_id, "idle") async def _send_event(self, connection: Any, event: str, **fields: Any) -> None: """Send a control event (attached, error, ...) to a single connection.""" diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index 4a34f894..bcfe16c8 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -130,6 +130,15 @@ async def cmd_stop(ctx: CommandContext) -> OutboundMessage: loop = ctx.loop msg = ctx.msg total = await loop._cancel_active_tasks(ctx.key) + # Also drain pending queue to prevent mid-turn injection deadlock + pending = loop._pending_queues.pop(ctx.key, None) + if pending is not None: + while not pending.empty(): + try: + pending.get_nowait() + total += 1 + except Exception: + break content = f"Stopped {total} task(s)." if total else "No active task to stop." return OutboundMessage( channel=msg.channel, chat_id=msg.chat_id, content=content, diff --git a/tests/channels/test_websocket_reconnect_idle.py b/tests/channels/test_websocket_reconnect_idle.py new file mode 100644 index 00000000..4cc1da8a --- /dev/null +++ b/tests/channels/test_websocket_reconnect_idle.py @@ -0,0 +1,61 @@ +"""Test websocket reconnect pushes idle status when no turn is active.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from nanobot.channels.websocket import WebSocketChannel + + +@pytest.mark.asyncio +async def test_hydrate_after_subscribe_pushes_idle_when_no_turn_active(): + """Reconnecting client should receive idle status when no turn is running.""" + channel = WebSocketChannel.__new__(WebSocketChannel) + channel.gateway = MagicMock() + channel.gateway.session_manager = MagicMock() + channel.gateway.session_manager.read_session_file = MagicMock(return_value={}) + + sent_events = [] + + async def mock_send_goal_state(chat_id, blob): + sent_events.append(("goal_state", chat_id, blob)) + + async def mock_send_goal_status(chat_id, status, **kwargs): + sent_events.append(("goal_status", chat_id, status, kwargs)) + + channel.send_goal_state = mock_send_goal_state + channel.send_goal_status = mock_send_goal_status + + with patch("nanobot.channels.websocket.websocket_turn_wall_started_at", return_value=None): + await channel._hydrate_after_subscribe("test-chat") + + # Should have pushed idle status + assert any(e[0] == "goal_status" and e[2] == "idle" for e in sent_events) + + +@pytest.mark.asyncio +async def test_hydrate_after_subscribe_pushes_running_when_turn_active(): + """Reconnecting client should receive running status when turn is active.""" + channel = WebSocketChannel.__new__(WebSocketChannel) + channel.gateway = MagicMock() + channel.gateway.session_manager = MagicMock() + channel.gateway.session_manager.read_session_file = MagicMock(return_value={}) + + sent_events = [] + + async def mock_send_goal_state(chat_id, blob): + sent_events.append(("goal_state", chat_id, blob)) + + async def mock_send_goal_status(chat_id, status, **kwargs): + sent_events.append(("goal_status", chat_id, status, kwargs)) + + channel.send_goal_state = mock_send_goal_state + channel.send_goal_status = mock_send_goal_status + + with patch("nanobot.channels.websocket.websocket_turn_wall_started_at", return_value=1234567890.0): + await channel._hydrate_after_subscribe("test-chat") + + # Should have pushed running status with started_at + running_events = [e for e in sent_events if e[0] == "goal_status" and e[2] == "running"] + assert len(running_events) == 1 + assert running_events[0][3]["started_at"] == 1234567890.0 diff --git a/tests/command/test_stop_pending_queue.py b/tests/command/test_stop_pending_queue.py new file mode 100644 index 00000000..c07dcedf --- /dev/null +++ b/tests/command/test_stop_pending_queue.py @@ -0,0 +1,81 @@ +"""Test cmd_stop drains pending queue to prevent mid-turn injection deadlock.""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from nanobot.bus.events import OutboundMessage +from nanobot.command.builtin import cmd_stop +from nanobot.command.router import CommandContext + + +@pytest.mark.asyncio +async def test_cmd_stop_drains_pending_queue(): + """cmd_stop should drain pending queue in addition to cancelling active tasks.""" + mock_loop = MagicMock() + mock_loop._cancel_active_tasks = AsyncMock(return_value=1) + mock_loop._pending_queues = {} + + pending = asyncio.Queue() + await pending.put("msg1") + await pending.put("msg2") + mock_loop._pending_queues["test-session"] = pending + + ctx = CommandContext( + msg=MagicMock(channel="websocket", chat_id="test-chat", metadata={}), + session=None, + key="test-session", + raw="/stop", + loop=mock_loop, + ) + + result = await cmd_stop(ctx) + + assert isinstance(result, OutboundMessage) + assert "Stopped 3 task(s)" in result.content # 1 cancelled + 2 drained + assert "test-session" not in mock_loop._pending_queues + + +@pytest.mark.asyncio +async def test_cmd_stop_with_empty_pending_queue(): + """cmd_stop should work correctly when pending queue is empty.""" + mock_loop = MagicMock() + mock_loop._cancel_active_tasks = AsyncMock(return_value=2) + mock_loop._pending_queues = {} + + pending = asyncio.Queue() + mock_loop._pending_queues["test-session"] = pending + + ctx = CommandContext( + msg=MagicMock(channel="websocket", chat_id="test-chat", metadata={}), + session=None, + key="test-session", + raw="/stop", + loop=mock_loop, + ) + + result = await cmd_stop(ctx) + + assert "Stopped 2 task(s)" in result.content + assert "test-session" not in mock_loop._pending_queues + + +@pytest.mark.asyncio +async def test_cmd_stop_no_pending_queue(): + """cmd_stop should work when no pending queue exists.""" + mock_loop = MagicMock() + mock_loop._cancel_active_tasks = AsyncMock(return_value=0) + mock_loop._pending_queues = {} + + ctx = CommandContext( + msg=MagicMock(channel="websocket", chat_id="test-chat", metadata={}), + session=None, + key="test-session", + raw="/stop", + loop=mock_loop, + ) + + result = await cmd_stop(ctx) + + assert "No active task to stop" in result.content