fix(extension): make tool cancellation cooperative

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<number>((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 });

apps/extension/src/session-manager/agent-window.ts

@@ -43,12 +43,11 @@ export const chromeAgentWindowApi: AgentWindowApi = {
     return win.id;
   },
   async remove(windowId: number): Promise<void> {
-    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<void> {
     const tabs = await chrome.tabs.query({ windowId });

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;
+    }
   }
 
   /**

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 () => {

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 <T = unknown>(
+          tabId: number,
+          method: string,
+          params?: object,
+        ) => Promise<T>,
+        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");

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 };
 }

apps/extension/src/tools/__tests__/tabs.test.ts

@@ -179,6 +179,7 @@ function makeWindowsApi(
     get: ReturnType<typeof vi.fn>;
     lastFocused: ReturnType<typeof vi.fn>;
     create: ReturnType<typeof vi.fn>;
+    remove: ReturnType<typeof vi.fn>;
   };
 } {
   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");