diff --git a/apps/extension/src/session-manager/__tests__/manager.test.ts b/apps/extension/src/session-manager/__tests__/manager.test.ts index c58b5d1..966693e 100644 --- a/apps/extension/src/session-manager/__tests__/manager.test.ts +++ b/apps/extension/src/session-manager/__tests__/manager.test.ts @@ -75,6 +75,52 @@ describe("SessionManager", () => { await expect(sm.start("aa11")).rejects.toThrow(/already exists/); }); + it("removes a newly created Agent Window when startup is aborted", async () => { + const aw = fakeAgentWindow(); + let resolveCreate: (windowId: number) => void = () => {}; + aw.createMock.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveCreate = resolve; + }), + ); + const sm = new SessionManager({ agentWindow: aw }); + const controller = new AbortController(); + const pending = sm.start("aa11", { signal: controller.signal }); + + controller.abort(); + resolveCreate(777); + + await expect(pending).rejects.toMatchObject({ name: "AbortError" }); + expect(aw.removeMock).toHaveBeenCalledWith(777); + expect(sm.has("aa11")).toBe(false); + }); + + it("removes an incomplete Agent Window when active-tab setup fails", async () => { + const aw = fakeAgentWindow(); + aw.ensureActiveTabMock.mockRejectedValueOnce(new Error("tab setup failed")); + const sm = new SessionManager({ agentWindow: aw }); + + await expect(sm.start("aa11")).rejects.toThrow("tab setup failed"); + + expect(aw.removeMock).toHaveBeenCalledWith(100); + expect(sm.has("aa11")).toBe(false); + }); + + it("surfaces the orphan Agent Window id when startup cleanup fails", async () => { + const aw = fakeAgentWindow(); + aw.ensureActiveTabMock.mockRejectedValueOnce(new Error("tab setup failed")); + aw.removeMock.mockRejectedValueOnce(new Error("window removal denied")); + const sm = new SessionManager({ agentWindow: aw }); + + await expect(sm.start("aa11")).rejects.toMatchObject({ + name: "SessionStartCleanupError", + windowId: 100, + message: expect.stringMatching(/cleanup of Agent Window 100 failed.*window removal denied/), + }); + expect(sm.has("aa11")).toBe(false); + }); + it("stop() closes the Agent Window and forgets the session", async () => { const aw = fakeAgentWindow(); const sm = new SessionManager({ agentWindow: aw }); diff --git a/apps/extension/src/session-manager/agent-window.ts b/apps/extension/src/session-manager/agent-window.ts index f345b32..73f634b 100644 --- a/apps/extension/src/session-manager/agent-window.ts +++ b/apps/extension/src/session-manager/agent-window.ts @@ -43,12 +43,11 @@ export const chromeAgentWindowApi: AgentWindowApi = { return win.id; }, async remove(windowId: number): Promise { - try { - await chrome.windows.remove(windowId); - } catch (err) { - // Window may have been closed by the user already; ignore. - console.debug("[bh] chrome.windows.remove failed", err); - } + // Callers decide whether a missing/failed removal is benign. In + // particular, transactional session-start cleanup must be able to + // surface a window it could not remove instead of reporting a false + // cancellation success while the Agent Window remains open. + await chrome.windows.remove(windowId); }, async ensureActiveTab(windowId: number, url: string): Promise { const tabs = await chrome.tabs.query({ windowId }); diff --git a/apps/extension/src/session-manager/manager.ts b/apps/extension/src/session-manager/manager.ts index d2c1022..2fef0d7 100644 --- a/apps/extension/src/session-manager/manager.ts +++ b/apps/extension/src/session-manager/manager.ts @@ -31,6 +31,38 @@ export interface SessionStartOptions { size?: { width: number; height: number }; /** Defaults to true so existing clients keep visible Agent Windows. */ focused?: boolean; + /** Cancellation for the transactional Agent Window startup sequence. */ + signal?: AbortSignal; +} + +export class SessionStartCleanupError extends Error { + readonly windowId: number; + readonly startupError: unknown; + readonly cleanupError: unknown; + + constructor(windowId: number, startupError: unknown, cleanupError: unknown) { + const startupMessage = + startupError instanceof Error ? startupError.message : String(startupError); + const cleanupMessage = + cleanupError instanceof Error ? cleanupError.message : String(cleanupError); + super( + `session_start failed (${startupMessage}) and cleanup of Agent Window ${windowId} failed: ${cleanupMessage}`, + ); + this.name = "SessionStartCleanupError"; + this.windowId = windowId; + this.startupError = startupError; + this.cleanupError = cleanupError; + } +} + +function sessionStartAbortError(): Error { + const error = new Error("session_start aborted"); + error.name = "AbortError"; + return error; +} + +function throwIfSessionStartAborted(signal: AbortSignal | undefined): void { + if (signal?.aborted) throw sessionStartAbortError(); } /** @@ -140,18 +172,36 @@ export class SessionManager { if (this.sessions.has(sessionId)) { throw new Error(`[bh] session ${sessionId} already exists`); } - const windowId = await this.agentWindow.create(AGENT_WINDOW_HOME, opts); - await this.agentWindow.ensureActiveTab(windowId, AGENT_WINDOW_HOME); - const ctx: SessionContext = { - sessionId, - agentWindowId: windowId, - refStore: new RefStore(), - borrowedTabs: new Map(), - createdAtMs: this.now(), - }; - this.sessions.set(sessionId, ctx); - this.windowIndex.set(windowId, sessionId); - return ctx; + throwIfSessionStartAborted(opts.signal); + + let windowId: number | null = null; + try { + const { signal: _signal, ...createOptions } = opts; + windowId = await this.agentWindow.create(AGENT_WINDOW_HOME, createOptions); + throwIfSessionStartAborted(opts.signal); + await this.agentWindow.ensureActiveTab(windowId, AGENT_WINDOW_HOME); + throwIfSessionStartAborted(opts.signal); + + const ctx: SessionContext = { + sessionId, + agentWindowId: windowId, + refStore: new RefStore(), + borrowedTabs: new Map(), + createdAtMs: this.now(), + }; + this.sessions.set(sessionId, ctx); + this.windowIndex.set(windowId, sessionId); + return ctx; + } catch (startupError) { + if (windowId !== null) { + try { + await this.agentWindow.remove(windowId); + } catch (cleanupError) { + throw new SessionStartCleanupError(windowId, startupError, cleanupError); + } + } + throw startupError; + } } /** diff --git a/apps/extension/src/tools/__tests__/dispatcher.test.ts b/apps/extension/src/tools/__tests__/dispatcher.test.ts index 2280b26..2c82a27 100644 --- a/apps/extension/src/tools/__tests__/dispatcher.test.ts +++ b/apps/extension/src/tools/__tests__/dispatcher.test.ts @@ -677,23 +677,31 @@ describe("ToolDispatcher", () => { await flushMicrotasks(); expect(ac?.signal.aborted).toBe(true); - // Cancel ack arrived synchronously; the slow tool replies with - // `cancelled` once the dispatcher's race observes the abort. + // Cancel ack arrives synchronously, but the original RPC must not reply + // until the in-progress window creation has completed and been rolled back. const ack = sent.find( (m) => typeof (m as { id?: string }).id === "string" && (m as { id: string }).id === "cancel-1", ); expect(ack).toEqual({ id: "cancel-1", result: { cancelled: true } }); + expect( + sent.find( + (m) => typeof (m as { id?: string }).id === "string" && (m as { id: string }).id === "r-1", + ), + ).toBeUndefined(); + expect(dispatcher.inflightAbortControllers.has("r-1")).toBe(true); + + resolveCreate(9999); + await flushMicrotasks(); + const slow = sent.find( (m) => typeof (m as { id?: string }).id === "string" && (m as { id: string }).id === "r-1", ); expect(slow).toMatchObject({ id: "r-1", error: { code: "cancelled" } }); expect(dispatcher.inflightAbortControllers.has("r-1")).toBe(false); - - // Drain the dangling create promise so vitest does not warn. - resolveCreate(9999); - await flushMicrotasks(); + expect(sessions.has("aa44")).toBe(false); + expect(sessions.list()).toEqual([]); }); it("cancel for an unknown rpc_id replies with cancelled=false", async () => { diff --git a/apps/extension/src/tools/__tests__/observation.test.ts b/apps/extension/src/tools/__tests__/observation.test.ts index bfbcb2f..603a7c4 100644 --- a/apps/extension/src/tools/__tests__/observation.test.ts +++ b/apps/extension/src/tools/__tests__/observation.test.ts @@ -2457,6 +2457,64 @@ describe("handleSnapshot", () => { expect(ctx.refStore.resolve("e1")).toBeNull(); }); + it("keeps the previous RefStore when cancellation lands during DOM capture", async () => { + const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) }); + const ctx = await sm.start("aa11"); + ctx.refStore.set("e1", 999, { tabId: 4 }); + const controller = new AbortController(); + let resolveCapture: (value: unknown) => void = () => {}; + const send = vi.fn(async (_tabId: number, method: string) => { + if (method === "Accessibility.enable") return {}; + if (method === "Accessibility.getFullAXTree") { + return { + nodes: [ + { + nodeId: "new", + role: { type: "role", value: "button" }, + name: { type: "computedString", value: "New" }, + backendDOMNodeId: 123, + }, + ], + }; + } + if (method === "Page.getLayoutMetrics") { + return { cssLayoutViewport: { clientWidth: 1000, clientHeight: 800 } }; + } + if (method === "DOMSnapshot.enable") return {}; + if (method === "DOMSnapshot.captureSnapshot") { + return new Promise((resolve) => { + resolveCapture = resolve; + }); + } + throw new Error(`unexpected CDP method ${method}`); + }); + const deps = { + cdp: { + send: send as unknown as ( + tabId: number, + method: string, + params?: object, + ) => Promise, + trackSessionTab: vi.fn(), + }, + tabsApi: { + get: vi.fn(), + query: vi.fn(async () => [{ id: 4, windowId: 100, active: true } as chrome.tabs.Tab]), + }, + }; + + const pending = handleSnapshot(sm, { session_id: "aa11" }, deps, controller.signal); + await vi.waitFor(() => + expect(send).toHaveBeenCalledWith(4, "DOMSnapshot.captureSnapshot", expect.any(Object)), + ); + controller.abort(); + resolveCapture({ strings: [], documents: [] }); + + await expect(pending).resolves.toMatchObject({ code: "cancelled" }); + expect(ctx.refStore.resolve("e1", { tabId: 4 })).toBe(999); + expect(ctx.refStore.resolve("e2")).toBeNull(); + }); + it("surfaces CDP failures as cdp_failed", async () => { const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) }); await sm.start("aa11"); diff --git a/apps/extension/src/tools/__tests__/session.test.ts b/apps/extension/src/tools/__tests__/session.test.ts index adecac4..1404d75 100644 --- a/apps/extension/src/tools/__tests__/session.test.ts +++ b/apps/extension/src/tools/__tests__/session.test.ts @@ -60,6 +60,7 @@ function makeApis( }), getLastFocused: vi.fn(async () => ({ id: 500 }) as chrome.windows.Window), create: vi.fn(async () => ({ id: 999 }) as chrome.windows.Window), + remove: vi.fn(async () => {}), }; return { tabs, windows }; } diff --git a/apps/extension/src/tools/__tests__/tabs.test.ts b/apps/extension/src/tools/__tests__/tabs.test.ts index e294f34..f8fbce2 100644 --- a/apps/extension/src/tools/__tests__/tabs.test.ts +++ b/apps/extension/src/tools/__tests__/tabs.test.ts @@ -179,6 +179,7 @@ function makeWindowsApi( get: ReturnType; lastFocused: ReturnType; create: ReturnType; + remove: ReturnType; }; } { const get = vi.fn(async (windowId: number) => { @@ -197,7 +198,13 @@ function makeWindowsApi( const id = opts?.createWindowId ?? 999; return { id } as chrome.windows.Window; }); - return { api: { get, getLastFocused: lastFocused, create }, spies: { get, lastFocused, create } }; + const remove = vi.fn(async (windowId: number) => { + state.windowsClosed.add(windowId); + }); + return { + api: { get, getLastFocused: lastFocused, create, remove }, + spies: { get, lastFocused, create, remove }, + }; } describe("handleTabCreate", () => { @@ -601,6 +608,86 @@ describe("handleTabReturn", () => { expect(winSpies.create).toHaveBeenCalledOnce(); }); + it("closes a newly created fallback window when cancellation wins before the move", 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([200]), + }; + const { api, spies } = makeTabMutationApi(state); + const { api: windowsApi, spies: winSpies } = makeWindowsApi(state, { + lastFocused: 100, + createWindowId: 777, + }); + let resolveCreate: (window: chrome.windows.Window) => void = () => {}; + winSpies.create.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveCreate = resolve; + }), + ); + const controller = new AbortController(); + + const pending = handleTabReturn( + sm, + { session_id: "aa11", tab_id: 7 }, + { tabs: api, windows: windowsApi, signal: controller.signal }, + ); + await vi.waitFor(() => expect(winSpies.create).toHaveBeenCalledOnce()); + controller.abort(); + resolveCreate({ id: 777 } as chrome.windows.Window); + + await expect(pending).resolves.toMatchObject({ code: "cancelled" }); + expect(winSpies.remove).toHaveBeenCalledWith(777); + expect(spies.move).not.toHaveBeenCalled(); + expect(ctx.borrowedTabs.has(7)).toBe(true); + }); + + it("surfaces the fallback window id when cancellation cleanup fails", 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([200]), + }; + const { api, spies } = makeTabMutationApi(state); + const { api: windowsApi, spies: winSpies } = makeWindowsApi(state, { + lastFocused: 100, + createWindowId: 777, + }); + let resolveCreate: (window: chrome.windows.Window) => void = () => {}; + winSpies.create.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveCreate = resolve; + }), + ); + winSpies.remove.mockRejectedValueOnce(new Error("window removal denied")); + const controller = new AbortController(); + + const pending = handleTabReturn( + sm, + { session_id: "aa11", tab_id: 7 }, + { tabs: api, windows: windowsApi, signal: controller.signal }, + ); + await vi.waitFor(() => expect(winSpies.create).toHaveBeenCalledOnce()); + controller.abort(); + resolveCreate({ id: 777 } as chrome.windows.Window); + + await expect(pending).resolves.toMatchObject({ + code: "protocol_error", + data: { reason: "cleanup_failed", resource_type: "window", resource_id: 777 }, + message: expect.stringMatching(/fallback window 777.*window removal denied/), + }); + expect(spies.move).not.toHaveBeenCalled(); + expect(ctx.borrowedTabs.has(7)).toBe(true); + }); + it("falls back to a new window when getLastFocused only returns the Agent Window", async () => { const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) }); const ctx = await sm.start("aa11"); diff --git a/apps/extension/src/tools/console.ts b/apps/extension/src/tools/console.ts index 68cc7e4..6d95e4e 100644 --- a/apps/extension/src/tools/console.ts +++ b/apps/extension/src/tools/console.ts @@ -27,19 +27,23 @@ export async function handleConsole( manager: SessionManager, params: ConsoleParams, deps: ConsoleDeps = defaultConsoleDeps(), + signal?: AbortSignal, ): Promise { + if (signal?.aborted) return { code: "cancelled", message: "console aborted" }; const ctxOrErr = lookupSession(manager, params, "console"); if (isRpcError(ctxOrErr)) return ctxOrErr; const bounds = parseBufferedReadBounds(params); if (isRpcError(bounds)) return bounds; const target = await resolveTargetTab(manager, ctxOrErr, params.tab_id, deps.tabsApi); if (isRpcError(target)) return target; + if (signal?.aborted) return { code: "cancelled", message: "console aborted" }; if (!deps.cdp.ensureConsoleCapture || !deps.cdp.consoleEntriesSince) { return { code: "cdp_failed", message: "console capture requires CDP console support" }; } try { await deps.cdp.ensureConsoleCapture(target.tabId); + if (signal?.aborted) return { code: "cancelled", message: "console aborted" }; deps.cdp.trackSessionTab?.(ctxOrErr.sessionId, target.tabId); return deps.cdp.consoleEntriesSince( target.tabId, diff --git a/apps/extension/src/tools/dispatcher.ts b/apps/extension/src/tools/dispatcher.ts index 26ec56d..25b5f3a 100644 --- a/apps/extension/src/tools/dispatcher.ts +++ b/apps/extension/src/tools/dispatcher.ts @@ -125,11 +125,10 @@ export interface DispatcherDeps { * `AbortController` keyed by its wire `id` in * [`inflightAbortControllers`]. When the daemon pushes a `cancel` * request the dispatcher trips the matching controller; tool - * handlers that already accept a `signal` (waits, navigation, - * interaction, evaluate, tabs) react in line, and the dispatcher - * additionally races the in-flight invocation against the abort - * promise so handlers without explicit signal plumbing still respond - * promptly with `cancelled`. + * handlers observe that signal between awaited operations. The + * original RPC remains pending until its handler has stopped or + * completed compensation; only the separate cancel acknowledgement + * takes the fast path. */ export class ToolDispatcher { private readonly transport: Transport; @@ -221,7 +220,7 @@ export class ToolDispatcher { try { const sessionId = sessionIdForBrowserControlMethod(req); if (sessionId) this.onBrowserControlResumed?.(sessionId); - const result = await Promise.race([this.invoke(req, ac.signal), abortPromise(ac.signal)]); + const result = await this.invoke(req, ac.signal); if (isRpcError(result)) { body = { id: req.id, error: result }; } else { @@ -286,7 +285,7 @@ export class ToolDispatcher { private async invoke(req: RequestFrame, signal: AbortSignal): Promise { switch (req.method) { case "tool.session_start": - return handleSessionStart(this.sessions, req.params as SessionStartParams); + return handleSessionStart(this.sessions, req.params as SessionStartParams, { signal }); case "tool.session_stop": { await this.releaseHoverLatch((req.params as SessionStopParams).session_id); return handleSessionStop(this.sessions, req.params as SessionStopParams, { @@ -294,29 +293,36 @@ export class ToolDispatcher { }); } case "tool.tab_list": - return handleTabList(this.sessions, req.params as TabListParams); + return handleTabList(this.sessions, req.params as TabListParams, chromeTabsApi, signal); case "tool.tab_create": - return handleTabCreate(this.sessions, req.params as TabCreateParams); + return handleTabCreate(this.sessions, req.params as TabCreateParams, { signal }); case "tool.tab_close": - return this.withHoverReleaseForRequest(req.params as TabCloseParams, () => - handleTabClose(this.sessions, req.params as TabCloseParams), + return this.withHoverReleaseForRequest( + req.params as TabCloseParams, + () => handleTabClose(this.sessions, req.params as TabCloseParams, { signal }), + signal, ); case "tool.tab_select": - return handleTabSelect(this.sessions, req.params as TabSelectParams); + return handleTabSelect(this.sessions, req.params as TabSelectParams, { signal }); case "tool.tab_borrow": return handleTabBorrow(this.sessions, req.params as TabBorrowParams, { signal, approveBorrow: this.approveBorrow, }); case "tool.tab_return": - return handleTabReturn(this.sessions, req.params as TabReturnParams); + return handleTabReturn(this.sessions, req.params as TabReturnParams, { signal }); case "tool.window_resize": - return handleWindowResize(this.sessions, req.params as WindowResizeParams); + return handleWindowResize( + this.sessions, + req.params as WindowResizeParams, + undefined, + signal, + ); case "tool.emulate": return handleEmulate( this.sessions, req.params as EmulateParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi } : undefined, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, ); case "tool.screenshot": return handleScreenshot( @@ -325,43 +331,57 @@ export class ToolDispatcher { this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsCaptureApi, captureApi: chromeTabsCaptureApi } : undefined, + signal, ); case "tool.console": return handleConsole( this.sessions, req.params as ConsoleParams, this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi } : undefined, + signal, ); case "tool.network": return handleNetwork( this.sessions, req.params as NetworkParams, this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi } : undefined, + signal, ); case "tool.snapshot": - return this.withHoverReassert(req.params as SnapshotParams, () => - handleSnapshot( - this.sessions, - req.params as SnapshotParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsCaptureApi } : undefined, - ), + return this.withHoverReassert( + req.params as SnapshotParams, + () => + handleSnapshot( + this.sessions, + req.params as SnapshotParams, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsCaptureApi } : undefined, + signal, + ), + {}, + signal, ); case "tool.observe": { const params = req.params as ObserveParams; const hoverScope = await this.resolveHoverLatchScope(params); - return this.withHoverReassert(params, () => - handleObserve( - this.sessions, - params, - this.cdp - ? { - cdp: this.cdp, - tabsApi: chromeTabsCaptureApi, - conditionalSurfaceProbe: !this.hasHoverLatchForScope(hoverScope), - hoverProbeBypassOverlay: bypassOverlay, - } - : undefined, - ), + throwIfDispatchAborted(signal); + return this.withHoverReassert( + params, + () => + handleObserve( + this.sessions, + params, + this.cdp + ? { + cdp: this.cdp, + tabsApi: chromeTabsCaptureApi, + conditionalSurfaceProbe: !this.hasHoverLatchForScope(hoverScope), + hoverProbeBypassOverlay: bypassOverlay, + } + : undefined, + signal, + ), + {}, + signal, ); } case "tool.get_html": @@ -369,38 +389,51 @@ export class ToolDispatcher { this.sessions, req.params as GetHtmlParams, this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsCaptureApi } : undefined, + signal, ); case "tool.navigate": - return this.withHoverReleaseForRequest(req.params as NavigateParams, () => - handleNavigate( - this.sessions, - req.params as NavigateParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, - ), + return this.withHoverReleaseForRequest( + req.params as NavigateParams, + () => + handleNavigate( + this.sessions, + req.params as NavigateParams, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, + ), + signal, ); case "tool.navigate_back": - return this.withHoverReleaseForRequest(req.params as NavigateBackParams, () => - handleNavigateBack( - this.sessions, - req.params as NavigateBackParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, - ), + return this.withHoverReleaseForRequest( + req.params as NavigateBackParams, + () => + handleNavigateBack( + this.sessions, + req.params as NavigateBackParams, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, + ), + signal, ); case "tool.navigate_forward": - return this.withHoverReleaseForRequest(req.params as NavigateForwardParams, () => - handleNavigateForward( - this.sessions, - req.params as NavigateForwardParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, - ), + return this.withHoverReleaseForRequest( + req.params as NavigateForwardParams, + () => + handleNavigateForward( + this.sessions, + req.params as NavigateForwardParams, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, + ), + signal, ); case "tool.reload": - return this.withHoverReleaseForRequest(req.params as ReloadParams, () => - handleReload( - this.sessions, - req.params as ReloadParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, - ), + return this.withHoverReleaseForRequest( + req.params as ReloadParams, + () => + handleReload( + this.sessions, + req.params as ReloadParams, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, + ), + signal, ); case "tool.click": return this.withHoverReassert( @@ -419,6 +452,7 @@ export class ToolDispatcher { : undefined, ), { releaseAfter: true }, + signal, ); case "tool.hover": { const result = await handleHover( @@ -438,28 +472,37 @@ export class ToolDispatcher { return this.rememberHover((req.params as HoverParams).session_id, result); } case "tool.fill": - return this.withHoverReleaseForRequest(req.params as FillParams, () => - handleFill( - this.sessions, - req.params as FillParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, - ), + return this.withHoverReleaseForRequest( + req.params as FillParams, + () => + handleFill( + this.sessions, + req.params as FillParams, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, + ), + signal, ); case "tool.press": - return this.withHoverReleaseForRequest(req.params as PressParams, () => - handlePress( - this.sessions, - req.params as PressParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, - ), + return this.withHoverReleaseForRequest( + req.params as PressParams, + () => + handlePress( + this.sessions, + req.params as PressParams, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, + ), + signal, ); case "tool.select": - return this.withHoverReleaseForRequest(req.params as SelectParams, () => - handleSelect( - this.sessions, - req.params as SelectParams, - this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, - ), + return this.withHoverReleaseForRequest( + req.params as SelectParams, + () => + handleSelect( + this.sessions, + req.params as SelectParams, + this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsApi, signal } : undefined, + ), + signal, ); case "tool.evaluate": return handleEvaluate( @@ -566,9 +609,13 @@ export class ToolDispatcher { params: { session_id: string; tab_id?: number }, work: () => Promise, options: { releaseAfter?: boolean } = {}, + signal?: AbortSignal, ): Promise { + throwIfDispatchAborted(signal); const scope = await this.resolveHoverLatchScope(params); + throwIfDispatchAborted(signal); await this.reassertHover(scope); + throwIfDispatchAborted(signal); try { return await work(); } finally { @@ -581,9 +628,13 @@ export class ToolDispatcher { private async withHoverReleaseForRequest( params: { session_id: string; tab_id?: number }, work: () => Promise, + signal?: AbortSignal, ): Promise { + throwIfDispatchAborted(signal); const scope = await this.resolveHoverLatchScope(params); + throwIfDispatchAborted(signal); await this.releaseHoverLatch(scope.session_id, scope.tab_id); + throwIfDispatchAborted(signal); return work(); } @@ -693,41 +744,14 @@ async function bypassOverlay(tabId: number, enabled: boolean): Promise { } } -/** - * Resolves never; rejects with `AbortLikeError` as soon as the signal - * fires (or immediately if it is already aborted). Used by the - * dispatcher to race the tool invocation so handlers without explicit - * signal plumbing still surface a `cancelled` reply promptly. - */ -function abortPromise(signal: AbortSignal): Promise { - return new Promise((_, reject) => { - if (signal.aborted) { - reject(new AbortLikeError()); - return; - } - signal.addEventListener( - "abort", - () => { - reject(new AbortLikeError()); - }, - { once: true }, - ); - }); -} - -/** - * Sentinel error class so [`isAbortLikeError`] can recognise our own - * race-rejection without confusing it with a real CDP failure. - */ -class AbortLikeError extends Error { - constructor() { - super("rpc aborted by daemon cancel"); - this.name = "BhAbortError"; - } +function throwIfDispatchAborted(signal: AbortSignal | undefined): void { + if (!signal?.aborted) return; + const error = new Error("rpc aborted by daemon cancel"); + error.name = "AbortError"; + throw error; } function isAbortLikeError(err: unknown): boolean { - if (err instanceof AbortLikeError) return true; if (err instanceof DOMException && err.name === "AbortError") return true; if (typeof err === "object" && err !== null && (err as { name?: string }).name === "AbortError") { return true; diff --git a/apps/extension/src/tools/emulate.ts b/apps/extension/src/tools/emulate.ts index 06e19b6..44a30a8 100644 --- a/apps/extension/src/tools/emulate.ts +++ b/apps/extension/src/tools/emulate.ts @@ -44,6 +44,7 @@ export interface EmulateCdpRunner { export interface EmulateDeps { cdp: EmulateCdpRunner; tabsApi: ChromeTabsApi; + signal?: AbortSignal; } function defaultEmulateDeps(): EmulateDeps { @@ -210,11 +211,13 @@ export async function handleEmulate( params: EmulateParams, deps: EmulateDeps = defaultEmulateDeps(), ): Promise { + if (deps.signal?.aborted) return { code: "cancelled", message: "emulate aborted" }; const ctxOrErr = lookupSession(manager, params, "emulate"); if (isRpcError(ctxOrErr)) return ctxOrErr; const ctx = ctxOrErr; const target = await resolveTargetTab(manager, ctx, params?.tab_id, deps.tabsApi); if (isRpcError(target)) return target; + if (deps.signal?.aborted) return { code: "cancelled", message: "emulate aborted" }; const denied = enforceAgentWindow(ctx, target, "emulate"); if (denied) return denied; @@ -223,9 +226,16 @@ export async function handleEmulate( return invalidParams("emulate off cannot be combined with overrides"); } try { + if (deps.signal?.aborted) return { code: "cancelled", message: "emulate aborted" }; deps.cdp.trackSessionTab?.(ctx.sessionId, target.tabId); await deps.cdp.clearDeviceMetricsOverride(target.tabId); + if (deps.signal?.aborted) { + return { code: "cancelled", message: "emulate aborted after clearing device metrics" }; + } await deps.cdp.setTouchEmulationEnabled(target.tabId, false); + if (deps.signal?.aborted) { + return { code: "cancelled", message: "emulate aborted after clearing touch emulation" }; + } // An empty userAgent string clears the override (CDP convention). await deps.cdp.setUserAgentOverride(target.tabId, { userAgent: "" }); } catch (err) { @@ -246,6 +256,7 @@ export async function handleEmulate( // values; the merged state is what gets applied (and echoed back). const merged = mergeEmulateOverrides(tabEmulationStates.get(target.tabId), overrides); try { + if (deps.signal?.aborted) return { code: "cancelled", message: "emulate aborted" }; deps.cdp.trackSessionTab?.(ctx.sessionId, target.tabId); if (merged.width !== undefined && merged.height !== undefined) { await deps.cdp.setDeviceMetricsOverride(target.tabId, { @@ -254,6 +265,9 @@ export async function handleEmulate( deviceScaleFactor: merged.device_scale_factor ?? 0, mobile: merged.mobile ?? false, }); + if (deps.signal?.aborted) { + return { code: "cancelled", message: "emulate aborted after applying device metrics" }; + } } if (merged.user_agent !== undefined) { await deps.cdp.setUserAgentOverride(target.tabId, { @@ -263,6 +277,9 @@ export async function handleEmulate( ? { userAgentMetadata: toCdpUserAgentMetadata(merged.user_agent_metadata) } : {}), }); + if (deps.signal?.aborted) { + return { code: "cancelled", message: "emulate aborted after applying user agent" }; + } } if (merged.touch !== undefined || merged.max_touch_points !== undefined) { await deps.cdp.setTouchEmulationEnabled( diff --git a/apps/extension/src/tools/network.ts b/apps/extension/src/tools/network.ts index 850f1b6..ea317ed 100644 --- a/apps/extension/src/tools/network.ts +++ b/apps/extension/src/tools/network.ts @@ -37,17 +37,22 @@ export async function handleNetwork( manager: SessionManager, params: NetworkParams, deps: NetworkDeps = defaultNetworkDeps(), + signal?: AbortSignal, ): Promise { + if (signal?.aborted) return { code: "cancelled", message: "network aborted" }; const ctxOrErr = lookupSession(manager, params, "network"); if (isRpcError(ctxOrErr)) return ctxOrErr; const bounds = parseBufferedReadBounds(params); if (isRpcError(bounds)) return bounds; const target = await resolveTargetTab(manager, ctxOrErr, params.tab_id, deps.tabsApi); if (isRpcError(target)) return target; + if (signal?.aborted) return { code: "cancelled", message: "network aborted" }; try { + if (signal?.aborted) return { code: "cancelled", message: "network aborted" }; deps.cdp.trackSessionTab?.(ctxOrErr.sessionId, target.tabId); await deps.cdp.ensureNetworkCapture(target.tabId); + if (signal?.aborted) return { code: "cancelled", message: "network aborted" }; return deps.cdp.networkEntriesSince( target.tabId, bounds.since, diff --git a/apps/extension/src/tools/observation.ts b/apps/extension/src/tools/observation.ts index 34c60f1..12c3efc 100644 --- a/apps/extension/src/tools/observation.ts +++ b/apps/extension/src/tools/observation.ts @@ -137,16 +137,42 @@ function defaultScreenshotDeps(): ScreenshotDeps { }; } +function cancelled(tool: string): RpcError { + return { code: "cancelled", message: `${tool} aborted` }; +} + +function abortError(tool: string): Error { + const error = new Error(`${tool} aborted`); + error.name = "AbortError"; + return error; +} + +function isAbortError(error: unknown): boolean { + return ( + typeof error === "object" && + error !== null && + (error as { name?: string }).name === "AbortError" + ); +} + +function throwIfAborted(signal: AbortSignal | undefined, tool: string): void { + if (signal?.aborted) throw abortError(tool); +} + async function captureElementScreenshot( cdp: SharedCdpRunner, tabId: number, backendNodeId: number, + signal?: AbortSignal, ): Promise<{ image_base64: string; width: number; height: number } | RpcError> { + if (signal?.aborted) return cancelled("screenshot"); const scrollErr = await scrollNodeIntoView(cdp, tabId, backendNodeId); if (scrollErr) return scrollErr; + if (signal?.aborted) return cancelled("screenshot"); const rectOrErr = await nodeBoundingRect(cdp, tabId, backendNodeId); if (isRpcError(rectOrErr)) return rectOrErr; + if (signal?.aborted) return cancelled("screenshot"); try { const shot = await cdp.send<{ data?: string }>(tabId, "Page.captureScreenshot", { @@ -159,6 +185,7 @@ async function captureElementScreenshot( scale: 1, }, }); + if (signal?.aborted) return cancelled("screenshot"); const image_base64 = shot.data ?? ""; if (!image_base64) { return { code: "cdp_failed", message: "Page.captureScreenshot returned no data" }; @@ -193,11 +220,15 @@ async function captureFullTabPng( deps: ScreenshotDeps, ctx: SessionContext, target: ResolvedTargetTab, + signal?: AbortSignal, ): Promise { + if (signal?.aborted) return cancelled("screenshot"); try { const dataUrl = await deps.captureApi.captureVisibleTab(target.windowId, { format: "png" }); + if (signal?.aborted) return cancelled("screenshot"); return stripDataUrlPrefix(dataUrl); } catch (primaryErr) { + if (signal?.aborted) return cancelled("screenshot"); const cdp = deps.cdp; if (!cdp) { return { @@ -207,12 +238,15 @@ async function captureFullTabPng( } let fallbackMsg: string; try { + if (signal?.aborted) return cancelled("screenshot"); cdp.trackSessionTab?.(ctx.sessionId, target.tabId); await cdp.ensureAttachedToUrl?.(target.tabId, target.url); + if (signal?.aborted) return cancelled("screenshot"); const shot = await cdp.send<{ data?: string }>(target.tabId, "Page.captureScreenshot", { format: "png", fromSurface: true, }); + if (signal?.aborted) return cancelled("screenshot"); if (shot.data) return shot.data; fallbackMsg = "Page.captureScreenshot returned no data"; } catch (fallbackErr) { @@ -231,7 +265,9 @@ export async function handleScreenshot( manager: SessionManager, params: ScreenshotParams, deps: ScreenshotDeps = defaultScreenshotDeps(), + signal?: AbortSignal, ): Promise { + if (signal?.aborted) return cancelled("screenshot"); const ctxOrErr = lookupSession(manager, params, "screenshot"); if (isRpcError(ctxOrErr)) return ctxOrErr; const ctx = ctxOrErr; @@ -243,6 +279,7 @@ export async function handleScreenshot( "screenshot", ); if (isRpcError(target)) return target; + if (signal?.aborted) return cancelled("screenshot"); const dialogCursor = deps.cdp ? markDialogCursor(deps.cdp, target.tabId) : 0; const withShotDialogs = (result: T) => deps.cdp ? attachDialogs(deps.cdp, target.tabId, dialogCursor, result) : result; @@ -254,15 +291,18 @@ export async function handleScreenshot( } const node = resolveSnapshotRef(ctx, ref, target.tabId); if (isRpcError(node)) return node; + if (signal?.aborted) return cancelled("screenshot"); deps.cdp.trackSessionTab?.(ctx.sessionId, target.tabId); await deps.cdp.ensureAttachedToUrl?.(target.tabId, target.url); + if (signal?.aborted) return cancelled("screenshot"); const cdp = deps.cdp; const captured = await withOverlaysHiddenForCapture( target.tabId, - () => captureElementScreenshot(cdp, target.tabId, node.backendNodeId), + () => captureElementScreenshot(cdp, target.tabId, node.backendNodeId, signal), deps.sendToTab, ); if (isRpcError(captured)) return captured; + if (signal?.aborted) return cancelled("screenshot"); return withShotDialogs({ image_base64: captured.image_base64, width: captured.width, @@ -282,10 +322,11 @@ export async function handleScreenshot( const captured = await withOverlaysHiddenForCapture( target.tabId, - () => captureFullTabPng(deps, ctx, target), + () => captureFullTabPng(deps, ctx, target, signal), deps.sendToTab, ); if (isRpcError(captured)) return captured; + if (signal?.aborted) return cancelled("screenshot"); const image_base64 = captured; const dims = parsePngDimensions(image_base64) ?? { width: 0, height: 0 }; return withShotDialogs({ @@ -1252,7 +1293,9 @@ export async function handleGetHtml( manager: SessionManager, params: GetHtmlParams, deps: SnapshotDeps = getDefaultDeps(), + signal?: AbortSignal, ): Promise { + if (signal?.aborted) return cancelled("get_html"); const ctxOrErr = lookupSession(manager, params, "get_html"); if (isRpcError(ctxOrErr)) return ctxOrErr; const ctx = ctxOrErr; @@ -1264,14 +1307,17 @@ export async function handleGetHtml( "get_html", ); if (isRpcError(target)) return target; + if (signal?.aborted) return cancelled("get_html"); const dialogCursor = markDialogCursor(deps.cdp, target.tabId); const maxBytes = params.max_bytes && params.max_bytes > 0 ? params.max_bytes : DEFAULT_GET_HTML_MAX_BYTES; try { + throwIfAborted(signal, "get_html"); deps.cdp.trackSessionTab?.(ctx.sessionId, target.tabId); await deps.cdp.ensureAttachedToUrl?.(target.tabId, target.url); + throwIfAborted(signal, "get_html"); let html: string; if (params.ref) { const resolved = resolveSnapshotRef(ctx, params.ref, target.tabId); @@ -1279,6 +1325,7 @@ export async function handleGetHtml( const resp = await deps.cdp.send<{ outerHTML?: string }>(target.tabId, "DOM.getOuterHTML", { backendNodeId: resolved.backendNodeId, }); + throwIfAborted(signal, "get_html"); html = resp.outerHTML ?? ""; } else { const doc = await deps.cdp.send<{ root?: { nodeId?: number } }>( @@ -1286,6 +1333,7 @@ export async function handleGetHtml( "DOM.getDocument", { depth: 0 }, ); + throwIfAborted(signal, "get_html"); const nodeId = doc.root?.nodeId; if (typeof nodeId !== "number") { return { @@ -1296,6 +1344,7 @@ export async function handleGetHtml( const resp = await deps.cdp.send<{ outerHTML?: string }>(target.tabId, "DOM.getOuterHTML", { nodeId, }); + throwIfAborted(signal, "get_html"); html = resp.outerHTML ?? ""; } const originalBytes = utf8ByteLength(html); @@ -1307,6 +1356,7 @@ export async function handleGetHtml( tab_id: target.tabId, }); } catch (err) { + if (isAbortError(err)) return cancelled("get_html"); return { code: "cdp_failed", message: err instanceof Error ? err.message : String(err), @@ -1326,19 +1376,24 @@ interface LayoutMetricsViewportReply { async function fallbackCapturedViewModel( cdp: CdpRunner, tabId: number, + signal?: AbortSignal, ): Promise { + throwIfAborted(signal, "observation"); let viewport = { width: 0, height: 0 }; try { const metrics = await cdp.send(tabId, "Page.getLayoutMetrics", {}); + throwIfAborted(signal, "observation"); const vpSrc = metrics.cssLayoutViewport ?? metrics.layoutViewport ?? {}; viewport = { width: vpSrc.clientWidth ?? 0, height: vpSrc.clientHeight ?? 0, }; - } catch { + } catch (err) { + if (isAbortError(err)) throw err; // viewport stays zero-sized } - const excludedBackendNodeIds = await collectOverlayExcludedBackendIds(cdp, tabId); + const excludedBackendNodeIds = await collectOverlayExcludedBackendIds(cdp, tabId, signal); + throwIfAborted(signal, "observation"); return { ...emptyCapturedViewModel(viewport), excludedBackendNodeIds }; } @@ -1347,14 +1402,17 @@ async function captureForVom( tabId: number, conditionalSurfaceProbe: boolean, hoverProbeBypassOverlay?: (tabId: number, enabled: boolean) => Promise, + signal?: AbortSignal, ): Promise { try { return await captureViewModel(cdp, tabId, { conditionalSurfaceProbe, hoverProbeBypassOverlay, + signal, }); - } catch { - return fallbackCapturedViewModel(cdp, tabId); + } catch (err) { + if (isAbortError(err)) throw err; + return fallbackCapturedViewModel(cdp, tabId, signal); } } @@ -1365,7 +1423,9 @@ async function handleVomObservation( effect: ToolEffect, conditionalSurfaceProbe: boolean, deps: SnapshotDeps = getDefaultDeps(), + signal?: AbortSignal, ): Promise { + if (signal?.aborted) return cancelled(toolName); const ctxOrErr = lookupSession(manager, params, toolName); if (isRpcError(ctxOrErr)) return ctxOrErr; const ctx = ctxOrErr; @@ -1377,19 +1437,24 @@ async function handleVomObservation( toolName, ); if (isRpcError(target)) return target; + if (signal?.aborted) return cancelled(toolName); const denied = enforceToolTargetScope(ctx, target, effect, toolName); if (denied) return denied; const dialogCursor = markDialogCursor(deps.cdp, target.tabId); try { + throwIfAborted(signal, toolName); deps.cdp.trackSessionTab?.(ctx.sessionId, target.tabId); await deps.cdp.ensureAttachedToUrl?.(target.tabId, target.url); + throwIfAborted(signal, toolName); await deps.cdp.send(target.tabId, "Accessibility.enable", {}); + throwIfAborted(signal, toolName); const result = await deps.cdp.send<{ nodes: CdpAxNode[] }>( target.tabId, "Accessibility.getFullAXTree", {}, ); + throwIfAborted(signal, toolName); const axNodes = result.nodes ?? []; const effectiveConditionalSurfaceProbe = deps.conditionalSurfaceProbe ?? conditionalSurfaceProbe; @@ -1398,13 +1463,16 @@ async function handleVomObservation( target.tabId, effectiveConditionalSurfaceProbe, deps.hoverProbeBypassOverlay, + signal, ); + throwIfAborted(signal, toolName); const scene = buildVomScene(axNodes, captured, { pageUrl: target.url }); const rendered = renderVom(scene, { maxDepth: params.max_depth, maxTokens: params.max_tokens, activeRegionPolicy: true, }); + throwIfAborted(signal, toolName); ctx.refStore.replace( rendered.refs.map( (r) => [r.ref, { backendNodeId: r.backendNodeId, tabId: target.tabId }] as const, @@ -1430,6 +1498,7 @@ async function handleVomObservation( : {}), }); } catch (err) { + if (isAbortError(err)) return cancelled(toolName); return { code: "cdp_failed", message: err instanceof Error ? err.message : String(err), @@ -1441,14 +1510,16 @@ export async function handleSnapshot( manager: SessionManager, params: SnapshotParams, deps: SnapshotDeps = getDefaultDeps(), + signal?: AbortSignal, ): Promise { - return handleVomObservation(manager, params, "snapshot", "passive_read", false, deps); + return handleVomObservation(manager, params, "snapshot", "passive_read", false, deps, signal); } export async function handleObserve( manager: SessionManager, params: ObserveParams, deps: SnapshotDeps = getDefaultDeps(), + signal?: AbortSignal, ): Promise { - return handleVomObservation(manager, params, "observe", "transient_input", true, deps); + return handleVomObservation(manager, params, "observe", "transient_input", true, deps, signal); } diff --git a/apps/extension/src/tools/session.ts b/apps/extension/src/tools/session.ts index edf564e..138b6f6 100644 --- a/apps/extension/src/tools/session.ts +++ b/apps/extension/src/tools/session.ts @@ -1,5 +1,6 @@ -import type { SessionManager } from "@/session-manager/manager"; +import { type SessionManager, SessionStartCleanupError } from "@/session-manager/manager"; import type { RpcError } from "@/transport/types"; +import { rpcError } from "./errors"; import { clearRecordingForSession } from "./record"; import { isRpcError } from "./shared"; import { returnBorrowedTab, type TabManagementDeps } from "./tabs"; @@ -58,6 +59,10 @@ export interface SessionStartResult { agent_window_id?: number; } +export interface SessionStartDeps { + signal?: AbortSignal; +} + export interface SessionStopParams { session_id: string; } @@ -91,6 +96,7 @@ export interface SessionStopDeps { export async function handleSessionStart( manager: SessionManager, params: SessionStartParams, + deps: SessionStartDeps = {}, ): Promise { if (!params?.session_id) { return { @@ -104,9 +110,23 @@ export async function handleSessionStart( const ctx = await manager.start(params.session_id, { size: sizeOrErr, focused: params.focused, + signal: deps.signal, }); return { agent_window_id: ctx.agentWindowId }; } catch (err) { + if (err instanceof SessionStartCleanupError) { + return rpcError("protocol_error", "cleanup_failed", err.message, { + resource_type: "agent_window", + resource_id: err.windowId, + }); + } + if ( + typeof err === "object" && + err !== null && + (err as { name?: string }).name === "AbortError" + ) { + return { code: "cancelled", message: "session_start aborted" }; + } // chrome.windows.create / SessionManager failures are not CDP // failures (§4.5 reserves cdp_failed for raw CDP errors). Surface // them as protocol_error so the CLI maps to the right exit code diff --git a/apps/extension/src/tools/tabs.ts b/apps/extension/src/tools/tabs.ts index 6bad610..04e46b8 100644 --- a/apps/extension/src/tools/tabs.ts +++ b/apps/extension/src/tools/tabs.ts @@ -157,6 +157,7 @@ export interface ChromeWindowsApi { get(windowId: number): Promise; getLastFocused(filters?: chrome.windows.QueryOptions): Promise; create(props: chrome.windows.CreateData): Promise; + remove(windowId: number): Promise; } export const chromeWindowsApi: ChromeWindowsApi = { @@ -164,6 +165,7 @@ export const chromeWindowsApi: ChromeWindowsApi = { getLastFocused: (filters) => filters ? chrome.windows.getLastFocused(filters) : chrome.windows.getLastFocused(), create: (p) => chrome.windows.create(p), + remove: (id) => chrome.windows.remove(id), }; export interface AgentOverlayResetApi { @@ -215,7 +217,9 @@ export async function handleTabList( manager: SessionManager, params: TabListParams, api: ChromeTabsApi = chromeTabsApi, + signal?: AbortSignal, ): Promise { + if (signal?.aborted) return { code: "cancelled", message: "tab_list aborted" }; if (!params || typeof params.session_id !== "string" || params.session_id.length === 0) { return { code: "invalid_params", @@ -249,6 +253,7 @@ export async function handleTabList( const myAgentWindowId = ctx.agentWindowId; const allTabs = await api.query({}); + if (signal?.aborted) return { code: "cancelled", message: "tab_list aborted" }; const tabs: TabInfo[] = []; for (const t of allTabs) { if (typeof t.id !== "number") continue; @@ -382,7 +387,12 @@ async function createTabAndCleanup( try { await getTabsApi(deps).remove(tab.id); } catch (cleanupErr) { - console.debug("[bsk tab_create] cleanup after abort failed", cleanupErr); + return rpcError( + "protocol_error", + "cleanup_failed", + `tab_create aborted but cleanup of tab ${tab.id} failed: ${describeError(cleanupErr)}`, + { resource_type: "tab", resource_id: tab.id }, + ); } } return { code: "cancelled", message: "tab_create aborted" }; @@ -693,7 +703,17 @@ async function moveTabForBorrow( try { await tabsApi.move(tabId, { windowId: originalWindowId, index: originalIndex }); } catch (rollbackErr) { - console.debug("[bsk tab_borrow] rollback move failed", rollbackErr); + return rpcError( + "protocol_error", + "cleanup_failed", + `tab_borrow aborted but rollback of tab ${tabId} to window ${originalWindowId} failed: ${describeError(rollbackErr)}`, + { + resource_type: "tab", + resource_id: tabId, + original_window_id: originalWindowId, + original_index: originalIndex, + }, + ); } return { code: "cancelled", message: "tab_borrow aborted" }; } @@ -825,6 +845,12 @@ export interface ReturnOutcome { fallback: boolean; } +interface FallbackWindowTarget { + windowId: number; + index: number; + created: boolean; +} + function describeError(err: unknown): string { return err instanceof Error ? err.message : String(err); } @@ -833,10 +859,13 @@ async function chooseFallbackWindow( ctx: SessionContext, windowsApi: ChromeWindowsApi, isAgentWindowId: (windowId: number) => boolean, -): Promise<{ windowId: number; index: number } | RpcError> { + signal?: AbortSignal, +): Promise { let lastFocusedError: unknown; try { const last = await windowsApi.getLastFocused({ windowTypes: ["normal"] }); + const cancelled = aborted(signal, "tab_return"); + if (cancelled) return cancelled; const lastId = typeof last?.id === "number" ? last.id : null; // Never relocate a user tab into *any* session's Agent Window: this // session's is excluded explicitly, and other sessions' are excluded @@ -846,12 +875,15 @@ async function chooseFallbackWindow( // session write to it and, worse, see it destroyed when that session // stops and closes its window. if (lastId !== null && lastId !== ctx.agentWindowId && !isAgentWindowId(lastId)) { - return { windowId: lastId, index: -1 }; + return { windowId: lastId, index: -1, created: false }; } } catch (err) { lastFocusedError = err; } + const cancelled = aborted(signal, "tab_return"); + if (cancelled) return cancelled; + try { const created = await windowsApi.create({ url: NEW_TAB_DEFAULT_URL, @@ -863,7 +895,20 @@ async function chooseFallbackWindow( message: "tab_return: failed to create fallback window", }; } - return { windowId: created.id, index: 0 }; + if (aborted(signal, "tab_return")) { + try { + await windowsApi.remove(created.id); + } catch (cleanupErr) { + return rpcError( + "protocol_error", + "cleanup_failed", + `tab_return aborted but cleanup of fallback window ${created.id} failed: ${describeError(cleanupErr)}`, + { resource_type: "window", resource_id: created.id }, + ); + } + return { code: "cancelled", message: "tab_return aborted" }; + } + return { windowId: created.id, index: 0, created: true }; } catch (err) { const suffix = lastFocusedError ? ` (after getLastFocused failed: ${describeError(lastFocusedError)})` @@ -875,6 +920,25 @@ async function chooseFallbackWindow( } } +async function cleanupUnusedFallbackWindow( + windowsApi: ChromeWindowsApi, + target: FallbackWindowTarget, + reason: string, +): Promise { + if (!target.created) return null; + try { + await windowsApi.remove(target.windowId); + return null; + } catch (cleanupErr) { + return rpcError( + "protocol_error", + "cleanup_failed", + `${reason}; cleanup of fallback window ${target.windowId} failed: ${describeError(cleanupErr)}`, + { resource_type: "window", resource_id: target.windowId }, + ); + } +} + function resetAgentOverlaysInReturnedTab( ctx: SessionContext, tabId: number, @@ -911,10 +975,13 @@ export async function returnBorrowedTab( const tabsApi = getTabsApi(deps); const windowsApi = getWindowsApi(deps); const isAgentWindowId = getIsAgentWindowId(deps); + const alreadyCancelled = aborted(deps.signal, "tab_return"); + if (alreadyCancelled) return alreadyCancelled; let targetWindowId = entry.originalWindowId; let targetIndex = entry.originalIndex; let fallback = false; + let fallbackTarget: FallbackWindowTarget | null = null; // Check the original window is still around. let originalAlive = true; @@ -924,14 +991,30 @@ export async function returnBorrowedTab( console.debug("[bsk tab_return] original window gone, falling back", err); originalAlive = false; } + const cancelledAfterLookup = aborted(deps.signal, "tab_return"); + if (cancelledAfterLookup) return cancelledAfterLookup; if (!originalAlive) { fallback = true; - const target = await chooseFallbackWindow(ctx, windowsApi, isAgentWindowId); + const target = await chooseFallbackWindow(ctx, windowsApi, isAgentWindowId, deps.signal); if ("code" in target) return target; + fallbackTarget = target; targetWindowId = target.windowId; targetIndex = target.index; } + const cancelledBeforeMove = aborted(deps.signal, "tab_return"); + if (cancelledBeforeMove) { + if (fallbackTarget) { + const cleanupError = await cleanupUnusedFallbackWindow( + windowsApi, + fallbackTarget, + "tab_return aborted before moving the borrowed tab", + ); + if (cleanupError) return cleanupError; + } + return cancelledBeforeMove; + } + try { const moved = await tabsApi.move(tabId, { windowId: targetWindowId, @@ -947,14 +1030,33 @@ export async function returnBorrowedTab( fallback, }; } catch (err) { + if (fallbackTarget) { + const cleanupError = await cleanupUnusedFallbackWindow( + windowsApi, + fallbackTarget, + `tab_return could not move tab ${tabId}`, + ); + if (cleanupError) return cleanupError; + } + const cancelledAfterMoveFailure = aborted(deps.signal, "tab_return"); + if (cancelledAfterMoveFailure) return cancelledAfterMoveFailure; if (!fallback) { - const target = await chooseFallbackWindow(ctx, windowsApi, isAgentWindowId); + const target = await chooseFallbackWindow(ctx, windowsApi, isAgentWindowId, deps.signal); if ("code" in target) { return { code: "cdp_failed", message: `tab_return: chrome.tabs.move failed: ${describeError(err)}; fallback failed: ${target.message}`, }; } + const cancelledBeforeFallbackMove = aborted(deps.signal, "tab_return"); + if (cancelledBeforeFallbackMove) { + const cleanupError = await cleanupUnusedFallbackWindow( + windowsApi, + target, + "tab_return aborted before fallback move", + ); + return cleanupError ?? cancelledBeforeFallbackMove; + } try { const moved = await tabsApi.move(tabId, { windowId: target.windowId, @@ -970,6 +1072,12 @@ export async function returnBorrowedTab( fallback: true, }; } catch (fallbackErr) { + const cleanupError = await cleanupUnusedFallbackWindow( + windowsApi, + target, + `tab_return fallback move for tab ${tabId} failed`, + ); + if (cleanupError) return cleanupError; return { code: "cdp_failed", message: `tab_return: chrome.tabs.move failed: ${describeError(err)}; fallback move failed: ${describeError(fallbackErr)}`, diff --git a/apps/extension/src/tools/vom/capture.ts b/apps/extension/src/tools/vom/capture.ts index 0f56577..af78382 100644 --- a/apps/extension/src/tools/vom/capture.ts +++ b/apps/extension/src/tools/vom/capture.ts @@ -65,6 +65,25 @@ export interface CapturedViewModel { export interface CaptureViewModelOptions { conditionalSurfaceProbe?: boolean; hoverProbeBypassOverlay?: (tabId: number, enabled: boolean) => Promise; + signal?: AbortSignal; +} + +function captureAbortError(): Error { + const error = new Error("observation aborted"); + error.name = "AbortError"; + return error; +} + +function isAbortError(error: unknown): boolean { + return ( + typeof error === "object" && + error !== null && + (error as { name?: string }).name === "AbortError" + ); +} + +function throwIfAborted(signal: AbortSignal | undefined): void { + if (signal?.aborted) throw captureAbortError(); } /** Sparse array format Chrome uses for infrequently-set per-node fields. */ @@ -203,21 +222,25 @@ async function enrichFormControlStates( cdp: CdpRunner, tabId: number, frameNodeGroups: CapturedNode[][], + signal?: AbortSignal, ): Promise { const hasControls = frameNodeGroups.some((nodes) => nodes.some((node) => isFormControlTag(node.tag)), ); if (!hasControls) return; try { + throwIfAborted(signal); const result = await cdp.send(tabId, "Runtime.evaluate", { expression: formStateBatchExpression(MAX_FORM_ENRICH_CONTROLS), returnByValue: true, }); + throwIfAborted(signal); const frameStates = flattenFrameFormStates(runtimeValue(result)); for (let i = 0; i < frameNodeGroups.length; i += 1) { applyFormStates(frameNodeGroups[i], frameStates[i] ?? []); } - } catch { + } catch (error) { + if (isAbortError(error)) throw error; // Best-effort enrichment. DOMSnapshot/AX data still carries the nodes. } } @@ -425,8 +448,20 @@ function hoverStateExpression(): string { })()`; } -async function wait(ms: number): Promise { - await new Promise((resolve) => setTimeout(resolve, ms)); +async function wait(ms: number, signal?: AbortSignal): Promise { + throwIfAborted(signal); + await new Promise((resolve, reject) => { + const timer = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(); + }, ms); + const onAbort = () => { + clearTimeout(timer); + signal?.removeEventListener("abort", onAbort); + reject(captureAbortError()); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + }); } async function clearHover(cdp: CdpRunner, tabId: number): Promise { @@ -598,41 +633,51 @@ async function probeHoverSurfaces( ): Promise { const started = Date.now(); try { + throwIfAborted(options.signal); const cssScan = await cdp.send(tabId, "Runtime.evaluate", { expression: hoverCssTriggerScanExpression(), returnByValue: true, }); + throwIfAborted(options.signal); const cssHoverPoints = runtimeValue>(cssScan) ?? []; const candidates = buildHoverCandidates(nodes, cssHoverPoints); if (candidates.length === 0) return []; const results: CapturedSurfaceProbe[] = []; const seen = new Set(); + throwIfAborted(options.signal); await options.hoverProbeBypassOverlay?.(tabId, true).catch(() => undefined); try { + throwIfAborted(options.signal); for (const candidate of candidates.slice(0, MAX_HOVER_TRIGGERS)) { + throwIfAborted(options.signal); if (Date.now() - started > MAX_HOVER_PROBE_MS) break; if (results.length >= MAX_HOVER_SURFACES) break; if (seen.has(candidate.backendNodeId)) continue; try { await clearHover(cdp, tabId); - await wait(HOVER_SETTLE_MS); + throwIfAborted(options.signal); + await wait(HOVER_SETTLE_MS, options.signal); const baselineReply = await cdp.send(tabId, "Runtime.evaluate", { expression: hoverStateExpression(), returnByValue: true, }); + throwIfAborted(options.signal); const baselineItems = runtimeValue(baselineReply) ?? []; + throwIfAborted(options.signal); await cdp.send(tabId, "Input.dispatchMouseEvent", { type: "mouseMoved", x: candidate.x, y: candidate.y, }); - await wait(HOVER_SETTLE_MS); + throwIfAborted(options.signal); + await wait(HOVER_SETTLE_MS, options.signal); const collected = await cdp.send(tabId, "Runtime.evaluate", { expression: hoverStateExpression(), returnByValue: true, }); + throwIfAborted(options.signal); const subItems = diffHoverItems( baselineItems, runtimeValue(collected) ?? [], @@ -646,7 +691,8 @@ async function probeHoverSurfaces( subItems, confidence: confidenceForHover(candidate, subItems), }); - } catch { + } catch (error) { + if (isAbortError(error)) throw error; continue; } finally { await clearHover(cdp, tabId); @@ -657,6 +703,7 @@ async function probeHoverSurfaces( } return results; } catch (err) { + if (isAbortError(err)) throw err; console.debug("[bsk capture] hover surface probe failed", err); return []; } @@ -669,13 +716,16 @@ async function probeHoverSurfaces( export async function collectOverlayExcludedBackendIds( cdp: CdpRunner, tabId: number, + signal?: AbortSignal, ): Promise> { const excluded = new Set(); try { + throwIfAborted(signal); const doc = await cdp.send<{ root?: { nodeId?: number } }>(tabId, "DOM.getDocument", { depth: 0, pierce: true, }); + throwIfAborted(signal); const rootNodeId = doc.root?.nodeId; if (typeof rootNodeId !== "number") return excluded; @@ -683,6 +733,7 @@ export async function collectOverlayExcludedBackendIds( nodeId: rootNodeId, selector: OVERLAY_HOST_SELECTOR, }); + throwIfAborted(signal); if (typeof found.nodeId !== "number" || found.nodeId === 0) return excluded; const described = await cdp.send<{ node?: CdpDomNode }>(tabId, "DOM.describeNode", { @@ -690,8 +741,10 @@ export async function collectOverlayExcludedBackendIds( depth: -1, pierce: true, }); + throwIfAborted(signal); collectBackendIdsFromDomNode(described.node, excluded); } catch (err) { + if (isAbortError(err)) throw err; console.debug("[bsk capture] overlay exclusion fallback failed", err); } return excluded; @@ -910,7 +963,9 @@ export async function captureViewModel( tabId: number, options: CaptureViewModelOptions = {}, ): Promise { + throwIfAborted(options.signal); const metrics = await cdp.send(tabId, "Page.getLayoutMetrics", {}); + throwIfAborted(options.signal); const dpr = devicePixelRatio(metrics); const vpSrc = metrics.cssLayoutViewport ?? metrics.layoutViewport ?? {}; const viewport: Viewport = { @@ -921,11 +976,13 @@ export async function captureViewModel( const scrollY = vpSrc.pageY ?? 0; await cdp.send(tabId, "DOMSnapshot.enable", {}); + throwIfAborted(options.signal); const snap = await cdp.send(tabId, "DOMSnapshot.captureSnapshot", { computedStyles: REQUESTED_STYLES, includePaintOrder: true, includeDOMRects: true, }); + throwIfAborted(options.signal); const strings = snap.strings ?? []; const documents = snap.documents ?? []; @@ -957,11 +1014,13 @@ export async function captureViewModel( excludedBackendNodeIds.add(id); } - await enrichFormControlStates(cdp, tabId, [nodes, ...iframeNodes.values()]); + await enrichFormControlStates(cdp, tabId, [nodes, ...iframeNodes.values()], options.signal); + throwIfAborted(options.signal); const surfaceProbes = options.conditionalSurfaceProbe ? await probeHoverSurfaces(cdp, tabId, nodes, options) : []; + throwIfAborted(options.signal); return { nodes, viewport, iframeNodes, surfaceProbes, excludedBackendNodeIds }; } diff --git a/apps/extension/src/tools/window.ts b/apps/extension/src/tools/window.ts index ad11ad7..ccec54e 100755 --- a/apps/extension/src/tools/window.ts +++ b/apps/extension/src/tools/window.ts @@ -47,7 +47,9 @@ export async function handleWindowResize( manager: SessionManager, params: WindowResizeParams, api: WindowResizeApi = chromeWindowResizeApi, + signal?: AbortSignal, ): Promise { + if (signal?.aborted) return { code: "cancelled", message: "window_resize aborted" }; const ctxOrErr = lookupSession(manager, params, "window_resize"); if (isRpcError(ctxOrErr)) return ctxOrErr; const ctx = ctxOrErr; @@ -64,6 +66,7 @@ export async function handleWindowResize( } try { + if (signal?.aborted) return { code: "cancelled", message: "window_resize aborted" }; await api.update(ctx.agentWindowId, { width: sizeOrErr.width, height: sizeOrErr.height, diff --git a/apps/extension/src/transport/types.ts b/apps/extension/src/transport/types.ts index 47de957..d53825a 100644 --- a/apps/extension/src/transport/types.ts +++ b/apps/extension/src/transport/types.ts @@ -32,7 +32,8 @@ export type RpcErrorReason = | "tab_not_active" | "restricted_tab_url" | "borrow_conflict" - | "screenshot_capture_failed"; + | "screenshot_capture_failed" + | "cleanup_failed"; export interface RpcErrorData { reason?: RpcErrorReason; diff --git a/crates/bsk-cli/src/daemon/ipc.rs b/crates/bsk-cli/src/daemon/ipc.rs index b464416..722342c 100644 --- a/crates/bsk-cli/src/daemon/ipc.rs +++ b/crates/bsk-cli/src/daemon/ipc.rs @@ -443,10 +443,10 @@ async fn handle_wait_ms( /// * **Queued** — the worker's pre-flight observes the cancelled /// token and short-circuits with `cancelled` before any WS /// frame leaves the daemon (review C2 fix). -/// * **Forwarded** — the worker's `tokio::select!` returns -/// `cancelled` immediately, AND we additionally push a WS-side -/// `cancel { rpc_id: ws_rpc_id }` frame so the extension's -/// dispatcher can trip its `AbortController`. +/// * **Forwarded** — we push a WS-side `cancel { rpc_id: ws_rpc_id }` +/// frame so the extension's dispatcher can trip its +/// `AbortController`; the worker keeps the session busy until the +/// extension returns its final result or the cleanup timeout expires. /// /// Returns `{ cancelled }` reflecting whether either surface /// matched. The RPC itself never errors — a cancelled tool surfaces diff --git a/crates/bsk-cli/src/daemon/queue.rs b/crates/bsk-cli/src/daemon/queue.rs index 3fac0eb..67a6042 100644 --- a/crates/bsk-cli/src/daemon/queue.rs +++ b/crates/bsk-cli/src/daemon/queue.rs @@ -58,6 +58,11 @@ pub const QUEUE_CAPACITY: usize = 64; /// so tests can default it. pub const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(30); +/// Once a forwarded request is cancelled, keep the session queue busy while +/// the extension finishes handler-side compensation. This is deliberately +/// bounded so a wedged extension cannot pin the session forever. +pub const CANCEL_CLEANUP_TIMEOUT: Duration = Duration::from_secs(2); + /// Job submitted into a session queue. Carries everything the worker /// needs to forward one RPC and one oneshot to deliver the answer. pub struct ToolJob { @@ -70,9 +75,9 @@ pub struct ToolJob { /// the queued/forwarded state machine; the worker checks it on /// pre-flight (so a cancel arriving while the job was still /// queued short-circuits without ever touching the extension) - /// and selects on `cancel.cancelled()` while awaiting the WS - /// response (so a cancel arriving mid-flight unblocks the worker - /// the same way the WS-side cancel frame already did). + /// and observes `cancel.cancelled()` while awaiting the WS response. + /// A forwarded cancel keeps the worker busy until the extension's + /// final response arrives or the bounded cleanup timeout expires. /// /// `None` for daemon-internal callers that do not flow through an /// IPC request id (e.g. `session.stop`'s queued teardown call — @@ -571,20 +576,41 @@ async fn forward_one( None } }; - let waited = await_with_optional_cancel(job.timeout, waiter, cancel_token.as_ref()).await; + let waited = await_with_optional_cancel( + job.timeout, + CANCEL_CLEANUP_TIMEOUT, + waiter, + cancel_token.as_ref(), + ) + .await; let response = match waited { WaitOutcome::Response(resp) => resp, - WaitOutcome::Cancelled => { - // Drop the WS waiter so a late extension reply is dropped - // cleanly (otherwise pending.resolve would log a stale - // entry). The CLI caller will get a synthesised - // `cancelled` here even if the extension never answers. - client.pending.lock().unwrap().cancel(&rpc_id); + WaitOutcome::CancelledAfterResponse(resp) => { + // Cancellation wins the external verdict, but only after the + // extension's original RPC has settled. Preserve any non-cancel + // error so compensation failures remain explicit instead of + // being hidden behind a generic cancelled result. + if let ResponseBody::Err(err) = resp.body + && !matches!(err.code, ErrorCode::Cancelled | ErrorCode::UserAborted) + { + return Err(err); + } return Err(cancelled_error( job.inflight.as_deref(), - "tool dispatch cancelled mid-flight", + "tool dispatch cancelled after extension cleanup", )); } + WaitOutcome::CleanupTimeout => { + client.pending.lock().unwrap().cancel(&rpc_id); + return Err(RpcError { + code: ErrorCode::Timeout, + message: format!( + "cancelled tool did not finish cleanup within {:?}", + CANCEL_CLEANUP_TIMEOUT + ), + data: Some(serde_json::json!({ "reason": "cancel_cleanup_timeout" })), + }); + } WaitOutcome::WaiterClosed => { client.pending.lock().unwrap().cancel(&rpc_id); return Err(RpcError { @@ -611,47 +637,42 @@ async fn forward_one( #[derive(Debug)] enum WaitOutcome { Response(bsk_protocol::ResponseFrame), - Cancelled, + CancelledAfterResponse(bsk_protocol::ResponseFrame), + CleanupTimeout, WaiterClosed, Timeout, } async fn await_with_optional_cancel( timeout: Duration, - waiter: oneshot::Receiver, + cleanup_timeout: Duration, + mut waiter: oneshot::Receiver, cancel: Option<&super::abort::AbortToken>, ) -> WaitOutcome { match cancel { - // Cancel wins if both ready: when a cancel notification and - // the extension's response are both observable on the same - // tokio tick, `biased;` polls `token.cancelled()` first and - // resolves the whole `select!` to `WaitOutcome::Cancelled`. - // The already-arrived tool result is intentionally dropped. - // - // This extends design §4.6 — which only specifies "CLI sends - // `cancel` on SIGINT and waits for the extension's - // `cancelled` reply (≤ 2s before forced exit)" — by pinning - // the same-tick race to a single, observable verdict. - // Without `biased;`, an extension racing a fast successful - // reply against a cancel notification could occasionally - // resolve as `ok` even though the caller had already moved on - // to compensation logic, leaving the agent trusting state the - // daemon had just been asked to roll back. Picking - // cancel-wins keeps the external observation rule simple - // ("once a cancel is in flight, the in-flight RPC's verdict - // is `cancelled`") and matches what the CLI / agent already - // assumes after firing SIGINT. Round 3 M1 / round 4 M2 - // nail-down. - Some(token) => tokio::select! { - biased; - _ = token.cancelled() => WaitOutcome::Cancelled, - outcome = tokio::time::timeout(timeout, waiter) => match outcome { - Ok(Ok(resp)) => WaitOutcome::Response(resp), - Ok(Err(_)) => WaitOutcome::WaiterClosed, - Err(_) => WaitOutcome::Timeout, - }, - }, - None => match tokio::time::timeout(timeout, waiter).await { + Some(token) => { + let deadline = tokio::time::sleep(timeout); + tokio::pin!(deadline); + tokio::select! { + // Cancel keeps same-tick priority, but it no longer drops the + // waiter. The worker remains busy until the extension replies + // after compensation or the cleanup deadline expires. + biased; + _ = token.cancelled() => { + match tokio::time::timeout(cleanup_timeout, &mut waiter).await { + Ok(Ok(resp)) => WaitOutcome::CancelledAfterResponse(resp), + Ok(Err(_)) => WaitOutcome::WaiterClosed, + Err(_) => WaitOutcome::CleanupTimeout, + } + }, + outcome = &mut waiter => match outcome { + Ok(resp) => WaitOutcome::Response(resp), + Err(_) => WaitOutcome::WaiterClosed, + }, + _ = &mut deadline => WaitOutcome::Timeout, + } + } + None => match tokio::time::timeout(timeout, &mut waiter).await { Ok(Ok(resp)) => WaitOutcome::Response(resp), Ok(Err(_)) => WaitOutcome::WaiterClosed, Err(_) => WaitOutcome::Timeout, @@ -733,10 +754,16 @@ mod await_with_optional_cancel_tests { let (tx, rx) = oneshot::channel(); tx.send(dummy_response()).unwrap(); - let outcome = await_with_optional_cancel(Duration::from_secs(10), rx, Some(&token)).await; + let outcome = await_with_optional_cancel( + Duration::from_secs(10), + Duration::from_secs(1), + rx, + Some(&token), + ) + .await; assert!( - matches!(outcome, WaitOutcome::Cancelled), - "expected Cancelled when both arms are ready under biased; (got something else)" + matches!(outcome, WaitOutcome::CancelledAfterResponse(_)), + "expected cancellation to win while retaining the ready response" ); } @@ -749,7 +776,13 @@ mod await_with_optional_cancel_tests { let (tx, rx) = oneshot::channel(); tx.send(dummy_response()).unwrap(); - let outcome = await_with_optional_cancel(Duration::from_secs(10), rx, Some(&token)).await; + let outcome = await_with_optional_cancel( + Duration::from_secs(10), + Duration::from_secs(1), + rx, + Some(&token), + ) + .await; match outcome { WaitOutcome::Response(frame) => { assert!(matches!(frame.body, ResponseBody::Ok(_))); @@ -762,9 +795,50 @@ mod await_with_optional_cancel_tests { async fn no_cancel_token_still_returns_response() { let (tx, rx) = oneshot::channel(); tx.send(dummy_response()).unwrap(); - let outcome = await_with_optional_cancel(Duration::from_secs(10), rx, None).await; + let outcome = + await_with_optional_cancel(Duration::from_secs(10), Duration::from_secs(1), rx, None) + .await; assert!(matches!(outcome, WaitOutcome::Response(_))); } + + #[tokio::test] + async fn cancelled_waiter_stays_pending_until_cleanup_response_arrives() { + let token = AbortToken::new(); + let (tx, rx) = oneshot::channel(); + token.cancel(); + let pending = tokio::spawn(async move { + await_with_optional_cancel( + Duration::from_secs(10), + Duration::from_secs(1), + rx, + Some(&token), + ) + .await + }); + + tokio::task::yield_now().await; + assert!(!pending.is_finished()); + tx.send(dummy_response()).unwrap(); + assert!(matches!( + pending.await.unwrap(), + WaitOutcome::CancelledAfterResponse(_) + )); + } + + #[tokio::test] + async fn cancelled_waiter_releases_after_cleanup_timeout() { + let token = AbortToken::new(); + let (_tx, rx) = oneshot::channel(); + token.cancel(); + let outcome = await_with_optional_cancel( + Duration::from_secs(10), + Duration::from_millis(10), + rx, + Some(&token), + ) + .await; + assert!(matches!(outcome, WaitOutcome::CleanupTimeout)); + } } #[cfg(test)] diff --git a/crates/bsk-cli/tests/cancel_forwarding.rs b/crates/bsk-cli/tests/cancel_forwarding.rs index e8a0aaf..c48113c 100644 --- a/crates/bsk-cli/tests/cancel_forwarding.rs +++ b/crates/bsk-cli/tests/cancel_forwarding.rs @@ -583,11 +583,11 @@ async fn cancel_for_unknown_rpc_id_returns_false() { } #[tokio::test] -async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() { - // With session busy fast-fail, a second tool RPC submitted while - // the worker is occupied returns `session_busy` immediately - // instead of queuing behind the slow call. Cancel still unblocks - // the in-flight RPC through the existing forwarded-cancel path. +async fn cancel_keeps_session_busy_until_delayed_extension_cleanup_finishes() { + // A forwarded cancel is acknowledged immediately, but the extension + // deliberately delays the original RPC's final response to model + // handler-side compensation. The daemon must keep the session busy + // throughout that delay and release it only after cleanup settles. use std::sync::atomic::{AtomicUsize, Ordering as AOrdering}; let (handle, sock) = spawn_daemon().await; let mut ws = connect_ext(handle.ws_addr()).await; @@ -596,6 +596,10 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() { let ws = Arc::new(tokio::sync::Mutex::new(ws)); let snapshots_seen = Arc::new(AtomicUsize::new(0)); let snapshots_seen_clone = Arc::clone(&snapshots_seen); + let cleanup_started = Arc::new(tokio::sync::Notify::new()); + let cleanup_started_clone = Arc::clone(&cleanup_started); + let cleanup_release = Arc::new(tokio::sync::Notify::new()); + let cleanup_release_clone = Arc::clone(&cleanup_release); let ws_clone = Arc::clone(&ws); let responder = tokio::spawn(async move { let mut pending_snapshot: Option = None; @@ -636,6 +640,16 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() { snapshots_seen_clone.fetch_add(1, AOrdering::SeqCst); pending_snapshot = Some(req.id); } + Method::ToolTabList => { + let reply = ResponseFrame { + id: req.id, + body: ResponseBody::Ok(json!({ "tabs": [] })), + }; + let mut g = ws_clone.lock().await; + g.send(Message::Text(serde_json::to_string(&reply).unwrap())) + .await + .unwrap(); + } Method::Cancel => { let target = req .params @@ -647,6 +661,8 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() { if let (Some(target), Some(snap)) = (target, snapshot_id) && target == snap { + cleanup_started_clone.notify_one(); + cleanup_release_clone.notified().await; let reply = ResponseFrame { id: snap, body: ResponseBody::Err(RpcError { @@ -753,6 +769,33 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() { .expect("cancel slow rpc succeeds"); assert!(cancel_slow.cancelled); + cleanup_started.notified().await; + let busy_after_cancel = tokio::time::timeout( + Duration::from_millis(200), + busy_ipc.call_with_id::<_, serde_json::Value>( + "snap-busy-after-cancel".into(), + Method::ToolSnapshot, + Some(json!({"session_id": session_id.clone()})), + Duration::from_secs(10), + ), + ) + .await + .expect("post-cancel RPC should fast-fail while cleanup is pending") + .unwrap(); + let busy_after_cancel_err = + busy_after_cancel.expect_err("session must remain busy during delayed compensation"); + assert_eq!(busy_after_cancel_err.code, ErrorCode::Timeout); + assert_eq!( + busy_after_cancel_err + .data + .as_ref() + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()), + Some(bsk::rpc_reason::SESSION_BUSY) + ); + + cleanup_release.notify_one(); + let slow_outcome = tokio::time::timeout(Duration::from_secs(5), slow_handle) .await .expect("slow snapshot did not resolve") @@ -761,6 +804,18 @@ async fn concurrent_tool_call_returns_session_busy_while_slow_rpc_inflight() { let slow_err = slow_outcome.expect_err("slow snapshot must surface cancelled"); assert_eq!(slow_err.code, ErrorCode::Cancelled); + let after_cleanup: serde_json::Value = busy_ipc + .call( + "tabs-after-cleanup", + Method::ToolTabList, + Some(json!({"session_id": session_id})), + Duration::from_secs(3), + ) + .await + .unwrap() + .expect("session queue should reopen after extension cleanup"); + assert_eq!(after_cleanup, json!({ "tabs": [] })); + assert_eq!( snapshots_seen.load(AOrdering::SeqCst), 1,