mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-31 08:13:11 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix(webui): make mutations reconnect-safe

Browse Source
This commit is contained in:
chengyongru
2026-08-16 21:40:14 +08:00
committed by chengyongru
parent dec89a49a3
commit e51ffc8978
4 changed files with 325 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/channels/websocket/runtime.py
+123 -36
View File
@@ -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()
nanobot/channels/websocket/tests/test_websocket_channel.py
+120
View File
@@ -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
+24 -11
View File
@@ -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.
* 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<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()) {
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;
}
webui/src/tests/nanobot-client.test.ts
+58
View File
@@ -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",