diff --git a/crates/bsk-cli/src/daemon/start.rs b/crates/bsk-cli/src/daemon/start.rs index 0e48fc6..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, } } } @@ -391,16 +410,17 @@ pub fn run_foreground(cfg: DaemonConfig) -> Result<()> { pub(crate) fn spawn_browser_liveness_reaper( state: Arc, ) -> tokio::task::JoinHandle<()> { - use crate::daemon::browsers::{BROWSER_LIVENESS_TICK, BROWSER_LIVENESS_TIMEOUT}; + let timeout = state.config.browser_liveness_timeout; + let tick = state.config.browser_liveness_tick; tokio::spawn(async move { - let mut ticker = tokio::time::interval(BROWSER_LIVENESS_TICK); + 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(BROWSER_LIVENESS_TIMEOUT) { + for client in state.browsers.stale_browsers(timeout) { if state .browsers .remove_if_generation_matches(&client.id, client.generation) diff --git a/crates/bsk-cli/tests/browser_liveness.rs b/crates/bsk-cli/tests/browser_liveness.rs new file mode 100644 index 0000000..444b7ab --- /dev/null +++ b/crates/bsk-cli/tests/browser_liveness.rs @@ -0,0 +1,111 @@ +//! Integration coverage for the daemon's browser-liveness reaper +//! (`spawn_browser_liveness_reaper`). +//! +//! These drive the *task* — not just the `stale_browsers` predicate — +//! by spinning up a real daemon with a sub-second liveness window and +//! observing that a silent, heartbeat-capable browser (and its sessions) +//! get dropped, while a legacy browser that never heartbeated is left +//! alone (backward-compat guard). + +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::time::{Duration, Instant}; + +use bsk::daemon::browsers::{ + BrowserClient, BrowserId, BrowserSink, Pending, next_browser_generation, +}; +use bsk::daemon::sessions::{Session, SessionId}; +use bsk::daemon::{self, DaemonConfig}; +use tokio::sync::mpsc; + +/// Build a registry client. `heartbeat_seen` opts it in to reaping and +/// `idle_secs` back-dates its `last_seen` so it looks silent immediately. +fn fake_client(id: &str, heartbeat_seen: bool, idle_secs: u64) -> 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; +}