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

Cache remote Bash environment snapshots

Browse Source
This commit is contained in:
Vaibhav Srivastav
2026-06-24 19:31:51 +01:00
parent a781761eda
commit e351f974f6
17 changed files with 779 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -2905,6 +2905,7 @@ dependencies = [
"serde",
"serde_json",
"serial_test",
"sha2 0.10.9",
"tempfile",
"test-case",
"thiserror 2.0.18",
codex-rs/core/src/unified_exec/process_manager.rs
+1
View File
@@ -183,6 +183,7 @@ fn exec_server_params_for_request(
sandbox: request.exec_server_sandbox.clone(),
enforce_managed_network: request.exec_server_enforce_managed_network,
managed_network: request.exec_server_managed_network.clone(),
bash_env_snapshot: None,
}
}
codex-rs/exec-server-protocol/src/protocol.rs
+27 -2
View File
@@ -144,6 +144,21 @@ pub struct ExecParams {
/// continue to fail closed. This preserves compatibility with older clients.
#[serde(default)]
pub managed_network: Option<ManagedNetworkSandboxContext>,
/// Optional request to reuse a workspace-scoped Bash environment snapshot on the executor.
/// Older exec-server peers ignore this field, so callers must treat reuse as opportunistic.
#[serde(default)]
pub bash_env_snapshot: Option<BashEnvSnapshotParams>,
}
/// Parameters for executor-owned Bash environment snapshot capture and replay.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BashEnvSnapshotParams {
/// Workspace root that owns the captured environment. The command cwd must remain inside it.
/// Direct `BASH_ENV` startup must not be command-dependent, trap-setting, or builtin-shadowing.
pub workspace_root: PathUri,
/// Process environment keys that must win after the snapshot is sourced.
pub preserve_env_keys: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -549,6 +564,7 @@ mod base64_bytes {
#[cfg(test)]
mod tests {
use super::BashEnvSnapshotParams;
use super::EnvironmentInfo;
use super::ExecParams;
use super::FsReadFileParams;
@@ -563,14 +579,14 @@ mod tests {
use std::collections::HashMap;
#[test]
fn exec_params_managed_network_context_round_trips_and_defaults_for_legacy_peers() {
fn exec_params_optional_context_round_trips_and_defaults_for_legacy_peers() {
let cwd =
PathUri::from_host_native_path(std::env::current_dir().expect("current directory"))
.expect("cwd URI");
let params = ExecParams {
process_id: ProcessId::from("managed-network"),
argv: vec!["true".to_string()],
cwd,
cwd: cwd.clone(),
env_policy: None,
env: HashMap::new(),
tty: false,
@@ -582,6 +598,10 @@ mod tests {
loopback_ports: vec![43123, 48081],
allow_local_binding: false,
}),
bash_env_snapshot: Some(BashEnvSnapshotParams {
workspace_root: cwd,
preserve_env_keys: vec!["PATH".to_string()],
}),
};
let mut serialized = serde_json::to_value(&params).expect("serialize exec params");
@@ -600,10 +620,15 @@ mod tests {
.as_object_mut()
.expect("exec params object")
.remove("managedNetwork");
serialized
.as_object_mut()
.expect("exec params object")
.remove("bashEnvSnapshot");
let legacy: ExecParams =
serde_json::from_value(serialized).expect("deserialize legacy exec params");
assert!(legacy.enforce_managed_network);
assert_eq!(legacy.managed_network, None);
assert_eq!(legacy.bash_env_snapshot, None);
}
#[test]
codex-rs/exec-server/Cargo.toml
+2 -1
View File
@@ -33,6 +33,8 @@ reqwest = { workspace = true, features = ["json", "rustls-tls", "stream"] }
prost = "0.14.3"
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
sha2 = { workspace = true }
tempfile = { workspace = true }
thiserror = { workspace = true }
toml = { workspace = true }
tokio = { workspace = true, features = [
@@ -67,6 +69,5 @@ ctor = { workspace = true }
http = { workspace = true }
pretty_assertions = { workspace = true }
serial_test = { workspace = true }
tempfile = { workspace = true }
test-case = "3.3.1"
wiremock = { workspace = true }
codex-rs/exec-server/src/bash_env_snapshot.rs
+30
View File
@@ -0,0 +1,30 @@
#[cfg(not(unix))]
use std::collections::HashMap;
use crate::ExecServerRuntimePaths;
use crate::protocol::ExecParams;
#[cfg(unix)]
#[path = "bash_env_snapshot_unix.rs"]
mod imp;
#[cfg(unix)]
pub(crate) use imp::BashEnvSnapshotCache;
#[cfg(not(unix))]
#[derive(Default)]
pub(crate) struct BashEnvSnapshotCache(());
#[cfg(not(unix))]
impl BashEnvSnapshotCache {
pub(crate) async fn prepare_launch(
&self,
params: &ExecParams,
environment: &HashMap<String, String>,
_runtime_paths: Option<&ExecServerRuntimePaths>,
) -> (Vec<String>, HashMap<String, String>) {
(params.argv.clone(), environment.clone())
}
pub(crate) fn clear(&self) {}
}
codex-rs/exec-server/src/bash_env_snapshot_unix.rs
+495
View File
@@ -0,0 +1,495 @@
use std::collections::HashMap;
use std::fs::Metadata;
use std::io::Write;
use std::os::unix::fs::MetadataExt;
use std::os::unix::fs::PermissionsExt;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use sha2::Digest;
use sha2::Sha256;
use tempfile::TempDir;
use tokio::sync::OnceCell;
use tokio::time::timeout;
use uuid::Uuid;
use super::ExecParams;
use super::ExecServerRuntimePaths;
use crate::process_sandbox::PreparedExecRequest;
use crate::process_sandbox::prepare_exec_request;
const CAPTURE_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_CACHE_ENTRIES: usize = 16;
const MAX_ENV_BYTES: usize = 256 * 1024;
const MAX_SCOPE_BYTES: usize = 384 * 1024;
const MAX_PRESERVE_KEYS: usize = 128;
const MAX_SNAPSHOT_BYTES: usize = 1024 * 1024;
type SnapshotCell = OnceCell<Option<Arc<BashEnvSnapshot>>>;
pub(crate) struct BashEnvSnapshotCache {
entries: Mutex<HashMap<[u8; 32], Arc<SnapshotCell>>>,
root: Option<TempDir>,
}
struct BashEnvSnapshotRequest {
key: [u8; 32],
environment: HashMap<String, String>,
preserve_env_keys: Vec<String>,
bash_env: String,
}
struct BashEnvSnapshot {
path: PathBuf,
bash_env: String,
checksum: String,
prefix: String,
}
impl Default for BashEnvSnapshotCache {
fn default() -> Self {
Self {
entries: Mutex::new(HashMap::new()),
root: tempfile::Builder::new()
.prefix("codex-exec-server-bash-env-")
.permissions(std::fs::Permissions::from_mode(0o700))
.tempdir()
.ok(),
}
}
}
impl BashEnvSnapshotCache {
pub(crate) async fn prepare_launch(
&self,
params: &ExecParams,
environment: &HashMap<String, String>,
runtime_paths: Option<&ExecServerRuntimePaths>,
) -> (Vec<String>, HashMap<String, String>) {
let Some(request) = BashEnvSnapshotRequest::new(params, environment) else {
return fallback(params, environment);
};
let preserve_env_keys = request.preserve_env_keys.clone();
let Some(snapshot) = self.snapshot(params, request, runtime_paths).await else {
return fallback(params, environment);
};
let Some(argv) = wrap_command(&params.argv, &snapshot, &preserve_env_keys) else {
return fallback(params, environment);
};
let mut environment = environment.clone();
environment.remove("BASH_ENV");
(argv, environment)
}
pub(crate) fn clear(&self) {
self.entries
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clear();
}
async fn snapshot(
&self,
params: &ExecParams,
request: BashEnvSnapshotRequest,
runtime_paths: Option<&ExecServerRuntimePaths>,
) -> Option<Arc<BashEnvSnapshot>> {
let key = request.key;
let cell = {
let mut entries = self
.entries
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if entries.len() >= MAX_CACHE_ENTRIES && !entries.contains_key(&key) {
return None;
}
Arc::clone(
entries
.entry(key)
.or_insert_with(|| Arc::new(OnceCell::new())),
)
};
cell.get_or_init(|| async {
self.capture(params, request, runtime_paths)
.await
.map(Arc::new)
})
.await
.clone()
}
async fn capture(
&self,
params: &ExecParams,
request: BashEnvSnapshotRequest,
runtime_paths: Option<&ExecServerRuntimePaths>,
) -> Option<BashEnvSnapshot> {
let nonce = Uuid::new_v4().simple().to_string();
let prefix = format!("__CODEX_SNAPSHOT_{}", nonce.to_ascii_uppercase());
let marker = format!("# Codex remote Bash environment snapshot {nonce}");
let stderr_marker = format!("# Codex remote Bash stderr complete {nonce}");
let wrapper_end = format!("# Codex remote Bash capture complete {nonce}");
let root = self.root.as_ref()?;
let mut wrapper = tempfile::NamedTempFile::new_in(root.path()).ok()?;
let script = capture_script(
&prefix,
&marker,
&stderr_marker,
&wrapper_end,
&request.bash_env,
&request.preserve_env_keys,
);
wrapper.write_all(script.as_bytes()).ok()?;
wrapper.as_file().sync_all().ok()?;
let mut capture = params.clone();
capture.bash_env_snapshot = None;
let mut environment = request.environment;
environment.insert("BASH_ENV".into(), wrapper.path().to_string_lossy().into());
let prepared = prepare_exec_request(&capture, environment, runtime_paths).ok()?;
let (output, stderr) = capture_output(prepared).await?;
let marker_offset = output
.windows(marker.len())
.position(|part| part == marker.as_bytes())?;
let stdout = &output[..marker_offset];
let snapshot =
output[marker_offset..].strip_suffix(format!("{wrapper_end}\n").as_bytes())?;
let stderr = &stderr[..stderr
.windows(stderr_marker.len())
.position(|part| part == stderr_marker.as_bytes())?];
let mut contents = snapshot.to_vec();
contents.extend(replay_output(stdout, /*fd*/ 3));
contents.extend(replay_output(stderr, /*fd*/ 4));
contents.extend_from_slice(b"builtin true\n");
if !contents.ends_with(b"\n") {
contents.push(b'\n');
}
if contents.len() > MAX_SNAPSHOT_BYTES {
return None;
}
let mut temp = tempfile::NamedTempFile::new_in(root.path()).ok()?;
temp.write_all(&contents).ok()?;
temp.as_file().sync_all().ok()?;
let path = root.path().join(format!("{nonce}.sh"));
temp.persist(&path).ok()?;
Some(BashEnvSnapshot {
path,
bash_env: request.bash_env,
checksum: format!("{:x}", Sha256::digest(&contents)),
prefix,
})
}
}
impl BashEnvSnapshotRequest {
fn new(params: &ExecParams, environment: &HashMap<String, String>) -> Option<Self> {
let snapshot = params.bash_env_snapshot.as_ref()?;
if params.tty
|| params.pipe_stdin
|| params.arg0.is_some()
|| params.argv.len() < 3
|| params.argv.get(1).map(String::as_str) != Some("-c")
|| !params.cwd.starts_with(&snapshot.workspace_root)
|| environment.iter().any(|(key, value)| {
(key == "builtin" || key.starts_with("BASH_FUNC_builtin"))
&& value.trim_start().starts_with("()")
})
{
return None;
}
let shell = Path::new(params.argv.first()?);
let bash_env = Path::new(environment.get("BASH_ENV")?);
if !shell.is_absolute()
|| shell.file_name()?.to_str()? != "bash"
|| !shell.is_file()
|| !bash_env.is_absolute()
|| !bash_env.is_file()
{
return None;
}
let bash_env_source = std::fs::read(bash_env).ok()?;
if bash_env_source.len() > MAX_SNAPSHOT_BYTES || unsupported_startup(&bash_env_source) {
return None;
}
let mut preserve_env_keys = snapshot.preserve_env_keys.clone();
if preserve_env_keys.len() > MAX_PRESERVE_KEYS
|| preserve_env_keys
.iter()
.any(|key| key.len() > 128 || !valid_variable(key))
{
return None;
}
preserve_env_keys.retain(|key| key != "BASH_ENV");
preserve_env_keys.sort_unstable();
preserve_env_keys.dedup();
let mut environment_scope = environment
.iter()
.map(|(key, value)| (key.clone(), value.clone()))
.collect::<Vec<_>>();
environment_scope.sort_unstable();
let environment_bytes = environment_scope
.iter()
.try_fold(0usize, |size, (key, value)| {
size.checked_add(key.len())?.checked_add(value.len())
})?;
if environment_bytes > MAX_ENV_BYTES {
return None;
}
let mut scope = serde_json::to_value((
unsafe { libc::geteuid() },
&snapshot.workspace_root,
&params.cwd,
shell.to_string_lossy(),
file_identity(&shell.metadata().ok()?),
bash_env.to_string_lossy(),
file_identity(&bash_env.metadata().ok()?),
environment_scope,
&params.env_policy,
&params.sandbox,
params.enforce_managed_network,
&params.managed_network,
&preserve_env_keys,
))
.ok()?;
sort_json(&mut scope);
let scope = serde_json::to_vec(&scope).ok()?;
if scope.len() > MAX_SCOPE_BYTES {
return None;
}
let key = Sha256::digest(scope).into();
Some(Self {
key,
environment: environment.clone(),
preserve_env_keys,
bash_env: bash_env.to_string_lossy().into_owned(),
})
}
}
impl Drop for BashEnvSnapshot {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.path);
}
}
async fn capture_output(prepared: PreparedExecRequest) -> Option<(Vec<u8>, Vec<u8>)> {
timeout(CAPTURE_TIMEOUT, async move {
let (program, args) = prepared.command.split_first()?;
let spawned = codex_utils_pty::spawn_pipe_process_no_stdin(
program,
args,
prepared.cwd.as_path(),
&prepared.env,
&prepared.arg0,
)
.await
.ok()?;
let codex_utils_pty::SpawnedProcess {
session: _session,
stdout_rx,
stderr_rx,
exit_rx,
} = spawned;
let (stdout, stderr, exit) = tokio::join!(
collect_output(stdout_rx),
collect_output(stderr_rx),
exit_rx
);
(exit.ok()? == 0).then_some((stdout?, stderr?))
})
.await
.ok()
.flatten()
}
async fn collect_output(mut receiver: tokio::sync::mpsc::Receiver<Vec<u8>>) -> Option<Vec<u8>> {
let mut output = Vec::new();
while let Some(chunk) = receiver.recv().await {
if output.len().checked_add(chunk.len())? > MAX_SNAPSHOT_BYTES {
return None;
}
output.extend_from_slice(&chunk);
}
Some(output)
}
fn capture_script(
prefix: &str,
marker: &str,
stderr_marker: &str,
wrapper_end: &str,
bash_env: &str,
preserve_env_keys: &[String],
) -> String {
let mut excluded = ":BASH:BASHOPTS:BASHPID:BASH_ALIASES:BASH_ARGC:BASH_ARGV:BASH_ARGV0:\
BASH_CMDS:BASH_COMMAND:BASH_ENV:BASH_EXECUTION_STRING:BASH_LINENO:BASH_REMATCH:BASH_SOURCE:BASH_SUBSHELL:\
BASH_VERSINFO:BASH_VERSION:COPROC:DIRSTACK:EPOCHREALTIME:EPOCHSECONDS:EUID:FUNCNAME:GROUPS:\
HISTCMD:HOSTTYPE:LINENO:MACHTYPE:OLDPWD:OSTYPE:PIPESTATUS:PPID:PWD:RANDOM:SECONDS:SHELLOPTS:SHLVL:SRANDOM:UID:_:"
.to_string();
for key in preserve_env_keys {
excluded.push_str(key);
excluded.push(':');
}
r#"if ! builtin type -P sha256sum >/dev/null 2>&1 && ! builtin type -P shasum >/dev/null 2>&1; then builtin exit 125; fi
builtin unset __CODEX_PREFIX___EXCLUDED __CODEX_PREFIX___P_NAMES __CODEX_PREFIX___NAME __CODEX_PREFIX___DECL __CODEX_PREFIX___ALIASES __CODEX_PREFIX___OPTS __CODEX_PREFIX___SHOPTS __CODEX_PREFIX___FUNC_ATTRS __CODEX_PREFIX___CWD __CODEX_PREFIX___UMASK
builtin export BASH_ENV='__CODEX_BASH_ENV__'
__CODEX_PREFIX___EXCLUDED='__CODEX_EXCLUDED__'; __CODEX_PREFIX___P_NAMES=; __CODEX_PREFIX___NAME=; __CODEX_PREFIX___DECL=
__CODEX_PREFIX___ALIASES=; __CODEX_PREFIX___OPTS=; __CODEX_PREFIX___SHOPTS=; __CODEX_PREFIX___FUNC_ATTRS=; __CODEX_PREFIX___CWD=; __CODEX_PREFIX___UMASK=
while IFS= builtin read -r __CODEX_PREFIX___NAME; do
[[ "$__CODEX_PREFIX___NAME" != __CODEX_PREFIX___* && "$__CODEX_PREFIX___EXCLUDED" != *":$__CODEX_PREFIX___NAME:"* && "$__CODEX_PREFIX___NAME" =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]] || continue
__CODEX_PREFIX___P_NAMES="${__CODEX_PREFIX___P_NAMES}${__CODEX_PREFIX___NAME}"$'\n'
done < <(builtin compgen -A variable)
builtin source "$BASH_ENV"
if [[ $- == *x* || $- == *v* ]]; then builtin exit 125; fi
builtin printf '%s\n' '__CODEX_MARKER__'; builtin printf '%s\n' '__CODEX_STDERR_MARKER__' >&2
__CODEX_PREFIX___ALIASES=$(builtin alias -p 2>/dev/null || builtin true)
__CODEX_PREFIX___OPTS=$(builtin set +o); if [[ $- == *e* ]]; then __CODEX_PREFIX___OPTS=${__CODEX_PREFIX___OPTS/set +o errexit/set -o errexit}; fi
__CODEX_PREFIX___SHOPTS=$(builtin shopt -p 2>/dev/null || builtin true)
builtin shopt -u extdebug 2>/dev/null || builtin true; __CODEX_PREFIX___FUNC_ATTRS=$(builtin declare -F)
__CODEX_PREFIX___CWD=$(builtin pwd -L); __CODEX_PREFIX___UMASK=$(builtin umask)
builtin set +e; builtin set +u; builtin unalias -a 2>/dev/null || builtin true
__CODEX_PREFIX___ALIASES=${__CODEX_PREFIX___ALIASES//$'\n'/$'\nbuiltin '}; __CODEX_PREFIX___OPTS=${__CODEX_PREFIX___OPTS//$'\n'/$'\nbuiltin '}; __CODEX_PREFIX___SHOPTS=${__CODEX_PREFIX___SHOPTS//$'\n'/$'\nbuiltin '}
[[ -z "$__CODEX_PREFIX___ALIASES" ]] || __CODEX_PREFIX___ALIASES="builtin $__CODEX_PREFIX___ALIASES"
[[ -z "$__CODEX_PREFIX___OPTS" ]] || __CODEX_PREFIX___OPTS="builtin $__CODEX_PREFIX___OPTS"
[[ -z "$__CODEX_PREFIX___SHOPTS" ]] || __CODEX_PREFIX___SHOPTS="builtin $__CODEX_PREFIX___SHOPTS"
builtin printf 'builtin cd -L -- %q || builtin return 1\n' "$__CODEX_PREFIX___CWD"
builtin declare -f
while IFS= builtin read -r __CODEX_PREFIX___DECL; do
[[ "$__CODEX_PREFIX___DECL" == "declare -f "* ]] && continue
builtin printf 'builtin %s\n' "$__CODEX_PREFIX___DECL"
done <<< "$__CODEX_PREFIX___FUNC_ATTRS"
builtin printf '\n'
while IFS= builtin read -r __CODEX_PREFIX___NAME; do
[[ "$__CODEX_PREFIX___NAME" != __CODEX_PREFIX___* && "$__CODEX_PREFIX___EXCLUDED" != *":$__CODEX_PREFIX___NAME:"* && "$__CODEX_PREFIX___NAME" =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]] || continue
__CODEX_PREFIX___DECL=$(builtin declare -p "$__CODEX_PREFIX___NAME" 2>/dev/null) || continue; builtin printf 'builtin unset -v %s\n' "$__CODEX_PREFIX___NAME"; __CODEX_PREFIX___NAME=${__CODEX_PREFIX___DECL#declare }; __CODEX_PREFIX___NAME=${__CODEX_PREFIX___NAME%% *}
if [[ "$__CODEX_PREFIX___NAME" == -*[aA]* ]]; then __CODEX_PREFIX___DECL=${__CODEX_PREFIX___DECL#declare }; __CODEX_PREFIX___DECL=${__CODEX_PREFIX___DECL#* }; if [[ "$__CODEX_PREFIX___NAME" == *A* ]]; then __CODEX_PREFIX___CWD=A; else __CODEX_PREFIX___CWD=a; fi
if [[ "${__CODEX_PREFIX___DECL#*=}" == \(* ]]; then builtin printf 'builtin declare -%s %q\n' "$__CODEX_PREFIX___CWD" "$__CODEX_PREFIX___DECL"; else builtin printf 'builtin declare -%s %s\n' "$__CODEX_PREFIX___CWD" "$__CODEX_PREFIX___DECL"; fi; builtin printf 'builtin declare %s %s\n' "$__CODEX_PREFIX___NAME" "${__CODEX_PREFIX___DECL%%=*}"; else builtin printf 'builtin %s\n' "$__CODEX_PREFIX___DECL"; fi
done < <(builtin compgen -A variable)
while IFS= builtin read -r __CODEX_PREFIX___NAME; do
[[ -z "$__CODEX_PREFIX___NAME" ]] && continue
builtin declare -p "$__CODEX_PREFIX___NAME" >/dev/null 2>&1 || builtin printf 'builtin unset -v %s\n' "$__CODEX_PREFIX___NAME"
done < <(builtin printf '%s' "$__CODEX_PREFIX___P_NAMES")
builtin printf '\n'
if __CODEX_PREFIX___DECL=$(builtin declare -p BASH_ENV 2>/dev/null); then builtin printf 'builtin %s\n' "$__CODEX_PREFIX___DECL"; else builtin printf 'builtin unset -v BASH_ENV\n'; fi
builtin printf '\nbuiltin umask %s\n' "$__CODEX_PREFIX___UMASK"
builtin printf '%s\n' "$__CODEX_PREFIX___OPTS"
builtin printf '\n%s\n\n' "$__CODEX_PREFIX___SHOPTS"
[[ -z "$__CODEX_PREFIX___ALIASES" ]] || builtin printf '%s\n' "$__CODEX_PREFIX___ALIASES"
builtin printf '%s\n' '__CODEX_WRAPPER_END__'; builtin exit 0
"#
.replace("__CODEX_PREFIX__", prefix)
.replace("__CODEX_MARKER__", marker)
.replace("__CODEX_STDERR_MARKER__", stderr_marker)
.replace("__CODEX_WRAPPER_END__", wrapper_end)
.replace("__CODEX_BASH_ENV__", &quote(bash_env))
.replace("__CODEX_EXCLUDED__", &excluded)
}
fn replay_output(output: &[u8], fd: u8) -> Vec<u8> {
if output.is_empty() {
return Vec::new();
}
let escaped = output
.iter()
.map(|byte| format!("\\{byte:03o}"))
.collect::<String>();
format!("builtin printf '%b' '{escaped}' >&{fd}\n").into_bytes()
}
fn wrap_command(
argv: &[String],
snapshot: &BashEnvSnapshot,
preserve_env_keys: &[String],
) -> Option<Vec<String>> {
let mut script = String::new();
for (index, key) in preserve_env_keys.iter().enumerate() {
script.push_str(&format!(
"{prefix}_{index}_SET=\"${{{key}+x}}\"\n{prefix}_{index}_VALUE=\"${{{key}-}}\"\n",
prefix = snapshot.prefix,
));
}
script.push_str("builtin unset BASH_ENV\n");
script.push_str(&format!(
"{prefix}_HASH=$(if builtin command -v sha256sum >/dev/null 2>&1; then builtin command sha256sum '{path}'; elif builtin command -v shasum >/dev/null 2>&1; then builtin command shasum -a 256 '{path}'; fi)\n\
{prefix}_HASH=${{{prefix}_HASH%% *}}\n\
if [[ \"${{{prefix}_HASH}}\" = '{checksum}' ]] && builtin source '{path}' 3>&1 4>&2 >/dev/null 2>&1; then builtin true; else builtin export BASH_ENV='{bash_env}'; builtin source \"$BASH_ENV\" || builtin true; fi\n\
builtin unset {prefix}_HASH\n",
prefix = snapshot.prefix,
path = quote(&snapshot.path.to_string_lossy()),
checksum = snapshot.checksum,
bash_env = quote(&snapshot.bash_env),
));
for (index, key) in preserve_env_keys.iter().enumerate() {
script.push_str(&format!(
"if [[ \"${{{prefix}_{index}_SET}}\" = x ]]; then builtin export {key}=\"${{{prefix}_{index}_VALUE}}\"; else builtin unset {key}; fi\nbuiltin unset {prefix}_{index}_SET {prefix}_{index}_VALUE\n",
prefix = snapshot.prefix,
));
}
script.push_str(&format!(
"BASH_EXECUTION_STRING='{}'\n",
quote(argv.get(2)?)
));
script.push_str(argv.get(2)?);
let mut wrapped = vec![argv.first()?.clone(), "-c".to_string(), script];
wrapped.extend(argv.get(3..)?.iter().cloned());
Some(wrapped)
}
fn file_identity(metadata: &Metadata) -> (u64, u64, u64, i64, i64, i64, i64, u32) {
(
metadata.dev(),
metadata.ino(),
metadata.len(),
metadata.mtime(),
metadata.mtime_nsec(),
metadata.ctime(),
metadata.ctime_nsec(),
metadata.mode(),
)
}
fn sort_json(value: &mut serde_json::Value) {
match value {
serde_json::Value::Array(values) => values.iter_mut().for_each(sort_json),
serde_json::Value::Object(map) => {
map.values_mut().for_each(sort_json);
map.sort_keys();
}
_ => {}
}
}
fn valid_variable(name: &str) -> bool {
let mut chars = name.chars();
matches!(chars.next(), Some('_') | Some('a'..='z') | Some('A'..='Z'))
&& chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
}
fn unsupported_startup(source: &[u8]) -> bool {
let positional = |byte| matches!(byte, b'0'..=b'9' | b'@' | b'*' | b'#');
source
.windows(b"BASH_EXECUTION_STRING".len())
.any(|part| part == b"BASH_EXECUTION_STRING")
|| source
.windows(2)
.any(|part| part[0] == b'$' && positional(part[1]))
|| source
.windows(3)
.any(|part| part.starts_with(b"${") && positional(part[2]))
|| source
.split(|byte| !byte.is_ascii_alphanumeric() && *byte != b'_')
.any(|word| word == b"trap" || word == b"builtin")
}
fn quote(value: &str) -> String {
value.replace('\'', "'\"'\"'")
}
fn fallback(
params: &ExecParams,
environment: &HashMap<String, String>,
) -> (Vec<String>, HashMap<String, String>) {
(params.argv.clone(), environment.clone())
}
codex-rs/exec-server/src/environment.rs
+2
View File
@@ -1155,6 +1155,7 @@ mod tests {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -1195,6 +1196,7 @@ mod tests {
pipe_stdin: false,
arg0: None,
sandbox: Some(sandbox),
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
codex-rs/exec-server/src/lib.rs
+2
View File
@@ -1,3 +1,4 @@
mod bash_env_snapshot;
mod client;
mod client_api;
mod client_transport;
@@ -88,6 +89,7 @@ pub use process::ExecProcessEvent;
pub use process::ExecProcessEventReceiver;
pub use process::ExecProcessFuture;
pub use process::StartedExecProcess;
pub use protocol::BashEnvSnapshotParams;
pub use protocol::ByteChunk;
pub use protocol::EnvironmentInfo;
pub use protocol::ExecClosedNotification;
codex-rs/exec-server/src/local_process.rs
+150 -18
View File
@@ -3,6 +3,7 @@ use std::collections::HashSet;
use std::collections::VecDeque;
use std::collections::hash_map::Entry;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::time::Duration;
@@ -33,6 +34,7 @@ use crate::ExecServerError;
use crate::ExecServerRuntimePaths;
use crate::ProcessId;
use crate::StartedExecProcess;
use crate::bash_env_snapshot::BashEnvSnapshotCache;
use crate::process::ExecProcessEventLog;
use crate::process_sandbox::prepare_exec_request;
use crate::protocol::EXEC_CLOSED_METHOD;
@@ -126,7 +128,10 @@ impl AcceptedStdinWriteIds {
}
}
struct ProcessStart;
#[derive(Default)]
struct ProcessStart {
cancelled: AtomicBool,
}
enum ProcessEntry {
Starting(Arc<ProcessStart>),
@@ -136,6 +141,27 @@ enum ProcessEntry {
struct Inner {
notifications: std::sync::RwLock<Option<RpcNotificationSender>>,
processes: Mutex<HashMap<ProcessId, ProcessEntry>>,
bash_env_snapshots: BashEnvSnapshotCache,
}
struct ProcessStartGuard(Option<(Arc<Inner>, ProcessId, Arc<ProcessStart>)>);
impl Drop for ProcessStartGuard {
fn drop(&mut self) {
let Some((inner, process_id, start)) = self.0.take() else {
return;
};
start.cancelled.store(true, Ordering::Release);
tokio::spawn(async move {
let mut process_map = inner.processes.lock().await;
if matches!(
process_map.get(&process_id),
Some(ProcessEntry::Starting(current)) if Arc::ptr_eq(current, &start)
) {
process_map.remove(&process_id);
}
});
}
}
#[derive(Clone)]
@@ -184,12 +210,14 @@ impl LocalProcess {
inner: Arc::new(Inner {
notifications: std::sync::RwLock::new(Some(notifications)),
processes: Mutex::new(HashMap::new()),
bash_env_snapshots: BashEnvSnapshotCache::default(),
}),
runtime_paths,
}
}
pub(crate) async fn shutdown(&self) {
self.clear_bash_env_snapshots();
let remaining = {
let mut processes = self.inner.processes.lock().await;
processes
@@ -214,22 +242,27 @@ impl LocalProcess {
*notification_sender = notifications;
}
pub(crate) fn clear_bash_env_snapshots(&self) {
self.inner.bash_env_snapshots.clear();
}
async fn start_process(
&self,
params: ExecParams,
mut params: ExecParams,
) -> Result<(ExecResponse, watch::Sender<u64>, ExecProcessEventLog), JSONRPCErrorError> {
if params.argv.is_empty() {
return Err(invalid_params("argv must not be empty".to_string()));
}
let process_id = params.process_id.clone();
let prepared =
prepare_exec_request(&params, child_env(&params), self.runtime_paths.as_ref())?;
let (program, args) = prepared
.command
.split_first()
.ok_or_else(|| invalid_params("argv must not be empty".to_string()))?;
let start = Arc::new(ProcessStart);
let start = Arc::new(ProcessStart::default());
{
let mut process_map = self.inner.processes.lock().await;
if process_map.contains_key(&process_id) {
let reusable = match process_map.get(&process_id) {
None => true,
Some(ProcessEntry::Starting(current)) => current.cancelled.load(Ordering::Acquire),
Some(ProcessEntry::Running(_)) => false,
};
if !reusable {
return Err(invalid_request(format!(
"process {process_id} already exists"
)));
@@ -239,6 +272,33 @@ impl LocalProcess {
ProcessEntry::Starting(Arc::clone(&start)),
);
}
let mut start_guard = ProcessStartGuard(Some((
Arc::clone(&self.inner),
process_id.clone(),
Arc::clone(&start),
)));
let environment = child_env(&params);
let (argv, environment) = self
.inner
.bash_env_snapshots
.prepare_launch(&params, &environment, self.runtime_paths.as_ref())
.await;
params.argv = argv;
let prepared = match prepare_exec_request(&params, environment, self.runtime_paths.as_ref())
{
Ok(prepared) => prepared,
Err(err) => {
self.remove_starting_process(&process_id, &start).await;
start_guard.0 = None;
return Err(err);
}
};
let Some((program, args)) = prepared.command.split_first() else {
self.remove_starting_process(&process_id, &start).await;
start_guard.0 = None;
return Err(invalid_params("argv must not be empty".to_string()));
};
let spawned_result = if params.tty {
codex_utils_pty::spawn_pty_process(
@@ -272,13 +332,8 @@ impl LocalProcess {
let spawned = match spawned_result {
Ok(spawned) => spawned,
Err(err) => {
let mut process_map = self.inner.processes.lock().await;
if matches!(
process_map.get(&process_id),
Some(ProcessEntry::Starting(current)) if Arc::ptr_eq(current, &start)
) {
process_map.remove(&process_id);
}
self.remove_starting_process(&process_id, &start).await;
start_guard.0 = None;
return Err(internal_error(err.to_string()));
}
};
@@ -324,6 +379,7 @@ impl LocalProcess {
})),
);
}
start_guard.0 = None;
tokio::spawn(stream_output(
process_id.clone(),
@@ -357,6 +413,16 @@ impl LocalProcess {
Ok((ExecResponse { process_id }, wake_tx, events))
}
async fn remove_starting_process(&self, process_id: &ProcessId, start: &Arc<ProcessStart>) {
let mut process_map = self.inner.processes.lock().await;
if matches!(
process_map.get(process_id),
Some(ProcessEntry::Starting(current)) if Arc::ptr_eq(current, start)
) {
process_map.remove(process_id);
}
}
pub(crate) async fn exec(&self, params: ExecParams) -> Result<ExecResponse, JSONRPCErrorError> {
self.start_process(params)
.await
@@ -964,6 +1030,7 @@ mod tests {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
}
@@ -993,6 +1060,71 @@ mod tests {
assert_eq!(error, expected);
}
#[cfg(unix)]
#[tokio::test]
async fn cancelled_snapshot_capture_allows_immediate_retry_and_kills_descendants() {
let temp = tempfile::tempdir().expect("tempdir");
let bash_env = temp.path().join("bash_env.sh");
let child_pid_file = temp.path().join("child.pid");
std::fs::write(
&bash_env,
format!(
"sleep 10 & printf '%s' \"$!\" > '{}'; wait\n",
child_pid_file.display()
),
)
.expect("write slow BASH_ENV");
let cwd = PathUri::from_host_native_path(temp.path()).expect("cwd URI");
let mut params = test_exec_params(HashMap::from([
("BASH_ENV".to_string(), bash_env.display().to_string()),
("PATH".to_string(), std::env::var("PATH").expect("PATH")),
]));
params.argv = vec![
"/bin/bash".to_string(),
"-c".to_string(),
"true".to_string(),
];
params.cwd = cwd.clone();
params.bash_env_snapshot = Some(crate::protocol::BashEnvSnapshotParams {
workspace_root: cwd,
preserve_env_keys: Vec::new(),
});
let backend = LocalProcess::default();
let task = tokio::spawn({
let (backend, params) = (backend.clone(), params.clone());
async move { backend.start_process(params).await }
});
let child_pid = timeout(Duration::from_secs(1), async {
loop {
if let Ok(contents) = std::fs::read_to_string(&child_pid_file)
&& let Ok(pid) = contents.parse::<libc::pid_t>()
{
break pid;
}
tokio::task::yield_now().await;
}
})
.await
.expect("capture should start a descendant");
task.abort();
let _ = task.await;
std::fs::write(&bash_env, ":\n").expect("replace slow BASH_ENV");
timeout(Duration::from_secs(2), backend.start_process(params))
.await
.expect("retry should not time out")
.expect("process id should be reusable");
timeout(Duration::from_secs(2), async {
while unsafe { libc::kill(child_pid, 0) } == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap_or_else(|_| {
unsafe { libc::kill(child_pid, libc::SIGKILL) };
panic!("capture descendant {child_pid} should be terminated");
});
}
#[test]
fn child_env_defaults_to_exact_env() {
let params = test_exec_params(HashMap::from([("ONLY_THIS".to_string(), "1".to_string())]));
codex-rs/exec-server/src/process_sandbox_tests.rs
+3
View File
@@ -45,6 +45,7 @@ fn sandbox_request_wraps_native_argv_on_executor() {
pipe_stdin: false,
arg0: None,
sandbox: Some(sandbox),
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
};
@@ -106,6 +107,7 @@ fn sandbox_request_allows_prepared_managed_proxy_port() {
pipe_stdin: false,
arg0: None,
sandbox: Some(sandbox),
bash_env_snapshot: None,
enforce_managed_network: true,
managed_network: Some(ManagedNetworkSandboxContext {
loopback_ports: vec![43123],
@@ -142,6 +144,7 @@ fn native_request_preserves_native_launch_fields() {
pipe_stdin: false,
arg0: Some("custom-arg0".to_string()),
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
};
codex-rs/exec-server/src/server/handler/tests.rs
+1
View File
@@ -35,6 +35,7 @@ fn exec_params_with_argv(process_id: &str, argv: Vec<String>) -> ExecParams {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
}
codex-rs/exec-server/src/server/process_handler.rs
+4
View File
@@ -33,6 +33,10 @@ impl ProcessHandler {
self.process.shutdown().await;
}
pub(crate) fn clear_bash_env_snapshots(&self) {
self.process.clear_bash_env_snapshots();
}
pub(crate) fn set_notification_sender(&self, notifications: Option<RpcNotificationSender>) {
self.process.set_notification_sender(notifications);
}
codex-rs/exec-server/src/server/processor.rs
+1
View File
@@ -439,6 +439,7 @@ mod tests {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
}
codex-rs/exec-server/src/server/session_registry.rs
+1
View File
@@ -248,6 +248,7 @@ impl SessionHandle {
if !self.entry.detach(self.connection_id) {
return;
}
self.entry.process.clear_bash_env_snapshots();
let registry = Arc::clone(&self.registry);
let session_id = self.entry.session_id.clone();
codex-rs/exec-server/tests/exec_process.rs
+57
View File
@@ -5,6 +5,8 @@ use std::sync::Arc;
use anyhow::Context;
use anyhow::Result;
#[cfg(unix)]
use codex_exec_server::BashEnvSnapshotParams;
use codex_exec_server::Environment;
use codex_exec_server::ExecBackend;
use codex_exec_server::ExecOutputStream;
@@ -82,6 +84,7 @@ async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -226,6 +229,7 @@ async fn assert_exec_process_streams_output(use_remote: bool) -> Result<()> {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -260,6 +264,7 @@ async fn assert_exec_process_pushes_events(use_remote: bool) -> Result<()> {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -310,6 +315,7 @@ async fn assert_exec_process_replays_events_after_close(use_remote: bool) -> Res
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -361,6 +367,7 @@ async fn assert_exec_process_retains_output_after_exit_until_streams_close(
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -437,6 +444,7 @@ async fn assert_exec_process_write_then_read(use_remote: bool) -> Result<()> {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -477,6 +485,7 @@ async fn assert_exec_process_write_then_read_without_tty(use_remote: bool) -> Re
pipe_stdin: true,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -513,6 +522,7 @@ async fn assert_exec_process_rejects_write_without_pipe_stdin(use_remote: bool)
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -550,6 +560,7 @@ async fn assert_exec_process_signal_interrupts_process(use_remote: bool) -> Resu
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -606,6 +617,7 @@ async fn assert_exec_process_signal_reports_unsupported_on_windows(use_remote: b
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -649,6 +661,7 @@ async fn assert_exec_process_preserves_queued_events_before_subscribe(
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -710,6 +723,7 @@ async fn remote_exec_process_recovers_after_transport_disconnect() -> Result<()>
pipe_stdin: true,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
@@ -816,6 +830,49 @@ async fn remote_exec_process_recovers_after_transport_disconnect() -> Result<()>
Ok(())
}
#[cfg(unix)]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial(remote_exec_server)]
async fn remote_bash_env_snapshot_is_reused_across_execs() -> Result<()> {
let server = exec_server().await?;
let backend =
Environment::create_for_tests(Some(server.websocket_url().to_string()))?.get_exec_backend();
let temp = tempfile::tempdir()?;
let bash_env = temp.path().join("bash_env.sh");
std::fs::write(&bash_env, "printf x >> counter\nexport SENTINEL=ready\n")?;
let workspace_uri = PathUri::from_host_native_path(temp.path())?;
for index in 0..2 {
let session = backend
.start(ExecParams {
process_id: format!("bash-snapshot-{index}").into(),
argv: vec![
"/bin/bash".into(),
"-c".into(),
"printf %s \"$SENTINEL\"".into(),
],
cwd: workspace_uri.clone(),
env_policy: None,
env: HashMap::from([("BASH_ENV".to_string(), bash_env.display().to_string())]),
tty: false,
pipe_stdin: false,
arg0: None,
sandbox: None,
enforce_managed_network: false,
managed_network: None,
bash_env_snapshot: Some(BashEnvSnapshotParams {
workspace_root: workspace_uri.clone(),
preserve_env_keys: Vec::new(),
}),
})
.await?;
let wake = session.process.subscribe_wake();
let (output, ..) = collect_process_output_from_reads(session.process, wake).await?;
assert_eq!(output, "ready");
}
assert_eq!(std::fs::read_to_string(temp.path().join("counter"))?, "x");
Ok(())
}
#[test_case(false ; "local")]
#[test_case(true ; "remote")]
#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")]
codex-rs/exec-server/tests/relay.rs
+1
View File
@@ -151,6 +151,7 @@ async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> {
pipe_stdin: false,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})
codex-rs/rmcp-client/src/stdio_server_launcher.rs
+1
View File
@@ -508,6 +508,7 @@ impl ExecutorStdioServerLauncher {
pipe_stdin: true,
arg0: None,
sandbox: None,
bash_env_snapshot: None,
enforce_managed_network: false,
managed_network: None,
})