fix(extension): make tool cancellation cooperative

apps/extension/src/session-manager/__tests__/manager.test.ts

@@ -57,6 +57,38 @@ 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("stop() closes the Agent Window and forgets the session", async () => {
     const aw = fakeAgentWindow();
     const sm = new SessionManager({ agentWindow: aw });

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

@@ -25,6 +25,16 @@ export interface SessionManagerOptions {
   now?: () => number;
 }
 
+export interface SessionStartOptions {
+  signal?: AbortSignal;
+}
+
+function sessionStartAbortError(): Error {
+  const err = new Error("session_start aborted");
+  err.name = "AbortError";
+  return err;
+}
+
 /**
  * Owner of all live agent sessions inside the extension.
  *
@@ -128,22 +138,38 @@ export class SessionManager {
    * Returns the created window id so callers can echo it back to the
    * daemon in the `tool.session_start` reply.
    */
-  async start(sessionId: string): Promise<SessionContext> {
+  async start(sessionId: string, options: SessionStartOptions = {}): Promise<SessionContext> {
     if (this.sessions.has(sessionId)) {
       throw new Error(`[bh] session ${sessionId} already exists`);
     }
-    const windowId = await this.agentWindow.create(AGENT_WINDOW_HOME);
-    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;
+    if (options.signal?.aborted) throw sessionStartAbortError();
+
+    let windowId: number | null = null;
+    try {
+      windowId = await this.agentWindow.create(AGENT_WINDOW_HOME);
+      if (options.signal?.aborted) throw sessionStartAbortError();
+      await this.agentWindow.ensureActiveTab(windowId, AGENT_WINDOW_HOME);
+      if (options.signal?.aborted) throw sessionStartAbortError();
+      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 (err) {
+      if (windowId !== null) {
+        try {
+          await this.agentWindow.remove(windowId);
+        } catch (cleanupErr) {
+          console.warn(`[bh] failed to remove incomplete Agent Window ${windowId}`, cleanupErr);
+        }
+      }
+      throw err;
+    }
   }
 
   /**

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

@@ -334,23 +334,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

@@ -524,6 +524,41 @@ describe("handleSnapshot", () => {
     expect(ctx.refStore.resolve("e2")).toBe(200);
   });
 
+  it("does not overwrite refs when cancellation arrives during CDP 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 resolveTree: (value: { nodes: CdpAxNode[] }) => void = () => {};
+    const deps = makeDeps([]);
+    deps.send.mockImplementation(async (_tabId: number, method: string) => {
+      if (method === "Accessibility.enable") return {};
+      if (method === "Accessibility.getFullAXTree") {
+        return new Promise((resolve) => {
+          resolveTree = resolve;
+        });
+      }
+      throw new Error(`unexpected CDP method ${method}`);
+    });
+
+    const pending = handleSnapshot(sm, { session_id: "aa11" }, deps, controller.signal);
+    await vi.waitFor(() => expect(deps.send).toHaveBeenCalledTimes(2));
+    controller.abort();
+    resolveTree({
+      nodes: [
+        {
+          nodeId: "new",
+          role: { type: "role", value: "button" },
+          name: { type: "computedString", value: "New" },
+          backendDOMNodeId: 123,
+        },
+      ],
+    });
+
+    await expect(pending).resolves.toMatchObject({ code: "cancelled" });
+    expect(ctx.refStore.resolve("e1")).toBe(999);
+  });
+
   it("resets the RefStore on every fresh snapshot", async () => {
     const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) });
     const ctx = await sm.start("aa11");

apps/extension/src/tools/console.ts

@@ -31,7 +31,9 @@ export async function handleConsole(
   manager: SessionManager,
   params: ConsoleParams,
   deps: ConsoleDeps = defaultConsoleDeps(),
+  signal?: AbortSignal,
 ): Promise<ConsoleResult | RpcError> {
+  if (signal?.aborted) return { code: "cancelled", message: "console aborted" };
   const ctxOrErr = lookupSession(manager, params, "console");
   if (isRpcError(ctxOrErr)) return ctxOrErr;
   const bounds = parseConsoleBounds(params);
@@ -44,6 +46,7 @@ export async function handleConsole(
 
   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,

apps/extension/src/tools/dispatcher.ts

@@ -186,7 +186,10 @@ export class ToolDispatcher {
     let body: ResponseFrame;
     let startedSession: string | null = null;
     try {
-      const result = await Promise.race([this.invoke(req, ac.signal), abortPromise(ac.signal)]);
+      // Await the handler itself. A cancel trips its AbortSignal immediately,
+      // but the RPC response is not sent until the handler has stopped or
+      // completed any required compensation for an already-started side effect.
+      const result = await this.invoke(req, ac.signal);
       if (isRpcError(result)) {
         body = { id: req.id, error: result };
       } else {
@@ -251,51 +254,55 @@ export class ToolDispatcher {
   private async invoke(req: RequestFrame, signal: AbortSignal): Promise<unknown | RpcError> {
     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":
         return handleSessionStop(this.sessions, req.params as SessionStopParams, {
           cdp: this.cdp,
         });
       case "tool.tab_list":
         return handleTabList(this.sessions, req.params as TabListParams);
       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 handleTabClose(this.sessions, req.params as TabCloseParams);
+        return handleTabClose(this.sessions, req.params as TabCloseParams, { 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.screenshot":
         return handleScreenshot(
           this.sessions,
           req.params as ScreenshotParams,
           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.snapshot":
         return handleSnapshot(
           this.sessions,
           req.params as SnapshotParams,
           this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsCaptureApi } : undefined,
+          signal,
         );
       case "tool.get_html":
         return handleGetHtml(
           this.sessions,
           req.params as GetHtmlParams,
           this.cdp ? { cdp: this.cdp, tabsApi: chromeTabsCaptureApi } : undefined,
+          signal,
         );
       case "tool.navigate":
         return handleNavigate(
@@ -406,41 +413,7 @@ function isRpcError(v: unknown): v is RpcError {
   );
 }
 
-/**
- * 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<never> {
-  return new Promise<never>((_, 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 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;

apps/extension/src/tools/observation.ts

@@ -150,7 +150,9 @@ export async function handleScreenshot(
   manager: SessionManager,
   params: ScreenshotParams,
   deps: ScreenshotDeps = defaultScreenshotDeps(),
+  signal?: AbortSignal,
 ): Promise<ScreenshotResult | RpcError> {
+  if (signal?.aborted) return { code: "cancelled", message: "screenshot aborted" };
   const ctxOrErr = lookupSession(manager, params, "screenshot");
   if (isRpcError(ctxOrErr)) return ctxOrErr;
   const ctx = ctxOrErr;
@@ -170,6 +172,7 @@ export async function handleScreenshot(
     deps.cdp.trackSessionTab?.(ctx.sessionId, target.tabId);
     const captured = await captureElementScreenshot(deps.cdp, target.tabId, node.backendNodeId);
     if (isRpcError(captured)) return captured;
+    if (signal?.aborted) return { code: "cancelled", message: "screenshot aborted" };
     return withShotDialogs({
       image_base64: captured.image_base64,
       width: captured.width,
@@ -189,6 +192,7 @@ export async function handleScreenshot(
 
   try {
     const dataUrl = await deps.captureApi.captureVisibleTab(target.windowId, { format: "png" });
+    if (signal?.aborted) return { code: "cancelled", message: "screenshot aborted" };
     const image_base64 = stripDataUrlPrefix(dataUrl);
     const dims = parsePngDimensions(image_base64) ?? { width: 0, height: 0 };
     return withShotDialogs({
@@ -459,7 +463,9 @@ export async function handleGetHtml(
   manager: SessionManager,
   params: GetHtmlParams,
   deps: SnapshotDeps = getDefaultDeps(),
+  signal?: AbortSignal,
 ): Promise<GetHtmlResult | RpcError> {
+  if (signal?.aborted) return { code: "cancelled", message: "get_html aborted" };
   const ctxOrErr = lookupSession(manager, params, "get_html");
   if (isRpcError(ctxOrErr)) return ctxOrErr;
   const ctx = ctxOrErr;
@@ -498,6 +504,7 @@ export async function handleGetHtml(
       });
       html = resp.outerHTML ?? "";
     }
+    if (signal?.aborted) return { code: "cancelled", message: "get_html aborted" };
     const originalBytes = utf8ByteLength(html);
     const { out, truncated } = truncateBytes(html, maxBytes);
     return attachDialogs(deps.cdp, target.tabId, dialogCursor, {
@@ -518,7 +525,9 @@ export async function handleSnapshot(
   manager: SessionManager,
   params: SnapshotParams,
   deps: SnapshotDeps = getDefaultDeps(),
+  signal?: AbortSignal,
 ): Promise<SnapshotResult | RpcError> {
+  if (signal?.aborted) return { code: "cancelled", message: "snapshot aborted" };
   const ctxOrErr = lookupSession(manager, params, "snapshot");
   if (isRpcError(ctxOrErr)) return ctxOrErr;
   const ctx = ctxOrErr;
@@ -534,6 +543,7 @@ export async function handleSnapshot(
       "Accessibility.getFullAXTree",
       {},
     );
+    if (signal?.aborted) return { code: "cancelled", message: "snapshot aborted" };
     const rendered = renderAxTree(result.nodes ?? [], {
       maxDepth: params.max_depth,
       maxTokens: params.max_tokens,