Ensure fork recovery before serving requests

Author: richiemcilroy

crates/sandchest-agent/src/service.rs

@@ -47,6 +47,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<ExecRequest>,
     ) -> Result<Response<Self::ExecStream>, Status> {
+        snapshot::ensure_recovered_if_needed(&self.session_manager).await;
         snapshot::note_user_activity();
         let stream = crate::exec::spawn_exec(request.into_inner());
         Ok(Response::new(stream))
@@ -56,6 +57,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<CreateSessionRequest>,
     ) -> Result<Response<SessionResponse>, Status> {
+        snapshot::ensure_recovered_if_needed(&self.session_manager).await;
         snapshot::note_user_activity();
         let req = request.into_inner();
         let session_id = self
@@ -71,6 +73,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<SessionExecRequest>,
     ) -> Result<Response<Self::SessionExecStream>, Status> {
+        snapshot::ensure_recovered_if_needed(&self.session_manager).await;
         snapshot::note_user_activity();
         let req = request.into_inner();
         let stream = self
@@ -84,6 +87,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<SessionInputRequest>,
     ) -> Result<Response<()>, Status> {
+        snapshot::ensure_recovered_if_needed(&self.session_manager).await;
         snapshot::note_user_activity();
         let req = request.into_inner();
         self.session_manager
@@ -96,6 +100,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<DestroySessionRequest>,
     ) -> Result<Response<()>, Status> {
+        snapshot::ensure_recovered_if_needed(&self.session_manager).await;
         snapshot::note_user_activity();
         let req = request.into_inner();
         self.session_manager
@@ -108,6 +113,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<Streaming<FileChunk>>,
     ) -> Result<Response<PutFileResponse>, Status> {
+        snapshot::ensure_recovered_if_needed(&self.session_manager).await;
         snapshot::note_user_activity();
         let response = crate::files::put_file(request.into_inner()).await?;
         Ok(Response::new(response))
@@ -119,6 +125,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<GetFileRequest>,
     ) -> Result<Response<Self::GetFileStream>, Status> {
+        snapshot::ensure_recovered_if_needed(&self.session_manager).await;
         snapshot::note_user_activity();
         let stream = crate::files::spawn_get_file(request.into_inner());
         Ok(Response::new(stream))
@@ -128,6 +135,7 @@ impl GuestAgent for GuestAgentService {
         &self,
         request: Request<ListFilesRequest>,
     ) -> Result<Response<ListFilesResponse>, Status> {
+        snapshot::ensure_recovered_if_needed(&self.session_manager).await;
         snapshot::note_user_activity();
         let response = crate::files::list_files(request.into_inner()).await?;
         Ok(Response::new(response))

crates/sandchest-agent/src/session.rs

@@ -487,15 +487,21 @@ async fn run_session_exec(
 /// Look for the sentinel pattern in the buffer. Returns (output_before_sentinel, exit_code).
 fn extract_sentinel(buf: &[u8], sentinel_marker: &str) -> Option<(Vec<u8>, i32)> {
     let buf_str = String::from_utf8_lossy(buf);
-    let marker_pos = buf_str.find(sentinel_marker)?;
-
-    let after_marker = &buf_str[marker_pos + sentinel_marker.len()..];
-    let suffix_pos = after_marker.find(SENTINEL_SUFFIX)?;
-    let exit_code_str = &after_marker[..suffix_pos];
-    let exit_code: i32 = exit_code_str.parse().unwrap_or(-1);
+    let mut search_from = 0;
+
+    while let Some(relative_pos) = buf_str[search_from..].find(sentinel_marker) {
+        let marker_pos = search_from + relative_pos;
+        let after_marker = &buf_str[marker_pos + sentinel_marker.len()..];
+        let suffix_pos = after_marker.find(SENTINEL_SUFFIX)?;
+        let exit_code_str = &after_marker[..suffix_pos];
+        if let Ok(exit_code) = exit_code_str.parse::<i32>() {
+            let output = buf[..marker_pos].to_vec();
+            return Some((output, exit_code));
+        }
+        search_from = marker_pos + sentinel_marker.len();
+    }
 
-    let output = buf[..marker_pos].to_vec();
-    Some((output, exit_code))
+    None
 }
 
 /// Strip the echoed command line from PTY output.
@@ -526,6 +532,20 @@ fn strip_command_echo(output: &[u8], cmd: &str) -> Vec<u8> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use tokio_stream::StreamExt;
+
+    async fn collect_session_events(
+        manager: &SessionManager,
+        session_id: &str,
+        cmd: &str,
+    ) -> Vec<Result<ExecEvent, Status>> {
+        manager
+            .spawn_session_exec(session_id, cmd.to_string(), 10)
+            .await
+            .unwrap()
+            .collect()
+            .await
+    }
 
     // ---- extract_sentinel tests ----
 
@@ -590,8 +610,7 @@ mod tests {
     fn extract_sentinel_invalid_exit_code() {
         let marker = "__SC_SENTINEL_999_";
         let buf = b"__SC_SENTINEL_999_notanumber__\n";
-        let (_, code) = extract_sentinel(buf, marker).unwrap();
-        assert_eq!(code, -1); // parse failure falls back to -1
+        assert!(extract_sentinel(buf, marker).is_none());
     }
 
     #[test]
@@ -611,6 +630,19 @@ mod tests {
         assert_eq!(code, 255);
     }
 
+    #[test]
+    fn extract_sentinel_skips_echoed_marker_literal() {
+        let marker = "__SC_SENTINEL_123_";
+        let buf =
+            b"echo \"__SC_SENTINEL_123_${__sc_exit}__\"\nreal output\n__SC_SENTINEL_123_0__\n";
+        let (output, code) = extract_sentinel(buf, marker).unwrap();
+        assert_eq!(
+            String::from_utf8_lossy(&output),
+            "echo \"__SC_SENTINEL_123_${__sc_exit}__\"\nreal output\n"
+        );
+        assert_eq!(code, 0);
+    }
+
     // ---- strip_command_echo tests ----
 
     #[test]
@@ -746,6 +778,46 @@ mod tests {
         manager.destroy_session(&id).await.unwrap();
     }
 
+    #[tokio::test]
+    async fn session_exec_runs_command_and_preserves_working_directory() {
+        let manager = SessionManager::new();
+        let env = HashMap::new();
+        let session_id = manager.create_session("", "", &env).await.unwrap();
+
+        let prime_events = collect_session_events(
+            &manager,
+            &session_id,
+            "cd /tmp && printf '%s' test-session > /tmp/session-test.txt",
+        )
+        .await;
+        let prime_exit = prime_events
+            .iter()
+            .find_map(|event| match event.as_ref().ok()?.event.as_ref()? {
+                exec_event::Event::Exit(exit) => Some(exit.exit_code),
+                _ => None,
+            });
+        assert_eq!(prime_exit, Some(0));
+
+        let persisted_events =
+            collect_session_events(&manager, &session_id, "pwd && cat /tmp/session-test.txt").await;
+        let mut stdout = Vec::new();
+        let mut exit = None;
+        for event in &persisted_events {
+            let event = event.as_ref().unwrap();
+            match event.event.as_ref() {
+                Some(exec_event::Event::Stdout(data)) => stdout.extend_from_slice(data),
+                Some(exec_event::Event::Exit(info)) => exit = Some(info.exit_code),
+                _ => {}
+            }
+        }
+
+        assert_eq!(exit, Some(0));
+        let normalized = String::from_utf8_lossy(&stdout).replace("\r\n", "\n");
+        assert_eq!(normalized.trim_end(), "/tmp\ntest-session");
+
+        manager.destroy_session(&session_id).await.unwrap();
+    }
+
     #[tokio::test]
     async fn session_manager_destroy_all() {
         let manager = SessionManager::new();

crates/sandchest-agent/src/snapshot.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{SystemTime, UNIX_EPOCH};
 
+use tokio::sync::Mutex;
 use tracing::{info, warn};
 
 use crate::session::SessionManager;
@@ -13,6 +14,7 @@ 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);
+static RECOVERY_LOCK: Mutex<()> = Mutex::const_new(());
 
 fn current_unix_secs() -> u64 {
     SystemTime::now()
@@ -229,6 +231,34 @@ async fn perform_fork_recovery(session_manager: &SessionManager) {
     info!("Fork recovery complete — agent ready");
 }
 
+pub async fn ensure_recovered_if_needed(session_manager: &SessionManager) {
+    let _guard = RECOVERY_LOCK.lock().await;
+    let heartbeat_path = Path::new(HEARTBEAT_PATH);
+
+    let Some(file_ts) = read_heartbeat_timestamp(heartbeat_path) else {
+        return;
+    };
+
+    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 before serving traffic — running fork recovery"
+        );
+        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;
+    }
+}
+
 /// Write current timestamp to the heartbeat file.
 async fn write_heartbeat() {
     let ts = SystemTime::now()
@@ -252,27 +282,7 @@ async fn write_heartbeat() {
 pub fn start_snapshot_watcher(session_manager: Arc<SessionManager>) {
     tokio::spawn(async move {
         loop {
-            // Check for stale heartbeat BEFORE writing a fresh one
-            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"
-                    );
-                }
-            }
+            ensure_recovered_if_needed(&session_manager).await;
 
             write_heartbeat().await;
 

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

@@ -355,7 +355,7 @@ export async function runSandboxSmokeTest(
           const persistedResult = await session.exec(`pwd && cat "${sessionPath}"`, { timeout: 60 })
           assertExecSuccess(persistedResult, 'session persisted exec')
 
-          const output = persistedResult.stdout.trimEnd()
+          const output = persistedResult.stdout.replace(/\r\n/g, '\n').trimEnd()
           assert(
             output === `/tmp\nsession:${runId}`,
             `Session state did not persist as expected: ${JSON.stringify(output)}`,