diff --git a/crates/bsk-cli/src/daemon/mod.rs b/crates/bsk-cli/src/daemon/mod.rs index 4d1d79c..99ee7e6 100644 --- a/crates/bsk-cli/src/daemon/mod.rs +++ b/crates/bsk-cli/src/daemon/mod.rs @@ -44,10 +44,16 @@ pub async fn run( Some(path) => Some(ipc::IpcServer::new(Arc::clone(&state)).bind(path).await?), None => None, }; + let session_idle_task = start::spawn_session_idle_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 // shutdown forever (hit by tests that connect/shutdown immediately). tokio::task::yield_now().await; - Ok(DaemonHandle::new(state, ws_handle, ipc_handle)) + Ok(DaemonHandle::new( + state, + ws_handle, + ipc_handle, + session_idle_task, + )) } diff --git a/crates/bsk-cli/src/daemon/queue.rs b/crates/bsk-cli/src/daemon/queue.rs index 3e79e5c..919c207 100644 --- a/crates/bsk-cli/src/daemon/queue.rs +++ b/crates/bsk-cli/src/daemon/queue.rs @@ -300,7 +300,8 @@ impl ToolQueueRegistry { queue_state.busy = true; (entry.sender.clone(), Arc::clone(&entry.state)) }; - dispatch_with_sender( + self.sessions.touch(sid); + let outcome = dispatch_with_sender( sender, state, sid.clone(), @@ -310,7 +311,13 @@ impl ToolQueueRegistry { false, inflight, ) - .await + .await; + // A long-running request may outlive the idle threshold. Touching + // after completion prevents the reaper from immediately closing it; + // while it is running, session.stop observes SessionBusy and retries + // on a later sweep rather than interrupting the tool. + self.sessions.touch(sid); + outcome } /// Stop accepting new jobs for `sid`, enqueue one final control RPC, diff --git a/crates/bsk-cli/src/daemon/sessions.rs b/crates/bsk-cli/src/daemon/sessions.rs index aff784f..60cf0fd 100644 --- a/crates/bsk-cli/src/daemon/sessions.rs +++ b/crates/bsk-cli/src/daemon/sessions.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; use bsk_protocol::system::{BrowserStatusEntry, SessionStatusEntry}; use bsk_protocol::tools::{ @@ -66,6 +66,9 @@ impl Session { #[derive(Debug, Default)] pub struct SessionRegistry { inner: Mutex>, + /// Operational metadata kept outside the public `Session` wire/domain + /// shape so idle enforcement does not break external struct users. + last_activity: Mutex>, } impl SessionRegistry { @@ -100,10 +103,13 @@ impl SessionRegistry { } pub fn insert(&self, session: Session) { - self.inner + let session_id = session.id.clone(); + let mut sessions = self.inner.lock().expect("session registry poisoned"); + sessions.insert(session_id.clone(), session); + self.last_activity .lock() - .expect("session registry poisoned") - .insert(session.id.clone(), session); + .expect("session activity registry poisoned") + .insert(session_id, Instant::now()); } /// Reserve a fresh, collision-free [`SessionId`] under the registry @@ -141,6 +147,10 @@ impl SessionRegistry { created_at_ms: now_ms_fn(), }, ); + self.last_activity + .lock() + .expect("session activity registry poisoned") + .insert(candidate.clone(), Instant::now()); return Some(candidate); } None @@ -168,13 +178,23 @@ impl SessionRegistry { .lock() .expect("session registry poisoned") .remove(session_id); + self.last_activity + .lock() + .expect("session activity registry poisoned") + .remove(session_id); } pub fn remove(&self, id: &SessionId) -> Option { - self.inner + let removed = self + .inner .lock() .expect("session registry poisoned") - .remove(id) + .remove(id); + self.last_activity + .lock() + .expect("session activity registry poisoned") + .remove(id); + removed } pub fn get(&self, id: &SessionId) -> Option { @@ -185,6 +205,42 @@ impl SessionRegistry { .cloned() } + /// Record accepted or completed tool activity for a live session. + pub fn touch(&self, id: &SessionId) -> bool { + if !self + .inner + .lock() + .expect("session registry poisoned") + .contains_key(id) + { + return false; + } + self.last_activity + .lock() + .expect("session activity registry poisoned") + .insert(id.clone(), Instant::now()); + true + } + + /// Return sessions whose last tool activity is at least `idle_for` + /// old. The caller supplies `now` to keep boundary tests deterministic. + pub fn idle_ids_at(&self, idle_for: Duration, now: Instant) -> Vec { + let sessions = self.inner.lock().expect("session registry poisoned"); + let activity = self + .last_activity + .lock() + .expect("session activity registry poisoned"); + sessions + .values() + .filter(|session| { + activity + .get(&session.id) + .is_some_and(|last| now.saturating_duration_since(*last) >= idle_for) + }) + .map(|session| session.id.clone()) + .collect() + } + /// Drop all sessions owned by `browser_id` (e.g. on disconnect). pub fn purge_browser(&self, browser_id: &BrowserId) -> Vec { let mut guard = self.inner.lock().expect("session registry poisoned"); @@ -196,6 +252,13 @@ impl SessionRegistry { for s in &drained { guard.remove(&s.id); } + let mut activity = self + .last_activity + .lock() + .expect("session activity registry poisoned"); + for session in &drained { + activity.remove(&session.id); + } drained } diff --git a/crates/bsk-cli/src/daemon/start.rs b/crates/bsk-cli/src/daemon/start.rs index d5f64e9..1d72e63 100644 --- a/crates/bsk-cli/src/daemon/start.rs +++ b/crates/bsk-cli/src/daemon/start.rs @@ -23,8 +23,11 @@ use tracing::{debug, info, warn}; use crate::cli::daemon::StartArgs; use crate::daemon::{ - browsers::EXTENSION_CONNECT_WAIT, info as daemon_info, ipc, lockfile, paths, - state::DaemonState, ws, + browsers::EXTENSION_CONNECT_WAIT, + info as daemon_info, ipc, lockfile, paths, + sessions::{StopSessionError, forget_session, stop_session}, + state::DaemonState, + ws, }; /// Internal env-var contract: the parent sets this on the spawned child @@ -196,6 +199,7 @@ pub fn run_foreground(cfg: DaemonConfig) -> Result<()> { .with_context(|| format!("bind IPC endpoint {}", sock_path.display()))?; let state = Arc::new(DaemonState::new(cfg.clone())); + let session_idle_task = spawn_session_idle_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) @@ -355,6 +359,8 @@ pub fn run_foreground(cfg: DaemonConfig) -> Result<()> { let _ = ipc_shutdown_tx.send(()); let _ = ipc_task.await; + session_idle_task.abort(); + let _ = session_idle_task.await; ws_handle.shutdown.notify_waiters(); let _ = ws_handle.task.await; @@ -367,6 +373,55 @@ pub fn run_foreground(cfg: DaemonConfig) -> Result<()> { Ok(()) } +/// 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<()> { + tokio::spawn(async move { + let session_idle = state.config.session_idle; + let tick = (session_idle / 4) + .max(Duration::from_millis(100)) + .min(Duration::from_secs(30)); + let mut ticker = tokio::time::interval(tick); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + // `interval`'s first tick is immediate. Consume it so a zero/very + // short test setting still gets one real inactivity window. + ticker.tick().await; + + loop { + ticker.tick().await; + let idle_ids = state.sessions.idle_ids_at(session_idle, Instant::now()); + for session_id in idle_ids { + match stop_session( + &state.browsers, + &state.sessions, + &state.tool_queues, + &state.session_interrupts, + &session_id, + Duration::from_secs(10), + ) + .await + { + Ok(_) => info!(session = %session_id, "idle session stopped"), + Err(StopSessionError::SessionBusy | StopSessionError::Stopping) => { + debug!(session = %session_id, "idle session still active; retrying later"); + } + Err(StopSessionError::NotFound | StopSessionError::BrowserGone) => { + forget_session( + &state.sessions, + &state.tool_queues, + &state.session_interrupts, + &session_id, + ); + } + Err(err) => { + warn!(session = %session_id, error = %err, "failed to stop idle session"); + } + } + } + } + }) +} + fn record_activity(activity: &Arc>) { if let Ok(mut a) = activity.lock() { *a = Instant::now(); diff --git a/crates/bsk-cli/src/daemon/state.rs b/crates/bsk-cli/src/daemon/state.rs index d5acc8e..203b3ac 100644 --- a/crates/bsk-cli/src/daemon/state.rs +++ b/crates/bsk-cli/src/daemon/state.rs @@ -81,11 +81,22 @@ pub struct DaemonHandle { state: Arc, ws: WsHandle, ipc: Option, + session_idle_task: JoinHandle<()>, } impl DaemonHandle { - pub(crate) fn new(state: Arc, ws: WsHandle, ipc: Option) -> Self { - Self { state, ws, ipc } + pub(crate) fn new( + state: Arc, + ws: WsHandle, + ipc: Option, + session_idle_task: JoinHandle<()>, + ) -> Self { + Self { + state, + ws, + ipc, + session_idle_task, + } } pub fn state(&self) -> Arc { @@ -103,6 +114,8 @@ impl DaemonHandle { /// Stop the WS server (and IPC if running). Returns once both join /// handles complete. pub async fn shutdown(self) { + self.session_idle_task.abort(); + let _ = await_join(self.session_idle_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/tests/sessions_ipc.rs b/crates/bsk-cli/tests/sessions_ipc.rs index 99b1535..29cff59 100644 --- a/crates/bsk-cli/tests/sessions_ipc.rs +++ b/crates/bsk-cli/tests/sessions_ipc.rs @@ -46,6 +46,16 @@ async fn spawn_daemon_with_connect_wait(connect_wait: Duration) -> (daemon::Daem (handle, sock) } +async fn spawn_daemon_with_session_idle(session_idle: Duration) -> (daemon::DaemonHandle, PathBuf) { + let port = 0; + + let mut config = DaemonConfig::new(port); + config.session_idle = session_idle; + let sock = tempfile_path("bsk-test-ipc"); + let handle = daemon::run(config, Some(sock.clone())).await.unwrap(); + (handle, sock) +} + async fn connect_ext( addr: std::net::SocketAddr, ) -> tokio_tungstenite::WebSocketStream> { @@ -237,6 +247,47 @@ async fn session_start_stop_round_trip_via_ipc() { handle.shutdown().await; } +#[tokio::test] +async fn session_idle_timeout_stops_and_unregisters_session() { + let (handle, _sock) = spawn_daemon_with_session_idle(Duration::from_millis(50)).await; + let mut ws = connect_ext(handle.ws_addr()).await; + let _ = handshake_as_ext(&mut ws).await; + let state = handle.state(); + let session_id = bsk::daemon::sessions::SessionId("idle".into()); + state.sessions.insert(bsk::daemon::sessions::Session { + id: session_id.clone(), + browser_id: bsk::daemon::browsers::BrowserId(TEST_EXT_ID.into()), + agent_window_id: Some(7), + created_at_ms: 0, + }); + state.tool_queues.spawn(session_id); + + let request = tokio::time::timeout(Duration::from_secs(1), ws.next()) + .await + .expect("idle reaper did not contact extension") + .expect("extension socket closed") + .expect("extension socket failed"); + let Message::Text(text) = request else { + panic!("expected text request"); + }; + let request: RequestFrame = serde_json::from_str(&text).unwrap(); + assert_eq!(request.method, Method::ToolSessionStop); + ws.send(Message::Text( + serde_json::to_string(&ResponseFrame { + id: request.id, + body: ResponseBody::Ok( + serde_json::to_value(bsk_protocol::tools::SessionStopResult::default()).unwrap(), + ), + }) + .unwrap(), + )) + .await + .unwrap(); + + wait_for_no_sessions(&state).await; + handle.shutdown().await; +} + #[tokio::test] async fn session_start_errors_without_browser() { let (handle, sock) = spawn_daemon_with_connect_wait(Duration::ZERO).await;