diff --git a/codex-rs/apply-patch/src/invocation.rs b/codex-rs/apply-patch/src/invocation.rs index abe223f1d9..1aa429e0d3 100644 --- a/codex-rs/apply-patch/src/invocation.rs +++ b/codex-rs/apply-patch/src/invocation.rs @@ -151,78 +151,87 @@ pub async fn maybe_parse_apply_patch_verified( } match maybe_parse_apply_patch(argv) { - MaybeApplyPatch::Body(ApplyPatchArgs { - patch, - hunks, - workdir, - }) => { - let effective_cwd = workdir - .as_ref() - .map(|dir| cwd.join(Path::new(dir))) - .unwrap_or_else(|| cwd.clone()); - let mut changes = HashMap::new(); - for hunk in hunks { - let path = hunk.resolve_path(&effective_cwd); - match hunk { - Hunk::AddFile { contents, .. } => { - changes.insert( - path.into_path_buf(), - ApplyPatchFileChange::Add { content: contents }, - ); - } - Hunk::DeleteFile { .. } => { - let content = match fs.read_file_text(&path, sandbox).await { - Ok(content) => content, - Err(e) => { - return MaybeApplyPatchVerified::CorrectnessError( - ApplyPatchError::IoError(IoError { - context: format!("Failed to read {}", path.display()), - source: e, - }), - ); - } - }; - changes.insert( - path.into_path_buf(), - ApplyPatchFileChange::Delete { content }, - ); - } - Hunk::UpdateFile { - move_path, chunks, .. - } => { - let ApplyPatchFileUpdate { - unified_diff, - content: contents, - .. - } = match unified_diff_from_chunks(&path, &chunks, fs, sandbox).await { - Ok(diff) => diff, - Err(e) => { - return MaybeApplyPatchVerified::CorrectnessError(e); - } - }; - changes.insert( - path.into_path_buf(), - ApplyPatchFileChange::Update { - unified_diff, - move_path: move_path.map(|p| effective_cwd.join(p).into_path_buf()), - new_content: contents, - }, - ); - } - } - } - MaybeApplyPatchVerified::Body(ApplyPatchAction { - changes, - patch, - cwd: effective_cwd, - }) - } + MaybeApplyPatch::Body(args) => verify_apply_patch_args(args, cwd, fs, sandbox).await, MaybeApplyPatch::ShellParseError(e) => MaybeApplyPatchVerified::ShellParseError(e), MaybeApplyPatch::PatchParseError(e) => MaybeApplyPatchVerified::CorrectnessError(e.into()), MaybeApplyPatch::NotApplyPatch => MaybeApplyPatchVerified::NotApplyPatch, } } +pub async fn verify_apply_patch_args( + args: ApplyPatchArgs, + cwd: &AbsolutePathBuf, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&codex_exec_server::FileSystemSandboxContext>, +) -> MaybeApplyPatchVerified { + let ApplyPatchArgs { + patch, + hunks, + workdir, + .. + } = args; + let effective_cwd = workdir + .as_ref() + .map(|dir| cwd.join(Path::new(dir))) + .unwrap_or_else(|| cwd.clone()); + let mut changes = HashMap::new(); + for hunk in hunks { + let path = hunk.resolve_path(&effective_cwd); + match hunk { + Hunk::AddFile { contents, .. } => { + changes.insert( + path.into_path_buf(), + ApplyPatchFileChange::Add { content: contents }, + ); + } + Hunk::DeleteFile { .. } => { + let content = match fs.read_file_text(&path, sandbox).await { + Ok(content) => content, + Err(e) => { + return MaybeApplyPatchVerified::CorrectnessError( + ApplyPatchError::IoError(IoError { + context: format!("Failed to read {}", path.display()), + source: e, + }), + ); + } + }; + changes.insert( + path.into_path_buf(), + ApplyPatchFileChange::Delete { content }, + ); + } + Hunk::UpdateFile { + move_path, chunks, .. + } => { + let ApplyPatchFileUpdate { + unified_diff, + content: contents, + .. + } = match unified_diff_from_chunks(&path, &chunks, fs, sandbox).await { + Ok(diff) => diff, + Err(e) => { + return MaybeApplyPatchVerified::CorrectnessError(e); + } + }; + changes.insert( + path.into_path_buf(), + ApplyPatchFileChange::Update { + unified_diff, + move_path: move_path.map(|p| effective_cwd.join(p).into_path_buf()), + new_content: contents, + }, + ); + } + } + } + MaybeApplyPatchVerified::Body(ApplyPatchAction { + changes, + patch, + cwd: effective_cwd, + }) +} + /// Extract the heredoc body (and optional `cd` workdir) from a `bash -lc` script /// that invokes the apply_patch tool using a heredoc. /// diff --git a/codex-rs/apply-patch/src/lib.rs b/codex-rs/apply-patch/src/lib.rs index 99a63e3ace..aa2326ab28 100644 --- a/codex-rs/apply-patch/src/lib.rs +++ b/codex-rs/apply-patch/src/lib.rs @@ -26,6 +26,7 @@ pub use streaming_parser::StreamingPatchParser; use thiserror::Error; pub use invocation::maybe_parse_apply_patch_verified; +pub use invocation::verify_apply_patch_args; pub use standalone_executable::main; use crate::invocation::ExtractHeredocError; @@ -97,6 +98,7 @@ pub struct ApplyPatchArgs { pub patch: String, pub hunks: Vec, pub workdir: Option, + pub environment_id: Option, } #[derive(Debug, PartialEq)] diff --git a/codex-rs/apply-patch/src/parser.rs b/codex-rs/apply-patch/src/parser.rs index 5403055272..b3b2337c39 100644 --- a/codex-rs/apply-patch/src/parser.rs +++ b/codex-rs/apply-patch/src/parser.rs @@ -3,8 +3,9 @@ //! //! The official Lark grammar for the apply-patch format is: //! -//! start: begin_patch hunk+ end_patch +//! start: begin_patch environment_id? hunk+ end_patch //! begin_patch: "*** Begin Patch" LF +//! environment_id: "*** Environment ID: " filename LF //! end_patch: "*** End Patch" LF? //! //! hunk: add_hunk | delete_hunk | update_hunk @@ -32,6 +33,7 @@ use std::path::PathBuf; use thiserror::Error; pub(crate) const BEGIN_PATCH_MARKER: &str = "*** Begin Patch"; +pub(crate) const ENVIRONMENT_ID_MARKER: &str = "*** Environment ID: "; pub(crate) const END_PATCH_MARKER: &str = "*** End Patch"; pub(crate) const ADD_FILE_MARKER: &str = "*** Add File: "; pub(crate) const DELETE_FILE_MARKER: &str = "*** Delete File: "; @@ -178,9 +180,9 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result check_patch_boundaries_lenient(&lines)?, }; + let (environment_id, mut remaining_lines, mut line_number) = + parse_environment_id_preamble(hunk_lines)?; let mut hunks: Vec = Vec::new(); - let mut remaining_lines = hunk_lines; - let mut line_number = 2; while !remaining_lines.is_empty() { let (hunk, hunk_lines) = parse_one_hunk(remaining_lines, line_number)?; hunks.push(hunk); @@ -192,9 +194,28 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result( + hunk_lines: &'a [&'a str], +) -> Result<(Option, &'a [&'a str], usize), ParseError> { + let Some(first_line) = hunk_lines.first() else { + return Ok((None, hunk_lines, 2)); + }; + let Some(environment_id) = first_line.trim_start().strip_prefix(ENVIRONMENT_ID_MARKER) else { + return Ok((None, hunk_lines, 2)); + }; + let environment_id = environment_id.trim(); + if environment_id.is_empty() { + return Err(InvalidPatchError( + "apply_patch environment_id cannot be empty".to_string(), + )); + } + Ok((Some(environment_id.to_string()), &hunk_lines[1..], 3)) +} + /// Checks the start and end lines of the patch text for `apply_patch`, /// returning an error if they do not match the expected markers. fn check_patch_boundaries_strict<'a>( @@ -837,6 +858,7 @@ fn test_parse_patch_lenient() { hunks: expected_patch.clone(), patch: patch_text.to_string(), workdir: None, + environment_id: None, }) ); @@ -851,6 +873,7 @@ fn test_parse_patch_lenient() { hunks: expected_patch.clone(), patch: patch_text.to_string(), workdir: None, + environment_id: None, }) ); @@ -865,6 +888,7 @@ fn test_parse_patch_lenient() { hunks: expected_patch, patch: patch_text.to_string(), workdir: None, + environment_id: None, }) ); @@ -891,3 +915,40 @@ fn test_parse_patch_lenient() { )) ); } + +#[test] +fn test_parse_patch_environment_id_preamble() { + assert_eq!( + parse_patch_text( + "*** Begin Patch\n\ + *** Environment ID: remote\n\ + *** Add File: hello.txt\n\ + +hello\n\ + *** End Patch", + ParseMode::Strict + ), + Ok(ApplyPatchArgs { + hunks: vec![AddFile { + path: PathBuf::from("hello.txt"), + contents: "hello\n".to_string(), + }], + patch: "*** Begin Patch\n*** Environment ID: remote\n*** Add File: hello.txt\n+hello\n*** End Patch".to_string(), + workdir: None, + environment_id: Some("remote".to_string()), + }) + ); + + assert_eq!( + parse_patch_text( + "*** Begin Patch\n\ + *** Environment ID: \n\ + *** Add File: hello.txt\n\ + +hello\n\ + *** End Patch", + ParseMode::Strict + ), + Err(InvalidPatchError( + "apply_patch environment_id cannot be empty".to_string() + )) + ); +} diff --git a/codex-rs/apply-patch/src/streaming_parser.rs b/codex-rs/apply-patch/src/streaming_parser.rs index 4acfad6720..86084af0b8 100644 --- a/codex-rs/apply-patch/src/streaming_parser.rs +++ b/codex-rs/apply-patch/src/streaming_parser.rs @@ -16,6 +16,8 @@ use crate::parser::UpdateFileChunk; use Hunk::*; use ParseError::*; +const ENVIRONMENT_ID_MARKER: &str = "*** Environment ID: "; + #[derive(Debug, Default, Clone)] pub struct StreamingPatchParser { line_buffer: String, @@ -29,7 +31,7 @@ struct StreamingParserState { hunks: Vec, } -#[derive(Debug, Default, Clone)] +#[derive(Debug, Default, Clone, Copy)] enum StreamingParserMode { #[default] NotStarted, @@ -43,6 +45,13 @@ enum StreamingParserMode { } impl StreamingPatchParser { + // The live streaming parser only needs to keep the patch preview flowing. + // Environment selection and validation happen on the final tool invocation, + // so here we just tolerate and skip the optional preamble line. + fn is_environment_id_preamble_line(&self, line: &str) -> bool { + line.starts_with(ENVIRONMENT_ID_MARKER) + } + fn ensure_update_hunk_is_not_empty(&self, line: &str) -> Result<(), ParseError> { if let Some(UpdateFile { path, chunks, .. }) = self.state.hunks.last() { if chunks.is_empty() @@ -150,7 +159,7 @@ impl StreamingPatchParser { fn process_line(&mut self, line: &str) -> Result<(), ParseError> { let trimmed = line.trim(); - match self.state.mode.clone() { + match self.state.mode { StreamingParserMode::NotStarted => { if trimmed == BEGIN_PATCH_MARKER { self.state.mode = StreamingParserMode::StartedPatch; @@ -161,6 +170,9 @@ impl StreamingPatchParser { )) } StreamingParserMode::StartedPatch => { + if self.is_environment_id_preamble_line(line) { + return Ok(()); + } if self.handle_hunk_headers_and_end_patch(trimmed)? { return Ok(()); } @@ -431,6 +443,32 @@ mod tests { ); } + #[test] + fn test_streaming_patch_parser_environment_id_mode() { + let patch = "\ +*** Begin Patch +*** Environment ID: remote +*** Add File: src/hello.txt ++hello +*** End Patch +"; + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta(patch), + Ok(vec![AddFile { + path: PathBuf::from("src/hello.txt"), + contents: "hello\n".to_string(), + }]) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Environment ID: \n"), + Ok(vec![]) + ); + } + #[test] fn test_streaming_patch_parser_large_patch_split_by_character() { let patch = "\ diff --git a/codex-rs/core/src/apply_patch.rs b/codex-rs/core/src/apply_patch.rs index 2463d69c2b..38dc45ec1f 100644 --- a/codex-rs/core/src/apply_patch.rs +++ b/codex-rs/core/src/apply_patch.rs @@ -40,7 +40,7 @@ pub(crate) async fn apply_patch( turn_context.approval_policy.value(), &turn_context.permission_profile(), file_system_sandbox_policy, - &turn_context.cwd, + &action.cwd, turn_context.windows_sandbox_level, ) { SafetyCheck::AutoApprove { diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 452bcb8d85..34e4663044 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -12,6 +12,7 @@ use crate::apply_patch::convert_apply_patch_to_protocol; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; +use crate::session::turn_context::TurnEnvironment; use crate::tools::context::ApplyPatchToolOutput; use crate::tools::context::FunctionToolOutput; use crate::tools::context::SharedTurnDiffTracker; @@ -22,6 +23,7 @@ use crate::tools::events::ToolEmitter; use crate::tools::events::ToolEventCtx; use crate::tools::handlers::apply_granted_turn_permissions; use crate::tools::handlers::apply_patch_spec::create_apply_patch_freeform_tool; +use crate::tools::handlers::resolve_tool_environment; use crate::tools::hook_names::HookToolName; use crate::tools::orchestrator::ToolOrchestrator; use crate::tools::registry::PostToolUsePayload; @@ -51,8 +53,18 @@ use codex_tools::ToolSpec; use codex_utils_absolute_path::AbsolutePathBuf; const APPLY_PATCH_ARGUMENT_DIFF_BUFFER_INTERVAL: Duration = Duration::from_millis(500); +/// Handles freeform `apply_patch` requests and routes verified patches to the +/// selected environment filesystem. +#[derive(Default)] +pub struct ApplyPatchHandler { + multi_environment: bool, +} -pub struct ApplyPatchHandler; +impl ApplyPatchHandler { + pub(crate) fn new(multi_environment: bool) -> Self { + Self { multi_environment } + } +} #[derive(Default)] struct ApplyPatchArgumentDiffConsumer { @@ -253,6 +265,7 @@ async fn effective_patch_permissions( session: &Session, turn: &TurnContext, action: &ApplyPatchAction, + cwd: &AbsolutePathBuf, ) -> ( Vec, crate::tools::handlers::EffectiveAdditionalPermissions, @@ -270,9 +283,9 @@ async fn effective_patch_permissions( ); let effective_additional_permissions = apply_granted_turn_permissions( session, - turn.cwd.as_path(), + cwd.as_path(), crate::sandboxing::SandboxPermissions::UseDefault, - write_permissions_for_paths(&file_paths, &file_system_sandbox_policy, &turn.cwd), + write_permissions_for_paths(&file_paths, &file_system_sandbox_policy, cwd), ) .await; @@ -291,7 +304,7 @@ impl ToolHandler for ApplyPatchHandler { } fn spec(&self) -> Option { - Some(create_apply_patch_freeform_tool()) + Some(create_apply_patch_freeform_tool(self.multi_environment)) } fn kind(&self) -> ToolKind { @@ -350,32 +363,36 @@ impl ToolHandler for ApplyPatchHandler { "apply_patch handler received unsupported payload".to_string(), )); }; + let args = match codex_apply_patch::parse_patch(&patch_input) { + Ok(args) => args, + Err(parse_error) => { + return Err(FunctionCallError::RespondToModel(format!( + "apply_patch verification failed: {parse_error}" + ))); + } + }; + let selected_environment_id = + require_environment_id(args.environment_id.as_deref(), self.multi_environment)?; - // Re-parse and verify the patch so we can compute changes and approval. - // Avoid building temporary ExecParams/command vectors; derive directly from inputs. - let cwd = turn.cwd.clone(); - let command = vec!["apply_patch".to_string(), patch_input.clone()]; - let Some(turn_environment) = turn.environments.primary() else { + // Verify the parsed patch against the selected environment filesystem. + let Some(turn_environment) = + resolve_tool_environment(turn.as_ref(), selected_environment_id.as_deref())? + else { return Err(FunctionCallError::RespondToModel( "apply_patch is unavailable in this session".to_string(), )); }; + let cwd = turn_environment.cwd.clone(); let fs = turn_environment.environment.get_filesystem(); - let sandbox = turn_environment - .environment - .is_remote() - .then(|| turn.file_system_sandbox_context(/*additional_permissions*/ None)); - match codex_apply_patch::maybe_parse_apply_patch_verified( - &command, - &cwd, - fs.as_ref(), - sandbox.as_ref(), - ) - .await + let mut sandbox = turn.file_system_sandbox_context(/*additional_permissions*/ None); + sandbox.cwd = Some(cwd.clone()); + match codex_apply_patch::verify_apply_patch_args(args, &cwd, fs.as_ref(), Some(&sandbox)) + .await { codex_apply_patch::MaybeApplyPatchVerified::Body(changes) => { let (file_paths, effective_additional_permissions, file_system_sandbox_policy) = - effective_patch_permissions(session.as_ref(), turn.as_ref(), &changes).await; + effective_patch_permissions(session.as_ref(), turn.as_ref(), &changes, &cwd) + .await; match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes) .await { @@ -396,6 +413,7 @@ impl ToolHandler for ApplyPatchHandler { emitter.begin(event_ctx).await; let req = ApplyPatchRequest { + turn_environment: turn_environment.clone(), action: apply.action, file_paths, changes, @@ -464,18 +482,16 @@ pub(crate) async fn intercept_apply_patch( command: &[String], cwd: &AbsolutePathBuf, fs: &dyn ExecutorFileSystem, + turn_environment: TurnEnvironment, session: Arc, turn: Arc, tracker: Option<&SharedTurnDiffTracker>, call_id: &str, tool_name: &str, ) -> Result, FunctionCallError> { - let sandbox = turn - .environments - .primary() - .filter(|env| env.environment.is_remote()) - .map(|_| turn.file_system_sandbox_context(/*additional_permissions*/ None)); - match codex_apply_patch::maybe_parse_apply_patch_verified(command, cwd, fs, sandbox.as_ref()) + let mut sandbox = turn.file_system_sandbox_context(/*additional_permissions*/ None); + sandbox.cwd = Some(cwd.clone()); + match codex_apply_patch::maybe_parse_apply_patch_verified(command, cwd, fs, Some(&sandbox)) .await { codex_apply_patch::MaybeApplyPatchVerified::Body(changes) => { @@ -488,7 +504,7 @@ pub(crate) async fn intercept_apply_patch( ) .await; let (approval_keys, effective_additional_permissions, file_system_sandbox_policy) = - effective_patch_permissions(session.as_ref(), turn.as_ref(), &changes).await; + effective_patch_permissions(session.as_ref(), turn.as_ref(), &changes, cwd).await; match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes) .await { @@ -508,6 +524,7 @@ pub(crate) async fn intercept_apply_patch( emitter.begin(event_ctx).await; let req = ApplyPatchRequest { + turn_environment, action: apply.action, file_paths: approval_keys, changes, @@ -564,6 +581,19 @@ pub(crate) async fn intercept_apply_patch( } } +fn require_environment_id( + parsed_environment_id: Option<&str>, + allow_environment_id: bool, +) -> Result, FunctionCallError> { + match parsed_environment_id { + Some(_) if !allow_environment_id => Err(FunctionCallError::RespondToModel( + "apply_patch environment selection is unavailable for this turn".to_string(), + )), + Some(environment_id) => Ok(Some(environment_id.to_string())), + None => Ok(None), + } +} + #[cfg(test)] #[path = "apply_patch_tests.rs"] mod tests; diff --git a/codex-rs/core/src/tools/handlers/apply_patch_spec.rs b/codex-rs/core/src/tools/handlers/apply_patch_spec.rs index fa1ffcb8c0..358f045db1 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_spec.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_spec.rs @@ -6,14 +6,22 @@ const APPLY_PATCH_LARK_GRAMMAR: &str = include_str!("apply_patch.lark"); /// Returns a custom tool that can be used to edit files. Well-suited for GPT-5 models /// https://platform.openai.com/docs/guides/function-calling#custom-tools -pub fn create_apply_patch_freeform_tool() -> ToolSpec { +pub fn create_apply_patch_freeform_tool(include_environment_id: bool) -> ToolSpec { + let definition = if include_environment_id { + APPLY_PATCH_LARK_GRAMMAR.replace( + "start: begin_patch hunk+ end_patch", + "start: begin_patch environment_id? hunk+ end_patch\nenvironment_id: \"*** Environment ID: \" filename LF", + ) + } else { + APPLY_PATCH_LARK_GRAMMAR.to_string() + }; ToolSpec::Freeform(FreeformTool { name: "apply_patch".to_string(), description: "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, so do not wrap the patch in JSON.".to_string(), format: FreeformToolFormat { r#type: "grammar".to_string(), syntax: "lark".to_string(), - definition: APPLY_PATCH_LARK_GRAMMAR.to_string(), + definition, }, }) } diff --git a/codex-rs/core/src/tools/handlers/apply_patch_spec_tests.rs b/codex-rs/core/src/tools/handlers/apply_patch_spec_tests.rs index 28bd57a991..7d172becae 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_spec_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_spec_tests.rs @@ -4,7 +4,7 @@ use pretty_assertions::assert_eq; #[test] fn create_apply_patch_freeform_tool_matches_expected_spec() { assert_eq!( - create_apply_patch_freeform_tool(), + create_apply_patch_freeform_tool(/*include_environment_id*/ false), ToolSpec::Freeform(FreeformTool { name: "apply_patch".to_string(), description: @@ -18,3 +18,19 @@ fn create_apply_patch_freeform_tool_matches_expected_spec() { }) ); } + +#[test] +fn create_apply_patch_freeform_tool_includes_environment_id_when_requested() { + let ToolSpec::Freeform(tool) = + create_apply_patch_freeform_tool(/*include_environment_id*/ true) + else { + panic!("expected freeform tool"); + }; + + assert!(tool.format.definition.contains("environment_id?")); + assert!( + tool.format + .definition + .contains("\"*** Environment ID: \" filename LF") + ); +} diff --git a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs index dfa642ade1..99854d23ea 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs @@ -49,7 +49,7 @@ async fn pre_tool_use_payload_uses_freeform_patch_input() { input: patch.to_string(), }; let invocation = invocation_for_payload(payload).await; - let handler = ApplyPatchHandler; + let handler = ApplyPatchHandler::default(); assert_eq!( handler.pre_tool_use_payload(&invocation), @@ -68,7 +68,7 @@ async fn post_tool_use_payload_uses_patch_input_and_tool_output() { }; let invocation = invocation_for_payload(payload).await; let output = ApplyPatchToolOutput::from_text("Success. Updated files.".to_string()); - let handler = ApplyPatchHandler; + let handler = ApplyPatchHandler::default(); assert_eq!( handler.post_tool_use_payload(&invocation, &output), @@ -135,6 +135,32 @@ fn diff_consumer_streams_apply_patch_changes() { ); } +#[test] +fn diff_consumer_streams_apply_patch_changes_with_environment_header() { + let mut consumer = ApplyPatchArgumentDiffConsumer::default(); + assert!( + consumer + .push_delta( + "call-1".to_string(), + "*** Begin Patch\n*** Environment ID: remote\n", + ) + .is_none() + ); + + let event = consumer + .push_delta("call-1".to_string(), "*** Add File: hello.txt\n+hello") + .expect("progress event"); + assert_eq!( + event.changes, + HashMap::from([( + PathBuf::from("hello.txt"), + FileChange::Add { + content: String::new(), + }, + )]) + ); +} + #[test] fn diff_consumer_sends_next_update_after_buffer_interval() { let mut consumer = ApplyPatchArgumentDiffConsumer::default(); @@ -168,6 +194,22 @@ fn diff_consumer_sends_next_update_after_buffer_interval() { ); } +#[test] +fn reconcile_environment_id_requires_selection_when_enabled() { + assert_eq!( + require_environment_id(Some("remote"), /*allow_environment_id*/ false), + Err(FunctionCallError::RespondToModel( + "apply_patch environment selection is unavailable for this turn".to_string(), + )) + ); + assert_eq!( + require_environment_id( + /*parsed_environment_id*/ None, /*allow_environment_id*/ true + ), + Ok(None) + ); +} + #[tokio::test] async fn approval_keys_include_move_destination() { let tmp = TempDir::new().expect("tmp"); diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 81a590e07d..8bffe13827 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -198,6 +198,7 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result, pub changes: std::collections::HashMap, @@ -108,10 +116,17 @@ impl Sandboxable for ApplyPatchRuntime { } impl Approvable for ApplyPatchRuntime { - type ApprovalKey = AbsolutePathBuf; + type ApprovalKey = ApplyPatchApprovalKey; fn approval_keys(&self, req: &ApplyPatchRequest) -> Vec { - req.file_paths.clone() + req.file_paths + .iter() + .cloned() + .map(|path| ApplyPatchApprovalKey { + environment_id: req.turn_environment.environment_id.clone(), + path, + }) + .collect() } fn start_approval_async<'a>( @@ -198,17 +213,18 @@ impl Approvable for ApplyPatchRuntime { } impl ToolRuntime for ApplyPatchRuntime { + fn sandbox_cwd<'a>(&self, req: &'a ApplyPatchRequest) -> Option<&'a AbsolutePathBuf> { + Some(&req.action.cwd) + } + async fn run( &mut self, req: &ApplyPatchRequest, attempt: &SandboxAttempt<'_>, - ctx: &ToolCtx, + _ctx: &ToolCtx, ) -> Result { - let turn_environment = ctx.turn.environments.primary().ok_or_else(|| { - ToolError::Rejected("apply_patch is unavailable in this session".to_string()) - })?; let started_at = Instant::now(); - let fs = turn_environment.environment.get_filesystem(); + let fs = req.turn_environment.environment.get_filesystem(); let sandbox = Self::file_system_sandbox_context_for_attempt(req, attempt); let mut stdout = Vec::new(); let mut stderr = Vec::new(); diff --git a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs index 173fa3e2a0..ff2c3d1e89 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs @@ -15,6 +15,14 @@ use codex_sandboxing::policy_transforms::effective_network_sandbox_policy; use core_test_support::PathBufExt; use pretty_assertions::assert_eq; use std::collections::HashMap; +fn test_turn_environment(environment_id: &str) -> crate::session::turn_context::TurnEnvironment { + crate::session::turn_context::TurnEnvironment { + environment_id: environment_id.to_string(), + environment: std::sync::Arc::new(codex_exec_server::Environment::default_for_tests()), + cwd: std::env::temp_dir().abs(), + shell: None, + } +} #[test] fn wants_no_sandbox_approval_granular_respects_sandbox_flag() { @@ -40,8 +48,8 @@ fn wants_no_sandbox_approval_granular_respects_sandbox_flag() { ); } -#[test] -fn guardian_review_request_includes_patch_context() { +#[tokio::test] +async fn guardian_review_request_includes_patch_context() { let path = std::env::temp_dir() .join("guardian-apply-patch-test.txt") .abs(); @@ -49,6 +57,7 @@ fn guardian_review_request_includes_patch_context() { let expected_cwd = action.cwd.clone(); let expected_patch = action.patch.clone(); let request = ApplyPatchRequest { + turn_environment: test_turn_environment(codex_exec_server::LOCAL_ENVIRONMENT_ID), action, file_paths: vec![path.clone()], changes: HashMap::from([( @@ -78,8 +87,8 @@ fn guardian_review_request_includes_patch_context() { ); } -#[test] -fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { +#[tokio::test] +async fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { let runtime = ApplyPatchRuntime::new(); let path = std::env::temp_dir() .join("apply-patch-permission-request-payload.txt") @@ -87,6 +96,7 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { let action = ApplyPatchAction::new_add_for_test(&path, "hello".to_string()); let expected_patch = action.patch.clone(); let req = ApplyPatchRequest { + turn_environment: test_turn_environment(codex_exec_server::LOCAL_ENVIRONMENT_ID), action, file_paths: vec![path], changes: HashMap::new(), @@ -113,8 +123,62 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { ); } -#[test] -fn file_system_sandbox_context_uses_active_attempt() { +#[tokio::test] +async fn approval_keys_include_environment_id() { + let runtime = ApplyPatchRuntime::new(); + let path = std::env::temp_dir() + .join("apply-patch-approval-key.txt") + .abs(); + let req = ApplyPatchRequest { + turn_environment: test_turn_environment("remote"), + action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), + file_paths: vec![path.clone()], + changes: HashMap::new(), + exec_approval_requirement: ExecApprovalRequirement::Skip { + bypass_sandbox: false, + proposed_execpolicy_amendment: None, + }, + additional_permissions: None, + permissions_preapproved: false, + }; + + let keys = runtime.approval_keys(&req); + + assert_eq!( + serde_json::to_value(&keys).expect("serialize approval keys"), + serde_json::json!([ + { + "environment_id": "remote", + "path": path, + } + ]) + ); +} + +#[tokio::test] +async fn sandbox_cwd_uses_patch_action_cwd() { + let runtime = ApplyPatchRuntime::new(); + let path = std::env::temp_dir() + .join("apply-patch-runtime-sandbox-cwd.txt") + .abs(); + let req = ApplyPatchRequest { + turn_environment: test_turn_environment(codex_exec_server::LOCAL_ENVIRONMENT_ID), + action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), + file_paths: vec![path.clone()], + changes: HashMap::new(), + exec_approval_requirement: ExecApprovalRequirement::Skip { + bypass_sandbox: false, + proposed_execpolicy_amendment: None, + }, + additional_permissions: None, + permissions_preapproved: false, + }; + + assert_eq!(runtime.sandbox_cwd(&req), Some(&req.action.cwd)); +} + +#[tokio::test] +async fn file_system_sandbox_context_uses_active_attempt() { let path = std::env::temp_dir() .join("apply-patch-runtime-attempt.txt") .abs(); @@ -126,6 +190,7 @@ fn file_system_sandbox_context_uses_active_attempt() { )), }; let req = ApplyPatchRequest { + turn_environment: test_turn_environment(codex_exec_server::LOCAL_ENVIRONMENT_ID), action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), file_paths: vec![path.clone()], changes: HashMap::new(), @@ -177,12 +242,13 @@ fn file_system_sandbox_context_uses_active_attempt() { assert_eq!(sandbox.use_legacy_landlock, true); } -#[test] -fn no_sandbox_attempt_has_no_file_system_context() { +#[tokio::test] +async fn no_sandbox_attempt_has_no_file_system_context() { let path = std::env::temp_dir() .join("apply-patch-runtime-none.txt") .abs(); let req = ApplyPatchRequest { + turn_environment: test_turn_environment(codex_exec_server::LOCAL_ENVIRONMENT_ID), action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), file_paths: vec![path.clone()], changes: HashMap::new(), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 13a79e5943..7fded61401 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -258,7 +258,9 @@ pub fn build_tool_registry_builder( } if config.environment_mode.has_environment() && config.apply_patch_tool_type.is_some() { - builder.register_handler(Arc::new(ApplyPatchHandler)); + let include_environment_id = + matches!(config.environment_mode, ToolEnvironmentMode::Multiple); + builder.register_handler(Arc::new(ApplyPatchHandler::new(include_environment_id))); } if config diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 7cd33a2aaf..fdae0452e2 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -190,7 +190,7 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() { create_write_stdin_tool(), create_update_plan_tool(), request_user_input_tool_spec(&request_user_input_available_modes(&features)), - create_apply_patch_freeform_tool(), + create_apply_patch_freeform_tool(/*include_environment_id*/ false), ToolSpec::WebSearch { external_web_access: Some(true), filters: None, @@ -301,6 +301,40 @@ fn exec_command_spec_includes_environment_id_only_for_multiple_selected_environm ); } +#[test] +fn apply_patch_spec_includes_environment_id_only_for_multiple_selected_environments() { + let model_info = model_info(); + let available_models = Vec::new(); + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &Features::with_defaults(), + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + permission_profile: &PermissionProfile::Disabled, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + + let (single_environment_tools, _) = build_specs( + &tools_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_apply_patch_environment_id(&single_environment_tools, /*expected_present*/ false); + + let multi_environment_config = + tools_config.with_environment_mode(ToolEnvironmentMode::Multiple); + let (multi_environment_tools, _) = build_specs( + &multi_environment_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_apply_patch_environment_id(&multi_environment_tools, /*expected_present*/ true); +} + #[test] fn test_build_specs_collab_tools_enabled() { let model_info = model_info(); @@ -2649,6 +2683,18 @@ fn assert_process_tool_environment_id( ); } +fn assert_apply_patch_environment_id(tools: &[ConfiguredToolSpec], expected_present: bool) { + let tool = find_tool(tools, "apply_patch"); + let ToolSpec::Freeform(FreeformTool { format, .. }) = &tool.spec else { + panic!("expected freeform apply_patch tool"); + }; + assert_eq!( + format.definition.contains("environment_id?"), + expected_present, + "apply_patch environment_id grammar presence" + ); +} + fn find_namespace_function_tool<'a>( tools: &'a [ConfiguredToolSpec], expected_namespace: &str, diff --git a/codex-rs/core/tests/suite/apply_patch_cli.rs b/codex-rs/core/tests/suite/apply_patch_cli.rs index f95fd1efef..82113ee494 100644 --- a/codex-rs/core/tests/suite/apply_patch_cli.rs +++ b/codex-rs/core/tests/suite/apply_patch_cli.rs @@ -139,6 +139,28 @@ fn workspace_write_with_read_only_root(read_only_root: AbsolutePathBuf) -> Permi ) } +#[cfg(unix)] +fn workspace_write_with_unreadable_path(unreadable_path: AbsolutePathBuf) -> PermissionProfile { + let file_system_sandbox_policy = FileSystemSandboxPolicy::restricted(vec![ + FileSystemSandboxEntry { + path: FileSystemPath::Path { + path: unreadable_path, + }, + access: FileSystemAccessMode::None, + }, + FileSystemSandboxEntry { + path: FileSystemPath::Special { + value: FileSystemSpecialPath::project_roots(/*subpath*/ None), + }, + access: FileSystemAccessMode::Write, + }, + ]); + PermissionProfile::from_runtime_permissions( + &file_system_sandbox_policy, + NetworkSandboxPolicy::Restricted, + ) +} + #[cfg(unix)] fn create_file_symlink(source: &std::path::Path, link: &std::path::Path) -> std::io::Result<()> { std::os::unix::fs::symlink(source, link) @@ -720,6 +742,59 @@ async fn apply_patch_cli_rejects_path_traversal_outside_workspace( Ok(()) } +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[test_case(ApplyPatchModelOutput::Shell ; "shell")] +#[test_case(ApplyPatchModelOutput::ShellViaHeredoc ; "shell_heredoc")] +#[test_case(ApplyPatchModelOutput::ShellCommandViaHeredoc ; "shell_command_heredoc")] +async fn intercepted_apply_patch_verification_uses_local_sandbox( + model_output: ApplyPatchModelOutput, +) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_remote!(Ok(()), "symlink setup needs local filesystem link creation"); + + let harness = apply_patch_harness().await?; + let denied_target = harness.path("denied-target.txt"); + std::fs::write(&denied_target, "outside content\n")?; + + let link_rel = "soft-link.txt"; + create_file_symlink(&denied_target, &harness.path(link_rel))?; + + let patch = format!( + r#"*** Begin Patch +*** Update File: {link_rel} +@@ +-outside content ++pwned +*** End Patch"# + ); + let call_id = "apply-sandboxed-read"; + mount_apply_patch(&harness, call_id, &patch, "fail", model_output).await; + + harness + .submit_with_permission_profile( + "attempt to read denied target via intercepted apply_patch", + workspace_write_with_unreadable_path(AbsolutePathBuf::try_from(denied_target.clone())?), + ) + .await?; + + let out = harness.apply_patch_output(call_id, model_output).await; + assert!( + out.contains("apply_patch verification failed"), + "expected sandboxed verification failure: {out}" + ); + assert!( + out.contains("Failed to read"), + "expected read failure: {out}" + ); + assert_eq!( + std::fs::read_to_string(&denied_target)?, + "outside content\n", + "verification failure should leave the denied target unchanged" + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[test_case(ApplyPatchModelOutput::Freeform ; "freeform")] #[test_case(ApplyPatchModelOutput::Shell ; "shell")] diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 0bd449188c..dc6d5cefc5 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -1,5 +1,7 @@ use anyhow::Context; use anyhow::Result; +use codex_config::types::ApprovalsReviewer; +use codex_core::config::Constrained; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::FileSystemSandboxContext; @@ -13,11 +15,19 @@ use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSandboxEntry; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::protocol::ApplyPatchApprovalRequestEvent; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::ReviewDecision; +use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::TurnEnvironmentSelection; +use codex_protocol::user_input::UserInput; use codex_utils_absolute_path::AbsolutePathBuf; use core_test_support::PathBufExt; use core_test_support::PathExt; use core_test_support::get_remote_test_env; +use core_test_support::responses::ev_apply_patch_custom_tool_call; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; @@ -29,6 +39,7 @@ use core_test_support::skip_if_no_network; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::test_codex::test_env; +use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -50,6 +61,76 @@ async fn unified_exec_test(server: &wiremock::MockServer) -> Result { builder.build_remote_aware(server).await } +async fn submit_turn_with_approval_and_environments( + test: &TestCodex, + prompt: &str, + environments: Vec, +) -> Result<()> { + test.codex + .submit(Op::UserTurn { + environments: Some(environments), + items: vec![UserInput::Text { + text: prompt.into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: test.cwd.path().to_path_buf(), + approval_policy: AskForApproval::OnRequest, + approvals_reviewer: Some(ApprovalsReviewer::User), + sandbox_policy: SandboxPolicy::new_workspace_write_policy(), + permission_profile: None, + model: test.session_configured.model.clone(), + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + + Ok(()) +} + +async fn expect_patch_approval( + test: &TestCodex, + expected_call_id: &str, +) -> ApplyPatchApprovalRequestEvent { + let event = wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::ApplyPatchApprovalRequest(_) | EventMsg::TurnComplete(_) + ) + }) + .await; + + match event { + EventMsg::ApplyPatchApprovalRequest(approval) => { + assert_eq!(approval.call_id, expected_call_id); + approval + } + EventMsg::TurnComplete(_) => panic!("expected patch approval request before completion"), + other => panic!("unexpected event: {other:?}"), + } +} + +async fn wait_for_completion_without_patch_approval(test: &TestCodex) { + let event = wait_for_event(&test.codex, |event| { + matches!( + event, + EventMsg::ApplyPatchApprovalRequest(_) | EventMsg::TurnComplete(_) + ) + }) + .await; + + match event { + EventMsg::TurnComplete(_) => {} + EventMsg::ApplyPatchApprovalRequest(event) => { + panic!("unexpected patch approval request: {:?}", event.call_id) + } + other => panic!("unexpected event: {other:?}"), + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_test_env_can_connect_and_use_filesystem() -> Result<()> { let Some(_remote_env) = get_remote_test_env() else { @@ -258,6 +339,375 @@ async fn exec_command_routes_to_selected_remote_environment() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn apply_patch_freeform_routes_to_selected_remote_environment() -> Result<()> { + skip_if_no_network!(Ok(())); + let Some(_remote_env) = get_remote_test_env() else { + return Ok(()); + }; + + let server = start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + config.include_apply_patch_tool = true; + }); + let test = builder.build_remote_aware(&server).await?; + let local_cwd = TempDir::new()?; + let file_name = "apply_patch_remote_freeform.txt"; + let remote_cwd = PathBuf::from(format!( + "/tmp/codex-remote-apply-patch-freeform-{}", + SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() + )) + .abs(); + test.fs() + .create_directory( + &remote_cwd, + CreateDirectoryOptions { recursive: true }, + /*sandbox*/ None, + ) + .await?; + + let patch = format!( + "*** Begin Patch\n*** Environment ID: {REMOTE_ENVIRONMENT_ID}\n*** Add File: {file_name}\n+patched remote freeform\n*** End Patch" + ); + let call_id = "apply-patch-remote-freeform"; + mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_apply_patch_custom_tool_call(call_id, &patch), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + test.submit_turn_with_environments( + "apply patch to remote environment", + Some(vec![ + TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: local_cwd.path().abs(), + }, + TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: remote_cwd.clone(), + }, + ]), + ) + .await?; + + let remote_contents = test + .fs() + .read_file_text(&remote_cwd.join(file_name), /*sandbox*/ None) + .await?; + assert_eq!(remote_contents, "patched remote freeform\n"); + assert!( + !local_cwd.path().join(file_name).exists(), + "freeform apply_patch should not create the file in the local environment" + ); + + test.fs() + .remove( + &remote_cwd, + RemoveOptions { + recursive: true, + force: true, + }, + /*sandbox*/ None, + ) + .await?; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn apply_patch_approvals_are_remembered_per_environment() -> Result<()> { + skip_if_no_network!(Ok(())); + let Some(_remote_env) = get_remote_test_env() else { + return Ok(()); + }; + + let server = start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + config.include_apply_patch_tool = true; + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config.approvals_reviewer = ApprovalsReviewer::User; + }); + let test = builder.build_remote_aware(&server).await?; + let local_cwd = TempDir::new()?; + let remote_cwd = PathBuf::from(format!( + "/tmp/codex-remote-apply-patch-approval-cwd-{}", + SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() + )) + .abs(); + test.fs() + .create_directory( + &remote_cwd, + CreateDirectoryOptions { recursive: true }, + /*sandbox*/ None, + ) + .await?; + + let target_path = PathBuf::from(format!( + "/tmp/codex-apply-patch-approval-scope-{}.txt", + SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() + )) + .abs(); + let _ = fs::remove_file(&target_path); + test.fs() + .remove( + &target_path, + RemoveOptions { + recursive: false, + force: true, + }, + /*sandbox*/ None, + ) + .await?; + + let environments = vec![ + TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: local_cwd.path().abs(), + }, + TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: remote_cwd.clone(), + }, + ]; + let local_patch = format!( + "*** Begin Patch\n*** Environment ID: {LOCAL_ENVIRONMENT_ID}\n*** Add File: {}\n+local\n*** End Patch", + target_path.display() + ); + let remote_patch = format!( + "*** Begin Patch\n*** Environment ID: {REMOTE_ENVIRONMENT_ID}\n*** Add File: {}\n+remote\n*** End Patch", + target_path.display() + ); + let remote_update_patch = format!( + "*** Begin Patch\n*** Environment ID: {REMOTE_ENVIRONMENT_ID}\n*** Update File: {}\n@@\n-remote\n+remote updated\n*** End Patch", + target_path.display() + ); + + mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-local-1"), + ev_apply_patch_custom_tool_call("call-local", &local_patch), + ev_completed("resp-local-1"), + ]), + sse(vec![ + ev_response_created("resp-local-2"), + ev_assistant_message("msg-local", "done"), + ev_completed("resp-local-2"), + ]), + sse(vec![ + ev_response_created("resp-remote-1"), + ev_apply_patch_custom_tool_call("call-remote", &remote_patch), + ev_completed("resp-remote-1"), + ]), + sse(vec![ + ev_response_created("resp-remote-2"), + ev_assistant_message("msg-remote", "done"), + ev_completed("resp-remote-2"), + ]), + sse(vec![ + ev_response_created("resp-remote-3"), + ev_apply_patch_custom_tool_call("call-remote-followup", &remote_update_patch), + ev_completed("resp-remote-3"), + ]), + sse(vec![ + ev_response_created("resp-remote-4"), + ev_assistant_message("msg-remote-followup", "done"), + ev_completed("resp-remote-4"), + ]), + ], + ) + .await; + + submit_turn_with_approval_and_environments( + &test, + "apply patch in local environment", + environments.clone(), + ) + .await?; + let approval = expect_patch_approval(&test, "call-local").await; + test.codex + .submit(Op::PatchApproval { + id: approval.call_id, + decision: ReviewDecision::ApprovedForSession, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!(fs::read_to_string(&target_path)?, "local\n"); + + submit_turn_with_approval_and_environments( + &test, + "apply patch in remote environment", + environments.clone(), + ) + .await?; + let approval = expect_patch_approval(&test, "call-remote").await; + test.codex + .submit(Op::PatchApproval { + id: approval.call_id, + decision: ReviewDecision::ApprovedForSession, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!( + test.fs() + .read_file_text(&target_path, /*sandbox*/ None) + .await?, + "remote\n" + ); + + submit_turn_with_approval_and_environments( + &test, + "apply patch again in remote environment", + environments, + ) + .await?; + wait_for_completion_without_patch_approval(&test).await; + assert_eq!( + test.fs() + .read_file_text(&target_path, /*sandbox*/ None) + .await?, + "remote updated\n" + ); + + let _ = fs::remove_file(&target_path); + test.fs() + .remove( + &target_path, + RemoveOptions { + recursive: false, + force: true, + }, + /*sandbox*/ None, + ) + .await?; + test.fs() + .remove( + &remote_cwd, + RemoveOptions { + recursive: true, + force: true, + }, + /*sandbox*/ None, + ) + .await?; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn apply_patch_intercepted_exec_command_routes_to_selected_remote_environment() -> Result<()> +{ + skip_if_no_network!(Ok(())); + let Some(_remote_env) = get_remote_test_env() else { + return Ok(()); + }; + + let server = start_mock_server().await; + let test = unified_exec_test(&server).await?; + let local_cwd = TempDir::new()?; + let file_name = "apply_patch_remote_exec.txt"; + let remote_cwd = PathBuf::from(format!( + "/tmp/codex-remote-apply-patch-exec-{}", + SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() + )) + .abs(); + test.fs() + .create_directory( + &remote_cwd, + CreateDirectoryOptions { recursive: true }, + /*sandbox*/ None, + ) + .await?; + + let patch = + format!("*** Begin Patch\n*** Add File: {file_name}\n+patched remote exec\n*** End Patch"); + let command = format!("apply_patch <<'EOF'\n{patch}\nEOF\n"); + let call_id = "apply-patch-remote-exec"; + mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + call_id, + "exec_command", + &serde_json::to_string(&json!({ + "shell": "/bin/sh", + "cmd": command, + "login": false, + "yield_time_ms": 5_000, + "environment_id": REMOTE_ENVIRONMENT_ID, + }))?, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + test.submit_turn_with_environments( + "apply patch through exec command to remote environment", + Some(vec![ + TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: local_cwd.path().abs(), + }, + TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: remote_cwd.clone(), + }, + ]), + ) + .await?; + + let remote_contents = test + .fs() + .read_file_text(&remote_cwd.join(file_name), /*sandbox*/ None) + .await?; + assert_eq!(remote_contents, "patched remote exec\n"); + assert!( + !local_cwd.path().join(file_name).exists(), + "intercepted apply_patch should not create the file in the local environment" + ); + + test.fs() + .remove( + &remote_cwd, + RemoveOptions { + recursive: true, + force: true, + }, + /*sandbox*/ None, + ) + .await?; + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_test_env_sandboxed_read_allows_readable_root() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/mcp-server/tests/suite/codex_tool.rs b/codex-rs/mcp-server/tests/suite/codex_tool.rs index 323d07264d..d9f290c3c5 100644 --- a/codex-rs/mcp-server/tests/suite/codex_tool.rs +++ b/codex-rs/mcp-server/tests/suite/codex_tool.rs @@ -261,6 +261,11 @@ async fn patch_approval_triggers_elicitation() -> anyhow::Result<()> { .send_codex_tool_call(CodexToolCallParam { cwd: Some(cwd.path().to_string_lossy().to_string()), prompt: "please modify the test file".to_string(), + // This test exercises patch approval elicitation, not local sandbox setup. + config: Some(HashMap::from([( + "sandbox_mode".to_string(), + json!("danger-full-access"), + )])), ..Default::default() }) .await?; @@ -296,7 +301,7 @@ async fn patch_approval_triggers_elicitation() -> anyhow::Result<()> { Some(create_expected_patch_approval_elicitation_request_params( expected_changes, /*grant_root*/ None, // No grant_root expected - /*reason*/ None, // No reason expected + /*reason*/ None, codex_request_id.to_string(), params.codex_event_id.clone(), params.thread_id,