fix(webui): make mutations reconnect-safe

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
+        # Delivery tasks are connection-bound, while operations are keyed only
-        # finish after a client-side deadline so an already-started mutation
+        # by request_id so reconnect retries join or replay the original work.
-        # isn't ambiguously cancelled halfway through.
         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)
+        payload_digest = hashlib.sha256(
-        if key in self._webui_request_tasks:
+            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(
+        if operation is None:
-            self._complete_webui_request(
+            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,
+                operation.task,
-                cast(dict[str, Any], payload),
             )
         )
-        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(
+                        return _WebUIRequestResult(
-                            connection,
-                            request_id,
                             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(
+                    return _WebUIRequestResult(result=result)
-                        connection,
+                return _WebUIRequestResult(
-                        request_id,
-                        result=result,
-                    )
-                    return
-                await self._send_webui_response(
-                    connection,
-                    request_id,
                     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(
+            return _WebUIRequestResult(
-                connection,
-                request_id,
                 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())
+        delivery_tasks = tuple(self._webui_request_tasks.values())
-        for task in mutation_tasks:
+        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:
+        if delivery_tasks:
-            await asyncio.gather(*mutation_tasks, return_exceptions=True)
+            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()

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)

webui/src/lib/nanobot-client.ts

@@ -109,6 +109,12 @@ interface PendingRequest<T> {
   timer: ReturnType<typeof setTimeout>;
 }
 
+type WebUIRequestFrame = Extract<Outbound, { type: "webui_request" }>;
+
+interface PendingWebUIRequest extends PendingRequest<unknown> {
+  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<string, PendingRequest<string>>();
   private pendingSystemCommands = new Map<string, PendingRequest<void>>();
-  private pendingWebUIRequests = new Map<string, PendingRequest<unknown>>();
+  private pendingWebUIRequests = new Map<string, PendingWebUIRequest>();
   // 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.
+   * Send one WebUI mutation over the authenticated socket. Pending requests are
-   * A client-side timeout only abandons the reply; the server may finish work
+   * replayed with the same request_id after reconnect so the gateway can join or
-   * that already started, so timed-out requests are never retried automatically.
+   * replay the original operation. A client-side timeout still ends all retries.
    */
   requestMutation<T>(
     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()) {
+    if (!willReconnect) {
-      clearTimeout(pending.timer);
+      for (const pending of this.pendingWebUIRequests.values()) {
-      pending.reject(
+        clearTimeout(pending.timer);
-        new WebUIMutationError(503, "Socket closed before WebUI response"),
+        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;
     }

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",