openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow larger shell snapshots when replaying from a file (#51350)

Browse Source 
## Why

File replay can accommodate shell state larger than the 512 KiB budget used for environment replay.

## What changed

Allow up to 4 MiB of shell state and 8 MiB of capture output for file replay. Keep captured environments capped at 512 KiB and preserve the 512 KiB combined state-and-environment limit for environment replay. Report distinct failures for oversized capture output, state, and environments.

## Testing

Add replay-specific size-limit tests and enlarge the Bash integration-test snapshot to exercise file replay beyond the environment replay budget.

GitOrigin-RevId: 0210848eb398d75426cb8b4b3bfaa0bce25264b8
This commit is contained in:
jif
2026-10-06 12:10:58 +00:00
committed by copyberry
parent 588f616e8b
commit e32365a2c6
3 changed files with 94 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/shell_snapshot.rs
+49 -10
View File
@@ -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<VecDeque<CachedShellSnapshot>>,
@@ -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,
})
}
codex-rs/exec-server/src/shell_snapshot_tests.rs
+44 -3
View File
@@ -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);
}
codex-rs/exec-server/tests/exec_process.rs
+1 -1
View File
@@ -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()