diff --git a/apps/extension/PRIVACY.md b/apps/extension/PRIVACY.md index 29527a4..9517a3c 100644 --- a/apps/extension/PRIVACY.md +++ b/apps/extension/PRIVACY.md @@ -48,6 +48,7 @@ The Extension requests the following Chrome permissions. Each is used solely for - **`tabs`** — Inspect, create, and close tabs in the Agent Window; query tab metadata. - **`windows`** — Create and manage the dedicated Agent Window that isolates agent activity from the user's normal browsing. - **`alarms`** — Periodically wake the service worker to keep the local WebSocket connection alive. +- **`idle`** — Detect when the device returns from idle/locked so the Extension can promptly re-establish the local WebSocket connection after the machine wakes. No idle data is stored or transmitted. - **`notifications`** — Show a system notification to obtain user approval before the agent borrows a user-owned tab. - **`storage`** — Persist a random instance ID and optional label in `chrome.storage.local`. - **Host permission ``** — Inject a small status overlay (showing "Agent Active") on pages controlled by the agent, and enable automation across whatever sites the user directs the agent to. The Extension does **not** read or transmit page content from sites the agent is not actively driving. diff --git a/apps/extension/package.json b/apps/extension/package.json index d651876..87a92db 100644 --- a/apps/extension/package.json +++ b/apps/extension/package.json @@ -1,6 +1,6 @@ { "name": "@browser-skill/extension", - "version": "0.1.4", + "version": "0.1.5", "private": true, "type": "module", "scripts": { diff --git a/apps/extension/src/entrypoints/background.ts b/apps/extension/src/entrypoints/background.ts index 26fd7b0..6192811 100644 --- a/apps/extension/src/entrypoints/background.ts +++ b/apps/extension/src/entrypoints/background.ts @@ -1,6 +1,7 @@ import { i18n } from "@browser-skill/i18n"; import { ChromiumCdp } from "@/browser-driver/chromium-cdp"; import { ConnectionController } from "@/lib/connection-controller"; +import { startHeartbeat } from "@/lib/heartbeat"; import { getConnectionEnabled, setConnectionEnabled as persistConnectionEnabled, @@ -139,6 +140,45 @@ export default defineBackground(() => { shouldConnect: () => controller.isConnectionEnabled, }); + // Application-level heartbeat (Chrome 116+): while the post-handshake + // link is live, beat every 20s so WebSocket activity keeps the service + // worker — and thus the daemon connection — alive during use, rather + // than depending on the boundary-hugging 30s keepalive alarm. The + // daemon also uses these beats to reap a silently-dead browser. + startHeartbeat({ + send: (frame) => transport.send(frame), + onActiveChange: (cb) => + controller.subscribe((snap) => + cb(snap.state === "connected" || snap.state === "version_skew"), + ), + }); + + // Wake-driven reconnect (best effort). An MV3 service worker is killed + // across OS sleep regardless of any keepalive, and the setTimeout-based + // transport backoff dies with it. These Chrome lifecycle events revive + // the worker and let us reconnect immediately instead of waiting for + // the next 30s alarm tick. They only help when the daemon is actually + // running; a cold daemon is (re)spawned by the next `bsk` command. + const reconnectIfNeeded = () => { + if (!controller.isConnectionEnabled) return; + if (transport.state === "connected") return; + void transport.connect().catch((err) => { + console.debug("[browser-skill] wake reconnect attempt failed", err); + }); + }; + if (typeof chrome.runtime?.onStartup?.addListener === "function") { + chrome.runtime.onStartup.addListener(reconnectIfNeeded); + } + if (typeof chrome.idle?.onStateChanged?.addListener === "function") { + chrome.idle.onStateChanged.addListener((state) => { + // "active" fires when the user returns from idle/locked — the most + // reliable "machine just woke" signal we get. + if (state === "active") reconnectIfNeeded(); + }); + // Treat >60s of no input as idle so the active transition is timely. + chrome.idle.setDetectionInterval?.(60); + } + void (async () => { const connectionEnabled = await getConnectionEnabled(); await controller.attach(transport, detectBrowserMeta(), connectionEnabled); diff --git a/apps/extension/src/lib/__tests__/heartbeat.test.ts b/apps/extension/src/lib/__tests__/heartbeat.test.ts new file mode 100644 index 0000000..12ddf74 --- /dev/null +++ b/apps/extension/src/lib/__tests__/heartbeat.test.ts @@ -0,0 +1,130 @@ +import { describe, expect, it, vi } from "vitest"; +import type { ProtocolFrame } from "@/transport/types"; +import { HEARTBEAT_EVENT, HEARTBEAT_INTERVAL_MS, startHeartbeat } from "../heartbeat"; + +function makeTimers() { + let nextId = 1; + const intervals = new Map void; ms: number }>(); + return { + intervals, + api: { + setInterval: (cb: () => void, ms: number) => { + const id = nextId++; + intervals.set(id, { cb, ms }); + return id; + }, + clearInterval: (id: number) => { + intervals.delete(id); + }, + }, + tick: (id: number) => intervals.get(id)?.cb(), + }; +} + +/** Drives `onActiveChange` subscribers so a test can flip the link state. */ +function makeActiveSource() { + const listeners = new Set<(active: boolean) => void>(); + return { + subscribe: (cb: (active: boolean) => void) => { + // Mirror the controller contract: fire immediately with the current + // value (starts inactive), then on every change. + cb(false); + listeners.add(cb); + return () => listeners.delete(cb); + }, + set: (active: boolean) => { + for (const cb of listeners) cb(active); + }, + listenerCount: () => listeners.size, + }; +} + +describe("startHeartbeat", () => { + it("does not beat while the link is inactive", () => { + const timers = makeTimers(); + const active = makeActiveSource(); + const send = vi.fn(); + startHeartbeat({ send, onActiveChange: active.subscribe, timers: timers.api }); + + expect(timers.intervals.size).toBe(0); + expect(send).not.toHaveBeenCalled(); + }); + + it("starts a 20s interval on activation and sends heartbeat events", () => { + const timers = makeTimers(); + const active = makeActiveSource(); + const send = vi.fn(); + startHeartbeat({ send, onActiveChange: active.subscribe, timers: timers.api }); + + active.set(true); + expect(timers.intervals.size).toBe(1); + const [id, info] = [...timers.intervals.entries()][0]; + expect(info.ms).toBe(HEARTBEAT_INTERVAL_MS); + + timers.tick(id); + expect(send).toHaveBeenCalledWith({ event: HEARTBEAT_EVENT, payload: {} }); + }); + + it("stops beating when the link goes inactive", () => { + const timers = makeTimers(); + const active = makeActiveSource(); + const send = vi.fn(); + startHeartbeat({ send, onActiveChange: active.subscribe, timers: timers.api }); + + active.set(true); + expect(timers.intervals.size).toBe(1); + active.set(false); + expect(timers.intervals.size).toBe(0); + }); + + it("does not stack intervals on repeated activation", () => { + const timers = makeTimers(); + const active = makeActiveSource(); + startHeartbeat({ send: vi.fn(), onActiveChange: active.subscribe, timers: timers.api }); + + active.set(true); + active.set(true); + expect(timers.intervals.size).toBe(1); + }); + + it("swallows send errors so the interval keeps running", () => { + const timers = makeTimers(); + const active = makeActiveSource(); + const send = vi.fn(() => { + throw new Error("socket closed"); + }); + const handle = startHeartbeat({ send, onActiveChange: active.subscribe, timers: timers.api }); + expect(() => handle.beatForTest()).not.toThrow(); + }); + + it("dispose clears the interval and unsubscribes", () => { + const timers = makeTimers(); + const active = makeActiveSource(); + const handle = startHeartbeat({ + send: vi.fn(), + onActiveChange: active.subscribe, + timers: timers.api, + }); + active.set(true); + expect(timers.intervals.size).toBe(1); + + handle.dispose(); + expect(timers.intervals.size).toBe(0); + expect(active.listenerCount()).toBe(0); + }); + + it("sends a well-formed protocol frame", () => { + const timers = makeTimers(); + const active = makeActiveSource(); + let captured: ProtocolFrame | null = null; + const handle = startHeartbeat({ + send: (f) => { + captured = f; + }, + onActiveChange: active.subscribe, + timers: timers.api, + }); + handle.beatForTest(); + expect(captured).toEqual({ event: HEARTBEAT_EVENT, payload: {} }); + }); +}); diff --git a/apps/extension/src/lib/heartbeat.ts b/apps/extension/src/lib/heartbeat.ts new file mode 100644 index 0000000..7869291 --- /dev/null +++ b/apps/extension/src/lib/heartbeat.ts @@ -0,0 +1,110 @@ +/** + * Application-level WebSocket heartbeat. + * + * Since Chrome 116, sending or receiving a message over a WebSocket + * resets the MV3 service-worker idle timer (see + * https://developer.chrome.com/docs/extensions/how-to/web-platform/websockets). + * By emitting a tiny `system.heartbeat` event every 20s — comfortably + * inside the 30s idle window — we keep the worker (and therefore the + * daemon link) alive for as long as the connection is up, instead of + * relying solely on the 30s keepalive alarm that sits right on the + * eviction boundary. + * + * The daemon also treats the heartbeat as a liveness signal, so a + * silently-dead browser can be reaped on its side. + * + * The heartbeat only runs while the *post-handshake* link is live: the + * daemon rejects any first frame that is not `system.handshake`, so + * beating before the handshake completes would get the connection + * kicked. Callers wire {@link HeartbeatOptions.onActiveChange} to the + * connection controller, which only reports `connected` / `version_skew` + * after the handshake has landed. + */ + +import type { ProtocolFrame } from "@/transport/types"; + +export const HEARTBEAT_EVENT = "system.heartbeat"; +/** 20s: the interval Chrome's own WebSocket keepalive sample recommends. */ +export const HEARTBEAT_INTERVAL_MS = 20_000; + +/** Subset of the timer API we use; lets tests inject fakes. */ +export interface HeartbeatTimers { + setInterval(cb: () => void, ms: number): number; + clearInterval(handle: number): void; +} + +export interface HeartbeatOptions { + /** Send a frame over the live transport. */ + send: (frame: ProtocolFrame) => void; + /** + * Subscribe to "is the post-handshake link live" changes. Must invoke + * the callback with the current value immediately and again on every + * transition, and return an unsubscribe function. + */ + onActiveChange: (cb: (active: boolean) => void) => () => void; + /** Override the beat interval (tests). */ + intervalMs?: number; + /** Inject fake timers (tests). */ + timers?: HeartbeatTimers; +} + +export interface HeartbeatHandle { + /** Stop beating and detach the active-state listener. */ + dispose(): void; + /** Test hook: emit one heartbeat frame right now. */ + beatForTest(): void; +} + +function defaultTimers(): HeartbeatTimers { + return { + setInterval: (cb, ms) => setInterval(cb, ms) as unknown as number, + clearInterval: (handle) => clearInterval(handle), + }; +} + +/** + * Start a heartbeat that beats every `intervalMs` while the link is + * active and pauses entirely while it is not. + */ +export function startHeartbeat(options: HeartbeatOptions): HeartbeatHandle { + const intervalMs = options.intervalMs ?? HEARTBEAT_INTERVAL_MS; + const timers = options.timers ?? defaultTimers(); + let timer: number | null = null; + + const beat = () => { + try { + options.send({ event: HEARTBEAT_EVENT, payload: {} }); + } catch (err) { + // The socket raced a close between the state change and this + // tick. The transport's own reconnect path recovers; swallow so + // the interval keeps running (it is cleared on the disconnect + // notification anyway). + console.debug("[bsk heartbeat] send failed", err); + } + }; + + const stop = () => { + if (timer !== null) { + timers.clearInterval(timer); + timer = null; + } + }; + + const start = () => { + if (timer !== null) return; + timer = timers.setInterval(beat, intervalMs); + }; + + const unsubscribe = options.onActiveChange((active) => { + if (active) start(); + else stop(); + }); + + return { + dispose: () => { + unsubscribe(); + stop(); + }, + beatForTest: beat, + }; +} diff --git a/apps/extension/wxt.config.ts b/apps/extension/wxt.config.ts index 4fd0ce8..34c72b6 100644 --- a/apps/extension/wxt.config.ts +++ b/apps/extension/wxt.config.ts @@ -25,9 +25,14 @@ export default defineConfig({ name: "BrowserSkill", description: "Let AI agents use your logged-in browser in a separate Agent Window—without interrupting your work. Powered by the bsk CLI.", + // Chrome 116 is the first version where WebSocket activity resets the + // service-worker idle timer, which the 20s heartbeat relies on to + // keep the daemon link alive. + minimum_chrome_version: "116", permissions: [ "alarms", "debugger", + "idle", "notifications", "tabs", "storage", diff --git a/crates/bsk-cli/src/cli/browser_wait.rs b/crates/bsk-cli/src/cli/browser_wait.rs index 56ef246..1ac4767 100644 --- a/crates/bsk-cli/src/cli/browser_wait.rs +++ b/crates/bsk-cli/src/cli/browser_wait.rs @@ -2,10 +2,16 @@ use std::time::Duration; -use crate::daemon::browsers::EXTENSION_CONNECT_WAIT; - /// Default wait passed to daemon status/list RPCs when the registry is empty. -pub const DEFAULT_BROWSER_CONNECT_WAIT: Duration = EXTENSION_CONNECT_WAIT; +/// +/// Deliberately shorter than the daemon's [`EXTENSION_CONNECT_WAIT`] +/// (used by `session.start`): `bsk status` / `bsk browsers` are quick +/// informational snapshots, so blocking them for a full keepalive-alarm +/// period would hurt more than it helps. The long wait is reserved for +/// the command path that actually needs the extension online. +/// +/// [`EXTENSION_CONNECT_WAIT`]: crate::daemon::browsers::EXTENSION_CONNECT_WAIT +pub const DEFAULT_BROWSER_CONNECT_WAIT: Duration = Duration::from_secs(5); /// Keep CLI-side waits aligned with the daemon's IPC clamp so a bad env value /// cannot inflate read timeouts after the daemon has already returned. const MAX_BROWSER_CONNECT_WAIT: Duration = Duration::from_secs(60); diff --git a/crates/bsk-cli/src/cli/session.rs b/crates/bsk-cli/src/cli/session.rs index 1fb1d36..198f568 100644 --- a/crates/bsk-cli/src/cli/session.rs +++ b/crates/bsk-cli/src/cli/session.rs @@ -6,6 +6,8 @@ //! structured JSON instead. use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use anyhow::Context; @@ -17,9 +19,18 @@ use serde::{Deserialize, Serialize}; use crate::cli::ensure_daemon::ensure_daemon; use crate::cli::error::{self, CliError, Format, RenderExtras}; +use crate::daemon::browsers::EXTENSION_CONNECT_WAIT; const SESSION_STOP_IPC_TIMEOUT: Duration = Duration::from_secs(60 * 60); +/// IPC read budget for `session.start`. The daemon holds this RPC open +/// while it polls for the extension to (re)connect for up to +/// [`EXTENSION_CONNECT_WAIT`], so the CLI-side timeout must exceed that +/// window plus scheduling slack — otherwise the CLI would give up and +/// surface a spurious timeout before the daemon ever answers. +const SESSION_START_IPC_TIMEOUT: Duration = + EXTENSION_CONNECT_WAIT.saturating_add(Duration::from_secs(10)); + /// `bsk session …` subcommand tree. #[derive(Debug, Clone, Args)] pub struct SessionCmd { @@ -113,7 +124,24 @@ pub fn dispatch(cmd: SessionCmd, format: Format) -> Result<(), CliError> { } fn run_start(sock: PathBuf, args: SessionStartArgs, format: Format) -> Result<(), CliError> { + // The daemon may hold `session.start` open for up to + // `EXTENSION_CONNECT_WAIT` while a just-woken service worker + // reconnects. Let the user know we are waiting rather than hung — + // but only if it actually takes a moment, so the common (already + // connected) fast path stays silent. Human output only; the hint + // goes to stderr so it never contaminates the session id on stdout. + let waited = Arc::new(AtomicBool::new(false)); + if matches!(format, Format::Human) { + let waited = Arc::clone(&waited); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(750)); + if !waited.swap(true, Ordering::SeqCst) { + eprintln!("waiting for browser extension to connect…"); + } + }); + } let result = start_session(sock, args.browser); + waited.store(true, Ordering::SeqCst); match result { Ok(reply) => match format { Format::Json => { @@ -144,7 +172,7 @@ pub fn start_session(sock: PathBuf, browser: Option) -> Result, + /// `true` once this browser has sent at least one `system.heartbeat`. + /// The liveness reaper only acts on heartbeat-capable browsers: a + /// pre-heartbeat extension never sets this, so it is never reaped on + /// silence and instead relies on socket-close detection exactly as + /// before — otherwise a new daemon paired with an old extension would + /// wrongly drop a live-but-idle connection. + pub heartbeat_seen: AtomicBool, } impl BrowserClient { + /// Record that a frame just arrived from this browser. + pub fn touch(&self) { + if let Ok(mut last) = self.last_seen.lock() { + *last = Instant::now(); + } + } + + /// Note that a `system.heartbeat` arrived, opting this browser in to + /// liveness reaping. + pub fn mark_heartbeat_seen(&self) { + self.heartbeat_seen.store(true, Ordering::Relaxed); + } + + /// Whether this browser has ever sent a heartbeat. + pub fn has_heartbeat(&self) -> bool { + self.heartbeat_seen.load(Ordering::Relaxed) + } + + /// How long since the last inbound frame from this browser. + pub fn idle_for(&self) -> Duration { + match self.last_seen.lock() { + Ok(last) => last.elapsed(), + Err(poisoned) => poisoned.into_inner().elapsed(), + } + } + pub fn status_entry(&self, session_count: u32) -> BrowserStatusEntry { BrowserStatusEntry { instance_id: self.id.0.clone(), @@ -206,6 +264,23 @@ impl BrowserRegistry { guard.values().cloned().collect() } + /// Return every heartbeat-capable browser whose last inbound frame is + /// older than `threshold`. Browsers that have never sent a heartbeat + /// (legacy extensions) are excluded so the reaper never drops a + /// live-but-idle connection it cannot probe. The caller is + /// responsible for dropping the returned entries (via + /// [`Self::remove_if_generation_matches`]) and purging their sessions + /// — done outside the lock so session teardown never blocks the + /// registry. + pub fn stale_browsers(&self, threshold: Duration) -> Vec> { + let guard = self.inner.lock().expect("browser registry poisoned"); + guard + .values() + .filter(|c| c.has_heartbeat() && c.idle_for() >= threshold) + .cloned() + .collect() + } + pub fn len(&self) -> usize { self.inner.lock().expect("browser registry poisoned").len() } @@ -348,6 +423,8 @@ mod tests { generation: next_browser_generation(), connected_at_ms: 0, version_skew: false, + last_seen: Mutex::new(Instant::now()), + heartbeat_seen: AtomicBool::new(false), }) } @@ -543,6 +620,51 @@ mod tests { assert_eq!(reg.len(), 1); } + #[test] + fn touch_resets_idle_for() { + let client = fake_client("a", ""); + // Force the last_seen well into the past, then touch. + *client.last_seen.lock().unwrap() = Instant::now() - Duration::from_secs(120); + assert!(client.idle_for() >= Duration::from_secs(120)); + client.touch(); + assert!(client.idle_for() < Duration::from_secs(1)); + } + + #[test] + fn stale_browsers_reports_only_silent_heartbeat_capable_connections() { + let reg = BrowserRegistry::new(); + let fresh = fake_client("fresh", ""); + fresh.mark_heartbeat_seen(); + let stale = fake_client("stale", ""); + stale.mark_heartbeat_seen(); + *stale.last_seen.lock().unwrap() = Instant::now() - Duration::from_secs(90); + reg.insert(fresh); + reg.insert(stale); + + let stale_ids: Vec = reg + .stale_browsers(Duration::from_secs(60)) + .into_iter() + .map(|c| c.id.0.clone()) + .collect(); + assert_eq!(stale_ids, vec!["stale".to_string()]); + } + + #[test] + fn stale_browsers_never_reaps_a_browser_that_never_sent_a_heartbeat() { + // A legacy extension that predates the heartbeat must fall back + // to socket-close detection, never liveness reaping — even when + // it has been silent well past the threshold. + let reg = BrowserRegistry::new(); + let legacy = fake_client("legacy", ""); + *legacy.last_seen.lock().unwrap() = Instant::now() - Duration::from_secs(600); + reg.insert(legacy); + + assert!( + reg.stale_browsers(Duration::from_secs(60)).is_empty(), + "a browser that never heartbeated must not be reaped" + ); + } + #[tokio::test] async fn wait_for_any_connected_skips_when_already_populated() { let reg = BrowserRegistry::new(); diff --git a/crates/bsk-cli/src/daemon/mod.rs b/crates/bsk-cli/src/daemon/mod.rs index 99ee7e6..c639f73 100644 --- a/crates/bsk-cli/src/daemon/mod.rs +++ b/crates/bsk-cli/src/daemon/mod.rs @@ -45,6 +45,7 @@ pub async fn run( None => None, }; let session_idle_task = start::spawn_session_idle_reaper(Arc::clone(&state)); + let browser_liveness_task = start::spawn_browser_liveness_reaper(Arc::clone(&state)); // Ensure WS/IPC accept loops have polled `shutdown.notified()` before any // caller can invoke `DaemonHandle::shutdown()` — `Notify::notify_waiters()` // drops wakeups when nothing is registered yet, which otherwise hangs @@ -55,5 +56,6 @@ pub async fn run( ws_handle, ipc_handle, session_idle_task, + browser_liveness_task, )) } diff --git a/crates/bsk-cli/src/daemon/start.rs b/crates/bsk-cli/src/daemon/start.rs index 1d72e63..33327a1 100644 --- a/crates/bsk-cli/src/daemon/start.rs +++ b/crates/bsk-cli/src/daemon/start.rs @@ -23,7 +23,7 @@ use tracing::{debug, info, warn}; use crate::cli::daemon::StartArgs; use crate::daemon::{ - browsers::EXTENSION_CONNECT_WAIT, + browsers::{BROWSER_LIVENESS_TICK, BROWSER_LIVENESS_TIMEOUT, EXTENSION_CONNECT_WAIT}, info as daemon_info, ipc, lockfile, paths, sessions::{StopSessionError, forget_session, stop_session}, state::DaemonState, @@ -45,6 +45,12 @@ pub struct DaemonConfig { /// How long `session.start` polls for an extension handshake before /// returning `no_browser_connected`. pub extension_connect_wait: Duration, + /// A heartbeat-capable browser silent for at least this long is + /// reaped by the liveness task. Defaults to [`BROWSER_LIVENESS_TIMEOUT`]. + pub browser_liveness_timeout: Duration, + /// How often the liveness task scans the registry. Defaults to + /// [`BROWSER_LIVENESS_TICK`]. + pub browser_liveness_tick: Duration, } impl DaemonConfig { @@ -58,6 +64,8 @@ impl DaemonConfig { daemon_idle: Duration::from_secs(60 * 30), allow_any_origin: false, extension_connect_wait: EXTENSION_CONNECT_WAIT, + browser_liveness_timeout: BROWSER_LIVENESS_TIMEOUT, + browser_liveness_tick: BROWSER_LIVENESS_TICK, } } @@ -65,6 +73,15 @@ impl DaemonConfig { self.extension_connect_wait = wait; self } + + /// Override the liveness reaper's silence threshold and scan cadence. + /// Primarily for tests that need the reaper to act within + /// sub-second windows instead of the production 60s/15s defaults. + pub fn with_browser_liveness(mut self, timeout: Duration, tick: Duration) -> Self { + self.browser_liveness_timeout = timeout; + self.browser_liveness_tick = tick; + self + } } impl From<&StartArgs> for DaemonConfig { @@ -75,6 +92,8 @@ impl From<&StartArgs> for DaemonConfig { daemon_idle: args.resolved_daemon_idle(), allow_any_origin: false, extension_connect_wait: EXTENSION_CONNECT_WAIT, + browser_liveness_timeout: BROWSER_LIVENESS_TIMEOUT, + browser_liveness_tick: BROWSER_LIVENESS_TICK, } } } @@ -200,6 +219,7 @@ pub fn run_foreground(cfg: DaemonConfig) -> Result<()> { let state = Arc::new(DaemonState::new(cfg.clone())); let session_idle_task = spawn_session_idle_reaper(Arc::clone(&state)); + let browser_liveness_task = spawn_browser_liveness_reaper(Arc::clone(&state)); let ws_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), cfg.ws_port); let ws_handle = ws::WsServer::new(Arc::clone(&state)) .bind(ws_addr) @@ -361,6 +381,8 @@ pub fn run_foreground(cfg: DaemonConfig) -> Result<()> { let _ = ipc_task.await; session_idle_task.abort(); let _ = session_idle_task.await; + browser_liveness_task.abort(); + let _ = browser_liveness_task.await; ws_handle.shutdown.notify_waiters(); let _ = ws_handle.task.await; @@ -373,6 +395,53 @@ pub fn run_foreground(cfg: DaemonConfig) -> Result<()> { Ok(()) } +/// Spawn the browser-liveness reaper shared by the production foreground +/// daemon and the test/embed entry point. +/// +/// A normal disconnect closes the WebSocket, which the per-connection +/// read loop observes and cleans up. But if the OS resumed from sleep +/// and killed the MV3 service worker, the socket can be left half-open +/// with no close frame ever delivered. Without the ~20s `system.heartbeat` +/// arriving, such a connection would otherwise sit in the registry +/// forever — pinning a phantom "connected" browser and preventing the +/// daemon from ever idle-exiting. This reaper drops any browser that has +/// gone silent past [`BROWSER_LIVENESS_TIMEOUT`] and purges its sessions, +/// mirroring the WS disconnect cleanup. +pub(crate) fn spawn_browser_liveness_reaper( + state: Arc, +) -> tokio::task::JoinHandle<()> { + let timeout = state.config.browser_liveness_timeout; + let tick = state.config.browser_liveness_tick; + tokio::spawn(async move { + let mut ticker = tokio::time::interval(tick); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + // The first tick is immediate; skip it so a freshly connected + // browser always gets at least one full window before a scan. + ticker.tick().await; + loop { + ticker.tick().await; + for client in state.browsers.stale_browsers(timeout) { + if state + .browsers + .remove_if_generation_matches(&client.id, client.generation) + .is_some() + { + warn!( + id = %client.id, + idle_secs = client.idle_for().as_secs(), + "reaping unresponsive browser (no heartbeat within liveness window)" + ); + for s in state.sessions.purge_browser(&client.id) { + state.tool_queues.remove(&s.id); + state.session_interrupts.drop_session(&s.id); + debug!(session = %s.id, "purged session on browser liveness timeout"); + } + } + } + } + }) +} + /// Spawn the cooperative session-idle reaper shared by the production /// foreground daemon and the test/embed daemon entry point. pub(crate) fn spawn_session_idle_reaper(state: Arc) -> tokio::task::JoinHandle<()> { diff --git a/crates/bsk-cli/src/daemon/state.rs b/crates/bsk-cli/src/daemon/state.rs index 203b3ac..d6632da 100644 --- a/crates/bsk-cli/src/daemon/state.rs +++ b/crates/bsk-cli/src/daemon/state.rs @@ -82,6 +82,7 @@ pub struct DaemonHandle { ws: WsHandle, ipc: Option, session_idle_task: JoinHandle<()>, + browser_liveness_task: JoinHandle<()>, } impl DaemonHandle { @@ -90,12 +91,14 @@ impl DaemonHandle { ws: WsHandle, ipc: Option, session_idle_task: JoinHandle<()>, + browser_liveness_task: JoinHandle<()>, ) -> Self { Self { state, ws, ipc, session_idle_task, + browser_liveness_task, } } @@ -116,6 +119,8 @@ impl DaemonHandle { pub async fn shutdown(self) { self.session_idle_task.abort(); let _ = await_join(self.session_idle_task).await; + self.browser_liveness_task.abort(); + let _ = await_join(self.browser_liveness_task).await; self.ws.shutdown.notify_waiters(); let _ = await_join(self.ws.task).await; if let Some(ipc) = self.ipc { diff --git a/crates/bsk-cli/src/daemon/ws.rs b/crates/bsk-cli/src/daemon/ws.rs index 4073771..7d34cc1 100644 --- a/crates/bsk-cli/src/daemon/ws.rs +++ b/crates/bsk-cli/src/daemon/ws.rs @@ -304,6 +304,8 @@ async fn drive_connection( generation, connected_at_ms, version_skew, + last_seen: std::sync::Mutex::new(std::time::Instant::now()), + heartbeat_seen: std::sync::atomic::AtomicBool::new(false), }); state.browsers.insert(Arc::clone(&client)); info!( @@ -352,15 +354,21 @@ async fn drive_connection( msg = reader.next() => { match msg { Some(Ok(Message::Text(t))) => { + // Any inbound frame — tool response, event, or the + // ~20s heartbeat — counts as proof of life. + pump_browser.touch(); handle_inbound_text(&pump_state, &pump_browser, &t).await; } Some(Ok(Message::Binary(_))) => { warn!("ignoring binary message"); } Some(Ok(Message::Ping(p))) => { + pump_browser.touch(); let _ = writer.send(Message::Pong(p)).await; } - Some(Ok(Message::Pong(_))) => {} + Some(Ok(Message::Pong(_))) => { + pump_browser.touch(); + } Some(Ok(Message::Frame(_))) => {} Some(Ok(Message::Close(_))) | None => break, Some(Err(err)) => { @@ -416,6 +424,13 @@ async fn handle_inbound_text(state: &Arc, client: &Arc match ev.event { + bsk_protocol::EventKind::SystemHeartbeat => { + // `touch()` already ran for this frame in the read loop; + // additionally opt this browser in to liveness reaping now + // that we know it speaks the heartbeat. + client.mark_heartbeat_seen(); + debug!(id = %client.id, "heartbeat"); + } bsk_protocol::EventKind::SessionWindowClosed => { handle_session_window_closed(state, &client.id, &ev.payload); } diff --git a/crates/bsk-cli/tests/browser_list_ordering.rs b/crates/bsk-cli/tests/browser_list_ordering.rs index f19f320..86031d7 100644 --- a/crates/bsk-cli/tests/browser_list_ordering.rs +++ b/crates/bsk-cli/tests/browser_list_ordering.rs @@ -51,6 +51,8 @@ fn fake_client(id: &str, label: &str, connected_at_ms: i64) -> Arc std::sync::Arc { + let (tx, _rx) = mpsc::unbounded_channel::(); + std::sync::Arc::new(BrowserClient { + id: BrowserId(id.into()), + browser_name: "chrome".into(), + browser_version: "131.0".into(), + extension_version: "0.1.4".into(), + extension_protocol_version: "1.0".into(), + label: String::new(), + sink: BrowserSink { tx }, + pending: Mutex::new(Pending::default()), + generation: next_browser_generation(), + connected_at_ms: 0, + version_skew: false, + last_seen: Mutex::new(Instant::now() - Duration::from_secs(idle_secs)), + heartbeat_seen: AtomicBool::new(heartbeat_seen), + }) +} + +fn fake_session(session_id: &str, browser_id: &str) -> Session { + Session { + id: SessionId(session_id.into()), + browser_id: BrowserId(browser_id.into()), + agent_window_id: None, + created_at_ms: 0, + } +} + +async fn spawn_reaping_daemon() -> daemon::DaemonHandle { + // Tiny window so the reaper acts within the test's lifetime. + let config = DaemonConfig::new(0) + .with_browser_liveness(Duration::from_millis(150), Duration::from_millis(25)); + daemon::run(config, None).await.unwrap() +} + +async fn wait_until bool>(pred: F, within: Duration) -> bool { + let deadline = Instant::now() + within; + while Instant::now() < deadline { + if pred() { + return true; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + pred() +} + +#[tokio::test] +async fn reaper_drops_silent_heartbeat_browser_and_purges_its_sessions() { + let handle = spawn_reaping_daemon().await; + let state = handle.state(); + + state.browsers.insert(fake_client("hb", true, 10)); + state.sessions.insert(fake_session("sess", "hb")); + assert_eq!(state.browsers.len(), 1); + assert_eq!(state.sessions.count_for_browser(&BrowserId("hb".into())), 1); + + let dropped = wait_until( + || state.browsers.get(&BrowserId("hb".into())).is_none(), + Duration::from_secs(2), + ) + .await; + assert!(dropped, "silent heartbeat-capable browser should be reaped"); + + // The reaper also tears down the browser's sessions. + assert!( + wait_until(|| state.sessions.is_empty(), Duration::from_secs(1)).await, + "reaping a browser must purge its sessions" + ); + + handle.shutdown().await; +} + +#[tokio::test] +async fn reaper_leaves_legacy_browser_without_heartbeat_alone() { + let handle = spawn_reaping_daemon().await; + let state = handle.state(); + + // Silent far past the window, but never heartbeated → must survive. + state.browsers.insert(fake_client("legacy", false, 600)); + + // Give the reaper several scan cycles to (wrongly) act. + tokio::time::sleep(Duration::from_millis(300)).await; + assert!( + state.browsers.get(&BrowserId("legacy".into())).is_some(), + "a browser that never heartbeated must not be reaped on silence" + ); + + handle.shutdown().await; +} diff --git a/crates/bsk-cli/tests/handshake_compat.rs b/crates/bsk-cli/tests/handshake_compat.rs index 596d35e..a5e91b0 100644 --- a/crates/bsk-cli/tests/handshake_compat.rs +++ b/crates/bsk-cli/tests/handshake_compat.rs @@ -242,6 +242,8 @@ async fn status_surfaces_version_skew_for_skewed_browser() { generation: bsk::daemon::browsers::next_browser_generation(), connected_at_ms: 0, version_skew: true, + last_seen: Mutex::new(std::time::Instant::now()), + heartbeat_seen: std::sync::atomic::AtomicBool::new(false), }); state.browsers.insert(client); diff --git a/crates/bsk-protocol/src/frame.rs b/crates/bsk-protocol/src/frame.rs index 019e58a..996d487 100644 --- a/crates/bsk-protocol/src/frame.rs +++ b/crates/bsk-protocol/src/frame.rs @@ -40,6 +40,14 @@ pub struct EventFrame { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum EventKind { + /// Application-level keepalive emitted by the extension roughly + /// every 20s while the WS link is up. Two purposes: (1) the + /// send/receive activity resets the MV3 service-worker idle timer + /// (Chrome 116+), keeping the worker — and therefore the socket — + /// alive during use; (2) the daemon treats it as a liveness signal + /// so a silently-dead browser can be reaped. Carries no payload. + #[serde(rename = "system.heartbeat")] + SystemHeartbeat, #[serde(rename = "session.activity")] SessionActivity, #[serde(rename = "session.window_closed")] @@ -230,6 +238,24 @@ impl<'de> Visitor<'de> for FrameVisitor { mod tests { use super::*; + #[test] + fn system_heartbeat_serialises_as_dotted_name() { + // The extension hardcodes the literal string "system.heartbeat" + // when it emits the keepalive event; this locks the daemon-side + // serde name to the same wire value so a rename cannot silently + // break liveness/keepalive. + let v = serde_json::to_value(EventKind::SystemHeartbeat).unwrap(); + assert_eq!(v, serde_json::json!("system.heartbeat")); + } + + #[test] + fn system_heartbeat_event_frame_round_trips_from_extension_shape() { + // Mirrors exactly what the extension sends: { event, payload: {} }. + let wire = serde_json::json!({ "event": "system.heartbeat", "payload": {} }); + let frame: EventFrame = serde_json::from_value(wire).unwrap(); + assert_eq!(frame.event, EventKind::SystemHeartbeat); + } + #[test] fn session_user_interrupt_serialises_as_snake_case() { let v = serde_json::to_value(EventKind::SessionUserInterrupt).unwrap();