Merge pull request #22 from NianJiuZst/codex/enforce-session-idle-timeout

fix(daemon): enforce session idle timeout

Author: BB-fat

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,
+    ))
 }

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,

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<HashMap<SessionId, Session>>,
+    /// Operational metadata kept outside the public `Session` wire/domain
+    /// shape so idle enforcement does not break external struct users.
+    last_activity: Mutex<HashMap<SessionId, Instant>>,
 }
 
 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<Session> {
-        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<Session> {
@@ -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<SessionId> {
+        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<Session> {
         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
     }
 

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<DaemonState>) -> 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<Mutex<Instant>>) {
     if let Ok(mut a) = activity.lock() {
         *a = Instant::now();

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

@@ -81,11 +81,22 @@ pub struct DaemonHandle {
     state: Arc<DaemonState>,
     ws: WsHandle,
     ipc: Option<IpcHandle>,
+    session_idle_task: JoinHandle<()>,
 }
 
 impl DaemonHandle {
-    pub(crate) fn new(state: Arc<DaemonState>, ws: WsHandle, ipc: Option<IpcHandle>) -> Self {
-        Self { state, ws, ipc }
+    pub(crate) fn new(
+        state: Arc<DaemonState>,
+        ws: WsHandle,
+        ipc: Option<IpcHandle>,
+        session_idle_task: JoinHandle<()>,
+    ) -> Self {
+        Self {
+            state,
+            ws,
+            ipc,
+            session_idle_task,
+        }
     }
 
     pub fn state(&self) -> Arc<DaemonState> {
@@ -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 {

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<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>> {
@@ -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;