From cef59137e4ce1b9dc007278791913d1103b842c3 Mon Sep 17 00:00:00 2001 From: drakezhang Date: Mon, 7 Sep 2026 11:44:11 +0800 Subject: [PATCH] fix(extension): handle tab return cleanup edge cases --- .../__tests__/chromium-cdp.test.ts | 63 +++++++++++++ .../src/browser-driver/chromium-cdp.ts | 30 ++++-- .../src/tools/__tests__/dispatcher.test.ts | 57 ++++++++++++ .../src/tools/__tests__/tabs.test.ts | 92 +++++++++++++++++++ apps/extension/src/tools/dispatcher.ts | 12 +-- apps/extension/src/tools/interaction.ts | 2 + apps/extension/src/tools/tabs.ts | 3 + 7 files changed, 244 insertions(+), 15 deletions(-) diff --git a/apps/extension/src/browser-driver/__tests__/chromium-cdp.test.ts b/apps/extension/src/browser-driver/__tests__/chromium-cdp.test.ts index a11c718..fccd4e2 100644 --- a/apps/extension/src/browser-driver/__tests__/chromium-cdp.test.ts +++ b/apps/extension/src/browser-driver/__tests__/chromium-cdp.test.ts @@ -254,6 +254,69 @@ describe("ChromiumCdp", () => { expect(cdp.isAttached(7)).toBe(false); }); + it.each([ + false, + true, + ])("waits for detach before reconnecting a returned tab (detach rejects: %s)", async (rejectDetach) => { + const { api } = fakeApi(); + const attached = new Set(); + let finishDetach!: () => void; + vi.mocked(api.attach).mockImplementation(async ({ tabId }) => { + if (attached.has(tabId!)) throw new Error("Another debugger is already attached"); + attached.add(tabId!); + }); + vi.mocked(api.detach).mockImplementation(async ({ tabId }) => { + attached.delete(tabId!); + }); + vi.mocked(api.detach).mockImplementationOnce( + ({ tabId }) => + new Promise((resolve, reject) => { + finishDetach = () => { + attached.delete(tabId!); + if (rejectDetach) reject(new Error("tab already detached")); + else resolve(); + }; + }), + ); + const cdp = new ChromiumCdp(api); + cdp.trackSessionTab("aa11", 7); + await cdp.ensureAttached(7); + const returning = cdp.releaseSessionTab("aa11", 7); + await vi.waitFor(() => expect(api.detach).toHaveBeenCalledOnce()); + let duplicateDetachFinished = false; + const duplicateDetach = cdp.detach(7).then(() => { + duplicateDetachFinished = true; + }); + + cdp.trackSessionTab("bb22", 7); + // Register rejection handlers immediately; the pre-fix driver rejects + // both commands while the browser still has the previous attachment. + const pending = Promise.allSettled([ + cdp.send(7, "Runtime.evaluate", { expression: "document.title" }), + cdp.send(7, "DOM.getDocument"), + ]); + await cdp.send(8, "DOM.getDocument"); + const finishedEarly = duplicateDetachFinished; + finishDetach(); + await Promise.all([returning, duplicateDetach]); + + expect(await pending).toEqual([ + { status: "fulfilled", value: { ok: true } }, + { status: "fulfilled", value: { ok: true } }, + ]); + expect(finishedEarly).toBe(false); + expect(vi.mocked(api.attach).mock.calls.filter(([target]) => target.tabId === 7)).toHaveLength( + 2, + ); + expect(cdp.isAttached(7)).toBe(true); + expect(cdp.isAttached(8)).toBe(true); + await cdp.detachSession("bb22"); + expect(api.detach).toHaveBeenCalledTimes(2); + expect(cdp.isAttached(7)).toBe(false); + expect(attached.has(7)).toBe(false); + expect(cdp.isAttached(8)).toBe(true); + }); + it("detachAll() iterates every cached tab", async () => { const { api } = fakeApi(); const cdp = new ChromiumCdp(api); diff --git a/apps/extension/src/browser-driver/chromium-cdp.ts b/apps/extension/src/browser-driver/chromium-cdp.ts index 4ebcad1..8af2abc 100644 --- a/apps/extension/src/browser-driver/chromium-cdp.ts +++ b/apps/extension/src/browser-driver/chromium-cdp.ts @@ -165,6 +165,7 @@ export class ChromiumCdp { private readonly api: CdpDebuggerApi; private readonly attachedTabs = new Set(); private readonly attachInFlight = new Map>(); + private readonly detachInFlight = new Map>(); private readonly tabOwners = new Map>(); private readonly dialogBuffers = new Map(); private readonly dialogSequences = new Map(); @@ -199,6 +200,10 @@ export class ChromiumCdp { /** Attach to `tabId` if we haven't already in this driver. */ async ensureAttached(tabId: number): Promise { + // Returning a tab clears the cache before Chrome finishes detaching. + // New observers must wait before opening the next connection to that tab. + const detaching = this.detachInFlight.get(tabId); + if (detaching) await detaching; if (this.attachedTabs.has(tabId)) return; const existing = this.attachInFlight.get(tabId); if (existing) { @@ -431,6 +436,11 @@ export class ChromiumCdp { /** Detach if attached; never throws. */ async detach(tabId: number): Promise { + const existing = this.detachInFlight.get(tabId); + if (existing) { + await existing; + return; + } this.attachInFlight.delete(tabId); if (!this.attachedTabs.has(tabId)) return; this.attachedTabs.delete(tabId); @@ -438,13 +448,19 @@ export class ChromiumCdp { this.clearConsoleState(tabId); this.clearNetworkState(tabId); this.clearFrameState(tabId); - try { - await this.api.detach({ tabId }); - } catch (err) { - // Tab may already be gone — Chrome auto-detaches on close. Log - // at debug so production builds aren't noisy. - console.debug("[bsk cdp] detach failed (likely tab already closed)", err); - } + const detach = (async () => { + try { + await this.api.detach({ tabId }); + } catch (err) { + // Tab may already be gone — Chrome auto-detaches on close. Log + // at debug so production builds aren't noisy. + console.debug("[bsk cdp] detach failed (likely tab already closed)", err); + } + })().finally(() => { + this.detachInFlight.delete(tabId); + }); + this.detachInFlight.set(tabId, detach); + await detach; } /** True iff `ensureAttached(tabId)` has succeeded since the last detach. */ diff --git a/apps/extension/src/tools/__tests__/dispatcher.test.ts b/apps/extension/src/tools/__tests__/dispatcher.test.ts index 94cf4ee..bb0216f 100644 --- a/apps/extension/src/tools/__tests__/dispatcher.test.ts +++ b/apps/extension/src/tools/__tests__/dispatcher.test.ts @@ -747,6 +747,63 @@ describe("ToolDispatcher", () => { ); }); + it.each([ + { tab_id: 7, code: "not_found" }, + { tab_id: undefined, code: "invalid_params" }, + ])("keeps hover when tab_return is rejected with $code", async ({ tab_id, code }) => { + const sendMessage = vi.fn(async () => undefined); + vi.stubGlobal("chrome", { + tabs: { + sendMessage, + get: vi.fn(async () => ({ id: 7, windowId: 100, active: true })), + query: vi.fn(async () => [{ id: 7, windowId: 100, active: true }]), + }, + }); + const { transport, sent, deliver } = fakeTransport(); + const sessions = new SessionManager({ + agentWindow: { + create: vi.fn(async () => 100), + remove: vi.fn(async () => {}), + ensureActiveTab: vi.fn(async () => 1), + }, + }); + await sessions.start("aa11"); + const cdp = { + send: vi.fn(async () => ({}) as T), + detachSession: vi.fn(async () => {}), + releaseSessionTab: vi.fn(async () => {}), + }; + const dispatcher = new ToolDispatcher({ + transport, + sessions, + cdp: cdp as unknown as TestDispatcherCdp, + }); + const helpers = dispatcher as unknown as { + rememberHover: (sessionId: string, result: unknown) => unknown; + setHoverBypass: (sessionId: string, tabId: number, enabled: boolean) => Promise; + withHoverReassert: ( + params: { session_id: string; tab_id?: number }, + work: () => Promise, + ) => Promise; + }; + helpers.rememberHover("aa11", { tab_id: 7, x: 10, y: 20 }); + await helpers.setHoverBypass("aa11", 7, true); + dispatcher.start(); + deliver(makeRequest("tool.tab_return", { session_id: "aa11", tab_id })); + await vi.waitFor(() => expect(sent).toHaveLength(1)); + + expect(sent[0]).toMatchObject({ error: { code } }); + expect(cdp.releaseSessionTab).not.toHaveBeenCalled(); + expect(sendMessage).not.toHaveBeenCalledWith(7, expect.objectContaining({ enabled: false })); + await helpers.withHoverReassert({ session_id: "aa11", tab_id: 7 }, async () => ({})); + expect(cdp.send).toHaveBeenCalledExactlyOnceWith(7, "Input.dispatchMouseEvent", { + type: "mouseMoved", + x: 10, + y: 20, + }); + dispatcher.stop(); + }); + it("releases the returned tab's hover before moving it and leaves other tabs alone", async () => { const sendMessage = vi.fn(async () => undefined); const move = vi.fn(async () => { diff --git a/apps/extension/src/tools/__tests__/tabs.test.ts b/apps/extension/src/tools/__tests__/tabs.test.ts index b6b3764..9e1270e 100644 --- a/apps/extension/src/tools/__tests__/tabs.test.ts +++ b/apps/extension/src/tools/__tests__/tabs.test.ts @@ -1,5 +1,7 @@ import { describe, expect, it, vi } from "vitest"; +import { type CdpDebuggerApi, ChromiumCdp } from "@/browser-driver/chromium-cdp"; import { SessionManager } from "@/session-manager/manager"; +import { handleHover } from "../interaction"; import { type AgentOverlayResetApi, type ChromeWindowsApi, @@ -547,6 +549,96 @@ describe("handleTabBorrow", () => { }); describe("handleTabReturn", () => { + it.each([ + "not_found", + "cdp_failed", + ])("detaches after selector resolution fails with %s", async (code) => { + const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) }); + const ctx = await sm.start("aa11"); + ctx.borrowedTabs.set(7, { tabId: 7, originalWindowId: 200, originalIndex: 4 }); + const state: FakeTabState = { + tabs: new Map([[7, { id: 7, windowId: 100 } as chrome.tabs.Tab]]), + nextTabId: 50, + windowsClosed: new Set(), + }; + const { api } = makeTabMutationApi(state); + const { api: windows } = makeWindowsApi(state); + const debuggerApi = { + attach: vi.fn(async () => {}), + detach: vi.fn(async () => {}), + sendCommand: vi.fn(async (_target: unknown, method: string) => { + if (method === "DOM.getDocument") return { root: { nodeId: 1 } }; + if (method === "DOM.querySelector") { + if (code === "cdp_failed") throw new Error("selector query failed"); + return { nodeId: 0 }; + } + return {}; + }), + onEvent: { addListener: vi.fn(), removeListener: vi.fn() }, + onDetach: { addListener: vi.fn(), removeListener: vi.fn() }, + }; + const cdp = new ChromiumCdp(debuggerApi as unknown as CdpDebuggerApi); + + const hover = await handleHover( + sm, + { session_id: "aa11", tab_id: 7, selector: "#missing" }, + { + cdp, + tabsApi: { get: api.get, query: vi.fn(async () => []) }, + }, + ); + expect(hover).toMatchObject({ code }); + expect(cdp.isAttached(7)).toBe(true); + + const returned = await handleTabReturn( + sm, + { session_id: "aa11", tab_id: 7 }, + { + tabs: api, + windows, + cdp, + agentOverlayReset: { resetAgentOverlays: vi.fn(async () => {}) }, + }, + ); + expect(returned).toMatchObject({ returned_to_window_id: 200 }); + expect(ctx.borrowedTabs.has(7)).toBe(false); + expect(cdp.isAttached(7)).toBe(false); + expect(debuggerApi.detach).toHaveBeenCalledExactlyOnceWith({ tabId: 7 }); + cdp.dispose(); + }); + + it("honours cancellation while preparing a validated return", async () => { + const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) }); + const ctx = await sm.start("aa11"); + ctx.borrowedTabs.set(7, { tabId: 7, originalWindowId: 200, originalIndex: 4 }); + const state: FakeTabState = { + tabs: new Map([[7, { id: 7, windowId: 100 } as chrome.tabs.Tab]]), + nextTabId: 50, + windowsClosed: new Set(), + }; + const { api, spies } = makeTabMutationApi(state); + const { api: windows } = makeWindowsApi(state); + const controller = new AbortController(); + const beforeReturn = vi.fn(async () => { + controller.abort(); + }); + const result = await handleTabReturn( + sm, + { session_id: "aa11", tab_id: 7 }, + { + tabs: api, + windows, + signal: controller.signal, + beforeReturn, + }, + ); + + expect(result).toMatchObject({ code: "cancelled" }); + expect(beforeReturn).toHaveBeenCalledExactlyOnceWith("aa11", 7); + expect(spies.move).not.toHaveBeenCalled(); + expect(ctx.borrowedTabs.has(7)).toBe(true); + }); + it.each([ "original", "fallback", diff --git a/apps/extension/src/tools/dispatcher.ts b/apps/extension/src/tools/dispatcher.ts index c396af8..64e451f 100644 --- a/apps/extension/src/tools/dispatcher.ts +++ b/apps/extension/src/tools/dispatcher.ts @@ -345,15 +345,11 @@ export class ToolDispatcher { return result; } case "tool.tab_return": - return this.withHoverReleaseForRequest( - req.params as TabReturnParams, - () => - handleTabReturn(this.sessions, req.params as TabReturnParams, { - signal, - cdp: this.cdp, - }), + return handleTabReturn(this.sessions, req.params as TabReturnParams, { signal, - ); + cdp: this.cdp, + beforeReturn: (sessionId, tabId) => this.releaseHoverLatch(sessionId, tabId), + }); case "tool.window_resize": return handleWindowResize( this.sessions, diff --git a/apps/extension/src/tools/interaction.ts b/apps/extension/src/tools/interaction.ts index 54646ed..18608ca 100644 --- a/apps/extension/src/tools/interaction.ts +++ b/apps/extension/src/tools/interaction.ts @@ -183,6 +183,8 @@ export async function resolveBackendNode( } // selector path try { + // Selector lookup itself attaches CDP, even when no element is found. + cdp.trackSessionTab?.(ctx.sessionId, target.tabId); const doc = await cdp.send<{ root?: { nodeId?: number } }>(target.tabId, "DOM.getDocument", { depth: 0, }); diff --git a/apps/extension/src/tools/tabs.ts b/apps/extension/src/tools/tabs.ts index 80fca96..feb9763 100644 --- a/apps/extension/src/tools/tabs.ts +++ b/apps/extension/src/tools/tabs.ts @@ -289,6 +289,8 @@ export interface TabManagementDeps { agentOverlayReset?: AgentOverlayResetApi; /** Releases this session's CDP claim after a borrowed tab is returned. */ cdp?: Pick; + /** Runs after tab_return validation, before moving the borrowed tab. */ + beforeReturn?: (sessionId: string, tabId: number) => Promise; /** * Reports whether `windowId` is any live session's Agent Window. * `tab_return`'s fallback window picker uses this to avoid moving a @@ -1124,6 +1126,7 @@ export async function handleTabReturn( message: `tab_return: tab ${params.tab_id} is not borrowed by this session`, }; } + await deps.beforeReturn?.(ctx.sessionId, params.tab_id); const outcome = await returnBorrowedTab(ctx, params.tab_id, { ...deps, isAgentWindowId: