fix(stagewise): scope browser tabs by agent

Author: cursoragent

apps/browser/src/backend/services/toolbox/index.ts

@@ -1198,13 +1198,22 @@ export class ToolboxService extends DisposableService {
     this.shellService?.killSession(sessionId);
   }
 
-  public getBrowserSnapshot(): BrowserSnapshot {
+  public getBrowserSnapshot(agentInstanceId?: string): BrowserSnapshot {
     const browser = this.uiKarton.state.browser;
-    const activeTab = browser.activeTabId
-      ? browser.tabs[browser.activeTabId]
+    const useAgentBrowserGroup =
+      !!agentInstanceId && agentInstanceId !== browser.activeAgentId;
+    const agentBrowserState = useAgentBrowserGroup
+      ? browser.tabsByAgent[agentInstanceId]
       : null;
-
-    const allTabs = Object.values(browser.tabs)
+    const browserTabs = useAgentBrowserGroup
+      ? (agentBrowserState?.tabs ?? {})
+      : browser.tabs;
+    const activeTabId = useAgentBrowserGroup
+      ? (agentBrowserState?.activeTabId ?? null)
+      : browser.activeTabId;
+    const activeTab = activeTabId ? browserTabs[activeTabId] : null;
+
+    const allTabs = Object.values(browserTabs)
       .sort((a, b) => b.lastFocusedAt - a.lastFocusedAt)
       .map((tab) => ({
         id: tab.id,
@@ -1231,14 +1240,14 @@ export class ToolboxService extends DisposableService {
           }
         : null,
       tabs: allTabs,
-      totalTabCount: Object.keys(browser.tabs).length,
+      totalTabCount: Object.keys(browserTabs).length,
     };
   }
 
   public async captureEnvironmentSnapshot(
     agentInstanceId: string,
   ): Promise<EnvironmentSnapshot> {
-    const browserState = this.getBrowserSnapshot();
+    const browserState = this.getBrowserSnapshot(agentInstanceId);
     const workspaceState =
       this.mountManagerService?.getWorkspaceSnapshot(agentInstanceId);
     const toolboxState = this.uiKarton.state.toolbox[agentInstanceId];

apps/browser/src/backend/services/window-layout/chat-state-controller.ts

@@ -1,14 +1,21 @@
 import type { KartonService } from '../karton';
 import type { TabController } from './tab-controller';
 import type { SelectedElement } from '@shared/selected-elements';
+import type { AppState } from '@shared/karton-contracts/ui';
 
 export class ChatStateController {
   private uiKarton: KartonService;
   private tabs: Record<string, TabController>;
+  private getActiveBrowserGroupId: () => string;
 
-  constructor(uiKarton: KartonService, tabs: Record<string, TabController>) {
+  constructor(
+    uiKarton: KartonService,
+    tabs: Record<string, TabController>,
+    getActiveBrowserGroupId: () => string,
+  ) {
     this.uiKarton = uiKarton;
     this.tabs = tabs;
+    this.getActiveBrowserGroupId = getActiveBrowserGroupId;
   }
 
   /**
@@ -32,6 +39,7 @@ export class ChatStateController {
       // Add if not exists
       if (!elements.some((e) => e.stagewiseId === element.stagewiseId))
         elements.push(element);
+      this.syncGroupSelection(draft.browser.tabsByAgent, elements);
     });
     this.broadcastSelectionUpdate();
   }
@@ -45,6 +53,10 @@ export class ChatStateController {
       draft.browser.selectedElements = draft.browser.selectedElements.filter(
         (e) => e.stagewiseId !== elementId,
       );
+      this.syncGroupSelection(
+        draft.browser.tabsByAgent,
+        draft.browser.selectedElements,
+      );
     });
     this.broadcastSelectionUpdate();
   }
@@ -55,6 +67,10 @@ export class ChatStateController {
   public clearElements(): void {
     this.uiKarton.setState((draft) => {
       draft.browser.selectedElements = [];
+      this.syncGroupSelection(
+        draft.browser.tabsByAgent,
+        draft.browser.selectedElements,
+      );
     });
     this.broadcastSelectionUpdate();
   }
@@ -66,6 +82,10 @@ export class ChatStateController {
   public restoreElements(elements: SelectedElement[]): void {
     this.uiKarton.setState((draft) => {
       draft.browser.selectedElements = [...elements];
+      this.syncGroupSelection(
+        draft.browser.tabsByAgent,
+        draft.browser.selectedElements,
+      );
     });
     this.broadcastSelectionUpdate();
   }
@@ -90,4 +110,14 @@ export class ChatStateController {
       tab.updateContextSelection(allSelectedElements);
     });
   }
+
+  private syncGroupSelection(
+    tabsByAgent: AppState['browser']['tabsByAgent'],
+    selectedElements: SelectedElement[],
+  ): void {
+    const group = tabsByAgent[this.getActiveBrowserGroupId()];
+    if (group) {
+      group.selectedElements = [...selectedElements];
+    }
+  }
 }