From baa02333779efe29a9386fcd97ad96f0e3195efb Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:17:34 +0800 Subject: [PATCH] fix(tui): preserve draft scope until first message --- nanobot/channels/websocket/runtime.py | 2 +- .../websocket/tests/test_websocket_channel.py | 43 +++++++++++++++++++ nanobot/webui/workspaces.py | 2 +- tui/src/protocol.test.ts | 32 +++++++++++--- tui/src/protocol.ts | 8 ++++ 5 files changed, 79 insertions(+), 8 deletions(-) diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index 225ba8a56..1d8ee8d55 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -1217,7 +1217,6 @@ class WebSocketChannel(BaseChannel): if session_mentions: metadata["session_mentions"] = session_mentions metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata() - self._workspaces.persist_scope(cid, scope) is_webui = metadata.get("webui") is True queued_owner = None if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content): @@ -1272,6 +1271,7 @@ class WebSocketChannel(BaseChannel): else False ), ) + self._workspaces.persist_scope(cid, scope) accepted = True finally: if not accepted and queued_owner is not None: diff --git a/nanobot/channels/websocket/tests/test_websocket_channel.py b/nanobot/channels/websocket/tests/test_websocket_channel.py index 4a185d8d3..c914e6620 100644 --- a/nanobot/channels/websocket/tests/test_websocket_channel.py +++ b/nanobot/channels/websocket/tests/test_websocket_channel.py @@ -1563,6 +1563,49 @@ async def test_new_chat_without_message_does_not_create_session( assert sessions.list_sessions() == [] +@pytest.mark.asyncio +async def test_failed_first_message_does_not_persist_draft_session( + bus: MagicMock, + tmp_path, +) -> None: + sessions = SessionManager(tmp_path / "sessions") + channel = WebSocketChannel( + {"enabled": True, "allowFrom": ["*"], "host": "127.0.0.1"}, + bus, + gateway=_basic_handler(bus, session_manager=sessions, workspace_path=tmp_path), + ) + conn = AsyncMock() + conn.remote_address = ("127.0.0.1", 50123) + + await channel._dispatch_envelope( + conn, + "tui-client", + { + "type": "new_chat", + "workspace_scope": { + "project_path": str(tmp_path), + "access_mode": "full", + }, + }, + ) + chat_id = json.loads(conn.send.await_args_list[0].args[0])["chat_id"] + bus.publish_inbound.side_effect = RuntimeError("queue unavailable") + + with pytest.raises(RuntimeError, match="queue unavailable"): + await channel._dispatch_envelope( + conn, + "tui-client", + { + "type": "message", + "chat_id": chat_id, + "content": "hello", + "webui": True, + }, + ) + + assert sessions.list_sessions() == [] + + @pytest.mark.asyncio async def test_workspace_scope_change_invalidates_other_attached_clients( bus: MagicMock, diff --git a/nanobot/webui/workspaces.py b/nanobot/webui/workspaces.py index dc035f985..4d0a853b1 100644 --- a/nanobot/webui/workspaces.py +++ b/nanobot/webui/workspaces.py @@ -336,12 +336,12 @@ class WebUIWorkspaceController: def persist_scope(self, chat_id: str, scope: WorkspaceScope) -> None: session_key = f"websocket:{chat_id}" - self._draft_scopes.pop(session_key, None) if self._sessions is not None: session = self._sessions.get_or_create(session_key) session.metadata["webui"] = True session.metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata() self._sessions.save(session) + self._draft_scopes.pop(session_key, None) def stage_scope(self, chat_id: str, scope: WorkspaceScope) -> None: """Keep a new chat's scope transient until its first accepted message.""" diff --git a/tui/src/protocol.test.ts b/tui/src/protocol.test.ts index 4ff1cbafc..5b7c2bd47 100644 --- a/tui/src/protocol.test.ts +++ b/tui/src/protocol.test.ts @@ -356,14 +356,34 @@ describe("gateway protocol", () => { socket.emit("message", { data: JSON.stringify({ event: "ready", chat_id: "", client_id: "client" }), }) + socket.emit("message", { + data: JSON.stringify({ event: "attached", chat_id: "draft-chat" }), + }) + // A gateway restart re-sends ready while the TUI keeps the draft chat alive. + socket.emit("message", { + data: JSON.stringify({ event: "ready", chat_id: "", client_id: "client" }), + }) + client.send("hello") - expect(socket.sent.map((value) => JSON.parse(value))).toEqual([{ - type: "new_chat", - workspace_scope: { - project_path: "/tmp/launch-project", - access_mode: "restricted", + expect(socket.sent.map((value) => JSON.parse(value))).toEqual([ + { + type: "new_chat", + workspace_scope: { + project_path: "/tmp/launch-project", + access_mode: "restricted", + }, }, - }]) + { type: "attach", chat_id: "draft-chat" }, + expect.objectContaining({ + type: "message", + chat_id: "draft-chat", + content: "hello", + workspace_scope: { + project_path: "/tmp/launch-project", + access_mode: "restricted", + }, + }), + ]) } finally { Object.defineProperty(globalThis, "WebSocket", { configurable: true, value: original }) } diff --git a/tui/src/protocol.ts b/tui/src/protocol.ts index 3d1237b71..885b5e96e 100644 --- a/tui/src/protocol.ts +++ b/tui/src/protocol.ts @@ -163,6 +163,7 @@ type OutboundEvent = content: string turn_id: string webui: true + workspace_scope?: WorkspaceScopePayload cli_apps?: Array<{ name: string }> mcp_presets?: Array<{ name: string }> session_mentions?: SessionMention[] @@ -903,6 +904,7 @@ export async function fetchGatewayConnection( export class NanobotClient { private socket: WebSocket | null = null private chatId = "" + private workspaceScope?: WorkspaceScopePayload private reconnectTimer: ReturnType | null = null private reconnectAttempt = 0 private closedByClient = false @@ -1002,6 +1004,7 @@ export class NanobotClient { content, turn_id: turnId, webui: true, + ...(this.workspaceScope ? { workspace_scope: this.workspaceScope } : {}), ...(options.userShell ? { user_shell: true } : {}), ...(options.cliApps?.length ? { cli_apps: options.cliApps } : {}), ...(options.mcpPresets?.length ? { mcp_presets: options.mcpPresets } : {}), @@ -1014,10 +1017,12 @@ export class NanobotClient { attach(chatId: string): void { if (!chatId) throw new Error("chat id is required") + this.workspaceScope = undefined this.write({ type: "attach", chat_id: chatId }) } newChat(scope?: WorkspaceScopePayload): void { + this.workspaceScope = scope this.write({ type: "new_chat", ...(scope ? { workspace_scope: scope } : {}) }) } @@ -1032,6 +1037,7 @@ export class NanobotClient { setWorkspaceScope(scope: WorkspaceScopePayload): void { if (!this.chatId) throw new Error("chat is not ready") + this.workspaceScope = scope this.write({ type: "set_workspace_scope", chat_id: this.chatId, workspace_scope: scope }) } @@ -1131,6 +1137,8 @@ export class NanobotClient { } } else if (event.event === "attached") { this.chatId = event.chat_id + } else if (event.event === "session_updated" && event.workspace_scope) { + this.workspaceScope = event.workspace_scope } this.options.onEvent(event) }