fix(tui): preserve draft scope until first message

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:

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,

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

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 })
     }

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<typeof setTimeout> | 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)
   }