From 1371c92721b8681261f9b39ab66076fe49585bcc Mon Sep 17 00:00:00 2001 From: NianJiu <3235467914@qq.com> Date: Wed, 12 Aug 2026 18:44:22 +0800 Subject: [PATCH] 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 | 23 +++++- .../__tests__/connection-controller.test.ts | 30 ++++++++ .../src/lib/connection-controller.ts | 51 ++++++++++++- .../__tests__/disconnect-cleanup.test.ts | 58 ++++++++++++++ .../src/session-manager/disconnect-cleanup.ts | 75 +++++++++++++++++++ crates/bsk-cli/src/daemon/sessions.rs | 7 +- crates/bsk-cli/src/daemon/ws.rs | 11 +++ crates/bsk-cli/tests/sessions_ipc.rs | 30 ++++---- 8 files changed, 261 insertions(+), 24 deletions(-) create mode 100644 apps/extension/src/session-manager/__tests__/disconnect-cleanup.test.ts create mode 100644 apps/extension/src/session-manager/disconnect-cleanup.ts diff --git a/apps/extension/src/entrypoints/background.ts b/apps/extension/src/entrypoints/background.ts index ce3586f..207b5b5 100644 --- a/apps/extension/src/entrypoints/background.ts +++ b/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); }); diff --git a/apps/extension/src/lib/__tests__/connection-controller.test.ts b/apps/extension/src/lib/__tests__/connection-controller.test.ts index b3eb9ae..408cd89 100644 --- a/apps/extension/src/lib/__tests__/connection-controller.test.ts +++ b/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(); diff --git a/apps/extension/src/lib/connection-controller.ts b/apps/extension/src/lib/connection-controller.ts index bc0187b..9f325b3 100644 --- a/apps/extension/src/lib/connection-controller.ts +++ b/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; + /** Best-effort teardown after an unexpected transport loss. */ + onDisconnected?: () => void | Promise; +} + /** * Orchestrates Transport + handshake lifecycle for the background SW. * @@ -37,6 +44,8 @@ export class ConnectionController { private connectionEnabled = true; private listeners = new Set(); private handshakeInFlight = false; + private lifecycleHooks: ConnectionLifecycleHooks = {}; + private disconnectRecovery: Promise | 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 { 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 { 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 { + 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; diff --git a/apps/extension/src/session-manager/__tests__/disconnect-cleanup.test.ts b/apps/extension/src/session-manager/__tests__/disconnect-cleanup.test.ts new file mode 100644 index 0000000..3ca893a --- /dev/null +++ b/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((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: [] }); + }); +}); diff --git a/apps/extension/src/session-manager/disconnect-cleanup.ts b/apps/extension/src/session-manager/disconnect-cleanup.ts new file mode 100644 index 0000000..59daafb --- /dev/null +++ b/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 | null = null; + + return (): Promise => { + if (inFlight) return inFlight; + inFlight = cleanupSessions(options).finally(() => { + inFlight = null; + }); + return inFlight; + }; +} + +async function cleanupSessions( + options: DisconnectCleanupOptions, +): Promise { + 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 }; +} diff --git a/crates/bsk-cli/src/daemon/sessions.rs b/crates/bsk-cli/src/daemon/sessions.rs index 53ba290..0e74e14 100644 --- a/crates/bsk-cli/src/daemon/sessions.rs +++ b/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` diff --git a/crates/bsk-cli/src/daemon/ws.rs b/crates/bsk-cli/src/daemon/ws.rs index 7d34cc1..435b0a9 100644 --- a/crates/bsk-cli/src/daemon/ws.rs +++ b/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::(); let generation = super::browsers::next_browser_generation(); let connected_at_ms = std::time::SystemTime::now() diff --git a/crates/bsk-cli/tests/sessions_ipc.rs b/crates/bsk-cli/tests/sessions_ipc.rs index 29cff59..9e6af21 100644 --- a/crates/bsk-cli/tests/sessions_ipc.rs +++ b/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;