Compare commits

...

1 Commits

Author SHA1 Message Date
celia-oai
9f678d256b changes 2026-03-16 23:22:32 -07:00
4 changed files with 213 additions and 49 deletions

View File

@@ -62,14 +62,56 @@ pub(crate) fn build_command_spec(
/// shell -lc "<script>"
/// => user_shell -c ". SNAPSHOT (best effort); exec shell -c <script>"
///
/// This wrapper script uses POSIX constructs (`if`, `.`, `exec`) so it can
/// be run by Bash/Zsh/sh. On non-matching commands, or when command cwd does
/// not match the snapshot cwd, this is a no-op.
/// This wrapper script uses POSIX constructs (`if`, `.`, `exec`) so it can be
/// run by Bash/Zsh/sh. On non-matching commands, or when command cwd does not
/// match the snapshot cwd, this is a no-op.
pub(crate) fn maybe_wrap_shell_lc_with_snapshot(
command: &[String],
session_shell: &Shell,
cwd: &Path,
explicit_env_overrides: &HashMap<String, String>,
) -> Vec<String> {
rewrite_shell_lc_with_snapshot(
command,
session_shell,
cwd,
explicit_env_overrides,
SnapshotWrapMode::ExecOriginalShell,
)
}
/// POSIX-only helper: equivalent to `maybe_wrap_shell_lc_with_snapshot`, but
/// keeps the original script in the same shell process after sourcing the
/// snapshot instead of `exec`-ing the original shell. This avoids an extra
/// shell re-exec, which matters for the zsh-fork path in restricted read-only
/// sandboxes.
pub(crate) fn maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
command: &[String],
session_shell: &Shell,
cwd: &Path,
explicit_env_overrides: &HashMap<String, String>,
) -> Vec<String> {
rewrite_shell_lc_with_snapshot(
command,
session_shell,
cwd,
explicit_env_overrides,
SnapshotWrapMode::RunInSessionShell,
)
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum SnapshotWrapMode {
ExecOriginalShell,
RunInSessionShell,
}
fn rewrite_shell_lc_with_snapshot(
command: &[String],
session_shell: &Shell,
cwd: &Path,
explicit_env_overrides: &HashMap<String, String>,
wrap_mode: SnapshotWrapMode,
) -> Vec<String> {
if cfg!(windows) {
return command.to_vec();
@@ -105,25 +147,46 @@ pub(crate) fn maybe_wrap_shell_lc_with_snapshot(
let snapshot_path = snapshot.path.to_string_lossy();
let shell_path = session_shell.shell_path.to_string_lossy();
let original_shell = shell_single_quote(&command[0]);
let original_script = shell_single_quote(&command[2]);
let snapshot_path = shell_single_quote(snapshot_path.as_ref());
let trailing_args = command[3..]
.iter()
.map(|arg| format!(" '{}'", shell_single_quote(arg)))
.collect::<String>();
let (override_captures, override_exports) = build_override_exports(explicit_env_overrides);
let rewritten_script = if override_exports.is_empty() {
format!(
"if . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\nexec '{original_shell}' -c '{original_script}'{trailing_args}"
)
} else {
format!(
"{override_captures}\n\nif . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\n{override_exports}\n\nexec '{original_shell}' -c '{original_script}'{trailing_args}"
)
let rewritten_script = match wrap_mode {
SnapshotWrapMode::ExecOriginalShell => {
let original_shell = shell_single_quote(&command[0]);
let original_script = shell_single_quote(&command[2]);
let trailing_args = command[3..]
.iter()
.map(|arg| format!(" '{}'", shell_single_quote(arg)))
.collect::<String>();
if override_exports.is_empty() {
format!(
"if . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\nexec '{original_shell}' -c '{original_script}'{trailing_args}"
)
} else {
format!(
"{override_captures}\n\nif . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\n{override_exports}\n\nexec '{original_shell}' -c '{original_script}'{trailing_args}"
)
}
}
SnapshotWrapMode::RunInSessionShell => {
if override_exports.is_empty() {
format!(
"if . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\n{}",
command[2]
)
} else {
format!(
"{override_captures}\n\nif . '{snapshot_path}' >/dev/null 2>&1; then :; fi\n\n{override_exports}\n\n{}",
command[2]
)
}
}
};
vec![shell_path.to_string(), "-c".to_string(), rewritten_script]
let mut rewritten = vec![shell_path.to_string(), "-c".to_string(), rewritten_script];
if wrap_mode == SnapshotWrapMode::RunInSessionShell {
rewritten.extend(command[3..].iter().cloned());
}
rewritten
}
fn build_override_exports(explicit_env_overrides: &HashMap<String, String>) -> (String, String) {

View File

@@ -33,7 +33,7 @@ fn maybe_wrap_shell_lc_with_snapshot_bootstraps_in_user_shell() {
let session_shell = shell_with_snapshot(
ShellType::Zsh,
"/bin/zsh",
snapshot_path,
snapshot_path.clone(),
dir.path().to_path_buf(),
);
let command = vec![
@@ -59,7 +59,7 @@ fn maybe_wrap_shell_lc_with_snapshot_escapes_single_quotes() {
let session_shell = shell_with_snapshot(
ShellType::Zsh,
"/bin/zsh",
snapshot_path,
snapshot_path.clone(),
dir.path().to_path_buf(),
);
let command = vec![
@@ -82,7 +82,7 @@ fn maybe_wrap_shell_lc_with_snapshot_uses_bash_bootstrap_shell() {
let session_shell = shell_with_snapshot(
ShellType::Bash,
"/bin/bash",
snapshot_path,
snapshot_path.clone(),
dir.path().to_path_buf(),
);
let command = vec![
@@ -134,7 +134,7 @@ fn maybe_wrap_shell_lc_with_snapshot_preserves_trailing_args() {
let session_shell = shell_with_snapshot(
ShellType::Zsh,
"/bin/zsh",
snapshot_path,
snapshot_path.clone(),
dir.path().to_path_buf(),
);
let command = vec![
@@ -148,6 +148,7 @@ fn maybe_wrap_shell_lc_with_snapshot_preserves_trailing_args() {
let rewritten =
maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path(), &HashMap::new());
assert_eq!(rewritten.len(), 3);
assert!(
rewritten[2]
.contains(r#"exec '/bin/bash' -c 'printf '"'"'%s %s'"'"' "$0" "$1"' 'arg0' 'arg1'"#)
@@ -195,13 +196,17 @@ fn maybe_wrap_shell_lc_with_snapshot_accepts_dot_alias_cwd() {
];
let command_cwd = dir.path().join(".");
let rewritten =
maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, &command_cwd, &HashMap::new());
let rewritten = maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&command,
&session_shell,
&command_cwd,
&HashMap::new(),
);
assert_eq!(rewritten[0], "/bin/zsh");
assert_eq!(rewritten[1], "-c");
assert!(rewritten[2].contains("if . '"));
assert!(rewritten[2].contains("exec '/bin/bash' -c 'echo hello'"));
assert!(rewritten[2].ends_with("\n\necho hello"));
}
#[test]
@@ -226,7 +231,7 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_explicit_override_precedence() {
];
let explicit_env_overrides =
HashMap::from([("TEST_ENV_SNAPSHOT".to_string(), "worktree".to_string())]);
let rewritten = maybe_wrap_shell_lc_with_snapshot(
let rewritten = maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&command,
&session_shell,
dir.path(),
@@ -265,8 +270,12 @@ fn maybe_wrap_shell_lc_with_snapshot_keeps_snapshot_path_without_override() {
"-lc".to_string(),
"printf '%s' \"$PATH\"".to_string(),
];
let rewritten =
maybe_wrap_shell_lc_with_snapshot(&command, &session_shell, dir.path(), &HashMap::new());
let rewritten = maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&command,
&session_shell,
dir.path(),
&HashMap::new(),
);
let output = Command::new(&rewritten[0])
.args(&rewritten[1..])
.output()
@@ -297,7 +306,7 @@ fn maybe_wrap_shell_lc_with_snapshot_applies_explicit_path_override() {
"printf '%s' \"$PATH\"".to_string(),
];
let explicit_env_overrides = HashMap::from([("PATH".to_string(), "/worktree/bin".to_string())]);
let rewritten = maybe_wrap_shell_lc_with_snapshot(
let rewritten = maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&command,
&session_shell,
dir.path(),
@@ -337,7 +346,7 @@ fn maybe_wrap_shell_lc_with_snapshot_does_not_embed_override_values_in_argv() {
"OPENAI_API_KEY".to_string(),
"super-secret-value".to_string(),
)]);
let rewritten = maybe_wrap_shell_lc_with_snapshot(
let rewritten = maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&command,
&session_shell,
dir.path(),
@@ -381,7 +390,7 @@ fn maybe_wrap_shell_lc_with_snapshot_preserves_unset_override_variables() {
"CODEX_TEST_UNSET_OVERRIDE".to_string(),
"worktree-value".to_string(),
)]);
let rewritten = maybe_wrap_shell_lc_with_snapshot(
let rewritten = maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&command,
&session_shell,
dir.path(),
@@ -396,3 +405,74 @@ fn maybe_wrap_shell_lc_with_snapshot_preserves_unset_override_variables() {
assert!(output.status.success(), "command failed: {output:?}");
assert_eq!(String::from_utf8_lossy(&output.stdout), "unset");
}
#[test]
fn maybe_wrap_shell_lc_with_snapshot_for_zsh_fork_bootstraps_in_user_shell() {
let dir = tempdir().expect("create temp dir");
let snapshot_path = dir.path().join("snapshot.sh");
std::fs::write(&snapshot_path, "# Snapshot file\n").expect("write snapshot");
let session_shell = shell_with_snapshot(
ShellType::Zsh,
"/bin/zsh",
snapshot_path.clone(),
dir.path().to_path_buf(),
);
let command = vec![
"/bin/bash".to_string(),
"-lc".to_string(),
"echo hello".to_string(),
];
let rewritten = maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&command,
&session_shell,
dir.path(),
&HashMap::new(),
);
assert_eq!(rewritten[0], "/bin/zsh");
assert_eq!(rewritten[1], "-c");
assert!(rewritten[2].contains("if . '"));
assert!(rewritten[2].ends_with("\n\necho hello"));
}
#[test]
fn maybe_wrap_shell_lc_with_snapshot_for_zsh_fork_preserves_trailing_args() {
let dir = tempdir().expect("create temp dir");
let snapshot_path = dir.path().join("snapshot.sh");
std::fs::write(&snapshot_path, "# Snapshot file\n").expect("write snapshot");
let session_shell = shell_with_snapshot(
ShellType::Zsh,
"/bin/zsh",
snapshot_path.clone(),
dir.path().to_path_buf(),
);
let command = vec![
"/bin/bash".to_string(),
"-lc".to_string(),
"printf '%s %s' \"$0\" \"$1\"".to_string(),
"arg0".to_string(),
"arg1".to_string(),
];
let rewritten = maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&command,
&session_shell,
dir.path(),
&HashMap::new(),
);
assert_eq!(
rewritten,
vec![
"/bin/zsh".to_string(),
"-c".to_string(),
format!(
"if . '{}' >/dev/null 2>&1; then :; fi\n\nprintf '%s %s' \"$0\" \"$1\"",
snapshot_path.display()
),
"arg0".to_string(),
"arg1".to_string(),
]
);
}

View File

@@ -22,6 +22,7 @@ use crate::tools::network_approval::NetworkApprovalMode;
use crate::tools::network_approval::NetworkApprovalSpec;
use crate::tools::runtimes::build_command_spec;
use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot;
use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot_for_zsh_fork;
use crate::tools::sandboxing::Approvable;
use crate::tools::sandboxing::ApprovalCtx;
use crate::tools::sandboxing::ExecApprovalRequirement;
@@ -221,12 +222,21 @@ impl ToolRuntime<ShellRequest, ExecToolCallOutput> for ShellRuntime {
ctx: &ToolCtx,
) -> Result<ExecToolCallOutput, ToolError> {
let session_shell = ctx.session.user_shell();
let command = maybe_wrap_shell_lc_with_snapshot(
&req.command,
session_shell.as_ref(),
&req.cwd,
&req.explicit_env_overrides,
);
let command = if self.backend == ShellRuntimeBackend::ShellCommandZshFork {
maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
&req.command,
session_shell.as_ref(),
&req.cwd,
&req.explicit_env_overrides,
)
} else {
maybe_wrap_shell_lc_with_snapshot(
&req.command,
session_shell.as_ref(),
&req.cwd,
&req.explicit_env_overrides,
)
};
let command = if matches!(session_shell.shell_type, ShellType::PowerShell)
&& ctx.session.features().enabled(Feature::PowershellUtf8)
{

View File

@@ -19,6 +19,7 @@ use crate::tools::network_approval::NetworkApprovalMode;
use crate::tools::network_approval::NetworkApprovalSpec;
use crate::tools::runtimes::build_command_spec;
use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot;
use crate::tools::runtimes::maybe_wrap_shell_lc_with_snapshot_for_zsh_fork;
use crate::tools::runtimes::shell::zsh_fork_backend;
use crate::tools::sandboxing::Approvable;
use crate::tools::sandboxing::ApprovalCtx;
@@ -194,18 +195,28 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecProcess> for UnifiedExecRunt
) -> Result<UnifiedExecProcess, ToolError> {
let base_command = &req.command;
let session_shell = ctx.session.user_shell();
let command = maybe_wrap_shell_lc_with_snapshot(
base_command,
session_shell.as_ref(),
&req.cwd,
&req.explicit_env_overrides,
);
let command = if matches!(session_shell.shell_type, ShellType::PowerShell)
&& ctx.session.features().enabled(Feature::PowershellUtf8)
{
prefix_powershell_script_with_utf8(&command)
} else {
command
let command = match &self.shell_mode {
UnifiedExecShellMode::Direct => {
let command = maybe_wrap_shell_lc_with_snapshot(
base_command,
session_shell.as_ref(),
&req.cwd,
&req.explicit_env_overrides,
);
if matches!(session_shell.shell_type, ShellType::PowerShell)
&& ctx.session.features().enabled(Feature::PowershellUtf8)
{
prefix_powershell_script_with_utf8(&command)
} else {
command
}
}
UnifiedExecShellMode::ZshFork(_) => maybe_wrap_shell_lc_with_snapshot_for_zsh_fork(
base_command,
session_shell.as_ref(),
&req.cwd,
&req.explicit_env_overrides,
),
};
let mut env = req.env.clone();