diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index b6a6d5938f..e4594f78b7 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -405,6 +405,11 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc) { } let _ = sess.conversation.shutdown().await; sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + let shell_snapshot_prewarm = sess.state.lock().await.shell_snapshot_prewarm.take(); + if let Some(shell_snapshot_prewarm) = shell_snapshot_prewarm { + shell_snapshot_prewarm.abort(); + let _ = shell_snapshot_prewarm.await; + } sess.hooks().shutdown().await; sess.async_hook_results.close(); while sess.async_hook_results.try_recv().is_ok() {} diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index f43ad3cc35..96ba5ffe03 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -269,6 +269,15 @@ pub(crate) async fn run_turn( return Ok(None); } + // Only speculate after hooks accept the turn, using its finalized tools and permissions. + { + let mut state = sess.state.lock().await; + if state.shell_snapshot_prewarm.is_none() { + state.shell_snapshot_prewarm = + sess.prewarm_shell_snapshots(first_step_context.as_ref()); + } + } + sess.merge_connector_selection(explicitly_enabled_connectors.clone()) .await; sess.set_previous_turn_settings(Some(PreviousTurnSettings { diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 768e8d9fe3..e9e42cf00a 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -27,6 +27,7 @@ use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TokenUsageRecord; use codex_protocol::protocol::TurnContextItem; use codex_utils_output_truncation::TruncationPolicy; +use tokio_util::task::AbortOnDropHandle; /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { @@ -47,6 +48,8 @@ pub(crate) struct SessionState { auto_compact_window: AutoCompactWindow, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, + /// Retained after completion so later turns do not repeat speculative captures. + pub(crate) shell_snapshot_prewarm: Option>, pub(crate) current_time_reminder: CurrentTimeReminderState, pub(crate) active_connector_selection: HashSet, pub(crate) pending_session_start_sources: VecDeque, @@ -81,6 +84,7 @@ impl SessionState { previous_turn_settings: None, auto_compact_window: AutoCompactWindow::new_with_ids(auto_compact_window_ids), startup_prewarm: None, + shell_snapshot_prewarm: None, current_time_reminder: CurrentTimeReminderState::default(), active_connector_selection: HashSet::new(), pending_session_start_sources: VecDeque::new(), diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 67c065faee..7de566390b 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -114,7 +114,7 @@ impl ToolCallRuntime { } let router = &self.step_context.tool_router; let supports_parallel = router.tool_supports_parallel(&call); - let tool_runtime = router.tool_runtime(&call); + let tool_runtime = router.tool_runtime(&call.tool_name); let router = Arc::clone(router); let session = Arc::clone(&self.session); let step_context = Arc::clone(&self.step_context); diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index f7d65f5185..8d5bde3085 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -238,8 +238,8 @@ impl ToolRouter { .unwrap_or(false) } - pub(crate) fn tool_runtime(&self, call: &ToolCall) -> Option> { - self.registry.tool(&call.tool_name) + pub(crate) fn tool_runtime(&self, tool_name: &ToolName) -> Option> { + self.registry.tool(tool_name) } #[instrument(level = "trace", skip_all, err)] diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 01935232b4..1d6b1d322f 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -190,14 +190,7 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow assert_eq!( router - .tool_runtime(&ToolCall { - tool_name: ToolName::plain(parallel_tool_name), - call_id: "call-local-tool".to_string(), - payload: ToolPayload::Function { - arguments: "{}".to_string(), - }, - encrypted_function_args: None, - }) + .tool_runtime(&ToolName::plain(parallel_tool_name)) .map(|runtime| runtime.tool_name()), Some(ToolName::plain(parallel_tool_name)) ); @@ -369,7 +362,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { assert!(router.tool_supports_parallel(&call)); assert_eq!( router - .tool_runtime(&call) + .tool_runtime(&call.tool_name) .map(|runtime| runtime.tool_name()), Some(call.tool_name.clone()) ); @@ -385,7 +378,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { assert!(!router.tool_supports_parallel(&different_server_call)); assert_eq!( router - .tool_runtime(&different_server_call) + .tool_runtime(&different_server_call.tool_name) .map(|runtime| runtime.tool_name()), Some(different_server_call.tool_name.clone()) ); @@ -399,7 +392,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { encrypted_function_args: None, }; assert!(!router.tool_supports_parallel(&hidden_call)); - assert!(router.tool_runtime(&hidden_call).is_some()); + assert!(router.tool_runtime(&hidden_call.tool_name).is_some()); let nested_only_call = ToolCall { tool_name: ToolName::namespaced("mcp__nested_echo__", "query_with_delay"), diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 1a0d829949..21218570af 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -121,14 +121,14 @@ fn should_use_deterministic_process_ids() -> bool { cfg!(test) || deterministic_process_ids_forced_for_tests() } -fn apply_unified_exec_env(mut env: HashMap) -> HashMap { +pub(super) fn apply_unified_exec_env(mut env: HashMap) -> HashMap { for (key, value) in UNIFIED_EXEC_ENV { env.insert(key.to_string(), value.to_string()); } env } -fn exec_env_policy_from_shell_policy( +pub(super) fn exec_env_policy_from_shell_policy( policy: &ShellEnvironmentPolicy, ) -> codex_exec_server::ExecEnvPolicy { let mut exclude = policy diff --git a/codex-rs/core/src/unified_exec/shell_snapshot.rs b/codex-rs/core/src/unified_exec/shell_snapshot.rs index ddd93ec8b5..43e95cc1cb 100644 --- a/codex-rs/core/src/unified_exec/shell_snapshot.rs +++ b/codex-rs/core/src/unified_exec/shell_snapshot.rs @@ -1,12 +1,175 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use codex_exec_server::ExecParams; +use codex_exec_server::FileSystemSandboxContext; use codex_exec_server::ShellInfo; use codex_exec_server::ShellSnapshotRequest; use codex_features::Feature; +use codex_protocol::protocol::AskForApproval; +use codex_sandboxing::SandboxManager; +use codex_sandboxing::SandboxablePreference; +use codex_tools::ToolName; use codex_tools::UnifiedExecShellMode; use codex_utils_path_uri::PathUri; +use tokio_util::task::AbortOnDropHandle; +use uuid::Uuid; use super::ExecCommandRequest; use super::UnifiedExecContext; +use super::process_manager::apply_unified_exec_env; +use super::process_manager::exec_env_policy_from_shell_policy; +use crate::config::NetworkProxySpec; +use crate::exec_env::CODEX_THREAD_ID_ENV_VAR; +use crate::exec_env::inject_apply_patch_env; +use crate::exec_env::inject_permission_profile_env; +use crate::exec_env::inject_session_id_env; +use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::shell::ShellType; +use crate::tools::sandboxing::executor_windows_sandbox_level; + +impl Session { + pub(crate) fn prewarm_shell_snapshots( + self: &Arc, + step_context: &StepContext, + ) -> Option> { + if !self.features().enabled(Feature::ShellSnapshotV2) + || step_context + .tool_router + .tool_runtime(&ToolName::plain("exec_command")) + .is_none() + { + return None; + } + + let config = Arc::clone(&step_context.turn.config); + let user_shell = self.user_shell(); + if config.active_project.is_untrusted() + || matches!( + step_context.settings.approval_policy(), + AskForApproval::UnlessTrusted + ) + || step_context.turn.network.is_some() + || config + .permissions + .network + .as_ref() + .is_some_and(NetworkProxySpec::enabled) + || !matches!( + step_context.turn.unified_exec_shell_mode, + UnifiedExecShellMode::Direct + ) + { + return None; + } + + let environments = step_context.environments.clone(); + let session = Arc::clone(self); + Some(AbortOnDropHandle::new(tokio::spawn(async move { + for environment in environments.turn_environments() { + // Remote starts cannot cancel an in-flight capture. Leave them, deferred + // environments, and managed networking on the normal lazy capture path. + if environment.environment.is_remote() + || !environment.shell_snapshot_v2_supported + || !environment.config().allow_login_shell + || environment.config().network_policy.is_some() + { + continue; + } + + let shell = environment.shell.as_ref().unwrap_or(user_shell.as_ref()); + if !matches!( + shell.shell_type, + ShellType::Bash | ShellType::Zsh | ShellType::Sh + ) { + continue; + } + let sandbox = SandboxManager::new() + .should_sandbox( + environment.permission_profile(), + SandboxablePreference::Auto, + /*has_managed_network_requirements*/ false, + ) + .then(|| FileSystemSandboxContext { + permissions: environment.permission_profile().clone().into(), + cwd: Some(environment.cwd().clone()), + workspace_roots: environment.workspace_roots().to_vec(), + user_home_dir: None, + temporary_directories: None, + windows_sandbox_level: executor_windows_sandbox_level( + environment.config().windows_sandbox_level, + environment.cwd(), + ), + windows_sandbox_private_desktop: environment + .config() + .windows_sandbox_private_desktop, + windows_sandbox_proxy_settings_mode: Some( + session.windows_sandbox_proxy_settings_mode, + ), + use_legacy_landlock: environment.config().use_legacy_landlock, + }); + let mut env = apply_unified_exec_env(HashMap::new()); + env.insert( + CODEX_THREAD_ID_ENV_VAR.to_string(), + session.thread_id().to_string(), + ); + inject_session_id_env(&mut env, session.session_id()); + inject_apply_patch_env(&mut env, &config.features); + inject_permission_profile_env( + &mut env, + environment.active_permission_profile().as_ref(), + ); + #[cfg(unix)] + if let Some(path) = environment.shell_environment_policy().r#set.get("PATH") { + env.insert("PATH".to_string(), path.clone()); + crate::tools::runtimes::apply_package_path_prepend( + &mut env, + &mut crate::tools::runtimes::RuntimePathPrepends::default(), + ); + } + let params = ExecParams { + process_id: format!("shell-snapshot-warmup-{}", Uuid::new_v4()).into(), + // This capture-only operation never starts the command or routes + // speculative work through tool approvals or exec-policy bypass. + argv: shell.derive_exec_args("", /*use_login_shell*/ true), + cwd: environment.cwd().clone(), + env_policy: Some(exec_env_policy_from_shell_policy( + environment.shell_environment_policy(), + )), + shell_snapshot: Some(ShellSnapshotRequest { + scope_id: format!( + "{}:{}", + session.thread_id(), + environment.selection.environment_id + ), + shell: ShellInfo { + name: shell.name().to_string(), + path: shell.shell_path.to_string_lossy().into_owned(), + }, + }), + env, + tty: false, + pipe_stdin: false, + arg0: None, + sandbox, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }; + + let result = environment + .environment + .get_exec_backend() + .prewarm_shell_snapshot(params) + .await; + if let Err(err) = result { + tracing::debug!("failed to warm shell snapshot: {err}"); + } + } + }))) + } +} pub(super) fn shell_snapshot_request( request: &ExecCommandRequest, diff --git a/codex-rs/core/tests/suite/shell_snapshot.rs b/codex-rs/core/tests/suite/shell_snapshot.rs index 21371a90c0..902b587bc7 100644 --- a/codex-rs/core/tests/suite/shell_snapshot.rs +++ b/codex-rs/core/tests/suite/shell_snapshot.rs @@ -1,27 +1,55 @@ use anyhow::Result; use codex_core::TurnInputRequest; #[cfg(unix)] +use codex_core::config::Constrained; +#[cfg(unix)] +use codex_core::config::NetworkProxySpec; +#[cfg(unix)] use codex_core::shell::get_shell_by_model_provided_path; +#[cfg(unix)] +use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_features::Feature; +#[cfg(unix)] +use codex_network_proxy::NetworkProxyConfig; +#[cfg(unix)] +use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; +#[cfg(unix)] +use codex_protocol::config_types::TrustLevel; +#[cfg(unix)] +use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::models::PermissionProfile; +#[cfg(unix)] +use codex_protocol::models::PermissionProfileSnapshot; +#[cfg(unix)] +use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::AskForApproval; +#[cfg(unix)] +use codex_protocol::protocol::EnvironmentConfig; +#[cfg(unix)] +use codex_protocol::protocol::EnvironmentConfigState; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ExecCommandBeginEvent; use codex_protocol::protocol::ExecCommandEndEvent; use codex_protocol::protocol::Op; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; +#[cfg(unix)] +use core_test_support::hooks::trust_discovered_hooks; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; +#[cfg(unix)] +use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; #[cfg(unix)] use core_test_support::skip_if_remote; +#[cfg(unix)] +use core_test_support::test_codex::TestCodexBuilder; use core_test_support::test_codex::TestCodexHarness; use core_test_support::test_codex::local_selections; use core_test_support::test_codex::test_codex; @@ -84,11 +112,12 @@ async fn wait_for_snapshot(codex_home: &Path) -> Result { } } -async fn wait_for_file_contents(path: &Path) -> Result { +async fn wait_for_file_contents(path: &Path, expected: &str) -> Result<()> { let deadline = Instant::now() + Duration::from_secs(15); loop { match fs::read_to_string(path).await { - Ok(contents) => return Ok(contents), + Ok(contents) if contents == expected => return Ok(()), + Ok(_) => {} Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} Err(err) => return Err(err.into()), } @@ -292,6 +321,437 @@ fn assert_posix_snapshot_sections(snapshot: &str) { ); } +#[cfg(unix)] +fn shell_snapshot_v2_prewarm_builder(profile_home: &Path) -> TestCodexBuilder { + let configured_home = profile_home.to_string_lossy().into_owned(); + let shell = get_shell_by_model_provided_path(&PathBuf::from("/bin/bash")); + test_codex() + .with_user_shell(shell) + .with_config(move |config| { + config + .features + .enable(Feature::UnifiedExec) + .expect("test config should enable unified exec"); + config + .features + .enable(Feature::ShellSnapshotV2) + .expect("test config should enable in-memory snapshots"); + config + .permissions + .set_permission_profile(PermissionProfile::Disabled) + .expect("test config should allow unrestricted permissions"); + config + .permissions + .shell_environment_policy + .ignore_default_excludes = false; + config.permissions.shell_environment_policy.r#set = + HashMap::from([("HOME".to_string(), configured_home)]); + }) +} + +#[cfg(unix)] +async fn run_no_shell_turn(harness: &TestCodexHarness) -> Result<()> { + let response = mount_sse_once( + harness.server(), + sse(vec![ + ev_response_created("no-shell"), + ev_assistant_message("done", "done"), + ev_completed("no-shell"), + ]), + ) + .await; + let codex = &harness.test().codex; + codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "hello".to_string(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(codex, |event| { + assert!( + !matches!(event, EventMsg::ExecApprovalRequest(_) | EventMsg::Error(_)), + "unexpected event: {event:?}" + ); + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + response.single_request(); + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_v2_warms_after_hooks_without_blocking_the_model() -> Result<()> { + skip_if_remote!(Ok(()), "profile fixture uses a host-local HOME directory"); + let profile_home = tempfile::tempdir()?; + fs::write( + profile_home.path().join(".bashrc"), + ". \"$HOME/prepared\"\nprintf x >> \"$HOME/captures\"\nwhile [ ! -f \"$HOME/ready\" ]; do /bin/sleep 0.01; done\nexport PROFILE_SECRET=secret\n", + ) + .await?; + let prepared = profile_home.path().join("prepared"); + let builder = shell_snapshot_v2_prewarm_builder(profile_home.path()) + .with_pre_build_hook(move |home| { + std::fs::write(home.join("hooks.json"), json!({ + "hooks": { "SessionStart": [{ "hooks": [{ + "type": "command", + "command": format!("printf 'profile_helper() {{ printf helper; }}\\n' > {}", shlex::try_quote(prepared.to_str().unwrap()).unwrap()), + }] }] } + }).to_string()).expect("write startup hook"); + }) + .with_config(trust_discovered_hooks); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + + // The model can finish a text-only turn while the profile is still blocked. + run_no_shell_turn(&harness).await?; + wait_for_file_contents(&profile_home.path().join("captures"), "x").await?; + fs::write(profile_home.path().join("ready"), "").await?; + + let end = run_tool_turn_on_harness( + &harness, + "use the prewarmed in-memory shell snapshot", + "shell-snapshot-v2-prewarm", + "exec_command", + json!({ + "cmd": "profile_helper; printf '|%s' \"${PROFILE_SECRET-missing}\"", + "yield_time_ms": 1_000, + }), + ) + .await?; + + assert_eq!(end.exit_code, 0); + assert!( + harness + .function_call_stdout("shell-snapshot-v2-prewarm") + .await + .trim() + .ends_with("helper|missing") + ); + assert_eq!( + fs::read_to_string(profile_home.path().join("captures")).await?, + "x" + ); + assert!(!harness.test().home.path().join("shell_snapshots").exists()); + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_v2_recovers_after_failed_prewarm() -> Result<()> { + skip_if_remote!(Ok(()), "prewarming is local-only"); + let profile_home = tempfile::tempdir()?; + let profile = profile_home.path().join(".bashrc"); + fs::write(&profile, "printf x > \"$HOME/failed\"\nexit 7\n").await?; + let bash_env = profile.to_string_lossy().into_owned(); + let builder = + shell_snapshot_v2_prewarm_builder(profile_home.path()).with_config(move |config| { + // Ordinary login-shell fallback must read the same profile, so the output + // distinguishes a recovered, filtered snapshot from an unfiltered fallback. + let policy = &mut config.permissions.shell_environment_policy; + policy.r#set.insert("BASH_ENV".to_string(), bash_env); + }); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + run_no_shell_turn(&harness).await?; + wait_for_file_contents(&profile_home.path().join("failed"), "x").await?; + + // Replace the file rather than modifying the failed capture's open script. + let repaired = profile_home.path().join("repaired"); + fs::write( + &repaired, + "printf x >> \"$HOME/captures\"\nprofile_helper() { printf recovered; }\nexport PROFILE_SECRET=secret\n", + ).await?; + fs::rename(repaired, profile).await?; + for call_id in ["after-failed-prewarm", "reuse-recovered-snapshot"] { + let end = run_tool_turn_on_harness( + &harness, + "use the recovered snapshot", + call_id, + "exec_command", + json!({"cmd": "profile_helper; printf '|%s' \"${PROFILE_SECRET-missing}\""}), + ) + .await?; + assert_eq!(end.exit_code, 0); + let output = harness.function_call_stdout(call_id).await; + assert!(output.trim().ends_with("recovered|missing")); + } + assert_eq!( + fs::read_to_string(profile_home.path().join("captures")).await?, + "x" + ); + harness.test().codex.shutdown_and_wait().await?; + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_v2_prewarm_preserves_the_sandbox() -> Result<()> { + skip_if_remote!(Ok(()), "profile fixture uses a host-local HOME directory"); + let profile_home = tempfile::tempdir()?; + let outside = tempfile::tempdir()?; + let forbidden = outside.path().join("profile-write"); + let writable_home = profile_home.path().to_path_buf().try_into()?; + let forbidden_path = forbidden.to_string_lossy().into_owned(); + fs::write( + profile_home.path().join(".bashrc"), + "(printf escaped > \"$OUTSIDE\") 2>/dev/null\nprintf done > \"$HOME/captures\"\n", + ) + .await?; + let builder = + shell_snapshot_v2_prewarm_builder(profile_home.path()).with_config(move |config| { + config + .permissions + .set_permission_profile(PermissionProfile::workspace_write_with( + std::slice::from_ref(&writable_home), + NetworkSandboxPolicy::Restricted, + /*exclude_tmpdir_env_var*/ true, + /*exclude_slash_tmp*/ true, + )) + .expect("set sandboxed permissions"); + config + .permissions + .shell_environment_policy + .r#set + .insert("OUTSIDE".to_string(), forbidden_path); + let rules = config.codex_home.join("rules"); + std::fs::create_dir_all(&rules).expect("create rules directory"); + std::fs::write( + rules.join("default.rules"), + "prefix_rule(pattern=[\"true\"], decision=\"allow\")\n", + ) + .expect("allow the former warm-up sentinel"); + }); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + run_no_shell_turn(&harness).await?; + wait_for_file_contents(&profile_home.path().join("captures"), "done").await?; + harness.test().codex.shutdown_and_wait().await?; + assert!( + !forbidden.exists(), + "profile startup must stay sandboxed even when true is allowed" + ); + Ok(()) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_v2_prewarm_stops_on_shutdown() -> Result<()> { + skip_if_remote!(Ok(()), "profile fixture uses a host-local HOME directory"); + let profile_home = tempfile::tempdir()?; + fs::write( + profile_home.path().join(".bashrc"), + "printf '%s' \"$$\" > \"$HOME/pid\"\nprintf x > \"$HOME/captures\"\nwhile :; do /bin/sleep 0.01; done\n", + ).await?; + let harness = TestCodexHarness::with_auto_env_builder(shell_snapshot_v2_prewarm_builder( + profile_home.path(), + )) + .await?; + run_no_shell_turn(&harness).await?; + wait_for_file_contents(&profile_home.path().join("captures"), "x").await?; + let pid = fs::read_to_string(profile_home.path().join("pid")).await?; + + let codex = &harness.test().codex; + codex.submit(Op::Shutdown {}).await?; + wait_for_event(codex, |event| matches!(event, EventMsg::ShutdownComplete)).await; + tokio::time::timeout(Duration::from_secs(2), async { + while tokio::process::Command::new("/bin/kill") + .args(["-0", &pid]) + .output() + .await? + .status + .success() + { + sleep(Duration::from_millis(25)).await; + } + anyhow::Ok(()) + }) + .await??; + Ok(()) +} + +#[cfg(unix)] +#[test_case::test_case(|config| { + config.active_project.trust_level = Some(TrustLevel::Untrusted); + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::UnlessTrusted); +}; "untrusted")] +#[test_case::test_case(|config| { + config.permissions.network = Some(NetworkProxySpec::from_config_and_constraints( + NetworkProxyConfig { enabled: true, allow_local_binding: true, ..Default::default() }, + /*requirements*/ None, + config.permissions.permission_profile(), + ).expect("configure managed network")); +}; "managed network")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_v2_prewarm_skips_ineligible_sessions( + configure: fn(&mut codex_core::config::Config), +) -> Result<()> { + skip_if_remote!(Ok(()), "profile fixture uses a host-local HOME directory"); + let profile_home = tempfile::tempdir()?; + fs::write( + profile_home.path().join(".bashrc"), + "printf x > \"$HOME/captures\"\n", + ) + .await?; + let builder = shell_snapshot_v2_prewarm_builder(profile_home.path()).with_config(configure); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + run_no_shell_turn(&harness).await?; + harness.test().codex.shutdown_and_wait().await?; + assert!(!profile_home.path().join("captures").exists()); + Ok(()) +} + +#[cfg(unix)] +#[test_case::test_case("SessionStart")] +#[test_case::test_case("UserPromptSubmit")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_v2_does_not_warm_a_hook_stopped_turn(hook_event: &str) -> Result<()> { + skip_if_remote!(Ok(()), "profile fixture uses a host-local HOME directory"); + let profile_home = tempfile::tempdir()?; + fs::write( + profile_home.path().join(".bashrc"), + "printf x > \"$HOME/captures\"\n", + ) + .await?; + let hook_event = hook_event.to_string(); + let builder = + shell_snapshot_v2_prewarm_builder(profile_home.path()) + .with_pre_build_hook(move |home| { + std::fs::write(home.join("hooks.json"), json!({ + "hooks": { (hook_event): [{ "hooks": [{ + "type": "command", + "command": "printf '%s' '{\"continue\":false,\"stopReason\":\"blocked\"}'", + }] }] } + }).to_string()).expect("write stopping hook"); + }) + .with_config(trust_discovered_hooks); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + harness.submit("do not start the model or shell").await?; + harness.test().codex.shutdown_and_wait().await?; + assert!(harness.request_bodies().await.is_empty()); + assert!(!profile_home.path().join("captures").exists()); + Ok(()) +} + +#[cfg(unix)] +#[test_case::test_case(PermissionProfile::workspace_write(); "read only reviewer")] +#[test_case::test_case(PermissionProfile::External { network: NetworkSandboxPolicy::Enabled }; "reviewer without shell tools")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shell_snapshot_v2_guardian_uses_its_resolved_permissions_and_tools( + parent_permissions: PermissionProfile, +) -> Result<()> { + skip_if_remote!(Ok(()), "profile fixture uses a host-local HOME directory"); + let profile_home = tempfile::tempdir()?; + fs::write( + profile_home.path().join(".bashrc"), + "printf capture > \"$HOME/$CODEX_THREAD_ID\"\n", + ) + .await?; + let builder = + shell_snapshot_v2_prewarm_builder(profile_home.path()).with_config(move |config| { + config + .permissions + .set_permission_profile(parent_permissions) + .expect("set parent permissions"); + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + let rules = config.codex_home.join("rules"); + std::fs::create_dir_all(&rules).expect("create rules directory"); + std::fs::write( + rules.join("default.rules"), + "prefix_rule(pattern=[\"printf\", \"parent\"], decision=\"prompt\")\n", + ) + .expect("require review even for an unrestricted owner environment"); + }); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + let test = harness.test(); + let responses = mount_sse_sequence( + harness.server(), + vec![ + sse(vec![ + ev_function_call( + "parent", + "exec_command", + &json!({ + "cmd": "printf parent", + "sandbox_permissions": "require_escalated", + "justification": "Exercise Guardian review", + }) + .to_string(), + ), + ev_completed("parent"), + ]), + sse(vec![ + ev_function_call("reviewer", "exec_command", r#"{"cmd":"printf reviewed"}"#), + ev_completed("reviewer"), + ]), + sse(vec![ + ev_assistant_message( + "assessment", + r#"{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"Harmless output."}"#, + ), + ev_completed("assessment"), + ]), + sse(vec![ev_completed("done")]), + ], + ) + .await; + let mut environments = local_selections(test.config.cwd.clone()); + environments.environments[0].config = EnvironmentConfigState::Ready(EnvironmentConfig { + allow_login_shell: true, + workspace_roots: environments.environments[0].workspace_roots.clone(), + permission_profile: PermissionProfileSnapshot::legacy(PermissionProfile::Disabled), + shell_environment_policy: test.config.permissions.shell_environment_policy.clone(), + windows_sandbox_level: WindowsSandboxLevel::from_config(&test.config), + windows_sandbox_private_desktop: test.config.permissions.windows_sandbox_private_desktop, + use_legacy_landlock: test.config.features.use_legacy_landlock(), + exec_policy: None, + mcp_policy: None, + network_policy: None, + selected_capability_roots: Vec::new(), + }); + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "run the reviewed command".to_string(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + environments: Some(environments), + ..Default::default() + }), + ) + .await?; + wait_for_event(&test.codex, |event| { + assert!( + !matches!(event, EventMsg::ExecApprovalRequest(_) | EventMsg::Error(_)), + "unexpected event: {event:?}" + ); + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.codex.shutdown_and_wait().await?; + let requests = responses.requests(); + let guardian_requests = requests + .iter() + .filter(|request| request.body_json()["client_metadata"]["x-openai-subagent"] == "guardian") + .collect::>(); + assert_eq!(guardian_requests.len(), 2); + let guardian_id = guardian_requests[0].body_json()["client_metadata"]["thread_id"] + .as_str() + .expect("Guardian thread ID") + .to_string(); + assert_ne!(guardian_id, test.session_configured.thread_id.to_string()); + assert!( + profile_home + .path() + .join(test.session_configured.thread_id.to_string()) + .exists() + ); + assert!( + !profile_home.path().join(guardian_id).exists(), + "Guardian profile must not inherit writable owner permissions" + ); + Ok(()) +} + #[cfg(unix)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn shell_snapshot_v2_filters_profile_secrets_without_creating_files() -> Result<()> { @@ -567,10 +1027,7 @@ async fn unified_exec_snapshot_still_intercepts_apply_patch() -> Result<()> { patch_end.stdout, patch_end.stderr, ); - assert_eq!( - wait_for_file_contents(&target).await?, - "hello from snapshot\n" - ); + wait_for_file_contents(&target, "hello from snapshot\n").await?; Ok(()) } diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index b03cc55d19..f45e7d472a 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -71,6 +71,8 @@ use crate::rpc::internal_error; use crate::rpc::invalid_params; use crate::rpc::invalid_request; use crate::rpc_server_requests::RpcServerRequestSender; +#[cfg(unix)] +use crate::shell_snapshot::CapturePurpose; use crate::telemetry::ExecServerTelemetry; use crate::telemetry::ProcessMetricGuard; @@ -355,7 +357,12 @@ impl LocalProcess { #[cfg(unix)] self.inner .shell_snapshots - .prepare(¶ms, &mut prepared, &self.inner.telemetry) + .prepare( + ¶ms, + &mut prepared, + &self.inner.telemetry, + CapturePurpose::Execution, + ) .await?; if prepared.command.is_empty() { return Err(invalid_params("argv must not be empty".to_string())); @@ -753,6 +760,39 @@ impl ExecBackend for LocalProcess { fn start(&self, params: ExecParams) -> ExecBackendFuture<'_> { Box::pin(LocalProcess::start(self, params)) } + + #[cfg(unix)] + fn prewarm_shell_snapshot(&self, params: ExecParams) -> ExecProcessFuture<'_, ()> { + Box::pin(async move { + if params.enforce_managed_network + || params.managed_network.is_some() + || params.network_proxy.is_some() + { + return Err(ExecServerError::Protocol( + "shell snapshot prewarming does not support managed networking".to_string(), + )); + } + let mut prepared = prepare_exec_request( + ¶ms, + child_env(¶ms), + self.runtime_paths.as_ref(), + /*network_policy_decider*/ None, + /*network_policy_audit_observer*/ None, + ) + .await + .map_err(map_handler_error)?; + self.inner + .shell_snapshots + .prepare( + ¶ms, + &mut prepared, + &self.inner.telemetry, + CapturePurpose::Prewarm, + ) + .await + .map_err(map_handler_error) + }) + } } impl LocalExecProcess { diff --git a/codex-rs/exec-server/src/process.rs b/codex-rs/exec-server/src/process.rs index bf3cc752f0..726c8a6392 100644 --- a/codex-rs/exec-server/src/process.rs +++ b/codex-rs/exec-server/src/process.rs @@ -222,6 +222,17 @@ pub type ExecProcessFuture<'a, T> = pub trait ExecBackend: Send + Sync { fn start(&self, params: ExecParams) -> ExecBackendFuture<'_>; + /// Captures a local shell snapshot without starting the requested command. + /// Failures must remain retryable by real commands. Remote backends do not + /// support this operation; callers should leave them on the lazy path. + fn prewarm_shell_snapshot(&self, _params: ExecParams) -> ExecProcessFuture<'_, ()> { + Box::pin(async { + Err(ExecServerError::Protocol( + "exec backend does not support shell snapshot prewarming".to_string(), + )) + }) + } + /// Starts a process with an authoritative controller-side policy decider. fn start_with_network_policy_decider( &self, diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs index 324580a66f..fe3e2a2a16 100644 --- a/codex-rs/exec-server/src/shell_snapshot.rs +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -56,12 +56,19 @@ struct ShellSnapshot { environment: HashMap, } +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) enum CapturePurpose { + Execution, + Prewarm, +} + impl ShellSnapshotCache { pub(crate) async fn prepare( &self, params: &ExecParams, prepared: &mut PreparedExecRequest, telemetry: &ExecServerTelemetry, + purpose: CapturePurpose, ) -> Result<(), JSONRPCErrorError> { let Some(request) = params.shell_snapshot.as_ref() else { return Ok(()); @@ -103,7 +110,8 @@ impl ShellSnapshotCache { let mut entry = entries.remove(position)?; // Share each failed attempt during backoff. After the retry // budget is exhausted, keep falling back until eviction. - if entry.attempts < MAX_SNAPSHOT_ATTEMPTS + if purpose == CapturePurpose::Execution + && entry.attempts < MAX_SNAPSHOT_ATTEMPTS && let Some(Err(retry_at)) = entry.snapshot.get() && Instant::now() >= *retry_at { @@ -134,23 +142,40 @@ impl ShellSnapshotCache { snapshot } }; - let Ok(snapshot) = snapshot - .get_or_init(|| async { - let started_at = std::time::Instant::now(); - let result = capture_snapshot(params, prepared, shell_type).await; - telemetry.shell_snapshot_captured( - started_at.elapsed(), - result.as_ref().map(|_| ()).map_err(|_| "capture_failed"), - ); - result.map_err(|err| { - tracing::warn!("failed to capture shell snapshot: {err:?}"); - Instant::now() + SNAPSHOT_RETRY_BACKOFF - }) - }) - .await - else { + let capture = async { + let started_at = std::time::Instant::now(); + let result = capture_snapshot(params, prepared, shell_type).await; + telemetry.shell_snapshot_captured( + started_at.elapsed(), + result.as_ref().map(|_| ()).map_err(|_| "capture_failed"), + ); + result + }; + let snapshot = match purpose { + CapturePurpose::Execution => { + snapshot + .get_or_init(|| async { + capture.await.map_err(|err| { + tracing::warn!("failed to capture shell snapshot: {err:?}"); + Instant::now() + SNAPSHOT_RETRY_BACKOFF + }) + }) + .await + } + CapturePurpose::Prewarm => { + // Leave the cell uninitialized on failure: a waiting real command + // can capture immediately, without spending its retry budget. + snapshot + .get_or_try_init(|| async { capture.await.map(Ok) }) + .await? + } + }; + let Ok(snapshot) = snapshot else { return Ok(()); }; + if purpose == CapturePurpose::Prewarm { + return Ok(()); + } let request_overrides = params .env diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs index f8a380a3f1..d0dcf1f88e 100644 --- a/codex-rs/exec-server/src/shell_snapshot_tests.rs +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -10,6 +10,7 @@ use opentelemetry_sdk::metrics::data::MetricData; use pretty_assertions::assert_eq; use test_case::test_case; +use super::CapturePurpose; use super::MAX_SNAPSHOT_ATTEMPTS; use super::SNAPSHOT_RETRY_BACKOFF; use super::ShellSnapshotCache; @@ -22,14 +23,19 @@ use crate::protocol::ShellInfo; use crate::protocol::ShellSnapshotRequest; use crate::telemetry::ExecServerTelemetry; -#[test_case(1; "succeeds_on_first_attempt")] -#[test_case(2; "recovers_on_second_attempt")] -#[test_case(3; "recovers_on_last_attempt")] -#[test_case(4; "stops_after_three_failures")] +#[test_case(1, CapturePurpose::Execution; "succeeds_on_first_attempt")] +#[test_case(2, CapturePurpose::Execution; "recovers_on_second_attempt")] +#[test_case(3, CapturePurpose::Execution; "recovers_on_last_attempt")] +#[test_case(4, CapturePurpose::Execution; "stops_after_three_failures")] +#[test_case(2, CapturePurpose::Prewarm; "failed_prewarm_releases_waiting_command")] +#[test_case(3, CapturePurpose::Prewarm; "failed_prewarm_preserves_last_attempt")] +#[test_case(4, CapturePurpose::Prewarm; "failed_prewarm_preserves_retry_limit")] #[tokio::test] async fn snapshot_failure_retries_are_bounded_and_single_flight( recovery_attempt: usize, + initial_purpose: CapturePurpose, ) -> anyhow::Result<()> { + let prewarm_fails_first = initial_purpose == CapturePurpose::Prewarm; let home = tempfile::TempDir::new()?; let profile = home.path().join(".bashrc"); std::fs::write(&profile, "printf x >> \"$HOME/captures\"\nexit 7\n")?; @@ -101,12 +107,23 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( ) .await .expect("prepare concurrent capture"); + let prewarming = prewarm_fails_first && attempt == 1; + let purpose = if prewarming { + CapturePurpose::Prewarm + } else { + CapturePurpose::Execution + }; let (first, second) = tokio::join!( - cache.prepare(¶ms, &mut prepared, &telemetry), - cache.prepare(¶ms, &mut concurrent, &telemetry), + biased; + cache.prepare(¶ms, &mut prepared, &telemetry, purpose), + cache.prepare(¶ms, &mut concurrent, &telemetry, CapturePurpose::Execution), ); - first.expect("capture failure must preserve command fallback"); - second.expect("concurrent request must share the capture attempt"); + if prewarming { + first.expect_err("prewarm must report failure without caching it"); + } else { + first.expect("capture failure must preserve command fallback"); + } + second.expect("waiting command must complete even when prewarm fails"); assert_eq!( (&prepared.command, &prepared.env), (&concurrent.command, &concurrent.env) @@ -115,7 +132,12 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( tokio::time::pause(); if attempt < recovery_attempt || recovery_attempt > MAX_SNAPSHOT_ATTEMPTS { cache - .prepare(¶ms, &mut prepared, &telemetry) + .prepare( + ¶ms, + &mut prepared, + &telemetry, + CapturePurpose::Execution, + ) .await .expect("capture must stay cached during backoff"); assert_eq!( @@ -127,7 +149,10 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( } assert_eq!( std::fs::read_to_string(home.path().join("captures"))?, - "x".repeat(attempt.min(recovery_attempt).min(MAX_SNAPSHOT_ATTEMPTS)) + "x".repeat( + attempt.min(recovery_attempt).min(MAX_SNAPSHOT_ATTEMPTS) + + usize::from(prewarm_fails_first) + ) ); tokio::time::advance(SNAPSHOT_RETRY_BACKOFF).await; tokio::time::resume(); @@ -170,7 +195,8 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( } let mut expected_counters = BTreeMap::new(); let mut expected_durations = BTreeMap::new(); - let failures = (recovery_attempt - 1).min(MAX_SNAPSHOT_ATTEMPTS) as u64; + let failures = + (recovery_attempt - 1).min(MAX_SNAPSHOT_ATTEMPTS) as u64 + u64::from(prewarm_fails_first); for (success, count) in [ ("false", failures), ("true", u64::from(recovery_attempt <= MAX_SNAPSHOT_ATTEMPTS)),