Track user activity to avoid false fork recovery

Author: richiemcilroy

crates/sandchest-agent/src/main.rs

@@ -22,6 +22,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         )
         .init();
 
+    snapshot::init_activity_clock();
+
     // Create service first so we have access to the session manager
     let service = service::GuestAgentService::new();
     let session_manager = service.session_manager();

crates/sandchest-agent/src/service.rs

@@ -9,6 +9,7 @@ use crate::proto::{
     GetFileRequest, HealthResponse, ListFilesRequest, ListFilesResponse, PutFileResponse,
     SessionExecRequest, SessionInputRequest, SessionResponse,
 };
+use crate::snapshot;
 use crate::session::SessionManager;
 
 pub struct GuestAgentService {
@@ -46,6 +47,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<ExecRequest>,
     ) -> Result<Response<Self::ExecStream>, Status> {
+        snapshot::note_user_activity();
         let stream = crate::exec::spawn_exec(request.into_inner());
         Ok(Response::new(stream))
     }
@@ -54,6 +56,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<CreateSessionRequest>,
     ) -> Result<Response<SessionResponse>, Status> {
+        snapshot::note_user_activity();
         let req = request.into_inner();
         let session_id = self
             .session_manager
@@ -68,6 +71,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<SessionExecRequest>,
     ) -> Result<Response<Self::SessionExecStream>, Status> {
+        snapshot::note_user_activity();
         let req = request.into_inner();
         let stream = self
             .session_manager
@@ -80,6 +84,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<SessionInputRequest>,
     ) -> Result<Response<()>, Status> {
+        snapshot::note_user_activity();
         let req = request.into_inner();
         self.session_manager
             .session_input(&req.session_id, &req.data)
@@ -91,6 +96,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<DestroySessionRequest>,
     ) -> Result<Response<()>, Status> {
+        snapshot::note_user_activity();
         let req = request.into_inner();
         self.session_manager
             .destroy_session(&req.session_id)
@@ -102,6 +108,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<Streaming<FileChunk>>,
     ) -> Result<Response<PutFileResponse>, Status> {
+        snapshot::note_user_activity();
         let response = crate::files::put_file(request.into_inner()).await?;
         Ok(Response::new(response))
     }
@@ -112,6 +119,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<GetFileRequest>,
     ) -> Result<Response<Self::GetFileStream>, Status> {
+        snapshot::note_user_activity();
         let stream = crate::files::spawn_get_file(request.into_inner());
         Ok(Response::new(stream))
     }
@@ -120,6 +128,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<ListFilesRequest>,
     ) -> Result<Response<ListFilesResponse>, Status> {
+        snapshot::note_user_activity();
         let response = crate::files::list_files(request.into_inner()).await?;
         Ok(Response::new(response))
     }

crates/sandchest-agent/src/snapshot.rs

@@ -1,5 +1,6 @@
 use std::path::Path;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use tracing::{info, warn};
@@ -11,34 +12,56 @@ const HEARTBEAT_INTERVAL_SECS: u64 = 1;
 const STALE_THRESHOLD_SECS: u64 = 5;
 #[cfg(target_os = "linux")]
 const URANDOM_SEED_BYTES: usize = 256;
+static LAST_USER_ACTIVITY_SECS: AtomicU64 = AtomicU64::new(0);
 
-/// Check if a heartbeat file at the given path is stale (indicating snapshot restore).
-fn is_heartbeat_stale(path: &Path) -> bool {
+fn current_unix_secs() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs()
+}
+
+pub fn init_activity_clock() {
+    note_user_activity();
+}
+
+pub fn note_user_activity() {
+    LAST_USER_ACTIVITY_SECS.store(current_unix_secs(), Ordering::Relaxed);
+}
+
+fn last_user_activity_secs() -> u64 {
+    LAST_USER_ACTIVITY_SECS.load(Ordering::Relaxed)
+}
+
+fn read_heartbeat_timestamp(path: &Path) -> Option<u64> {
     if !path.exists() {
-        return false;
+        return None;
     }
 
-    let contents = match std::fs::read_to_string(path) {
-        Ok(c) => c,
-        Err(_) => return false,
-    };
+    let contents = std::fs::read_to_string(path).ok()?;
+    contents.trim().parse().ok()
+}
 
-    let file_ts: u64 = match contents.trim().parse() {
-        Ok(ts) => ts,
-        Err(_) => return false,
-    };
+fn heartbeat_is_stale(file_ts: u64, now: u64) -> bool {
+    now > file_ts && (now - file_ts) > STALE_THRESHOLD_SECS
+}
 
-    let now = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap_or_default()
-        .as_secs();
+fn should_perform_fork_recovery(file_ts: u64, now: u64, last_activity: u64) -> bool {
+    heartbeat_is_stale(file_ts, now) && last_activity <= file_ts
+}
 
-    if now > file_ts && (now - file_ts) > STALE_THRESHOLD_SECS {
-        info!(
-            stale_secs = now - file_ts,
-            "Stale heartbeat detected — snapshot restore likely"
-        );
-        return true;
+/// Check if a heartbeat file at the given path is stale (indicating snapshot restore).
+fn is_heartbeat_stale(path: &Path) -> bool {
+    if let Some(file_ts) = read_heartbeat_timestamp(path) {
+        let now = current_unix_secs();
+
+        if heartbeat_is_stale(file_ts, now) {
+            info!(
+                stale_secs = now - file_ts,
+                "Stale heartbeat detected — snapshot restore likely"
+            );
+            return true;
+        }
     }
 
     false
@@ -230,8 +253,25 @@ pub fn start_snapshot_watcher(session_manager: Arc<SessionManager>) {
     tokio::spawn(async move {
         loop {
             // Check for stale heartbeat BEFORE writing a fresh one
-            if is_heartbeat_stale(Path::new(HEARTBEAT_PATH)) {
-                perform_fork_recovery(&session_manager).await;
+            let heartbeat_path = Path::new(HEARTBEAT_PATH);
+            if let Some(file_ts) = read_heartbeat_timestamp(heartbeat_path) {
+                let now = current_unix_secs();
+                let last_activity = last_user_activity_secs();
+
+                if should_perform_fork_recovery(file_ts, now, last_activity) {
+                    info!(
+                        stale_secs = now - file_ts,
+                        "Stale heartbeat detected — snapshot restore likely"
+                    );
+                    perform_fork_recovery(&session_manager).await;
+                } else if heartbeat_is_stale(file_ts, now) {
+                    warn!(
+                        stale_secs = now - file_ts,
+                        last_activity_secs = last_activity,
+                        heartbeat_secs = file_ts,
+                        "Skipping fork recovery because the agent has already served post-resume traffic"
+                    );
+                }
             }
 
             write_heartbeat().await;
@@ -335,6 +375,15 @@ mod tests {
         assert!(!is_heartbeat_stale(&path));
     }
 
+    #[test]
+    fn fork_recovery_requires_no_post_resume_activity() {
+        assert!(should_perform_fork_recovery(100, 110, 100));
+        assert!(should_perform_fork_recovery(100, 110, 95));
+        assert!(!should_perform_fork_recovery(100, 110, 101));
+        assert!(!should_perform_fork_recovery(100, 110, 120));
+        assert!(!should_perform_fork_recovery(100, 103, 100));
+    }
+
     // ---- fork recovery tests ----
 
     #[tokio::test]

packages/admin-cli/src/sandbox-smoke.ts

@@ -343,7 +343,7 @@ export async function runSandboxSmokeTest(
     checks.push(
       await measureCheck('session lifecycle', logger, async () => {
         try {
-          const session = await sandbox.session.create({ shell: '/bin/sh' })
+          const session = await sandbox.session.create()
           tracker.trackSession('root-shell', session)
 
           const primeResult = await session.exec(