From e51ffc8978cbb391cf35027109613430912d1bbe Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Sun, 16 Aug 2026 02:49:00 +0800 Subject: [PATCH] fix(webui): make mutations reconnect-safe --- nanobot/channels/websocket/runtime.py | 159 ++++++++++++++---- .../websocket/tests/test_websocket_channel.py | 120 +++++++++++++ webui/src/lib/nanobot-client.ts | 35 ++-- webui/src/tests/nanobot-client.test.ts | 58 +++++++ 4 files changed, 325 insertions(+), 47 deletions(-) diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index 6534799cc..a84d41cb5 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -3,14 +3,17 @@ from __future__ import annotations import asyncio +import hashlib import hmac import ipaddress import json import re import ssl +import time import uuid from collections.abc import Callable from contextlib import suppress +from dataclasses import dataclass from pathlib import Path from typing import Any, Self, TypeGuard, cast from urllib.parse import urlsplit, urlunsplit @@ -92,6 +95,8 @@ from nanobot.webui.websocket_logging import websockets_server_logger # Plain HTTP WebUI routes also run through websockets.process_request. _WEBUI_HTTP_OPEN_TIMEOUT_S = 360.0 +_WEBUI_REQUEST_CACHE_TTL_S = 5 * 60.0 +_WEBUI_REQUEST_CACHE_MAX = 256 _ROUTING_ASSERTION_HEADERS = frozenset( @@ -348,6 +353,21 @@ def _is_websocket_upgrade(request: WsRequest) -> bool: return True +@dataclass(frozen=True) +class _WebUIRequestResult: + result: Any = None + status: int | None = None + message: str | None = None + + +@dataclass +class _WebUIRequestOperation: + action: str + payload_digest: bytes + task: asyncio.Task[_WebUIRequestResult] + completed_at: float | None = None + + class WebSocketChannel(BaseChannel): """Run a local WebSocket server; forward text/JSON messages to the message bus.""" @@ -373,14 +393,14 @@ class WebSocketChannel(BaseChannel): self._conn_default: dict[ServerConnection, str] = {} # Connections authenticated with a one-time token from /webui/bootstrap. self._webui_connections: set[ServerConnection] = set() - # Request/reply mutations aren't replayed across reconnects. Tasks may - # finish after a client-side deadline so an already-started mutation - # isn't ambiguously cancelled halfway through. + # Delivery tasks are connection-bound, while operations are keyed only + # by request_id so reconnect retries join or replay the original work. self._webui_request_tasks: dict[ tuple[ServerConnection, str], asyncio.Task[None], ] = {} - # Preserve request/response order for non-replayable mutations from one + self._webui_request_operations: dict[str, _WebUIRequestOperation] = {} + # Preserve request/response order for mutations from one # UI. Without this, an earlier slow settings response can overwrite a # newer settings snapshot in the client. self._webui_request_locks: dict[ServerConnection, asyncio.Lock] = {} @@ -1182,32 +1202,107 @@ class WebSocketChannel(BaseChannel): ) return - key = (connection, request_id) - if key in self._webui_request_tasks: + payload_digest = hashlib.sha256( + json.dumps( + payload, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).digest() + self._prune_webui_request_operations() + operation = self._webui_request_operations.get(request_id) + if operation is not None and ( + operation.action != action or operation.payload_digest != payload_digest + ): await self._send_webui_response( connection, request_id, status=409, - message="duplicate WebUI request_id", + message="request_id was already used for a different WebUI mutation", ) return - task = asyncio.create_task( - self._complete_webui_request( + if operation is None: + operation_task = asyncio.create_task( + self._execute_webui_request( + connection, + action, + cast(dict[str, Any], payload), + ) + ) + new_operation = _WebUIRequestOperation( + action=action, + payload_digest=payload_digest, + task=operation_task, + ) + operation = new_operation + self._webui_request_operations[request_id] = new_operation + + def mark_complete(_task: asyncio.Task[_WebUIRequestResult]) -> None: + current = self._webui_request_operations.get(request_id) + if current is not new_operation: + return + new_operation.completed_at = time.monotonic() + self._prune_webui_request_operations() + + operation_task.add_done_callback(mark_complete) + + key = (connection, request_id) + if key in self._webui_request_tasks: + return + delivery_task = asyncio.create_task( + self._deliver_webui_request( connection, request_id, - action, - cast(dict[str, Any], payload), + operation.task, ) ) - self._webui_request_tasks[key] = task + self._webui_request_tasks[key] = delivery_task - async def _complete_webui_request( + def _prune_webui_request_operations(self) -> None: + now = time.monotonic() + for request_id, operation in tuple(self._webui_request_operations.items()): + if ( + operation.completed_at is not None + and now - operation.completed_at >= _WEBUI_REQUEST_CACHE_TTL_S + ): + self._webui_request_operations.pop(request_id, None) + + completed = sorted( + ( + (operation.completed_at, request_id) + for request_id, operation in self._webui_request_operations.items() + if operation.completed_at is not None + ), + key=lambda item: item[0], + ) + for _, request_id in completed[:-_WEBUI_REQUEST_CACHE_MAX]: + self._webui_request_operations.pop(request_id, None) + + async def _deliver_webui_request( self, connection: ServerConnection, request_id: str, + operation_task: asyncio.Task[_WebUIRequestResult], + ) -> None: + try: + result = await asyncio.shield(operation_task) + await self._send_webui_response( + connection, + request_id, + result=result.result, + status=result.status, + message=result.message, + ) + finally: + self._webui_request_tasks.pop((connection, request_id), None) + + async def _execute_webui_request( + self, + connection: ServerConnection, action: str, payload: dict[str, Any], - ) -> None: + ) -> _WebUIRequestResult: try: lock = self._webui_request_locks.setdefault(connection, asyncio.Lock()) async with lock: @@ -1222,27 +1317,17 @@ class WebSocketChannel(BaseChannel): try: result = json.loads(body) except json.JSONDecodeError: - await self._send_webui_response( - connection, - request_id, + return _WebUIRequestResult( status=502, message="WebUI mutation returned an invalid response", ) - return if action == "sidebar.update" and isinstance(result, dict): await self._broadcast_webui_event( "sidebar_state_updated", state=result, ) - await self._send_webui_response( - connection, - request_id, - result=result, - ) - return - await self._send_webui_response( - connection, - request_id, + return _WebUIRequestResult(result=result) + return _WebUIRequestResult( status=status, message=body or response.reason_phrase, ) @@ -1250,14 +1335,10 @@ class WebSocketChannel(BaseChannel): raise except Exception: self.logger.exception("WebUI mutation '{}' failed", action) - await self._send_webui_response( - connection, - request_id, + return _WebUIRequestResult( status=500, message="WebUI mutation failed", ) - finally: - self._webui_request_tasks.pop((connection, request_id), None) async def _send_webui_response( self, @@ -1328,13 +1409,19 @@ class WebSocketChannel(BaseChannel): except Exception as e: self.logger.warning("server task error during shutdown: {}", e) self._server_task = None - mutation_tasks = tuple(self._webui_request_tasks.values()) - for task in mutation_tasks: + delivery_tasks = tuple(self._webui_request_tasks.values()) + operation_tasks = tuple( + operation.task for operation in self._webui_request_operations.values() + ) + for task in (*delivery_tasks, *operation_tasks): task.cancel() - if mutation_tasks: - await asyncio.gather(*mutation_tasks, return_exceptions=True) + if delivery_tasks: + await asyncio.gather(*delivery_tasks, return_exceptions=True) + if operation_tasks: + await asyncio.gather(*operation_tasks, return_exceptions=True) self._webui_request_tasks.clear() self._webui_request_locks.clear() + self._webui_request_operations.clear() self._subs.clear() self._conn_chats.clear() self._conn_default.clear() diff --git a/nanobot/channels/websocket/tests/test_websocket_channel.py b/nanobot/channels/websocket/tests/test_websocket_channel.py index c4bf7efb3..b737e5ff3 100644 --- a/nanobot/channels/websocket/tests/test_websocket_channel.py +++ b/nanobot/channels/websocket/tests/test_websocket_channel.py @@ -1002,6 +1002,126 @@ async def test_webui_mutations_preserve_request_and_response_order(bus: MagicMoc ] +@pytest.mark.asyncio +async def test_webui_request_reuses_inflight_operation_after_reconnect(bus: MagicMock) -> None: + channel = _ch(bus) + first_conn = AsyncMock() + retry_conn = AsyncMock() + channel._webui_connections.update({first_conn, retry_conn}) + started = asyncio.Event() + release = asyncio.Event() + + async def mutate(*_args: Any) -> Any: + started.set() + await release.wait() + return _http_json_response({"ran": True}) + + channel.gateway.http.dispatch_webui_mutation = AsyncMock(side_effect=mutate) + envelope = { + "type": "webui_request", + "request_id": "request-retry", + "action": "automation.run", + "payload": {"id": "daily-summary"}, + } + + await channel._dispatch_envelope(first_conn, "webui-client", envelope) + await started.wait() + await channel._dispatch_envelope(retry_conn, "webui-client", envelope) + 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(first_conn.send.await_args.args[0]) == expected + assert json.loads(retry_conn.send.await_args.args[0]) == expected + + +@pytest.mark.asyncio +async def test_webui_request_replays_completed_result_after_reconnect(bus: MagicMock) -> None: + channel = _ch(bus) + first_conn = AsyncMock() + retry_conn = AsyncMock() + channel._webui_connections.update({first_conn, retry_conn}) + channel.gateway.http.dispatch_webui_mutation = AsyncMock( + return_value=_http_json_response({"installed": True}) + ) + envelope = { + "type": "webui_request", + "request_id": "request-completed", + "action": "skill.install", + "payload": {"skill": "demo"}, + } + + await channel._dispatch_envelope(first_conn, "webui-client", envelope) + await asyncio.gather(*tuple(channel._webui_request_tasks.values())) + await channel._dispatch_envelope(retry_conn, "webui-client", envelope) + await asyncio.gather(*tuple(channel._webui_request_tasks.values())) + + channel.gateway.http.dispatch_webui_mutation.assert_awaited_once() + assert json.loads(retry_conn.send.await_args.args[0]) == { + "event": "webui_response", + "request_id": "request-completed", + "ok": True, + "result": {"installed": True}, + } + + +@pytest.mark.asyncio +async def test_webui_request_rejects_request_id_reuse_with_different_payload( + bus: MagicMock, +) -> None: + channel = _ch(bus) + first_conn = AsyncMock() + retry_conn = AsyncMock() + channel._webui_connections.update({first_conn, retry_conn}) + channel.gateway.http.dispatch_webui_mutation = AsyncMock( + return_value=_http_json_response({"ran": True}) + ) + + await channel._dispatch_envelope( + first_conn, + "webui-client", + { + "type": "webui_request", + "request_id": "request-conflict", + "action": "automation.run", + "payload": {"id": "job-a"}, + }, + ) + await asyncio.gather(*tuple(channel._webui_request_tasks.values())) + await channel._dispatch_envelope( + retry_conn, + "webui-client", + { + "type": "webui_request", + "request_id": "request-conflict", + "action": "automation.run", + "payload": {"id": "job-b"}, + }, + ) + + channel.gateway.http.dispatch_webui_mutation.assert_awaited_once() + assert json.loads(retry_conn.send.await_args.args[0]) == { + "event": "webui_response", + "request_id": "request-conflict", + "ok": False, + "error": { + "status": 409, + "message": "request_id was already used for a different WebUI mutation", + }, + } + + @pytest.mark.asyncio async def test_webui_request_returns_correlated_route_error(bus: MagicMock) -> None: channel = _ch(bus) diff --git a/webui/src/lib/nanobot-client.ts b/webui/src/lib/nanobot-client.ts index fed2f4e5b..4b4d9d079 100644 --- a/webui/src/lib/nanobot-client.ts +++ b/webui/src/lib/nanobot-client.ts @@ -109,6 +109,12 @@ interface PendingRequest { timer: ReturnType; } +type WebUIRequestFrame = Extract; + +interface PendingWebUIRequest extends PendingRequest { + frame: WebUIRequestFrame; +} + export class WebUIMutationError extends Error { status: number; @@ -215,7 +221,7 @@ export class NanobotClient { private pendingNewChat: PendingChatRequest | null = null; private pendingTranscriptions = new Map>(); private pendingSystemCommands = new Map>(); - private pendingWebUIRequests = new Map>(); + private pendingWebUIRequests = new Map(); // Frames queued while the socket is not yet OPEN private sendQueue: Outbound[] = []; private reconnectAttempts = 0; @@ -833,9 +839,9 @@ export class NanobotClient { } /** - * Send one non-replayable WebUI mutation over the authenticated socket. - * A client-side timeout only abandons the reply; the server may finish work - * that already started, so timed-out requests are never retried automatically. + * Send one WebUI mutation over the authenticated socket. Pending requests are + * replayed with the same request_id after reconnect so the gateway can join or + * replay the original operation. A client-side timeout still ends all retries. */ requestMutation( action: string, @@ -875,6 +881,7 @@ export class NanobotClient { resolve: (value) => resolve(value as T), reject, timer, + frame, }); try { socket.send(JSON.stringify(frame)); @@ -1020,6 +1027,9 @@ export class NanobotClient { for (const chatId of this.knownChats) { this.rawSend({ type: "attach", chat_id: chatId }); } + for (const pending of this.pendingWebUIRequests.values()) { + this.rawSend(pending.frame); + } // Flush anything queued during reconnect. const queued = this.sendQueue.splice(0); for (const frame of queued) this.rawSend(frame); @@ -1252,19 +1262,22 @@ export class NanobotClient { private handleClose(event?: { code?: number }): void { this.socket = null; this.clearTemporaryChats(); + const willReconnect = !this.intentionallyClosed && this.shouldReconnect; if (this.pendingNewChat) { clearTimeout(this.pendingNewChat.timer); this.pendingNewChat.reject(new Error("socket closed")); this.pendingNewChat = null; } this.rejectAllTranscriptions("socket closed"); - for (const pending of this.pendingWebUIRequests.values()) { - clearTimeout(pending.timer); - pending.reject( - new WebUIMutationError(503, "Socket closed before WebUI response"), - ); + if (!willReconnect) { + for (const pending of this.pendingWebUIRequests.values()) { + clearTimeout(pending.timer); + pending.reject( + new WebUIMutationError(503, "Socket closed before WebUI response"), + ); + } + this.pendingWebUIRequests.clear(); } - this.pendingWebUIRequests.clear(); for (const pending of this.pendingSystemCommands.values()) { clearTimeout(pending.timer); pending.reject(new Error("socket closed")); @@ -1312,7 +1325,7 @@ export class NanobotClient { } this.socketPendingMessageSendKeys.clear(); this.lastSocketMessageSendKey = null; - if (this.intentionallyClosed || !this.shouldReconnect) { + if (!willReconnect) { this.setStatus("closed"); return; } diff --git a/webui/src/tests/nanobot-client.test.ts b/webui/src/tests/nanobot-client.test.ts index 7fd3de3e2..6256e904a 100644 --- a/webui/src/tests/nanobot-client.test.ts +++ b/webui/src/tests/nanobot-client.test.ts @@ -172,6 +172,64 @@ describe("NanobotClient", () => { }); }); + it("retries in-flight WebUI mutations with the same request id after reconnect", async () => { + const client = new NanobotClient({ + url: "ws://test", + reconnect: true, + maxBackoffMs: 10, + socketFactory: (url) => new FakeSocket(url) as unknown as WebSocket, + }); + client.connect(); + const firstSocket = lastSocket(); + firstSocket.fakeOpen(); + + const pending = client.requestMutation<{ ran: boolean }>( + "automation.run", + { id: "daily-summary" }, + ); + const frame = firstSocket.sent.at(-1) as string; + const requestId = JSON.parse(frame).request_id; + const settled = expect(pending).resolves.toEqual({ ran: true }); + firstSocket.fakeCloseWithCode(1006); + + await vi.advanceTimersByTimeAsync(20); + const retrySocket = lastSocket(); + retrySocket.fakeOpen(); + expect(retrySocket.sent).toEqual([frame]); + retrySocket.fakeMessage({ + event: "webui_response", + request_id: requestId, + ok: true, + result: { ran: true }, + }); + + await settled; + }); + + it("does not retry a WebUI mutation after its timeout expires", async () => { + const client = new NanobotClient({ + url: "ws://test", + reconnect: true, + maxBackoffMs: 1_000, + socketFactory: (url) => new FakeSocket(url) as unknown as WebSocket, + }); + client.connect(); + const firstSocket = lastSocket(); + firstSocket.fakeOpen(); + + const pending = expect( + client.requestMutation("skill.install", { skill: "docs" }, 25), + ).rejects.toMatchObject({ status: 504 }); + firstSocket.fakeCloseWithCode(1006); + await vi.advanceTimersByTimeAsync(25); + await pending; + + await vi.advanceTimersByTimeAsync(1_000); + const retrySocket = lastSocket(); + retrySocket.fakeOpen(); + expect(retrySocket.sent).toEqual([]); + }); + it("does not queue WebUI mutations before the authenticated socket opens", async () => { const client = new NanobotClient({ url: "ws://test",