fix(runtime): clean up sessions on disconnect (#17)

* fix(runtime): clean up sessions on disconnect

* chore(ci): match Biome 2.4 formatting

* ci: retry stalled Rust job

* test: bind daemon ports without probe races

apps/extension/src/entrypoints/background.ts

@@ -22,6 +22,7 @@ import {
 } from "@/lib/overlay-bridge";
 import { POPUP_PORT_NAME, type PopupInbound, type PopupOutbound } from "@/lib/popup-bridge";
 import { attachSessionsLiveFlag } from "@/lib/sessions-live-flag";
+import { createDisconnectCleanup } from "@/session-manager/disconnect-cleanup";
 import { attachSessionEventHandler } from "@/session-manager/event-handler";
 import { SessionManager } from "@/session-manager/manager";
 import {
@@ -137,6 +138,13 @@ export default defineBackground(() => {
   // stale `true` does not keep waking us on every page load until the
   // first mutation (review M4/M5 round 3 m-R3-1).
   void sessionsLive.refresh();
+  const cleanupAfterDisconnect = createDisconnectCleanup({
+    manager: sessions,
+    sessionStopDeps: { cdp },
+    onSessionsChanged: () => {
+      void sessionsLive.syncFromManager();
+    },
+  });
   const dispatcher = new ToolDispatcher({
     transport,
     sessions,
@@ -267,7 +275,20 @@ export default defineBackground(() => {
 
   void (async () => {
     const connectionEnabled = await getConnectionEnabled();
-    await controller.attach(transport, detectBrowserMeta(), connectionEnabled);
+    await controller.attach(transport, detectBrowserMeta(), connectionEnabled, {
+      beforeDisconnect: async () => {
+        const report = await cleanupAfterDisconnect();
+        if (report.failures.length > 0) {
+          console.warn("[browser-skill] session cleanup before disconnect was incomplete", report);
+        }
+      },
+      onDisconnected: async () => {
+        const report = await cleanupAfterDisconnect();
+        if (report.failures.length > 0) {
+          console.warn("[browser-skill] session cleanup after disconnect was incomplete", report);
+        }
+      },
+    });
   })().catch((err) => {
     console.error("[browser-skill] controller failed to attach", err);
   });

apps/extension/src/lib/__tests__/connection-controller.test.ts

@@ -152,6 +152,36 @@ describe("ConnectionController connectionEnabled", () => {
     expect(controller.snapshot().lastError).toBeNull();
   });
 
+  it("runs safe session cleanup before an intentional disconnect", async () => {
+    const controller = new ConnectionController();
+    const transport = makeMockTransport();
+    const order: string[] = [];
+    transport.disconnect.mockImplementation(async () => {
+      order.push("disconnect");
+    });
+    await controller.attach(transport, { name: "Chrome", version: "120" }, true, {
+      beforeDisconnect: async () => {
+        order.push("cleanup");
+      },
+    });
+
+    await controller.setConnectionEnabled(false);
+
+    expect(order).toEqual(["cleanup", "disconnect"]);
+  });
+
+  it("runs session cleanup when the transport disconnects unexpectedly", async () => {
+    const controller = new ConnectionController();
+    const transport = makeMockTransport();
+    const onDisconnected = vi.fn(async () => {});
+    await controller.attach(transport, { name: "Chrome", version: "120" }, true, {
+      onDisconnected,
+    });
+
+    transport.emitState("disconnected");
+    await vi.waitFor(() => expect(onDisconnected).toHaveBeenCalledTimes(1));
+  });
+
   it("reconnects when connection is re-enabled", async () => {
     const controller = new ConnectionController();
     const transport = makeMockTransport();

apps/extension/src/lib/connection-controller.ts

@@ -21,6 +21,13 @@ export interface SnapshotInfo {
 
 type Listener = (s: SnapshotInfo) => void;
 
+export interface ConnectionLifecycleHooks {
+  /** Safe local teardown that must finish before an intentional disconnect. */
+  beforeDisconnect?: () => void | Promise<void>;
+  /** Best-effort teardown after an unexpected transport loss. */
+  onDisconnected?: () => void | Promise<void>;
+}
+
 /**
  * Orchestrates Transport + handshake lifecycle for the background SW.
  *
@@ -37,6 +44,8 @@ export class ConnectionController {
   private connectionEnabled = true;
   private listeners = new Set<Listener>();
   private handshakeInFlight = false;
+  private lifecycleHooks: ConnectionLifecycleHooks = {};
+  private disconnectRecovery: Promise<void> | null = null;
 
   get isConnectionEnabled(): boolean {
     return this.connectionEnabled;
@@ -68,21 +77,28 @@ export class ConnectionController {
     transport: Transport,
     browser: { name: string; version: string },
     connectionEnabled = true,
+    lifecycleHooks: ConnectionLifecycleHooks = {},
   ): Promise<void> {
     this.transport = transport;
     this.connectionEnabled = connectionEnabled;
+    this.lifecycleHooks = lifecycleHooks;
     this.instanceId = await getOrCreateInstanceId();
     this.label = await getLabel();
 
     transport.onConnectionStateChange((s) => {
+      if (s === "disconnected") {
+        this.handshake = null;
+        if (this.connectionEnabled) {
+          this.setState("disconnected");
+          void this.recoverFromDisconnect();
+        }
+        return;
+      }
       if (!this.connectionEnabled) return;
       if (s === "connected") {
         void this.runHandshake(browser);
         return;
       }
-      if (s === "disconnected") {
-        this.handshake = null;
-      }
       this.setState(s);
     });
 
@@ -161,6 +177,7 @@ export class ConnectionController {
   private async applyDisabledState(): Promise<void> {
     this.handshake = null;
     this.lastError = null;
+    await this.lifecycleHooks.beforeDisconnect?.();
     if (this.transport) {
       await this.transport.disconnect().catch(() => {});
     }
@@ -169,6 +186,34 @@ export class ConnectionController {
     if (was === "disconnected") this.fire();
   }
 
+  private recoverFromDisconnect(): Promise<void> {
+    if (this.disconnectRecovery) return this.disconnectRecovery;
+    const recovery = (async () => {
+      // WSTransport schedules its reconnect timer immediately after notifying
+      // state listeners. Yield once, then disconnect explicitly so that timer
+      // is cancelled before local session teardown begins.
+      await Promise.resolve();
+      if (!this.connectionEnabled || !this.transport) return;
+      await this.transport.disconnect().catch(() => {});
+      try {
+        await this.lifecycleHooks.onDisconnected?.();
+      } catch (err) {
+        console.warn("[browser-skill] session cleanup after disconnect failed", err);
+      }
+      if (!this.connectionEnabled) return;
+      try {
+        await this.transport.connect();
+      } catch (err) {
+        this.lastError = err instanceof Error ? err.message : String(err);
+        this.setState("disconnected");
+      }
+    })();
+    this.disconnectRecovery = recovery.finally(() => {
+      this.disconnectRecovery = null;
+    });
+    return this.disconnectRecovery;
+  }
+
   private setState(next: ConnectionState): void {
     if (this.currentState === next) return;
     this.currentState = next;

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

@@ -0,0 +1,58 @@
+import { describe, expect, it, vi } from "vitest";
+import { createDisconnectCleanup } from "../disconnect-cleanup";
+import { SessionManager } from "../manager";
+
+describe("disconnect session cleanup", () => {
+  it("stops every local session through the safe session-stop path", async () => {
+    const remove = vi.fn(async () => {});
+    let nextWindowId = 100;
+    const manager = new SessionManager({
+      agentWindow: {
+        create: vi.fn(async () => nextWindowId++),
+        ensureActiveTab: vi.fn(async () => {}),
+        remove,
+      },
+    });
+    await manager.start("aa11");
+    await manager.start("bb22");
+    const detachSession = vi.fn(async () => {});
+    const onSessionsChanged = vi.fn();
+    const cleanup = createDisconnectCleanup({
+      manager,
+      sessionStopDeps: { cdp: { detachSession } },
+      onSessionsChanged,
+    });
+
+    const report = await cleanup();
+
+    expect(report).toEqual({ stoppedSessionIds: ["aa11", "bb22"], failures: [] });
+    expect(manager.list()).toEqual([]);
+    expect(remove).toHaveBeenCalledTimes(2);
+    expect(detachSession).toHaveBeenCalledWith("aa11");
+    expect(detachSession).toHaveBeenCalledWith("bb22");
+    expect(onSessionsChanged).toHaveBeenCalledTimes(1);
+  });
+
+  it("coalesces overlapping disconnect notifications", async () => {
+    let releaseRemove: () => void = () => {};
+    const removeGate = new Promise<void>((resolve) => {
+      releaseRemove = resolve;
+    });
+    const manager = new SessionManager({
+      agentWindow: {
+        create: vi.fn(async () => 100),
+        ensureActiveTab: vi.fn(async () => {}),
+        remove: vi.fn(() => removeGate),
+      },
+    });
+    await manager.start("aa11");
+    const cleanup = createDisconnectCleanup({ manager });
+
+    const first = cleanup();
+    const second = cleanup();
+    expect(second).toBe(first);
+    releaseRemove();
+
+    await expect(first).resolves.toEqual({ stoppedSessionIds: ["aa11"], failures: [] });
+  });
+});

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

@@ -0,0 +1,75 @@
+import { handleSessionStop, type SessionStopDeps } from "@/tools/session";
+import type { SessionManager } from "./manager";
+
+export interface DisconnectCleanupFailure {
+  sessionId: string;
+  message: string;
+}
+
+export interface DisconnectCleanupReport {
+  stoppedSessionIds: string[];
+  failures: DisconnectCleanupFailure[];
+}
+
+export interface DisconnectCleanupOptions {
+  manager: SessionManager;
+  sessionStopDeps?: SessionStopDeps;
+  onSessionsChanged?: () => void;
+}
+
+/**
+ * Build a coalesced cleanup operation for transport loss.
+ *
+ * The daemon purges its registry when the extension socket disappears. To
+ * keep the extension-side mirror consistent, every local session must follow
+ * the normal safe stop path as well: return borrowed tabs, clear refs, detach
+ * CDP, and only then close the Agent Window.
+ */
+export function createDisconnectCleanup(options: DisconnectCleanupOptions) {
+  let inFlight: Promise<DisconnectCleanupReport> | null = null;
+
+  return (): Promise<DisconnectCleanupReport> => {
+    if (inFlight) return inFlight;
+    inFlight = cleanupSessions(options).finally(() => {
+      inFlight = null;
+    });
+    return inFlight;
+  };
+}
+
+async function cleanupSessions(
+  options: DisconnectCleanupOptions,
+): Promise<DisconnectCleanupReport> {
+  const stoppedSessionIds: string[] = [];
+  const failures: DisconnectCleanupFailure[] = [];
+
+  for (const ctx of options.manager.list()) {
+    try {
+      const result = await handleSessionStop(
+        options.manager,
+        { session_id: ctx.sessionId },
+        options.sessionStopDeps,
+      );
+      if ("code" in result) {
+        failures.push({ sessionId: ctx.sessionId, message: result.message });
+        continue;
+      }
+      if (result.return_failures?.length) {
+        failures.push({
+          sessionId: ctx.sessionId,
+          message: result.return_failures.map((failure) => failure.message).join("; "),
+        });
+        continue;
+      }
+      stoppedSessionIds.push(ctx.sessionId);
+    } catch (err) {
+      failures.push({
+        sessionId: ctx.sessionId,
+        message: err instanceof Error ? err.message : String(err),
+      });
+    }
+  }
+
+  options.onSessionsChanged?.();
+  return { stoppedSessionIds, failures };
+}

crates/bsk-cli/src/daemon/sessions.rs

@@ -539,10 +539,9 @@ pub async fn stop_session(
             Ok(result)
         }
         Err(DispatchError::Rpc(err)) => {
-            // After an extension SW restart the daemon still owns the
-            // session entry (the generation-guard ensures we do not
-            // purge across reconnects), but the extension's in-memory
-            // SessionManager is reset and now answers `not_found`.
+            // A legacy extension or an unusual reconnect race may still
+            // leave a daemon session after the extension's in-memory
+            // SessionManager has reset and now answers `not_found`.
             // Treat that as an authoritative signal that the session is
             // already gone and reconcile our local state instead of
             // leaving an orphan row visible to `bsk session list`

crates/bsk-cli/src/daemon/ws.rs

@@ -286,6 +286,17 @@ async fn drive_connection(
     // socket's cleanup path uses `remove_if_generation_matches` to
     // avoid clobbering the newer entry (review M4/M5 round 2 #1).
     let browser_id = BrowserId(params.instance_id.clone());
+    // A fresh socket represents a fresh extension control plane. The
+    // extension tears down its local sessions before reconnecting, so any
+    // daemon-side sessions left under the same instance id are stale. Purge
+    // them before replacing the browser registration; otherwise an old WS
+    // cleanup racing with this handshake can preserve rows that no longer
+    // have an Agent Window on the extension side.
+    for session in state.sessions.purge_browser(&browser_id) {
+        state.tool_queues.remove(&session.id);
+        state.session_interrupts.drop_session(&session.id);
+        debug!(session = %session.id, "purged stale session before browser reconnect");
+    }
     let (tx, mut rx) = mpsc::unbounded_channel::<Frame>();
     let generation = super::browsers::next_browser_generation();
     let connected_at_ms = std::time::SystemTime::now()

crates/bsk-cli/tests/sessions_ipc.rs

@@ -19,7 +19,7 @@ use tokio_tungstenite::tungstenite::handshake::client::generate_key;
 use tokio_tungstenite::tungstenite::http::Request;
 use tokio_tungstenite::tungstenite::protocol::Message;
 
-use support::{wait_for_browser_count, wait_for_no_sessions, wait_for_session_count};
+use support::{wait_for_browser_count, wait_for_no_sessions};
 
 const TEST_EXT_ID: &str = "abcdefghijklmnopabcdefghijklmnop";
 
@@ -728,11 +728,10 @@ async fn browser_disconnect_purges_sessions() {
 
 #[tokio::test]
 async fn session_stop_self_heals_when_extension_reports_not_found() {
-    // After an extension SW restart the daemon still owns the session
-    // entry (Round 2 generation-guard) but the extension's
-    // SessionManager is reset and answers `not_found`. Stop must
-    // reconcile local state so `session.list` does not show an orphan
-    // (review M4/M5 round 3 I-R3-2).
+    // Legacy extension versions or unusual reconnect races can leave a
+    // daemon session after the extension's SessionManager resets and answers
+    // `not_found`. Stop must still reconcile local state so `session.list`
+    // does not show an orphan (review M4/M5 round 3 I-R3-2).
 
     let (handle, sock) = spawn_daemon().await;
     let mut ws = connect_ext(handle.ws_addr()).await;
@@ -826,14 +825,13 @@ async fn session_stop_self_heals_when_extension_reports_not_found() {
 }
 
 #[tokio::test]
-async fn reconnect_with_same_instance_id_does_not_clobber_new_browser() {
+async fn reconnect_with_same_instance_id_purges_stale_sessions_but_keeps_new_browser() {
     // Spawn a daemon, connect ext A, register a session bound to it,
     // then connect ext B reusing the same instance_id (mimics a SW
-    // restart / WS reconnect under MV3). When the OLD WS task tears
-    // down, the generation guard MUST keep the new entry + its
-    // session in the registry. Without the guard the stale cleanup
-    // path would call browsers.remove() + sessions.purge_browser()
-    // and the daemon would silently drop the live session.
+    // restart / WS reconnect under MV3). The extension now safely stops
+    // its local sessions before reconnecting, so the new handshake must
+    // purge the matching daemon rows. The generation guard must still keep
+    // the NEW browser registration when the old WS task later tears down.
 
     let (handle, _sock) = spawn_daemon().await;
     let mut ws_a = connect_ext(handle.ws_addr()).await;
@@ -854,14 +852,14 @@ async fn reconnect_with_same_instance_id_does_not_clobber_new_browser() {
     let _ = handshake_as_ext(&mut ws_b).await;
     // Registry now holds the newer generation under the same id.
     assert_eq!(state.browsers.len(), 1);
+    wait_for_no_sessions(&state).await;
 
     // Tear down ext A only; the cleanup path will run but must
     // observe that the registered generation no longer matches and
-    // leave the new entry + session alone.
+    // leave the new browser entry alone.
     let _ = ws_a.close(None).await;
     drop(ws_a);
     wait_for_browser_count(&state, 1).await;
-    wait_for_session_count(&state, 1).await;
 
     assert_eq!(
         state.browsers.len(),
@@ -870,8 +868,8 @@ async fn reconnect_with_same_instance_id_does_not_clobber_new_browser() {
     );
     assert_eq!(
         state.sessions.len(),
-        1,
-        "session must survive old browser cleanup"
+        0,
+        "stale sessions must not survive reconnect"
     );
 
     let _ = ws_b.close(None).await;