diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 5b4aead69d..d9bcb9335e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3369,6 +3369,7 @@ dependencies = [ "codex-config", "codex-context-fragments", "codex-exec-server-protocol", + "codex-file-system", "codex-history", "codex-mcp", "codex-protocol", diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index a90db60abe..bdf6877389 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -289,7 +289,6 @@ impl ExecCommandHandler { ))); } } - let process_id = manager.allocate_process_id().await; let resolved_command = get_command( &args, shell, @@ -347,7 +346,6 @@ impl ExecCommandHandler { && !effective_additional_permissions.permissions_preapproved && prompt_is_rejected_by_policy(approval_policy, /*prompt_is_rule*/ false).is_some() { - manager.release_process_id(process_id).await; return Err(FunctionCallError::RespondToModel(format!( "approval policy is {approval_policy:?}; reject command — you cannot ask for escalated permissions if the approval policy is {approval_policy:?}" ))); @@ -373,7 +371,6 @@ impl ExecCommandHandler { ) { Ok(normalized) => normalized, Err(err) => { - manager.release_process_id(process_id).await; return Err(FunctionCallError::RespondToModel(err)); } }; @@ -391,13 +388,7 @@ impl ExecCommandHandler { "exec_command", ) .await; - // Keep the reservation when interception returns `Ok(None)`: the normal command below - // still needs this process ID. - if intercepted_patch.is_err() { - manager.release_process_id(process_id).await; - } if let Some(output) = intercepted_patch? { - manager.release_process_id(process_id).await; return Ok(boxed_tool_output(ExecCommandToolOutput { event_call_id: String::new(), chunk_id: String::new(), @@ -414,6 +405,17 @@ impl ExecCommandHandler { } emit_unified_exec_tty_metric(&step_context.session_telemetry, tty); + crate::tools::lifecycle::notify_command_start( + context.session.as_ref(), + context.step_context.turn.as_ref(), + &context.call_id, + &command, + &cwd, + fs.as_ref(), + ) + .await; + // Preparation can be cancelled. Reserve a process only once it is done. + let process_id = manager.allocate_process_id().await; let request = ExecCommandRequest { command, shell_type, diff --git a/codex-rs/core/src/tools/lifecycle.rs b/codex-rs/core/src/tools/lifecycle.rs index 55288f7559..82ba4a6c6b 100644 --- a/codex-rs/core/src/tools/lifecycle.rs +++ b/codex-rs/core/src/tools/lifecycle.rs @@ -1,13 +1,16 @@ use std::sync::Arc; +use codex_extension_api::CommandStartInput; use codex_extension_api::McpToolContext; use codex_extension_api::McpToolResultInput; use codex_extension_api::ToolCallOutcome; use codex_extension_api::ToolCallSource as ExtensionToolCallSource; use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolStartInput; +use codex_file_system::ExecutorFileSystem; use codex_protocol::mcp::CallToolResult; use codex_tools::ToolName; +use codex_utils_path_uri::PathUri; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -54,6 +57,30 @@ pub(crate) async fn notify_tool_start( } } +pub(crate) async fn notify_command_start( + session: &Session, + turn: &TurnContext, + call_id: &str, + command: &[String], + cwd: &PathUri, + file_system: &dyn ExecutorFileSystem, +) { + for contributor in session.services.extensions.tool_lifecycle_contributors() { + contributor + .on_command_start(CommandStartInput { + session_store: &session.services.session_extension_data, + thread_store: &session.services.thread_extension_data, + turn_store: turn.extension_data.as_ref(), + turn_id: &turn.sub_id, + call_id, + command, + cwd, + file_system, + }) + .await; + } +} + pub(crate) async fn process_mcp_tool_result( session: &Session, turn: &TurnContext, diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 8f6c61fdee..3645306c27 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -160,7 +160,9 @@ pub(crate) struct ProcessStore { impl ProcessStore { fn remove(&mut self, process_id: i32) -> Option { - self.reserved_process_ids.remove(&process_id); + if !process_manager::should_use_deterministic_process_ids() { + self.reserved_process_ids.remove(&process_id); + } self.processes.remove(&process_id) } } diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index ef18be3447..f2e9afddab 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -121,7 +121,7 @@ fn deterministic_process_ids_forced_for_tests() -> bool { FORCE_DETERMINISTIC_PROCESS_IDS.load(Ordering::Relaxed) } -fn should_use_deterministic_process_ids() -> bool { +pub(super) fn should_use_deterministic_process_ids() -> bool { cfg!(test) || deterministic_process_ids_forced_for_tests() } diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index bb928ad8f9..e72cf039ea 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -47,6 +47,16 @@ async fn deterministic_process_ids_are_not_reused_after_release() { assert_eq!((first, second), (1000, 1001)); } +#[tokio::test] +async fn deterministic_process_ids_are_not_reused_after_removal() { + let manager = UnifiedExecProcessManager::default(); + let first = manager.allocate_process_id().await; + let _removed = manager.process_store.lock().await.remove(first); + let second = manager.allocate_process_id().await; + + assert_eq!((first, second), (1000, 1001)); +} + #[test] fn env_overlay_for_exec_server_keeps_runtime_changes_only() { let local_policy_env = HashMap::from([ diff --git a/codex-rs/core/tests/suite/command_lifecycle_tests.rs b/codex-rs/core/tests/suite/command_lifecycle_tests.rs new file mode 100644 index 0000000000..3d0545df4c --- /dev/null +++ b/codex-rs/core/tests/suite/command_lifecycle_tests.rs @@ -0,0 +1,377 @@ +//! Command lifecycle callbacks observe the command and filesystem selected for execution. + +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use anyhow::Result; +use codex_config::Constrained; +use codex_core::TurnInputRequest; +use codex_core::config::Config; +use codex_extension_api::CommandStartInput; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ToolLifecycleContributor; +use codex_extension_api::ToolLifecycleFuture; +use codex_features::Feature; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use codex_utils_path_uri::PathUri; +use core_test_support::TestTargetOs; +use core_test_support::hooks::trust_discovered_hooks; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::skip_if_wine_exec; +use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::TestCodexHarness; +use core_test_support::test_codex::test_codex; +use core_test_support::test_target_os; +use core_test_support::wait_for_event; +use core_test_support::wait_for_event_with_timeout; +use pretty_assertions::assert_eq; +use serde_json::json; +use test_case::test_case; +use tokio::sync::Notify; +use tokio::sync::oneshot; +use tokio::time::timeout; + +#[derive(Debug, PartialEq)] +struct RecordedCommand { + call_id: String, + command: Vec, + cwd: PathUri, + fixture: String, +} + +#[derive(Default)] +struct CommandRecorder { + commands: Mutex>, +} + +impl ToolLifecycleContributor for CommandRecorder { + fn on_command_start<'a>(&'a self, input: CommandStartInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(async move { + let fixture_path = input + .cwd + .join("executor-fixture.txt") + .expect("fixture path"); + let fixture = input + .file_system + .read_file_text(&fixture_path, Default::default(), /*sandbox*/ None) + .await + .expect("callback must use the selected executor's filesystem"); + let prepared_path = input.cwd.join("prepared.txt").expect("prepared path"); + input + .file_system + .write_file( + &prepared_path, + fixture.as_bytes().to_vec(), + Default::default(), + /*sandbox*/ None, + ) + .await + .expect("prepare the command's executor before it starts"); + self.commands + .lock() + .expect("command records") + .push(RecordedCommand { + call_id: input.call_id.to_owned(), + command: input.command.to_vec(), + cwd: input.cwd.clone(), + fixture, + }); + }) + } +} + +#[derive(Clone, Copy)] +enum CommandCallMode { + Direct, + CodeMode, +} + +impl CommandCallMode { + fn call(self, call_id: &str, arguments: serde_json::Value) -> serde_json::Value { + match self { + Self::Direct => { + responses::ev_function_call(call_id, "exec_command", &arguments.to_string()) + } + Self::CodeMode => responses::ev_custom_tool_call( + call_id, + "exec", + &format!("text(await tools.exec_command({arguments}));"), + ), + } + } +} + +fn configure_command_test_permissions(config: &mut Config) { + // Match TestCodex::submit_turn: these fixtures exercise execution callbacks + // without waiting for interactive approval on hosts lacking a sandbox backend. + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::Never); + config + .permissions + .set_permission_profile(PermissionProfile::Disabled) + .expect("set command test permissions"); +} + +async fn start_command_turn(test: &TestCodex) -> Result<()> { + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Run the command.".into(), + text_elements: Vec::new(), + }])) + .await?; + Ok(()) +} + +#[test_case(CommandCallMode::Direct; "direct")] +#[test_case(CommandCallMode::CodeMode; "code_mode")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn command_start_receives_rewritten_command_and_executor_workdir( + mode: CommandCallMode, +) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!(Ok(()), "command hooks require a host-native executor"); + + let rewritten_command = match test_target_os() { + TestTargetOs::Linux | TestTargetOs::MacOs => "cat prepared.txt", + TestTargetOs::Windows => "Get-Content -Raw prepared.txt", + }; + let hook_output = json!({ + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": { "command": rewritten_command }, + } + }); + let recorder = Arc::new(CommandRecorder::default()); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.tool_lifecycle_contributor(recorder.clone()); + let builder = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(configure_command_test_permissions) + .with_pre_build_hook(move |home| { + super::write_pre_tool_hook(home, "^Bash$", &hook_output) + .expect("write command-rewriting hook"); + }) + .with_config(trust_discovered_hooks) + .with_config(move |config| { + if matches!(mode, CommandCallMode::CodeMode) { + let _ = config.features.enable(Feature::CodeMode); + } + }); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + let (test, server) = (harness.test(), harness.server()); + + let cwd = test.workspace_path_uri("command-workdir")?; + let fixture = "selected-executor-command-marker"; + harness + .write_file("command-workdir/executor-fixture.txt", fixture) + .await?; + + let call_id = "command-lifecycle-call"; + let arguments = json!({ + "cmd": "echo original-command-must-not-run", + "workdir": "command-workdir", + "yield_time_ms": 5_000, + }); + let call = mode.call(call_id, arguments); + responses::mount_sse_once( + server, + responses::sse(vec![call, responses::ev_completed("first-response")]), + ) + .await; + let follow_up = responses::mount_sse_once( + server, + responses::sse(vec![responses::ev_completed("second-response")]), + ) + .await; + start_command_turn(test).await?; + + let mut begin = None; + wait_for_event_with_timeout( + &test.codex, + |event| { + assert!( + !matches!(event, EventMsg::ExecApprovalRequest(_)), + "command fixture unexpectedly requested approval: {event:?}" + ); + if let EventMsg::ExecCommandBegin(command) = event { + begin = Some(command.clone()); + } + matches!(event, EventMsg::TurnComplete(_)) + }, + Duration::from_secs(30), + ) + .await; + let request = follow_up.single_request(); + let output = match mode { + CommandCallMode::Direct => request.function_call_output(call_id), + CommandCallMode::CodeMode => request.custom_tool_call_output(call_id), + }; + let begin = begin.unwrap_or_else(|| panic!("command did not start: {output}")); + + assert_eq!(begin.cwd, cwd); + assert!( + begin + .command + .last() + .is_some_and(|argument| argument.contains(rewritten_command)), + "the executor must receive the post-hook shell command" + ); + assert_eq!( + *recorder.commands.lock().expect("command records"), + vec![RecordedCommand { + call_id: begin.call_id, + command: begin.command, + cwd, + fixture: fixture.to_owned(), + }] + ); + assert!( + output["output"].to_string().contains(fixture), + "execution must wait for the callback to prepare its selected filesystem" + ); + Ok(()) +} + +struct BlockingCommandContributor { + entered: Mutex>>, + completed: Mutex>>, + release: Notify, +} + +impl ToolLifecycleContributor for BlockingCommandContributor { + fn on_command_start<'a>(&'a self, _input: CommandStartInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(async move { + let Some(entered) = self.entered.lock().expect("entry sender").take() else { + return; + }; + // Closing this channel tells the test that cancellation dropped the callback. + let _completion = self.completed.lock().expect("completion sender").take(); + entered.send(()).expect("test awaits command preparation"); + self.release.notified().await; + }) + } +} + +#[test_case(CommandCallMode::Direct; "direct")] +#[test_case(CommandCallMode::CodeMode; "code_mode")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn interrupting_command_preparation_does_not_start_the_command( + mode: CommandCallMode, +) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!( + Ok(()), + "basic PowerShell execution through Wine is not passing yet" + ); + + let (entered_tx, entered_rx) = oneshot::channel(); + let (completed_tx, completed_rx) = oneshot::channel(); + let contributor = Arc::new(BlockingCommandContributor { + entered: Mutex::new(Some(entered_tx)), + completed: Mutex::new(Some(completed_tx)), + release: Notify::new(), + }); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.tool_lifecycle_contributor(contributor.clone()); + let builder = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(configure_command_test_permissions) + .with_config(move |config| { + if matches!(mode, CommandCallMode::CodeMode) { + let _ = config.features.enable(Feature::CodeMode); + let _ = config.features.enable(Feature::CodeModeInterrupt); + } + }); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + let (test, server) = (harness.test(), harness.server()); + let command = match test_target_os() { + TestTargetOs::Linux | TestTargetOs::MacOs => "printf executed > must-not-run.txt", + TestTargetOs::Windows => "Set-Content must-not-run.txt executed", + }; + responses::mount_sse_once( + server, + responses::sse(vec![ + mode.call("interrupted-command", json!({ "cmd": command })), + responses::ev_completed("first-response"), + ]), + ) + .await; + start_command_turn(test).await?; + timeout(Duration::from_secs(30), entered_rx).await??; + test.codex.submit(Op::Interrupt).await?; + wait_for_event(&test.codex, |event| { + assert!( + !matches!(event, EventMsg::ExecApprovalRequest(_)), + "command fixture unexpectedly requested approval: {event:?}" + ); + matches!(event, EventMsg::TurnAborted(_)) + }) + .await; + timeout(Duration::from_secs(30), completed_rx) + .await? + .expect_err("cancellation must drop the callback's completion sender"); + contributor.release.notify_one(); + + assert!( + !harness.path_exists("must-not-run.txt").await?, + "the command must not start after its preparation is cancelled" + ); + + let call_id = "command-after-interruption"; + responses::mount_sse_once( + server, + responses::sse(vec![ + mode.call(call_id, json!({ "cmd": "echo command-after-interruption" })), + responses::ev_completed("second-response"), + ]), + ) + .await; + let follow_up = responses::mount_sse_once( + server, + responses::sse(vec![responses::ev_completed("third-response")]), + ) + .await; + start_command_turn(test).await?; + let mut process_id = None; + let mut exit_code = None; + wait_for_event_with_timeout( + &test.codex, + |event| { + assert!( + !matches!(event, EventMsg::ExecApprovalRequest(_)), + "command fixture unexpectedly requested approval: {event:?}" + ); + match event { + EventMsg::ExecCommandBegin(begin) => { + process_id = begin.process_id.clone(); + } + EventMsg::ExecCommandEnd(end) => { + exit_code = Some(end.exit_code); + } + _ => {} + } + matches!(event, EventMsg::TurnComplete(_)) + }, + Duration::from_secs(30), + ) + .await; + assert_eq!(process_id.as_deref(), Some("1000")); + assert_eq!(exit_code, Some(0)); + let request = follow_up.single_request(); + let output = match mode { + CommandCallMode::Direct => request.function_call_output(call_id), + CommandCallMode::CodeMode => request.custom_tool_call_output(call_id), + }; + assert!( + output["output"] + .to_string() + .contains("command-after-interruption") + ); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/tool_lifecycle.rs b/codex-rs/core/tests/suite/tool_lifecycle.rs index 4e093fa70e..911d9d00f7 100644 --- a/codex-rs/core/tests/suite/tool_lifecycle.rs +++ b/codex-rs/core/tests/suite/tool_lifecycle.rs @@ -51,6 +51,9 @@ use wiremock::Request; use wiremock::ResponseTemplate; use wiremock::matchers::body_partial_json; +#[path = "command_lifecycle_tests.rs"] +mod command_lifecycle; + struct RecordedHistory { call_id: String, arguments: String, diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml index a0d4fe1f4c..7769b788a5 100644 --- a/codex-rs/ext/extension-api/Cargo.toml +++ b/codex-rs/ext/extension-api/Cargo.toml @@ -18,6 +18,7 @@ codex-history = { workspace = true } codex-config = { workspace = true } codex-context-fragments = { workspace = true } codex-exec-server-protocol = { workspace = true } +codex-file-system = { workspace = true } codex-mcp = { workspace = true } codex-protocol = { workspace = true } codex-tools = { workspace = true } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index ad102973ba..58a9909ce1 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -42,6 +42,7 @@ pub use thread_lifecycle::ThreadReadyInput; pub use thread_lifecycle::ThreadResumeInput; pub use thread_lifecycle::ThreadStartInput; pub use thread_lifecycle::ThreadStopInput; +pub use tool_lifecycle::CommandStartInput; pub use tool_lifecycle::McpToolContext; pub use tool_lifecycle::McpToolResultInput; pub use tool_lifecycle::McpToolSource; @@ -342,6 +343,14 @@ pub trait ToolLifecycleContributor: Send + Sync { Box::pin(std::future::ready(())) } + /// Called for a resolved builtin command before attribution and execution. + /// + /// This includes commands issued by code mode. It does not imply that + /// subsequent approval or process creation will succeed. + fn on_command_start<'a>(&'a self, _input: CommandStartInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(std::future::ready(())) + } + /// Runs before the MCP result is sent to the client and model. fn on_mcp_tool_result<'a>(&'a self, _input: McpToolResultInput<'a>) -> ToolLifecycleFuture<'a> { Box::pin(std::future::ready(())) diff --git a/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs index eb8c8154d4..79e20c8ab5 100644 --- a/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs @@ -3,6 +3,7 @@ use std::pin::Pin; use std::sync::Arc; use codex_config::McpServerConfig; +use codex_file_system::ExecutorFileSystem; use codex_mcp::McpServerSource; use codex_mcp::PreparedMcpCall; use codex_mcp::ResolvedMcpServer; @@ -144,6 +145,31 @@ pub struct ToolStartInput<'a> { pub source: ToolCallSource, } +/// Finalized builtin command and the executor that will run it. +/// +/// This runs after command hooks and environment resolution, before command +/// attribution and measurement selection. Filesystem paths belong to the selected +/// executor and must not be interpreted as host-local paths. +pub struct CommandStartInput<'a> { + /// Store scoped to the host session runtime. + pub session_store: &'a ExtensionData, + /// Store scoped to this thread runtime. + pub thread_store: &'a ExtensionData, + /// Store scoped to this turn runtime. + pub turn_store: &'a ExtensionData, + /// Current turn submission id. + pub turn_id: &'a str, + /// Host tool call id. + pub call_id: &'a str, + /// Resolved argv, including the selected executor's shell wrapper. + /// Commands can contain sensitive plaintext and must not be logged. + pub command: &'a [String], + /// Effective working directory in the selected executor. + pub cwd: &'a PathUri, + /// Filesystem of the selected executor. + pub file_system: &'a dyn ExecutorFileSystem, +} + /// Input supplied after an MCP server responds, before the host reports completion. pub struct McpToolResultInput<'a> { /// Store scoped to the host session runtime. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 72c56f26f5..ea2d0eb046 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -46,6 +46,7 @@ pub use codex_tools::parse_tool_input_schema_without_compaction; pub use contributors::ApprovalDecision; pub use contributors::ApprovalDecisionInput; pub use contributors::ApprovalReviewContributor; +pub use contributors::CommandStartInput; pub use contributors::ConfigContributor; pub use contributors::ContextContributor; pub use contributors::ExtensionFuture;