diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs index 511ad9520f..6619bbb733 100644 --- a/codex-rs/exec-server/src/shell_snapshot.rs +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -1,6 +1,7 @@ //! Cache captured shell state in executor memory. Replay through a per-launch //! unnamed reader when the capture sandbox permits it; otherwise retain env replay. //! Command metrics exclude prewarm and measure lookup/capture wait before replay. +//! File replay admits larger shell state while environment replay keeps its launch budget. use std::collections::HashMap; use std::collections::VecDeque; @@ -40,14 +41,22 @@ use crate::telemetry::ExecServerTelemetry; const MAX_CACHED_SNAPSHOTS: usize = 16; const MAX_SNAPSHOT_BYTES: usize = 512 * 1024; +const MAX_FILE_SNAPSHOT_BYTES: usize = 4 * 1024 * 1024; // Capture also includes quoted export records and an optional pre-startup environment. const MAX_SNAPSHOT_CAPTURE_BYTES: usize = 8 * MAX_SNAPSHOT_BYTES; +const MAX_FILE_SNAPSHOT_CAPTURE_BYTES: usize = 8 * 1024 * 1024; const MAX_SNAPSHOT_ENV_VALUE_BYTES: usize = 60 * 1024; const MAX_SNAPSHOT_SCOPE_BYTES: usize = 256; const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); const SNAPSHOT_RETRY_BACKOFF: Duration = Duration::from_secs(1); const MAX_SNAPSHOT_ATTEMPTS: usize = 3; +#[derive(Clone, Copy)] +enum SnapshotReplay { + File, + Environment, +} + #[derive(Default)] pub(crate) struct ShellSnapshotCache { entries: Mutex>, @@ -355,6 +364,11 @@ async fn capture_snapshot( .ok() .and_then(Result::ok) }; + let capture_limit = if probe.is_some() { + MAX_FILE_SNAPSHOT_CAPTURE_BYTES + } else { + MAX_SNAPSHOT_CAPTURE_BYTES + }; const SOURCE_MARKER: &str = "CODEX_SNAPSHOT_SOURCE"; script = if let Some(probe) = &probe { let fd = probe.as_raw_fd(); @@ -417,7 +431,7 @@ async fn capture_snapshot( let capture = async { let mut output = Vec::new(); stdout - .take((MAX_SNAPSHOT_CAPTURE_BYTES + 1) as u64) + .take((capture_limit + 1) as u64) .read_to_end(&mut output) .await .map_err(|err| { @@ -426,11 +440,11 @@ async fn capture_snapshot( internal_error(format!("cannot read shell snapshot: {err}")), ) })?; - if output.len() > MAX_SNAPSHOT_CAPTURE_BYTES { + if output.len() > capture_limit { return Err(( - "too_large", + "capture_too_large", internal_error(format!( - "shell snapshot capture exceeds {MAX_SNAPSHOT_CAPTURE_BYTES} bytes" + "shell snapshot capture exceeds {capture_limit} bytes" )), )); } @@ -474,8 +488,20 @@ async fn capture_snapshot( internal_error("incomplete snapshot source probe".to_string()), ) })?; - let mut snapshot = parse_snapshot(shell_type, captured, params.env_policy.as_ref())?; - snapshot.file_source = flag == b"1\0"; + let replay = if flag == b"1\0" { + SnapshotReplay::File + } else { + if output.len() > MAX_SNAPSHOT_CAPTURE_BYTES { + return Err(( + "capture_too_large", + internal_error(format!( + "shell snapshot capture exceeds {MAX_SNAPSHOT_CAPTURE_BYTES} bytes" + )), + )); + } + SnapshotReplay::Environment + }; + let snapshot = parse_snapshot(shell_type, captured, params.env_policy.as_ref(), replay)?; child.preserve_helpers(); Ok(snapshot) } @@ -484,6 +510,7 @@ fn parse_snapshot( shell_type: ShellType, output: &[u8], env_policy: Option<&ExecEnvPolicy>, + replay: SnapshotReplay, ) -> CaptureResult { let captured = CapturedSnapshot::parse(shell_type, output).ok_or_else(|| { ( @@ -492,10 +519,22 @@ fn parse_snapshot( ) })?; let state = captured.render_state(); - if state.len().saturating_add(captured.environment.len()) > MAX_SNAPSHOT_BYTES { + if captured.environment.len() > MAX_SNAPSHOT_BYTES { return Err(( - "too_large", - internal_error(format!("shell snapshot exceeds {MAX_SNAPSHOT_BYTES} bytes")), + "environment_too_large", + internal_error(format!( + "shell snapshot environment exceeds {MAX_SNAPSHOT_BYTES} bytes" + )), + )); + } + let state_limit = match replay { + SnapshotReplay::File => MAX_FILE_SNAPSHOT_BYTES, + SnapshotReplay::Environment => MAX_SNAPSHOT_BYTES - captured.environment.len(), + }; + if state.len() > state_limit { + return Err(( + "state_too_large", + internal_error(format!("shell snapshot state exceeds {state_limit} bytes")), )); } @@ -526,7 +565,7 @@ fn parse_snapshot( Ok(ShellSnapshot { prewarmed: false, state, - file_source: false, + file_source: matches!(replay, SnapshotReplay::File), environment, }) } diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs index b3852d548a..adc0ceb141 100644 --- a/codex-rs/exec-server/src/shell_snapshot_tests.rs +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -11,10 +11,12 @@ use pretty_assertions::assert_eq; use test_case::test_case; use super::CapturePurpose; +use super::MAX_FILE_SNAPSHOT_BYTES; use super::MAX_SNAPSHOT_ATTEMPTS; use super::MAX_SNAPSHOT_BYTES; use super::SNAPSHOT_RETRY_BACKOFF; use super::ShellSnapshotCache; +use super::SnapshotReplay; use super::parse_snapshot; use crate::process_sandbox::prepare_exec_request_with_telemetry; use crate::process_telemetry::ProcessTelemetry; @@ -291,7 +293,15 @@ fn snapshot_size_limit_counts_state_and_environment_before_filtering() { r#set: HashMap::new(), include_only: Vec::new(), }; - assert!(parse_snapshot(ShellType::Bash, oversized.as_bytes(), Some(&policy)).is_err()); + assert!( + parse_snapshot( + ShellType::Bash, + oversized.as_bytes(), + Some(&policy), + SnapshotReplay::Environment + ) + .is_err() + ); } #[test] @@ -307,6 +317,7 @@ fn snapshot_filters_profile_exports_after_capture() { ShellType::Bash, b"profile \xff noise\n# Snapshot file\nfunction profile_helper() { :; }\n\0alias profile_alias='profile_helper'\n\0PROFILE_DENIED\0export PROFILE_DENIED=denied\n\0NON_UTF8\0export NON_UTF8='\xff'\n\0\0PROFILE_ALLOWED=profile\0PROFILE_DENIED=denied\0PROFILE_SECRET=secret\0PWD=/tmp\0NON_UTF8=\xff\0", Some(&policy), + SnapshotReplay::Environment, ) .expect("snapshot should parse"); @@ -337,6 +348,7 @@ fn snapshot_preserves_profile_exports_with_restrictive_inheritance() { ShellType::Bash, b"# Snapshot file\n\0\0\0PROFILE_ALLOWED=profile\0SDKROOT=/sdk\0PROFILE_SECRET=secret\0PROFILE_DENIED=denied\0", Some(&policy), + SnapshotReplay::Environment, ) .expect("snapshot should parse"); @@ -369,10 +381,39 @@ fn snapshot_caches_only_unmanaged_proxy_state() { ), ] { let output = format!("# Snapshot file\n\0\0\0{exports}"); - let snapshot = parse_snapshot(ShellType::Bash, output.as_bytes(), /*env_policy*/ None) - .expect("snapshot should parse"); + let snapshot = parse_snapshot( + ShellType::Bash, + output.as_bytes(), + /*env_policy*/ None, + SnapshotReplay::Environment, + ) + .expect("snapshot should parse"); assert_eq!(snapshot.environment, expected); } } use codex_shell_command::shell_detect::ShellType; + +#[test_case(SnapshotReplay::File, 1024 * 1024, 0, None; "file_accepts_large_state")] +#[test_case(SnapshotReplay::Environment, 1024 * 1024, 0, Some("state_too_large"); "fallback_rejects_large_state")] +#[test_case(SnapshotReplay::File, MAX_FILE_SNAPSHOT_BYTES, 0, Some("state_too_large"); "file_rejects_oversized_state")] +#[test_case(SnapshotReplay::File, 0, MAX_SNAPSHOT_BYTES, Some("environment_too_large"); "file_keeps_environment_limit")] +fn snapshot_replay_size_limits( + replay: SnapshotReplay, + state_bytes: usize, + environment_bytes: usize, + expected_failure: Option<&str>, +) { + let output = format!( + "# Snapshot file\n# {}\n\0\0\0PADDING={}\0", + "x".repeat(state_bytes), + "x".repeat(environment_bytes), + ); + let result = parse_snapshot( + ShellType::Bash, + output.as_bytes(), + /*env_policy*/ None, + replay, + ); + assert_eq!(result.err().map(|(reason, _)| reason), expected_failure); +} diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index dba1c7c640..8cf2a5d76f 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -244,7 +244,7 @@ async fn shell_snapshot_v2_filters_profile_exports_and_stays_in_memory( let padding = if !use_remote && !tty && shell_name == "bash" { format!( "snapshot_padding() {{ printf '%s' '{}'; }}\n", - "🦀".repeat(20_000) + "🦀".repeat(200_000) ) } else { String::new()