diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index a84d41cb5..b27fd18ff 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -500,7 +500,7 @@ class WebSocketChannel(BaseChannel): await self._discard_connection_owned_chat(connection, cid) self._conn_default.pop(connection, None) self._webui_connections.discard(connection) - self._webui_request_locks.pop(connection, None) + self._discard_webui_request_lock_if_idle(connection) async def _maybe_push_active_goal_state(self, chat_id: str) -> None: """Replay an active sustained goal from session metadata after *chat_id* is subscribed. @@ -1212,6 +1212,7 @@ class WebSocketChannel(BaseChannel): ).digest() self._prune_webui_request_operations() operation = self._webui_request_operations.get(request_id) + is_replay = operation is not None if operation is not None and ( operation.action != action or operation.payload_digest != payload_digest ): @@ -1255,6 +1256,7 @@ class WebSocketChannel(BaseChannel): connection, request_id, operation.task, + sequence=is_replay, ) ) self._webui_request_tasks[key] = delivery_task @@ -1279,13 +1281,36 @@ class WebSocketChannel(BaseChannel): for _, request_id in completed[:-_WEBUI_REQUEST_CACHE_MAX]: self._webui_request_operations.pop(request_id, None) + def _discard_webui_request_lock_if_idle(self, connection: ServerConnection) -> None: + if connection in self._webui_connections: + return + if any(task_connection is connection for task_connection, _ in self._webui_request_tasks): + return + self._webui_request_locks.pop(connection, None) + async def _deliver_webui_request( self, connection: ServerConnection, request_id: str, operation_task: asyncio.Task[_WebUIRequestResult], + *, + sequence: bool = False, ) -> None: try: + if sequence: + # Make replayed work the predecessor for subsequent mutations on + # this connection without blocking its receive loop. + lock = self._webui_request_locks.setdefault(connection, asyncio.Lock()) + async with lock: + result = await asyncio.shield(operation_task) + await self._send_webui_response( + connection, + request_id, + result=result.result, + status=result.status, + message=result.message, + ) + return result = await asyncio.shield(operation_task) await self._send_webui_response( connection, @@ -1296,6 +1321,7 @@ class WebSocketChannel(BaseChannel): ) finally: self._webui_request_tasks.pop((connection, request_id), None) + self._discard_webui_request_lock_if_idle(connection) async def _execute_webui_request( self, diff --git a/nanobot/channels/websocket/tests/test_websocket_channel.py b/nanobot/channels/websocket/tests/test_websocket_channel.py index db8568eef..e89181116 100644 --- a/nanobot/channels/websocket/tests/test_websocket_channel.py +++ b/nanobot/channels/websocket/tests/test_websocket_channel.py @@ -1003,7 +1003,7 @@ async def test_webui_mutations_preserve_request_and_response_order(bus: MagicMoc @pytest.mark.asyncio -async def test_webui_request_survives_disconnect_and_reuses_inflight_operation( +async def test_webui_request_survives_disconnect_and_preserves_reconnect_order( bus: MagicMock, ) -> None: channel = _ch(bus) @@ -1013,11 +1013,14 @@ async def test_webui_request_survives_disconnect_and_reuses_inflight_operation( channel._webui_connections.update({first_conn, retry_conn}) started = asyncio.Event() release = asyncio.Event() + dispatch_order: list[str] = [] - async def mutate(*_args: Any) -> Any: - started.set() - await release.wait() - return _http_json_response({"ran": True}) + async def mutate(_connection: object, action: str, _payload: dict[str, Any]) -> Any: + dispatch_order.append(action) + if action == "automation.run": + started.set() + await release.wait() + return _http_json_response({"action": action}) channel.gateway.http.dispatch_webui_mutation = AsyncMock(side_effect=mutate) envelope = { @@ -1026,28 +1029,47 @@ async def test_webui_request_survives_disconnect_and_reuses_inflight_operation( "action": "automation.run", "payload": {"id": "daily-summary"}, } + queued_envelope = { + "type": "webui_request", + "request_id": "request-queued", + "action": "automation.update", + "payload": {"id": "daily-summary"}, + } await channel._dispatch_envelope(first_conn, "webui-client", envelope) await started.wait() + await channel._dispatch_envelope(first_conn, "webui-client", queued_envelope) + assert dispatch_order == ["automation.run"] + await channel._cleanup_connection(first_conn) assert first_conn not in channel._webui_connections await channel._dispatch_envelope(retry_conn, "webui-client", envelope) + await channel._dispatch_envelope(retry_conn, "webui-client", queued_envelope) + await channel._dispatch_envelope( + retry_conn, + "webui-client", + { + "type": "webui_request", + "request_id": "request-next", + "action": "automation.delete", + "payload": {"id": "daily-summary"}, + }, + ) + await asyncio.sleep(0) + assert dispatch_order == ["automation.run"] + pending = tuple(channel._webui_request_tasks.values()) release.set() await asyncio.gather(*pending) - channel.gateway.http.dispatch_webui_mutation.assert_awaited_once_with( - first_conn, - "automation.run", - {"id": "daily-summary"}, - ) - expected = { - "event": "webui_response", - "request_id": "request-retry", - "ok": True, - "result": {"ran": True}, - } - assert json.loads(retry_conn.send.await_args.args[0]) == expected + assert dispatch_order == ["automation.run", "automation.update", "automation.delete"] + responses = [json.loads(call.args[0]) for call in retry_conn.send.await_args_list] + assert [response["request_id"] for response in responses] == [ + "request-retry", + "request-queued", + "request-next", + ] + assert first_conn not in channel._webui_request_locks @pytest.mark.asyncio