fix(runtime): clean up sessions on disconnect

apps/extension/src/entrypoints/background.ts

@@ -16,6 +16,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 {
@@ -41,6 +42,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,
@@ -115,7 +123,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 };
+}

apps/extension/vitest.config.ts

@@ -21,9 +21,7 @@ export default defineConfig({
   },
   define: {
     __BSK_EXT_VERSION__: JSON.stringify(pkg.version),
-    __BSK_DAEMON_WS_URL__: JSON.stringify(
-      process.env.BSK_DAEMON_WS_URL ?? "ws://127.0.0.1:52800",
-    ),
+    __BSK_DAEMON_WS_URL__: JSON.stringify(process.env.BSK_DAEMON_WS_URL ?? "ws://127.0.0.1:52800"),
   },
   resolve: {
     alias: {

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

@@ -472,10 +472,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/browser_list_ordering.rs

@@ -16,7 +16,6 @@ use bsk_protocol::system::BrowserStatusEntry;
 use bsk_protocol::{ErrorCode, Method};
 use rand::Rng;
 use serde::Deserialize;
-use tokio::net::TcpListener;
 use tokio::sync::mpsc;
 
 fn tempfile_path(prefix: &str) -> PathBuf {
@@ -30,9 +29,7 @@ fn tempfile_path(prefix: &str) -> PathBuf {
 }
 
 async fn spawn_daemon() -> (daemon::DaemonHandle, PathBuf) {
-    let probe = TcpListener::bind("127.0.0.1:0").await.unwrap();
-    let port = probe.local_addr().unwrap().port();
-    drop(probe);
+    let port = 0;
 
     let config = DaemonConfig::new(port);
     let sock = tempfile_path("bsk-test-browser-list");

crates/bsk-cli/tests/browser_wait.rs

@@ -15,7 +15,6 @@ use bsk_protocol::system::{BrowserListParams, HandshakeParams, HandshakeResult,
 use bsk_protocol::{BrowserPeerInfo, Method, RequestFrame, ResponseBody, ResponseFrame};
 use futures_util::{SinkExt, StreamExt};
 use rand::Rng;
-use tokio::net::TcpListener;
 use tokio_tungstenite::tungstenite::handshake::client::generate_key;
 use tokio_tungstenite::tungstenite::http::Request;
 use tokio_tungstenite::tungstenite::protocol::Message;
@@ -77,9 +76,7 @@ fn tempfile_path(prefix: &str) -> PathBuf {
 }
 
 async fn spawn_daemon() -> (daemon::DaemonHandle, PathBuf) {
-    let probe = TcpListener::bind("127.0.0.1:0").await.unwrap();
-    let port = probe.local_addr().unwrap().port();
-    drop(probe);
+    let port = 0;
 
     let config = DaemonConfig::new(port);
     let sock = tempfile_path("bsk-test-browser-wait");