diff --git a/src/main/browser/agent-browser-bridge-tab-routing.test.ts b/src/main/browser/agent-browser-bridge-tab-routing.test.ts index e64e2ee85cc..103de60f4c7 100644 --- a/src/main/browser/agent-browser-bridge-tab-routing.test.ts +++ b/src/main/browser/agent-browser-bridge-tab-routing.test.ts @@ -234,6 +234,44 @@ describe('AgentBrowserBridge', () => { expect(bridge.getActiveWebContentsId()).toBeNull() }) + it('resolves repeated guest lookups without enumerating registered tabs', () => { + const tabCount = 1_000 + const tabs = new Map() + const tabIdByWebContentsId = new Map() + for (let index = 0; index < tabCount; index += 1) { + const tabId = `tab-${index}` + const webContentsId = 20_000 + index + tabs.set(tabId, webContentsId) + tabIdByWebContentsId.set(webContentsId, tabId) + } + + const getRegisteredTabs = vi.fn(() => { + throw new Error('unexpected registered-tab enumeration') + }) + const getTabIdForWebContentsId = vi.fn( + (webContentsId: number) => tabIdByWebContentsId.get(webContentsId) ?? null + ) + const b = new AgentBrowserBridge( + mockBrowserManager(tabs, new Map(), { + getWebContentsIdByTabId: getRegisteredTabs, + getTabIdForWebContentsId + }) + ) + const resolveTabIdSafe = ( + b as unknown as { resolveTabIdSafe: (webContentsId: number) => string | null } + ).resolveTabIdSafe.bind(b) + const targetWebContentsId = 20_000 + tabCount - 1 + + // The old forward scan inspected 1,000,000 entries for this repeated last-tab lookup. + for (let lookup = 0; lookup < tabCount; lookup += 1) { + expect(resolveTabIdSafe(targetWebContentsId)).toBe(`tab-${tabCount - 1}`) + } + + expect(getTabIdForWebContentsId).toHaveBeenCalledTimes(tabCount) + expect(getRegisteredTabs).not.toHaveBeenCalled() + expect(resolveTabIdSafe(99_999)).toBeNull() + }) + it('closes the named agent-browser session when a tab closes', async () => { succeedWith({ snapshot: 'tree' }) await bridge.snapshot() diff --git a/src/main/browser/agent-browser-bridge-test-harness.ts b/src/main/browser/agent-browser-bridge-test-harness.ts index 72ed978ac48..0e614600174 100644 --- a/src/main/browser/agent-browser-bridge-test-harness.ts +++ b/src/main/browser/agent-browser-bridge-test-harness.ts @@ -18,6 +18,14 @@ export function mockBrowserManager( ): BrowserManager { return { getWebContentsIdByTabId: () => tabs, + getTabIdForWebContentsId: (webContentsId: number) => { + for (const [tabId, tabWebContentsId] of tabs) { + if (tabWebContentsId === webContentsId) { + return tabId + } + } + return null + }, getWorktreeIdForTab: (tabId: string) => worktrees.get(tabId), getGuestWebContentsId: vi.fn(() => null), getBrowserPageLoadError: vi.fn(() => null), diff --git a/src/main/browser/agent-browser-bridge.ts b/src/main/browser/agent-browser-bridge.ts index 54611bffa8f..01f77b45451 100644 --- a/src/main/browser/agent-browser-bridge.ts +++ b/src/main/browser/agent-browser-bridge.ts @@ -2861,13 +2861,7 @@ export class AgentBrowserBridge { } private resolveTabIdSafe(webContentsId: number): string | null { - const tabs = this.browserManager.getWebContentsIdByTabId() - for (const [tabId, wcId] of tabs) { - if (wcId === webContentsId) { - return tabId - } - } - return null + return this.browserManager.getTabIdForWebContentsId(webContentsId) } private requireTargetWebContents(target: ResolvedBrowserCommandTarget): WebContents { diff --git a/src/main/browser/browser-manager-guest-lifecycle.test.ts b/src/main/browser/browser-manager-guest-lifecycle.test.ts index 6d71db94585..7b806c9d2d1 100644 --- a/src/main/browser/browser-manager-guest-lifecycle.test.ts +++ b/src/main/browser/browser-manager-guest-lifecycle.test.ts @@ -541,6 +541,11 @@ describe('browserManager', () => { }) expect(oldGuestOffMock).toHaveBeenCalled() expect(browserManager.getGuestWebContentsId('browser-1')).toBe(newGuest.id) + expect(browserManager.getTabIdForWebContentsId(oldGuest.id)).toBeNull() + expect(browserManager.getTabIdForWebContentsId(newGuest.id)).toBe('browser-1') + + browserManager.unregisterGuest('browser-1') + expect(browserManager.getTabIdForWebContentsId(newGuest.id)).toBeNull() }) it('cleans up prior guest listeners before re-registering the same tab', () => { diff --git a/src/main/browser/browser-manager.ts b/src/main/browser/browser-manager.ts index 2b43e17a371..fbd667d4119 100644 --- a/src/main/browser/browser-manager.ts +++ b/src/main/browser/browser-manager.ts @@ -1575,6 +1575,10 @@ export class BrowserManager { return this.webContentsIdByTabId } + getTabIdForWebContentsId(webContentsId: number): string | null { + return this.tabIdByWebContentsId.get(webContentsId) ?? null + } + getWorktreeIdForTab(browserTabId: string): string | undefined { return this.worktreeIdByTabId.get(browserTabId) } diff --git a/src/main/browser/cdp-bridge-state.test.ts b/src/main/browser/cdp-bridge-state.test.ts new file mode 100644 index 00000000000..1267201b39d --- /dev/null +++ b/src/main/browser/cdp-bridge-state.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, it, vi } from 'vitest' + +const { webContentsFromIdMock } = vi.hoisted(() => ({ + webContentsFromIdMock: vi.fn() +})) + +vi.mock('electron', () => ({ + webContents: { fromId: webContentsFromIdMock } +})) + +import type { CdpTabState } from './cdp-auxiliary-commands' +import { CdpBridgeState } from './cdp-bridge-state' +import { CdpTabCommands } from './cdp-tab-commands' + +function createBridgeState( + getRegisteredTabs: () => Map, + getTabIdForWebContentsId: (webContentsId: number) => string | null +) { + let activeWebContentsId: number | null = null + const bindings = { + getActiveWebContentsId: () => activeWebContentsId, + setActiveWebContentsId: (webContentsId: number | null) => { + activeWebContentsId = webContentsId + }, + getRegisteredTabs, + getTabIdForWebContentsId, + tabState: new Map(), + commandQueues: new Map(), + processingQueues: new Set() + } + + return new CdpBridgeState(bindings) +} + +describe('CdpBridgeState reverse tab lookup', () => { + it('resolves repeated lookups from the reverse map without enumerating tabs', () => { + const tabCount = 1_000 + const lookupCount = 1_000 + const registeredTabs = new Map() + const tabIdByWebContentsId = new Map() + for (let index = 0; index < tabCount; index += 1) { + const tabId = `tab-${index}` + const webContentsId = 10_000 + index + registeredTabs.set(tabId, webContentsId) + tabIdByWebContentsId.set(webContentsId, tabId) + } + + const getRegisteredTabs = vi.fn(() => { + throw new Error('unexpected registered-tab enumeration') + }) + const getTabIdForWebContentsId = vi.fn( + (webContentsId: number) => tabIdByWebContentsId.get(webContentsId) ?? null + ) + const state = createBridgeState(getRegisteredTabs, getTabIdForWebContentsId) + const targetWebContentsId = 10_000 + tabCount - 1 + + // The old forward scan inspected 1,000,000 entries for this repeated last-tab lookup. + for (let lookup = 0; lookup < lookupCount; lookup += 1) { + expect(state.resolveTabIdSafe(targetWebContentsId)).toBe(`tab-${tabCount - 1}`) + } + + expect(getTabIdForWebContentsId).toHaveBeenCalledTimes(lookupCount) + expect(getRegisteredTabs).not.toHaveBeenCalled() + }) + + it('keeps strict and safe resolution semantics when a registration changes', () => { + const registeredTabs = new Map([['tab-a', 101]]) + const tabIdByWebContentsId = new Map([[101, 'tab-a']]) + const state = createBridgeState( + () => registeredTabs, + (webContentsId) => tabIdByWebContentsId.get(webContentsId) ?? null + ) + + expect(state.resolveTabId(101)).toBe('tab-a') + expect(state.resolveTabIdSafe(999)).toBeNull() + expect(() => state.resolveTabId(999)).toThrow('Tab is no longer registered.') + + tabIdByWebContentsId.delete(101) + tabIdByWebContentsId.set(202, 'tab-a') + expect(state.resolveTabIdSafe(101)).toBeNull() + expect(state.resolveTabId(202)).toBe('tab-a') + }) + + it('lets CdpTabCommands read the active page id through the reverse map', () => { + const tabIdByWebContentsId = new Map([[101, 'tab-a']]) + const getRegisteredTabs = vi.fn(() => { + throw new Error('unexpected registered-tab enumeration') + }) + const state = createBridgeState( + getRegisteredTabs, + (webContentsId) => tabIdByWebContentsId.get(webContentsId) ?? null + ) + const commands = new CdpTabCommands(state, {} as never, {} as never, {} as never) + + state.activeWebContentsId = 101 + expect(commands.getActivePageId()).toBe('tab-a') + state.activeWebContentsId = 999 + expect(commands.getActivePageId()).toBeNull() + expect(getRegisteredTabs).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/browser/cdp-bridge-state.ts b/src/main/browser/cdp-bridge-state.ts index 9eed291b04d..837b44fb8f2 100644 --- a/src/main/browser/cdp-bridge-state.ts +++ b/src/main/browser/cdp-bridge-state.ts @@ -12,6 +12,7 @@ export type CdpBridgeStateBindings = { getActiveWebContentsId: () => number | null setActiveWebContentsId: (webContentsId: number | null) => void getRegisteredTabs: () => Map + getTabIdForWebContentsId: (webContentsId: number) => string | null tabState: Map commandQueues: Map processingQueues: Set @@ -84,21 +85,15 @@ export class CdpBridgeState { } resolveTabId(webContentsId: number): string { - for (const [tabId, wcId] of this.getRegisteredTabs()) { - if (wcId === webContentsId) { - return tabId - } + const tabId = this.bindings.getTabIdForWebContentsId(webContentsId) + if (tabId !== null) { + return tabId } throw new BrowserError('browser_debugger_detached', 'Tab is no longer registered.') } resolveTabIdSafe(webContentsId: number): string | null { - for (const [tabId, wcId] of this.getRegisteredTabs()) { - if (wcId === webContentsId) { - return tabId - } - } - return null + return this.bindings.getTabIdForWebContentsId(webContentsId) } getOrCreateTabState(tabId: string): CdpTabState { diff --git a/src/main/browser/cdp-bridge.ts b/src/main/browser/cdp-bridge.ts index dbe6dadf74d..77aba3ba6d5 100644 --- a/src/main/browser/cdp-bridge.ts +++ b/src/main/browser/cdp-bridge.ts @@ -60,6 +60,8 @@ export class CdpBridge { this.activeWebContentsId = webContentsId }, getRegisteredTabs: () => this.getRegisteredTabs(), + getTabIdForWebContentsId: (webContentsId) => + this.browserManager.getTabIdForWebContentsId(webContentsId), tabState: this.tabState, commandQueues: this.commandQueues, processingQueues: this.processingQueues diff --git a/src/main/browser/cdp-tab-commands.ts b/src/main/browser/cdp-tab-commands.ts index 2a13f58173e..ae13bf4d994 100644 --- a/src/main/browser/cdp-tab-commands.ts +++ b/src/main/browser/cdp-tab-commands.ts @@ -20,12 +20,7 @@ export class CdpTabCommands extends CdpBridgeCommandModule { if (!this.activeWebContentsId) { return null } - for (const [tabId, wcId] of this.getRegisteredTabs()) { - if (wcId === this.activeWebContentsId) { - return tabId - } - } - return null + return this.resolveTabIdSafe(this.activeWebContentsId) } getPageInfo(