Merge pull request #97 from NianJiuZst/codex/cancel-session-lifecycle-rpcs

fix(session): cancel lifecycle RPCs cooperatively

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

@@ -203,6 +203,55 @@ describe("ToolDispatcher", () => {
     expect(sent[0]).toEqual({ id: "r-1", result: {} });
   });
 
+  it("forwards daemon cancellation into an in-flight session stop", async () => {
+    const { transport, sent, deliver } = fakeTransport();
+    const remove = vi.fn(async () => {});
+    const sessions = new SessionManager({
+      agentWindow: {
+        create: vi.fn(async () => 4242),
+        remove,
+        ensureActiveTab: vi.fn(async () => {}),
+      },
+    });
+    await sessions.start("aa11");
+    let finishDetach: () => void = () => {};
+    const detach = new Promise<void>((resolve) => {
+      finishDetach = resolve;
+    });
+    const cdp = {
+      send: vi.fn(),
+      detachSession: vi.fn(() => detach),
+      ensureNetworkCapture: vi.fn(async () => {}),
+      networkEntriesSince: vi.fn(() => ({
+        tab_id: 7,
+        entries: [],
+        next_since: 0,
+        truncated: false,
+      })),
+      setDeviceMetricsOverride: vi.fn(async () => {}),
+      clearDeviceMetricsOverride: vi.fn(async () => {}),
+      setUserAgentOverride: vi.fn(async () => {}),
+      setTouchEmulationEnabled: vi.fn(async () => {}),
+    };
+    const dispatcher = new ToolDispatcher({ transport, sessions, cdp: cdp as TestDispatcherCdp });
+    dispatcher.start();
+
+    deliver(makeRequest("tool.session_stop", { session_id: "aa11" }));
+    await vi.waitFor(() => expect(cdp.detachSession).toHaveBeenCalledWith("aa11"));
+    deliver({ id: "cancel-stop", method: "cancel", params: { rpc_id: "r-1" } });
+    await flushMicrotasks();
+    finishDetach();
+    await flushMicrotasks();
+
+    expect(sent).toContainEqual({ id: "cancel-stop", result: { cancelled: true } });
+    expect(sent).toContainEqual({
+      id: "r-1",
+      error: expect.objectContaining({ code: "cancelled" }),
+    });
+    expect(sessions.has("aa11")).toBe(true);
+    expect(remove).not.toHaveBeenCalled();
+  });
+
   it("returns not_found when stopping an unknown session", async () => {
     const { transport, sent, deliver } = fakeTransport();
     const sessions = new SessionManager({

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

@@ -66,6 +66,52 @@ function makeApis(
 }
 
 describe("handleSessionStop with auto-return", () => {
+  it("leaves the session untouched when cancellation arrived before teardown", async () => {
+    const aw = fakeAgentWindow([100]);
+    const sm = new SessionManager({ agentWindow: aw });
+    await sm.start("aa11");
+    const cdp = { detachSession: vi.fn(async () => {}) };
+    const controller = new AbortController();
+    controller.abort();
+
+    const res = await handleSessionStop(
+      sm,
+      { session_id: "aa11" },
+      { cdp, signal: controller.signal },
+    );
+
+    expect(res).toMatchObject({ code: "cancelled" });
+    expect(sm.has("aa11")).toBe(true);
+    expect(cdp.detachSession).not.toHaveBeenCalled();
+    expect(aw.remove).not.toHaveBeenCalled();
+  });
+
+  it("does not close the Agent Window when cancellation lands during CDP detach", async () => {
+    const aw = fakeAgentWindow([100]);
+    const sm = new SessionManager({ agentWindow: aw });
+    await sm.start("aa11");
+    const controller = new AbortController();
+    let finishDetach: () => void = () => {};
+    const detach = new Promise<void>((resolve) => {
+      finishDetach = resolve;
+    });
+    const cdp = { detachSession: vi.fn(() => detach) };
+
+    const stopping = handleSessionStop(
+      sm,
+      { session_id: "aa11" },
+      { cdp, signal: controller.signal },
+    );
+    await vi.waitFor(() => expect(cdp.detachSession).toHaveBeenCalledWith("aa11"));
+    controller.abort();
+    finishDetach();
+
+    const res = await stopping;
+    expect(res).toMatchObject({ code: "cancelled" });
+    expect(sm.has("aa11")).toBe(true);
+    expect(aw.remove).not.toHaveBeenCalled();
+  });
+
   it("returns every borrowed tab and closes the Agent Window in the right order", async () => {
     const aw = fakeAgentWindow([100]);
     const sm = new SessionManager({ agentWindow: aw });

apps/extension/src/tools/dispatcher.ts

@@ -290,6 +290,7 @@ export class ToolDispatcher {
         await this.releaseHoverLatch((req.params as SessionStopParams).session_id);
         return handleSessionStop(this.sessions, req.params as SessionStopParams, {
           cdp: this.cdp,
+          signal,
         });
       }
       case "tool.tab_list":

apps/extension/src/tools/session.ts

@@ -76,6 +76,7 @@ export interface SessionStopResult {
 }
 
 export interface SessionStopDeps {
+  signal?: AbortSignal;
   cdp?: {
     detachSession(sessionId: string): Promise<void>;
   };
@@ -177,6 +178,9 @@ export async function handleSessionStop(
       message: `session ${params.session_id} unknown`,
     };
   }
+  if (deps.signal?.aborted) {
+    return { code: "cancelled", message: "session_stop aborted before teardown" };
+  }
 
   // Step 1: auto-return borrowed tabs. Iterate over a snapshot of the
   // ids so deletions during iteration do not break the Map iterator.
@@ -187,6 +191,7 @@ export async function handleSessionStop(
     try {
       const tabManagement = {
         ...(deps.tabManagement ?? {}),
+        signal: deps.signal,
         isAgentWindowId:
           deps.tabManagement?.isAgentWindowId ??
           ((windowId: number) => manager.findByWindowId(windowId) !== null),
@@ -214,6 +219,17 @@ export async function handleSessionStop(
     }
   }
 
+  if (deps.signal?.aborted) {
+    const nonCancelFailures = returnFailures?.filter((failure) => failure.code !== "cancelled");
+    if (nonCancelFailures && nonCancelFailures.length > 0) {
+      return {
+        ...(returnedTabIds.length > 0 ? { returned_tab_ids: returnedTabIds } : {}),
+        return_failures: nonCancelFailures,
+      };
+    }
+    return { code: "cancelled", message: "session_stop aborted during tab return" };
+  }
+
   const result: SessionStopResult = {};
   if (returnedTabIds.length > 0) result.returned_tab_ids = returnedTabIds;
   if (returnFailures && returnFailures.length > 0) {
@@ -233,6 +249,10 @@ export async function handleSessionStop(
   // Step 3: detach CDP sessions this session opened (no-op if none).
   await deps.cdp?.detachSession(params.session_id);
 
+  if (deps.signal?.aborted) {
+    return { code: "cancelled", message: "session_stop aborted before window close" };
+  }
+
   // Step 4: close the Agent Window and drop the context.
   await manager.stop(params.session_id);