From e831d65db3b7fcb3267159baa101cf5248aa2e61 Mon Sep 17 00:00:00 2001 From: drakezhang Date: Mon, 7 Sep 2026 11:12:44 +0800 Subject: [PATCH] fix(extension): prevent closed borrowed tabs from blocking session stop --- apps/extension/src/entrypoints/background.ts | 4 +- .../__tests__/event-handler.test.ts | 1 + .../session-manager/__tests__/manager.test.ts | 37 ++++++++++- apps/extension/src/session-manager/manager.ts | 10 ++- .../src/tools/__tests__/session.test.ts | 65 +++++++++++++++++++ apps/extension/src/tools/session.ts | 17 ++++- 6 files changed, 126 insertions(+), 8 deletions(-) diff --git a/apps/extension/src/entrypoints/background.ts b/apps/extension/src/entrypoints/background.ts index c70ae6a..34d67a3 100644 --- a/apps/extension/src/entrypoints/background.ts +++ b/apps/extension/src/entrypoints/background.ts @@ -169,8 +169,8 @@ export default defineBackground(() => { if (!sessions.findByWindowId(tab.windowId)) return; void pushOverlayStateForTab(tab.id, tab.windowId); }); - chrome.tabs.onRemoved.addListener((tabId) => { - sessions.forgetAgentCreatedTab(tabId); + chrome.tabs.onRemoved.addListener((tabId, removeInfo) => { + sessions.forgetClosedTab(tabId, { isWindowClosing: removeInfo.isWindowClosing }); }); // Re-sync the storage.session flag on SW startup so a previous SW's // stale `true` does not keep waking us on every page load until the diff --git a/apps/extension/src/session-manager/__tests__/event-handler.test.ts b/apps/extension/src/session-manager/__tests__/event-handler.test.ts index 5aaf338..7d4f33e 100644 --- a/apps/extension/src/session-manager/__tests__/event-handler.test.ts +++ b/apps/extension/src/session-manager/__tests__/event-handler.test.ts @@ -87,6 +87,7 @@ describe("attachSessionEventHandler", () => { windowEvents: events.api, }); + manager.forgetClosedTab(7, { isWindowClosing: true }); events.emit(4242); for (let i = 0; i < 4; i += 1) await Promise.resolve(); diff --git a/apps/extension/src/session-manager/__tests__/manager.test.ts b/apps/extension/src/session-manager/__tests__/manager.test.ts index ab231af..5abd2a7 100644 --- a/apps/extension/src/session-manager/__tests__/manager.test.ts +++ b/apps/extension/src/session-manager/__tests__/manager.test.ts @@ -180,10 +180,45 @@ describe("SessionManager", () => { const sm = new SessionManager({ agentWindow: fakeAgentWindow() }); const ctx = await sm.start("aa11"); ctx.agentCreatedTabs.add(42); - sm.forgetAgentCreatedTab(42); + sm.forgetClosedTab(42); expect(ctx.agentCreatedTabs.has(42)).toBe(false); expect(isAgentControlledTab(ctx, 0)).toBe(true); }); + + it("forgets closed borrows idempotently without affecting other tabs or sessions", async () => { + const sm = new SessionManager({ agentWindow: fakeAgentWindow() }); + const a = await sm.start("aa11"); + const b = await sm.start("bb22"); + a.borrowedTabs.set(42, { tabId: 42, originalWindowId: 7, originalIndex: 0 }); + a.borrowedTabs.set(43, { tabId: 43, originalWindowId: 7, originalIndex: 1 }); + b.borrowedTabs.set(44, { tabId: 44, originalWindowId: 8, originalIndex: 0 }); + + sm.forgetClosedTab(42); + sm.forgetClosedTab(42); + sm.forgetClosedTab(999); + + expect(isAgentControlledTab(a, 42)).toBe(false); + expect(sm.findBorrowingSession(42, null)).toBeNull(); + expect([...a.borrowedTabs.keys()]).toEqual([43]); + expect([...b.borrowedTabs.keys()]).toEqual([44]); + expect(sm.list()).toHaveLength(2); + }); + + it("prevents an in-flight borrow from reclaiming a closed tab", async () => { + const sm = new SessionManager({ agentWindow: fakeAgentWindow() }); + const ctx = await sm.start("aa11"); + const reservation = sm.tryReserveBorrow(42, ctx.sessionId); + if ("borrowedBy" in reservation) throw new Error("unexpected borrow conflict"); + + sm.forgetClosedTab(42); + + expect(() => + reservation.commit({ tabId: 42, originalWindowId: 7, originalIndex: 0 }), + ).toThrow(/reservation disappeared/); + reservation.release(); + expect(sm.findBorrowingSession(42, null)).toBeNull(); + expect(ctx.borrowedTabs.has(42)).toBe(false); + }); }); describe("findBorrowingSession", () => { diff --git a/apps/extension/src/session-manager/manager.ts b/apps/extension/src/session-manager/manager.ts index 4c0df05..97982e3 100644 --- a/apps/extension/src/session-manager/manager.ts +++ b/apps/extension/src/session-manager/manager.ts @@ -117,10 +117,16 @@ export class SessionManager { return Array.from(this.sessions.values()); } - /** Remove a closed tab from agent-created ownership tracking. */ - forgetAgentCreatedTab(tabId: number): void { + /** + * Forget a tab Chrome has removed, including any uncommitted borrow. + * Whole-window closures keep committed borrows until the window-removed + * handler reports which user tabs could not be returned. + */ + forgetClosedTab(tabId: number, { isWindowClosing = false } = {}): void { + this.borrowReservations.delete(tabId); for (const ctx of this.sessions.values()) { ctx.agentCreatedTabs.delete(tabId); + if (!isWindowClosing) ctx.borrowedTabs.delete(tabId); } } diff --git a/apps/extension/src/tools/__tests__/session.test.ts b/apps/extension/src/tools/__tests__/session.test.ts index d4cd548..f575c93 100644 --- a/apps/extension/src/tools/__tests__/session.test.ts +++ b/apps/extension/src/tools/__tests__/session.test.ts @@ -287,6 +287,71 @@ describe("handleSessionStop with auto-return", () => { await handleSessionStop(sm, { session_id: "aa11" }, { tabManagement: { tabs, windows } }); expect(ctx.refStore.isEmpty()).toBe(true); }); + + it.each([ + "before teardown", + "during its return", + "after its return failed", + "before its turn in the return loop", + ])("finishes stopping when a borrowed tab closes %s", async (timing) => { + const aw = fakeAgentWindow([100]); + const sm = new SessionManager({ agentWindow: aw }); + const ctx = await sm.start("aa11"); + ctx.borrowedTabs.set(1, { tabId: 1, originalWindowId: 200, originalIndex: 0 }); + ctx.borrowedTabs.set(2, { tabId: 2, originalWindowId: 200, originalIndex: 1 }); + const state: FakeState = { + tabs: new Map([ + [1, { id: 1, windowId: 100 } as chrome.tabs.Tab], + [2, { id: 2, windowId: 100 } as chrome.tabs.Tab], + ]), + windowsClosed: new Set(), + moves: [], + }; + const { tabs, windows } = makeApis(state); + const cdp = { detachSession: vi.fn(async () => {}) }; + const closeTab = (tabId: number) => { + state.tabs.delete(tabId); + sm.forgetClosedTab(tabId); + }; + const baseMove = tabs.move; + tabs.move = vi.fn(async (tabId, props) => { + if (tabId === 1 && timing === "during its return") closeTab(1); + if (tabId === 1 && timing === "after its return failed") { + throw new Error("simulated move failure"); + } + if (tabId === 2 && timing === "after its return failed") closeTab(1); + if (tabId === 1 && timing === "before its turn in the return loop") closeTab(2); + if (!state.tabs.has(tabId)) throw new Error(`tab ${tabId} not found`); + return baseMove(tabId, props); + }); + if (timing === "before teardown") closeTab(1); + + const res = await handleSessionStop( + sm, + { session_id: "aa11" }, + { + cdp, + tabManagement: { + tabs, + windows, + agentOverlayReset: { resetAgentOverlays: vi.fn(async () => {}) }, + }, + tabsQuery: makeQuery(state), + }, + ); + + if ("code" in res) throw new Error(`unexpected error: ${JSON.stringify(res)}`); + const closedTabId = timing === "before its turn in the return loop" ? 2 : 1; + expect(res.returned_tab_ids).toEqual([closedTabId === 1 ? 2 : 1]); + expect(res.return_failures).toBeUndefined(); + expect(ctx.borrowedTabs.size).toBe(0); + expect(sm.has("aa11")).toBe(false); + expect(cdp.detachSession).toHaveBeenCalledWith("aa11"); + expect(aw.remove).toHaveBeenCalledWith(100); + if (timing === "before teardown" || timing === "before its turn in the return loop") { + expect(tabs.move).not.toHaveBeenCalledWith(closedTabId, expect.anything()); + } + }); }); describe("handleSessionStop window release (issue #57)", () => { diff --git a/apps/extension/src/tools/session.ts b/apps/extension/src/tools/session.ts index b3af268..338e5c2 100644 --- a/apps/extension/src/tools/session.ts +++ b/apps/extension/src/tools/session.ts @@ -204,6 +204,8 @@ export async function handleSessionStop( const returnFailures: SessionStopResult["return_failures"] = []; const borrowedIds = Array.from(ctx.borrowedTabs.keys()); for (const tabId of borrowedIds) { + // A close event may have removed this borrow while another tab returned. + if (!ctx.borrowedTabs.has(tabId)) continue; try { const tabManagement = { ...(deps.tabManagement ?? {}), @@ -213,6 +215,7 @@ export async function handleSessionStop( ((windowId: number) => manager.findByWindowId(windowId) !== null), }; const outcome = await returnBorrowedTab(ctx, tabId, tabManagement); + if (!ctx.borrowedTabs.has(tabId)) continue; if (typeof outcome === "object" && "code" in outcome) { console.warn(`[bsk session_stop] auto-return failed for tab ${tabId}`, outcome); returnFailures?.push({ @@ -225,6 +228,7 @@ export async function handleSessionStop( ctx.borrowedTabs.delete(tabId); } } catch (err) { + if (!ctx.borrowedTabs.has(tabId)) continue; const message = err instanceof Error ? err.message : String(err); console.warn(`[bsk session_stop] auto-return threw for tab ${tabId}: ${message}`); returnFailures?.push({ @@ -235,8 +239,15 @@ export async function handleSessionStop( } } + // A previously failed tab may have closed during a later return. Only + // confirmed close events remove borrows; ordinary move failures stay retryable. + const pendingReturnFailures = returnFailures.filter((failure) => + ctx.borrowedTabs.has(failure.tab_id), + ); if (deps.signal?.aborted) { - const nonCancelFailures = returnFailures?.filter((failure) => failure.code !== "cancelled"); + const nonCancelFailures = pendingReturnFailures.filter( + (failure) => failure.code !== "cancelled", + ); if (nonCancelFailures && nonCancelFailures.length > 0) { return { ...(returnedTabIds.length > 0 ? { returned_tab_ids: returnedTabIds } : {}), @@ -248,8 +259,8 @@ export async function handleSessionStop( const result: SessionStopResult = {}; if (returnedTabIds.length > 0) result.returned_tab_ids = returnedTabIds; - if (returnFailures && returnFailures.length > 0) { - result.return_failures = returnFailures; + if (pendingReturnFailures.length > 0) { + result.return_failures = pendingReturnFailures; // A failed return means at least one borrowed user tab may still be // inside the Agent Window. Keep the session/window alive so the user // can retry `bsk session stop` or explicitly `bsk tab return` after the