From ebf66ae577d3fec239663d3bb77ae43d3b379beb Mon Sep 17 00:00:00 2001 From: starr-openai Date: Thu, 7 May 2026 16:43:32 -0700 Subject: [PATCH] Support multi-environment apply_patch selection Co-authored-by: Codex --- .../src/protocol/thread_history.rs | 4 + codex-rs/apply-patch/src/invocation.rs | 52 ++- codex-rs/apply-patch/src/lib.rs | 21 ++ codex-rs/apply-patch/src/parser.rs | 160 +++++++-- codex-rs/apply-patch/src/streaming_parser.rs | 19 + codex-rs/core/src/codex_delegate.rs | 20 +- .../core/src/guardian/approval_request.rs | 3 + codex-rs/core/src/guardian/tests.rs | 5 + codex-rs/core/src/session/mod.rs | 24 +- codex-rs/core/src/session/turn_context.rs | 10 +- .../core/src/tools/handlers/apply_patch.lark | 3 +- .../core/src/tools/handlers/apply_patch.rs | 336 ++++++++++++------ .../src/tools/handlers/apply_patch_spec.rs | 39 +- .../tools/handlers/apply_patch_spec_tests.rs | 51 ++- .../core/src/tools/runtimes/apply_patch.rs | 63 +++- .../src/tools/runtimes/apply_patch_tests.rs | 40 ++- codex-rs/core/src/tools/spec_plan.rs | 2 + codex-rs/core/src/tools/spec_plan_tests.rs | 5 +- codex-rs/core/tests/suite/tool_harness.rs | 156 ++++++++ codex-rs/mcp-server/src/codex_tool_runner.rs | 2 + codex-rs/protocol/src/approvals.rs | 4 + codex-rs/tui/src/app/thread_routing.rs | 1 + codex-rs/tui/src/approval_events.rs | 3 + .../tui/src/bottom_pane/approval_overlay.rs | 12 + codex-rs/tui/src/chatwidget.rs | 11 +- codex-rs/tui/src/chatwidget/slash_dispatch.rs | 6 + .../tui/src/chatwidget/tests/exec_flow.rs | 32 ++ 27 files changed, 871 insertions(+), 213 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 1121d3a35b..b053cf7e76 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -2532,6 +2532,10 @@ mod tests { call_id: "patch-call".into(), turn_id: turn_id.to_string(), started_at_ms: 0, + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes: [( PathBuf::from("README.md"), codex_protocol::protocol::FileChange::Add { diff --git a/codex-rs/apply-patch/src/invocation.rs b/codex-rs/apply-patch/src/invocation.rs index abe223f1d9..db5ab10edf 100644 --- a/codex-rs/apply-patch/src/invocation.rs +++ b/codex-rs/apply-patch/src/invocation.rs @@ -14,7 +14,6 @@ use crate::ApplyPatchAction; use crate::ApplyPatchArgs; use crate::ApplyPatchError; use crate::ApplyPatchFileChange; -use crate::ApplyPatchFileUpdate; use crate::IoError; use crate::MaybeApplyPatchVerified; use crate::parser::Hunk; @@ -154,6 +153,7 @@ pub async fn maybe_parse_apply_patch_verified( MaybeApplyPatch::Body(ApplyPatchArgs { patch, hunks, + environment_id: _, workdir, }) => { let effective_cwd = workdir @@ -163,53 +163,44 @@ pub async fn maybe_parse_apply_patch_verified( let mut changes = HashMap::new(); for hunk in hunks { let path = hunk.resolve_path(&effective_cwd); - match hunk { + let change = match hunk { Hunk::AddFile { contents, .. } => { - changes.insert( - path.into_path_buf(), - ApplyPatchFileChange::Add { content: contents }, - ); + ApplyPatchFileChange::Add { content: contents } } Hunk::DeleteFile { .. } => { let content = match fs.read_file_text(&path, sandbox).await { Ok(content) => content, - Err(e) => { + Err(source) => { return MaybeApplyPatchVerified::CorrectnessError( ApplyPatchError::IoError(IoError { context: format!("Failed to read {}", path.display()), - source: e, + source, }), ); } }; - changes.insert( - path.into_path_buf(), - ApplyPatchFileChange::Delete { content }, - ); + ApplyPatchFileChange::Delete { content } } Hunk::UpdateFile { move_path, chunks, .. } => { - let ApplyPatchFileUpdate { + let update = + match unified_diff_from_chunks(&path, &chunks, fs, sandbox).await { + Ok(update) => update, + Err(err) => { + return MaybeApplyPatchVerified::CorrectnessError(err); + } + }; + let (unified_diff, new_content) = update.into_parts(); + ApplyPatchFileChange::Update { 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, - }, - ); + move_path: move_path + .map(|move_path| effective_cwd.join(move_path).into_path_buf()), + new_content, + } } - } + }; + changes.insert(path.into_path_buf(), change); } MaybeApplyPatchVerified::Body(ApplyPatchAction { changes, @@ -378,6 +369,7 @@ fn extract_apply_patch_from_bash( #[cfg(test)] mod tests { use super::*; + use crate::ApplyPatchFileUpdate; use crate::unified_diff_from_chunks; use assert_matches::assert_matches; use codex_exec_server::LOCAL_FS; diff --git a/codex-rs/apply-patch/src/lib.rs b/codex-rs/apply-patch/src/lib.rs index 99a63e3ace..b496b9f1b5 100644 --- a/codex-rs/apply-patch/src/lib.rs +++ b/codex-rs/apply-patch/src/lib.rs @@ -19,8 +19,10 @@ use codex_utils_absolute_path::AbsolutePathBuf; pub use parser::Hunk; pub use parser::ParseError; use parser::ParseError::*; +pub use parser::ParseOptions; pub use parser::UpdateFileChunk; pub use parser::parse_patch; +pub use parser::parse_patch_with_options; use similar::TextDiff; pub use streaming_parser::StreamingPatchParser; use thiserror::Error; @@ -96,6 +98,7 @@ impl PartialEq for IoError { pub struct ApplyPatchArgs { pub patch: String, pub hunks: Vec, + pub environment_id: Option, pub workdir: Option, } @@ -146,6 +149,18 @@ pub struct ApplyPatchAction { } impl ApplyPatchAction { + pub fn new( + changes: HashMap, + patch: String, + cwd: AbsolutePathBuf, + ) -> Self { + Self { + changes, + patch, + cwd, + } + } + pub fn is_empty(&self) -> bool { self.changes.is_empty() } @@ -815,6 +830,12 @@ pub struct ApplyPatchFileUpdate { content: String, } +impl ApplyPatchFileUpdate { + pub fn into_parts(self) -> (String, String) { + (self.unified_diff, self.content) + } +} + pub async fn unified_diff_from_chunks( path_abs: &AbsolutePathBuf, chunks: &[UpdateFileChunk], diff --git a/codex-rs/apply-patch/src/parser.rs b/codex-rs/apply-patch/src/parser.rs index 5403055272..bfe73e972b 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: " /(.+)/ 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: "; @@ -124,12 +126,24 @@ pub struct UpdateFileChunk { } pub fn parse_patch(patch: &str) -> Result { + parse_patch_with_options(patch, ParseOptions::default()) +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ParseOptions { + pub allow_environment_id: bool, +} + +pub fn parse_patch_with_options( + patch: &str, + options: ParseOptions, +) -> Result { let mode = if PARSE_IN_STRICT_MODE { ParseMode::Strict } else { ParseMode::Lenient }; - parse_patch_text(patch, mode) + parse_patch_text_with_options(patch, mode, options) } enum ParseMode { @@ -171,30 +185,63 @@ enum ParseMode { Lenient, } -fn parse_patch_text(patch: &str, mode: ParseMode) -> Result { +fn parse_patch_text_with_options( + patch: &str, + mode: ParseMode, + options: ParseOptions, +) -> Result { let lines: Vec<&str> = patch.trim().lines().collect(); let (patch_lines, hunk_lines) = match mode { ParseMode::Strict => check_patch_boundaries_strict(&lines)?, ParseMode::Lenient => check_patch_boundaries_lenient(&lines)?, }; + let (environment_id, hunk_lines, hunk_start_line_number) = + parse_environment_id(hunk_lines, options.allow_environment_id)?; let mut hunks: Vec = Vec::new(); let mut remaining_lines = hunk_lines; - let mut line_number = 2; + let mut line_number = hunk_start_line_number; while !remaining_lines.is_empty() { let (hunk, hunk_lines) = parse_one_hunk(remaining_lines, line_number)?; hunks.push(hunk); line_number += hunk_lines; remaining_lines = &remaining_lines[hunk_lines..] } - let patch = patch_lines.join("\n"); + let mut patch_lines_without_metadata = Vec::with_capacity(patch_lines.len()); + patch_lines_without_metadata.push(patch_lines[0]); + patch_lines_without_metadata.extend_from_slice(hunk_lines); + patch_lines_without_metadata.push(patch_lines[patch_lines.len() - 1]); + let patch = patch_lines_without_metadata.join("\n"); Ok(ApplyPatchArgs { hunks, patch, workdir: None, + environment_id, }) } +fn parse_environment_id<'a>( + lines: &'a [&'a str], + allow_environment_id: bool, +) -> Result<(Option, &'a [&'a str], usize), ParseError> { + let Some(first_line) = lines.first().map(|line| line.trim()) else { + return Ok((None, lines, 2)); + }; + if !allow_environment_id { + return Ok((None, lines, 2)); + } + let Some(environment_id) = first_line.strip_prefix(ENVIRONMENT_ID_MARKER) else { + return Ok((None, lines, 2)); + }; + let environment_id = environment_id.trim(); + if environment_id.is_empty() { + return Err(InvalidPatchError( + "Environment ID cannot be empty".to_string(), + )); + } + Ok((Some(environment_id.to_string()), &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>( @@ -557,20 +604,24 @@ fn test_update_file_chunk() { #[test] fn test_parse_patch() { assert_eq!( - parse_patch_text("bad", ParseMode::Strict), + parse_patch_text_with_options("bad", ParseMode::Strict, ParseOptions::default()), Err(InvalidPatchError( "The first line of the patch must be '*** Begin Patch'".to_string() )) ); assert_eq!( - parse_patch_text("*** Begin Patch\nbad", ParseMode::Strict), + parse_patch_text_with_options( + "*** Begin Patch\nbad", + ParseMode::Strict, + ParseOptions::default() + ), Err(InvalidPatchError( "The last line of the patch must be '*** End Patch'".to_string() )) ); assert_eq!( - parse_patch_text( + parse_patch_text_with_options( concat!( "*** Begin Patch", " ", @@ -578,7 +629,8 @@ fn test_parse_patch() { " ", "*** End Patch" ), - ParseMode::Strict + ParseMode::Strict, + ParseOptions::default() ) .unwrap() .hunks, @@ -588,11 +640,12 @@ fn test_parse_patch() { }] ); assert_eq!( - parse_patch_text( + parse_patch_text_with_options( "*** Begin Patch\n\ *** Update File: test.py\n\ *** End Patch", - ParseMode::Strict + ParseMode::Strict, + ParseOptions::default() ), Err(InvalidHunkError { message: "Update file hunk for path 'test.py' is empty".to_string(), @@ -600,17 +653,18 @@ fn test_parse_patch() { }) ); assert_eq!( - parse_patch_text( + parse_patch_text_with_options( "*** Begin Patch\n\ *** End Patch", - ParseMode::Strict + ParseMode::Strict, + ParseOptions::default() ) .unwrap() .hunks, Vec::new() ); assert_eq!( - parse_patch_text( + parse_patch_text_with_options( "*** Begin Patch\n\ *** Add File: path/add.py\n\ +abc\n\ @@ -622,7 +676,8 @@ fn test_parse_patch() { - pass\n\ + return 123\n\ *** End Patch", - ParseMode::Strict + ParseMode::Strict, + ParseOptions::default() ) .unwrap() .hunks, @@ -648,7 +703,7 @@ fn test_parse_patch() { ); // Update hunk followed by another hunk (Add File). assert_eq!( - parse_patch_text( + parse_patch_text_with_options( "*** Begin Patch\n\ *** Update File: file.py\n\ @@\n\ @@ -656,7 +711,8 @@ fn test_parse_patch() { *** Add File: other.py\n\ +content\n\ *** End Patch", - ParseMode::Strict + ParseMode::Strict, + ParseOptions::default() ) .unwrap() .hunks, @@ -681,13 +737,14 @@ fn test_parse_patch() { // Update hunk without an explicit @@ header for the first chunk should parse. // Use a raw string to preserve the leading space diff marker on the context line. assert_eq!( - parse_patch_text( + parse_patch_text_with_options( r#"*** Begin Patch *** Update File: file2.py import foo +bar *** End Patch"#, - ParseMode::Strict + ParseMode::Strict, + ParseOptions::default() ) .unwrap() .hunks, @@ -724,7 +781,7 @@ fn test_parse_patch_accepts_relative_and_absolute_hunk_paths() { ); assert_eq!( - parse_patch_text(&patch_text, ParseMode::Strict) + parse_patch_text_with_options(&patch_text, ParseMode::Strict, ParseOptions::default()) .unwrap() .hunks, vec![ @@ -828,64 +885,107 @@ fn test_parse_patch_lenient() { let patch_text_in_heredoc = format!("< Result<(), ParseError> { if let Some(UpdateFile { path, chunks, .. }) = self.state.hunks.last() { if chunks.is_empty() @@ -161,6 +168,18 @@ impl StreamingPatchParser { )) } StreamingParserMode::StartedPatch => { + if self.allow_environment_id + && !self.saw_environment_id + && let Some(environment_id) = trimmed.strip_prefix(ENVIRONMENT_ID_MARKER) + { + if environment_id.trim().is_empty() { + return Err(InvalidPatchError( + "Environment ID cannot be empty".to_string(), + )); + } + self.saw_environment_id = true; + return Ok(()); + } if self.handle_hunk_headers_and_end_patch(trimmed)? { return Ok(()); } diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index a89d8fc973..dbecd2b57f 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -44,6 +44,7 @@ use crate::mcp_tool_call::lookup_mcp_tool_metadata; use crate::session::Codex; use crate::session::CodexSpawnArgs; use crate::session::CodexSpawnOk; +use crate::session::PatchApprovalRequest; use crate::session::SUBMISSION_CHANNEL_CAPACITY; use crate::session::emit_subagent_session_started; use crate::session::session::Session; @@ -520,6 +521,8 @@ async fn handle_patch_approval( ) { let ApplyPatchApprovalRequestEvent { call_id, + environment_id, + cwd, changes, reason, grant_root, @@ -529,7 +532,7 @@ async fn handle_patch_approval( let guardian_decision = if routes_approval_to_guardian(parent_ctx) { let files = changes .keys() - .map(|path| parent_ctx.cwd.join(path)) + .map(|path| cwd.join(path)) .collect::>(); let review_cancel = cancel_token.child_token(); let patch = changes @@ -565,7 +568,8 @@ async fn handle_patch_approval( new_guardian_review_id(), GuardianApprovalRequest::ApplyPatch { id: approval_id.clone(), - cwd: parent_ctx.cwd.clone(), + environment_id: environment_id.clone(), + cwd: cwd.clone(), files, patch, }, @@ -590,7 +594,17 @@ async fn handle_patch_approval( decision } else { let decision_rx = parent_session - .request_patch_approval(parent_ctx, call_id, changes, reason, grant_root) + .request_patch_approval( + parent_ctx, + PatchApprovalRequest { + call_id, + environment_id, + cwd, + changes, + reason, + grant_root, + }, + ) .await; await_approval_with_cancel( async move { decision_rx.await.unwrap_or_default() }, diff --git a/codex-rs/core/src/guardian/approval_request.rs b/codex-rs/core/src/guardian/approval_request.rs index fba227834a..f015edad51 100644 --- a/codex-rs/core/src/guardian/approval_request.rs +++ b/codex-rs/core/src/guardian/approval_request.rs @@ -43,6 +43,7 @@ pub(crate) enum GuardianApprovalRequest { }, ApplyPatch { id: String, + environment_id: String, cwd: AbsolutePathBuf, files: Vec, patch: String, @@ -310,11 +311,13 @@ pub(crate) fn guardian_approval_request_to_json( }), GuardianApprovalRequest::ApplyPatch { id: _, + environment_id, cwd, files, patch, } => Ok(serde_json::json!({ "tool": "apply_patch", + "environment_id": environment_id, "cwd": cwd, "files": files, "patch": patch, diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 6d03dfa813..1d1f34b59b 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -686,6 +686,7 @@ fn format_guardian_action_pretty_truncates_large_string_fields() -> serde_json:: let patch = "line\n".repeat(100_000); let action = GuardianApprovalRequest::ApplyPatch { id: "patch-1".to_string(), + environment_id: "local".to_string(), cwd: test_path_buf("/tmp").abs(), files: Vec::new(), patch: patch.clone(), @@ -705,6 +706,7 @@ fn format_guardian_action_pretty_reports_no_truncation_for_small_payload() -> se { let action = GuardianApprovalRequest::ApplyPatch { id: "patch-1".to_string(), + environment_id: "local".to_string(), cwd: test_path_buf("/tmp").abs(), files: Vec::new(), patch: "line\n".to_string(), @@ -877,6 +879,7 @@ fn guardian_assessment_action_redacts_apply_patch_patch_text() { let file = test_path_buf("/tmp/guardian.txt").abs(); let action = GuardianApprovalRequest::ApplyPatch { id: "patch-1".to_string(), + environment_id: "local".to_string(), cwd: cwd.clone(), files: vec![file.clone()], patch: "*** Begin Patch\n*** Update File: guardian.txt\n@@\n+secret\n*** End Patch" @@ -906,6 +909,7 @@ fn guardian_request_turn_id_prefers_network_access_owner_turn() { }; let apply_patch = GuardianApprovalRequest::ApplyPatch { id: "patch-1".to_string(), + environment_id: "local".to_string(), cwd: test_path_buf("/tmp").abs(), files: vec![test_path_buf("/tmp/guardian.txt").abs()], patch: "*** Begin Patch\n*** Update File: guardian.txt\n@@\n+hello\n*** End Patch" @@ -958,6 +962,7 @@ async fn cancelled_guardian_review_emits_terminal_abort_without_warning() { "review-cancelled-guardian".to_string(), GuardianApprovalRequest::ApplyPatch { id: "patch-1".to_string(), + environment_id: "local".to_string(), cwd: test_path_buf("/tmp").abs(), files: vec![test_path_buf("/tmp/guardian.txt").abs()], patch: "*** Begin Patch\n*** Update File: guardian.txt\n@@\n+hello\n*** End Patch" diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 89c12aaf81..0ce6f7ee1c 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -414,6 +414,15 @@ pub(crate) struct CodexSpawnArgs { pub(crate) thread_store: Arc, } +pub(crate) struct PatchApprovalRequest { + pub(crate) call_id: String, + pub(crate) environment_id: String, + pub(crate) cwd: AbsolutePathBuf, + pub(crate) changes: HashMap, + pub(crate) reason: Option, + pub(crate) grant_root: Option, +} + pub(crate) const INITIAL_SUBMIT_ID: &str = ""; pub(crate) const SUBMISSION_CHANNEL_CAPACITY: usize = 512; const CYBER_VERIFY_URL: &str = "https://chatgpt.com/cyber"; @@ -1978,11 +1987,16 @@ impl Session { pub async fn request_patch_approval( &self, turn_context: &TurnContext, - call_id: String, - changes: HashMap, - reason: Option, - grant_root: Option, + request: PatchApprovalRequest, ) -> oneshot::Receiver { + let PatchApprovalRequest { + call_id, + environment_id, + cwd, + changes, + reason, + grant_root, + } = request; // Add the tx_approve callback to the map before sending the request. let (tx_approve, rx_approve) = oneshot::channel(); let approval_id = call_id.clone(); @@ -2004,6 +2018,8 @@ impl Session { call_id, turn_id: turn_context.sub_id.clone(), started_at_ms: now_unix_timestamp_ms(), + environment_id, + cwd, changes, reason, grant_root, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index d4fe30063f..a06394905e 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -295,6 +295,14 @@ impl TurnContext { pub(crate) fn file_system_sandbox_context( &self, additional_permissions: Option, + ) -> FileSystemSandboxContext { + self.file_system_sandbox_context_for_cwd(self.cwd.clone(), additional_permissions) + } + + pub(crate) fn file_system_sandbox_context_for_cwd( + &self, + cwd: AbsolutePathBuf, + additional_permissions: Option, ) -> FileSystemSandboxContext { let (base_file_system_sandbox_policy, base_network_sandbox_policy) = self.permission_profile.to_runtime_permissions(); @@ -313,7 +321,7 @@ impl TurnContext { ); FileSystemSandboxContext { permissions, - cwd: Some(self.cwd.clone()), + cwd: Some(cwd), windows_sandbox_level: self.windows_sandbox_level, windows_sandbox_private_desktop: self .config diff --git a/codex-rs/core/src/tools/handlers/apply_patch.lark b/codex-rs/core/src/tools/handlers/apply_patch.lark index 5aa41b0af7..c61987d02a 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.lark +++ b/codex-rs/core/src/tools/handlers/apply_patch.lark @@ -1,5 +1,6 @@ -start: begin_patch hunk+ end_patch +start: begin_patch environment_id? hunk+ end_patch begin_patch: "*** Begin Patch" LF +environment_id: "*** Environment ID: " /(.+)/ LF end_patch: "*** End Patch" LF? hunk: add_hunk | delete_hunk | update_hunk diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 2b63c1cb17..f57dd015e4 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -38,6 +38,7 @@ use crate::tools::sandboxing::ToolCtx; use codex_apply_patch::ApplyPatchAction; use codex_apply_patch::ApplyPatchFileChange; use codex_apply_patch::Hunk; +use codex_apply_patch::ParseOptions; use codex_apply_patch::StreamingPatchParser; use codex_exec_server::ExecutorFileSystem; use codex_features::Feature; @@ -50,6 +51,7 @@ use codex_protocol::protocol::PatchApplyUpdatedEvent; use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy; use codex_sandboxing::policy_transforms::merge_permission_profiles; use codex_sandboxing::policy_transforms::normalize_additional_permissions; +use codex_tools::ToolEnvironmentMode; use codex_tools::ToolName; use codex_tools::ToolSpec; use codex_utils_absolute_path::AbsolutePathBuf; @@ -57,21 +59,24 @@ use codex_utils_absolute_path::AbsolutePathBuf; const APPLY_PATCH_ARGUMENT_DIFF_BUFFER_INTERVAL: Duration = Duration::from_millis(500); pub struct ApplyPatchHandler { - options: ApplyPatchToolType, + tool_type: ApplyPatchToolType, + multi_environment: bool, } impl Default for ApplyPatchHandler { fn default() -> Self { Self { - options: ApplyPatchToolType::Freeform, + tool_type: ApplyPatchToolType::Freeform, + multi_environment: false, } } } impl ApplyPatchHandler { - pub(crate) fn new(apply_patch_tool_type: ApplyPatchToolType) -> Self { + pub(crate) fn new(apply_patch_tool_type: ApplyPatchToolType, multi_environment: bool) -> Self { Self { - options: apply_patch_tool_type, + tool_type: apply_patch_tool_type, + multi_environment, } } } @@ -93,6 +98,11 @@ impl ToolArgumentDiffConsumer for ApplyPatchArgumentDiffConsumer { if !turn.features.enabled(Feature::ApplyPatchStreamingEvents) { return None; } + let allow_environment_id = matches!( + turn.tools_config.environment_mode, + ToolEnvironmentMode::Multiple + ); + self.parser.set_allow_environment_id(allow_environment_id); self.push_delta(call_id, diff) .map(EventMsg::PatchApplyUpdated) @@ -278,10 +288,67 @@ fn apply_patch_payload_command(payload: &ToolPayload) -> Option { } } +fn reconcile_environment_id( + parsed_environment_id: Option<&str>, + function_environment_id: Option<&str>, + allow_environment_id: bool, +) -> Result, FunctionCallError> { + if let Some(environment_id) = function_environment_id + && environment_id.trim().is_empty() + { + return Err(FunctionCallError::RespondToModel( + "apply_patch environment_id cannot be empty".to_string(), + )); + } + if !allow_environment_id + && (parsed_environment_id.is_some() || function_environment_id.is_some()) + { + return Err(FunctionCallError::RespondToModel( + "apply_patch environment selection is unavailable for this turn".to_string(), + )); + } + + match (parsed_environment_id, function_environment_id) { + (Some(parsed), Some(function)) if parsed != function => { + Err(FunctionCallError::RespondToModel(format!( + "apply_patch environment_id mismatch: patch header requested `{parsed}` but tool arguments requested `{function}`" + ))) + } + (Some(environment_id), _) | (_, Some(environment_id)) => { + Ok(Some(environment_id.to_string())) + } + (None, None) => Ok(None), + } +} + +fn select_patch_environment<'a>( + turn: &'a TurnContext, + requested_environment_id: Option<&str>, +) -> Result<&'a crate::session::turn_context::TurnEnvironment, FunctionCallError> { + match requested_environment_id { + Some(environment_id) => turn + .environments + .turn_environments + .iter() + .find(|environment| environment.environment_id == environment_id) + .ok_or_else(|| { + FunctionCallError::RespondToModel(format!( + "unknown turn environment id `{environment_id}`" + )) + }), + None => turn.environments.primary().ok_or_else(|| { + FunctionCallError::RespondToModel( + "apply_patch is unavailable in this session".to_string(), + ) + }), + } +} + async fn effective_patch_permissions( session: &Session, turn: &TurnContext, action: &ApplyPatchAction, + sandbox_cwd: &AbsolutePathBuf, ) -> ( Vec, crate::tools::handlers::EffectiveAdditionalPermissions, @@ -299,9 +366,9 @@ async fn effective_patch_permissions( ); let effective_additional_permissions = apply_granted_turn_permissions( session, - turn.cwd.as_path(), + sandbox_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, sandbox_cwd), ) .await; @@ -320,9 +387,11 @@ impl ToolHandler for ApplyPatchHandler { } fn spec(&self) -> Option { - Some(match self.options { - ApplyPatchToolType::Freeform => create_apply_patch_freeform_tool(), - ApplyPatchToolType::Function => create_apply_patch_json_tool(), + Some(match self.tool_type { + ApplyPatchToolType::Freeform => { + create_apply_patch_freeform_tool(self.multi_environment) + } + ApplyPatchToolType::Function => create_apply_patch_json_tool(self.multi_environment), }) } @@ -380,122 +449,154 @@ impl ToolHandler for ApplyPatchHandler { .. } = invocation; - let patch_input = match payload { + let (patch_input, function_environment_id) = match payload { ToolPayload::Function { arguments } => { let args: ApplyPatchToolArgs = parse_arguments(&arguments)?; - args.input + (args.input, args.environment_id) } - ToolPayload::Custom { input } => input, + ToolPayload::Custom { input } => (input, None), _ => { return Err(FunctionCallError::RespondToModel( "apply_patch handler received unsupported payload".to_string(), )); } }; - - // 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 { - return Err(FunctionCallError::RespondToModel( - "apply_patch is unavailable in this session".to_string(), - )); - }; - 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(), + let allow_environment_id = matches!( + turn.tools_config.environment_mode, + ToolEnvironmentMode::Multiple + ); + let args = codex_apply_patch::parse_patch_with_options( + &patch_input, + ParseOptions { + allow_environment_id, + }, ) - .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; - match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes) - .await - { - InternalApplyPatchInvocation::Output(item) => { - let content = item?; - Ok(ApplyPatchToolOutput::from_text(content)) - } - InternalApplyPatchInvocation::DelegateToRuntime(apply) => { - let changes = convert_apply_patch_to_protocol(&apply.action); - let emitter = - ToolEmitter::apply_patch(changes.clone(), apply.auto_approved); - let event_ctx = ToolEventCtx::new( - session.as_ref(), - turn.as_ref(), - &call_id, - Some(&tracker), - ); - emitter.begin(event_ctx).await; - - let req = ApplyPatchRequest { - action: apply.action, - file_paths, - changes, - exec_approval_requirement: apply.exec_approval_requirement, - additional_permissions: effective_additional_permissions - .additional_permissions, - permissions_preapproved: effective_additional_permissions - .permissions_preapproved, - }; - - let mut orchestrator = ToolOrchestrator::new(); - let mut runtime = ApplyPatchRuntime::new(); - let tool_ctx = ToolCtx { - session: session.clone(), - turn: turn.clone(), - call_id: call_id.clone(), - tool_name: tool_name.display(), - }; - let out = orchestrator - .run( - &mut runtime, - &req, - &tool_ctx, - turn.as_ref(), - turn.approval_policy.value(), - ) + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "apply_patch verification failed: {}", + codex_apply_patch::ApplyPatchError::from(err) + )) + })?; + let requested_environment_id = reconcile_environment_id( + args.environment_id.as_deref(), + function_environment_id.as_deref(), + allow_environment_id, + )?; + let selected_environment = + select_patch_environment(turn.as_ref(), requested_environment_id.as_deref())?; + let fs = selected_environment.environment.get_filesystem(); + let sandbox = selected_environment.environment.is_remote().then(|| { + turn.file_system_sandbox_context_for_cwd( + selected_environment.cwd.clone(), + /*additional_permissions*/ None, + ) + }); + let codex_apply_patch::ApplyPatchArgs { + patch, + hunks, + environment_id: _, + workdir, + } = args; + let effective_cwd = workdir + .as_ref() + .map(|dir| selected_environment.cwd.join(Path::new(dir))) + .unwrap_or_else(|| selected_environment.cwd.clone()); + let mut change_map = HashMap::new(); + for hunk in hunks { + let path = hunk.resolve_path(&effective_cwd); + let change = match hunk { + Hunk::AddFile { contents, .. } => ApplyPatchFileChange::Add { content: contents }, + Hunk::DeleteFile { .. } => { + let content = + fs.read_file_text(&path, sandbox.as_ref()) .await - .map(|result| result.output); - let (out, delta) = match out { - Ok(output) => (Ok(output.exec_output), Some(output.delta)), - Err(error) => (Err(error), Some(runtime.committed_delta().clone())), - }; - let event_ctx = ToolEventCtx::new( - session.as_ref(), - turn.as_ref(), - &call_id, - Some(&tracker), - ); - let content = emitter.finish(event_ctx, out, delta.as_ref()).await?; - Ok(ApplyPatchToolOutput::from_text(content)) + .map_err(|source| { + FunctionCallError::RespondToModel(format!( + "apply_patch verification failed: Failed to read {}: {source}", + path.display() + )) + })?; + ApplyPatchFileChange::Delete { content } + } + Hunk::UpdateFile { + move_path, chunks, .. + } => { + let update = codex_apply_patch::unified_diff_from_chunks( + &path, + &chunks, + fs.as_ref(), + sandbox.as_ref(), + ) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "apply_patch verification failed: {err}" + )) + })?; + let (unified_diff, new_content) = update.into_parts(); + ApplyPatchFileChange::Update { + unified_diff, + move_path: move_path + .map(|move_path| effective_cwd.join(move_path).into_path_buf()), + new_content, } } + }; + change_map.insert(path.into_path_buf(), change); + } + let changes = ApplyPatchAction::new(change_map, patch, effective_cwd); + let (file_paths, effective_additional_permissions, file_system_sandbox_policy) = + effective_patch_permissions(session.as_ref(), turn.as_ref(), &changes, &changes.cwd) + .await; + match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes).await { + InternalApplyPatchInvocation::Output(item) => { + let content = item?; + Ok(ApplyPatchToolOutput::from_text(content)) } - codex_apply_patch::MaybeApplyPatchVerified::CorrectnessError(parse_error) => { - Err(FunctionCallError::RespondToModel(format!( - "apply_patch verification failed: {parse_error}" - ))) - } - codex_apply_patch::MaybeApplyPatchVerified::ShellParseError(error) => { - tracing::trace!("Failed to parse apply_patch input, {error:?}"); - Err(FunctionCallError::RespondToModel( - "apply_patch handler received invalid patch input".to_string(), - )) - } - codex_apply_patch::MaybeApplyPatchVerified::NotApplyPatch => { - Err(FunctionCallError::RespondToModel( - "apply_patch handler received non-apply_patch input".to_string(), - )) + InternalApplyPatchInvocation::DelegateToRuntime(apply) => { + let changes = convert_apply_patch_to_protocol(&apply.action); + let emitter = ToolEmitter::apply_patch(changes.clone(), apply.auto_approved); + let event_ctx = + ToolEventCtx::new(session.as_ref(), turn.as_ref(), &call_id, Some(&tracker)); + emitter.begin(event_ctx).await; + + let req = ApplyPatchRequest { + action: apply.action, + environment_id: selected_environment.environment_id.clone(), + file_paths, + changes, + exec_approval_requirement: apply.exec_approval_requirement, + additional_permissions: effective_additional_permissions.additional_permissions, + permissions_preapproved: effective_additional_permissions + .permissions_preapproved, + }; + + let mut orchestrator = ToolOrchestrator::new(); + let mut runtime = ApplyPatchRuntime::new(); + let tool_ctx = ToolCtx { + session: session.clone(), + turn: turn.clone(), + call_id: call_id.clone(), + tool_name: tool_name.display(), + }; + let out = orchestrator + .run( + &mut runtime, + &req, + &tool_ctx, + turn.as_ref(), + turn.approval_policy.value(), + ) + .await + .map(|result| result.output); + let event_ctx = + ToolEventCtx::new(session.as_ref(), turn.as_ref(), &call_id, Some(&tracker)); + let (out, delta) = match out { + Ok(output) => (Ok(output.exec_output), Some(output.delta)), + Err(error) => (Err(error), Some(runtime.committed_delta().clone())), + }; + let content = emitter.finish(event_ctx, out, delta.as_ref()).await?; + Ok(ApplyPatchToolOutput::from_text(content)) } } } @@ -530,7 +631,13 @@ 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, + &changes.cwd, + ) + .await; match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes) .await { @@ -551,6 +658,11 @@ pub(crate) async fn intercept_apply_patch( let req = ApplyPatchRequest { action: apply.action, + environment_id: turn + .environments + .primary() + .map(|environment| environment.environment_id.clone()) + .unwrap_or_default(), file_paths: approval_keys, changes, exec_approval_requirement: apply.exec_approval_requirement, 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 93a3ce4aac..d021fbdff4 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_spec.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_spec.rs @@ -8,6 +8,9 @@ use serde::Serialize; use std::collections::BTreeMap; const APPLY_PATCH_LARK_GRAMMAR: &str = include_str!("apply_patch.lark"); +const APPLY_PATCH_WITH_ENV_START: &str = "start: begin_patch environment_id? hunk+ end_patch"; +const APPLY_PATCH_STRICT_START: &str = "start: begin_patch hunk+ end_patch"; +const APPLY_PATCH_ENVIRONMENT_RULE: &str = "environment_id: \"*** Environment ID: \" /(.+)/ LF\n"; const APPLY_PATCH_JSON_TOOL_DESCRIPTION: &str = r#"Use the `apply_patch` tool to edit files. Your patch language is a stripped‑down, file‑oriented diff format designed to be easy to parse and safe to apply. You can think of it as a high‑level envelope: @@ -78,38 +81,64 @@ It is important to remember: - File references can only be relative, NEVER ABSOLUTE. "#; +const APPLY_PATCH_JSON_TOOL_MULTI_ENVIRONMENT_SUFFIX: &str = r#" + +When multiple environments are available, you may target a non-default environment by setting the `environment_id` function parameter. If the patch body also includes a `*** Environment ID: ...` header, it must match the `environment_id` parameter. +"#; + /// TODO(dylan): deprecate once we get rid of json tool #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ApplyPatchToolArgs { pub input: String, + pub environment_id: Option, } /// 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(multi_environment: bool) -> ToolSpec { 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: if multi_environment { + APPLY_PATCH_LARK_GRAMMAR.to_string() + } else { + APPLY_PATCH_LARK_GRAMMAR + .replace(APPLY_PATCH_WITH_ENV_START, APPLY_PATCH_STRICT_START) + .replace(APPLY_PATCH_ENVIRONMENT_RULE, "") + }, }, }) } /// Returns a json tool that can be used to edit files. Should only be used with gpt-oss models -pub fn create_apply_patch_json_tool() -> ToolSpec { - let properties = BTreeMap::from([( +pub fn create_apply_patch_json_tool(multi_environment: bool) -> ToolSpec { + let mut properties = BTreeMap::from([( "input".to_string(), JsonSchema::string(Some( "The entire contents of the apply_patch command".to_string(), )), )]); + if multi_environment { + properties.insert( + "environment_id".to_string(), + JsonSchema::string(Some( + "Optional environment id to target for this patch".to_string(), + )), + ); + } ToolSpec::Function(ResponsesApiTool { name: "apply_patch".to_string(), - description: APPLY_PATCH_JSON_TOOL_DESCRIPTION.to_string(), + description: if multi_environment { + format!( + "{APPLY_PATCH_JSON_TOOL_DESCRIPTION}{APPLY_PATCH_JSON_TOOL_MULTI_ENVIRONMENT_SUFFIX}" + ) + } else { + APPLY_PATCH_JSON_TOOL_DESCRIPTION.to_string() + }, strict: false, defer_loading: None, parameters: JsonSchema::object( 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 beda5cc916..80ff086b73 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 @@ -6,7 +6,7 @@ use std::collections::BTreeMap; #[test] fn create_apply_patch_freeform_tool_matches_expected_spec() { assert_eq!( - create_apply_patch_freeform_tool(), + create_apply_patch_freeform_tool(/*multi_environment*/ true), ToolSpec::Freeform(FreeformTool { name: "apply_patch".to_string(), description: @@ -24,19 +24,29 @@ fn create_apply_patch_freeform_tool_matches_expected_spec() { #[test] fn create_apply_patch_json_tool_matches_expected_spec() { assert_eq!( - create_apply_patch_json_tool(), + create_apply_patch_json_tool(/*multi_environment*/ true), ToolSpec::Function(ResponsesApiTool { name: "apply_patch".to_string(), - description: APPLY_PATCH_JSON_TOOL_DESCRIPTION.to_string(), + description: format!( + "{APPLY_PATCH_JSON_TOOL_DESCRIPTION}{APPLY_PATCH_JSON_TOOL_MULTI_ENVIRONMENT_SUFFIX}" + ), strict: false, defer_loading: None, parameters: JsonSchema::object( - BTreeMap::from([( - "input".to_string(), - JsonSchema::string(Some( - "The entire contents of the apply_patch command".to_string(), - ),), - )]), + BTreeMap::from([ + ( + "input".to_string(), + JsonSchema::string(Some( + "The entire contents of the apply_patch command".to_string(), + ),), + ), + ( + "environment_id".to_string(), + JsonSchema::string(Some( + "Optional environment id to target for this patch".to_string(), + ),), + ) + ]), Some(vec!["input".to_string()]), Some(false.into()) ), @@ -44,3 +54,26 @@ fn create_apply_patch_json_tool_matches_expected_spec() { }) ); } + +#[test] +fn strict_apply_patch_tools_do_not_advertise_multi_environment() { + let freeform = create_apply_patch_freeform_tool(/*multi_environment*/ false); + let ToolSpec::Freeform(freeform) = freeform else { + panic!("expected freeform tool"); + }; + let expected_strict = APPLY_PATCH_LARK_GRAMMAR + .replace(APPLY_PATCH_WITH_ENV_START, APPLY_PATCH_STRICT_START) + .replace(APPLY_PATCH_ENVIRONMENT_RULE, ""); + assert_eq!(freeform.format.definition, expected_strict); + + let json_tool = create_apply_patch_json_tool(/*multi_environment*/ false); + let ToolSpec::Function(json_tool) = json_tool else { + panic!("expected function tool"); + }; + assert_eq!(json_tool.description, APPLY_PATCH_JSON_TOOL_DESCRIPTION); + let properties = json_tool + .parameters + .properties + .expect("object properties should exist"); + assert!(!properties.contains_key("environment_id")); +} diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index 1ab249fb0a..c7a7c2cd74 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -6,6 +6,7 @@ use crate::exec::is_likely_sandbox_denied; use crate::guardian::GuardianApprovalRequest; use crate::guardian::review_approval_request; +use crate::session::PatchApprovalRequest; use crate::tools::hook_names::HookToolName; use crate::tools::sandboxing::Approvable; use crate::tools::sandboxing::ApprovalCtx; @@ -36,9 +37,16 @@ use futures::future::BoxFuture; use std::path::PathBuf; use std::time::Instant; +#[derive(serde::Serialize, Clone, Debug, Eq, PartialEq, Hash)] +pub(crate) struct ApplyPatchApprovalKey { + environment_id: String, + file_path: AbsolutePathBuf, +} + #[derive(Debug)] pub struct ApplyPatchRequest { pub action: ApplyPatchAction, + pub environment_id: String, pub file_paths: Vec, pub changes: std::collections::HashMap, pub exec_approval_requirement: ExecApprovalRequirement, @@ -72,6 +80,7 @@ impl ApplyPatchRuntime { ) -> GuardianApprovalRequest { GuardianApprovalRequest::ApplyPatch { id: call_id.to_string(), + environment_id: req.environment_id.clone(), cwd: req.action.cwd.clone(), files: req.file_paths.clone(), patch: req.action.patch.clone(), @@ -90,7 +99,7 @@ impl ApplyPatchRuntime { effective_permission_profile(attempt.permissions, req.additional_permissions.as_ref()); Some(FileSystemSandboxContext { permissions, - cwd: Some(attempt.sandbox_cwd.clone()), + cwd: Some(req.action.cwd.clone()), windows_sandbox_level: attempt.windows_sandbox_level, windows_sandbox_private_desktop: attempt.windows_sandbox_private_desktop, use_legacy_landlock: attempt.use_legacy_landlock, @@ -108,10 +117,16 @@ 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() + .map(|file_path| ApplyPatchApprovalKey { + environment_id: req.environment_id.clone(), + file_path: file_path.clone(), + }) + .collect() } fn start_approval_async<'a>( @@ -139,10 +154,14 @@ impl Approvable for ApplyPatchRuntime { let rx_approve = session .request_patch_approval( turn, - call_id, - changes.clone(), - Some(reason), - /*grant_root*/ None, + PatchApprovalRequest { + call_id, + environment_id: req.environment_id.clone(), + cwd: req.action.cwd.clone(), + changes: changes.clone(), + reason: Some(reason), + grant_root: None, + }, ) .await; return rx_approve.await.unwrap_or_default(); @@ -155,7 +174,15 @@ impl Approvable for ApplyPatchRuntime { || async move { let rx_approve = session .request_patch_approval( - turn, call_id, changes, /*reason*/ None, /*grant_root*/ None, + turn, + PatchApprovalRequest { + call_id, + environment_id: req.environment_id.clone(), + cwd: req.action.cwd.clone(), + changes, + reason: None, + grant_root: None, + }, ) .await; rx_approve.await.unwrap_or_default() @@ -192,7 +219,10 @@ impl Approvable for ApplyPatchRuntime { ) -> Option { Some(PermissionRequestPayload { tool_name: HookToolName::apply_patch(), - tool_input: serde_json::json!({ "command": req.action.patch }), + tool_input: serde_json::json!({ + "command": req.action.patch.clone(), + "environment_id": req.environment_id.clone(), + }), }) } } @@ -204,9 +234,18 @@ impl ToolRuntime for ApplyPatchRunti attempt: &SandboxAttempt<'_>, 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 turn_environment = ctx + .turn + .environments + .turn_environments + .iter() + .find(|environment| environment.environment_id == req.environment_id) + .ok_or_else(|| { + ToolError::Rejected(format!( + "apply_patch environment `{}` is unavailable in this session", + req.environment_id + )) + })?; let started_at = Instant::now(); let fs = turn_environment.environment.get_filesystem(); let sandbox = Self::file_system_sandbox_context_for_attempt(req, attempt); 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..66afbbb217 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs @@ -50,6 +50,7 @@ fn guardian_review_request_includes_patch_context() { let expected_patch = action.patch.clone(); let request = ApplyPatchRequest { action, + environment_id: "local".to_string(), file_paths: vec![path.clone()], changes: HashMap::from([( path.to_path_buf(), @@ -71,6 +72,7 @@ fn guardian_review_request_includes_patch_context() { guardian_request, GuardianApprovalRequest::ApplyPatch { id: "call-1".to_string(), + environment_id: "local".to_string(), cwd: expected_cwd, files: request.file_paths, patch: expected_patch, @@ -78,6 +80,34 @@ fn guardian_review_request_includes_patch_context() { ); } +#[test] +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 { + action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), + environment_id: "remote".to_string(), + file_paths: vec![path.clone()], + changes: HashMap::new(), + exec_approval_requirement: ExecApprovalRequirement::NeedsApproval { + reason: None, + proposed_execpolicy_amendment: None, + }, + additional_permissions: None, + permissions_preapproved: false, + }; + + assert_eq!( + runtime.approval_keys(&req), + vec![ApplyPatchApprovalKey { + environment_id: "remote".to_string(), + file_path: path, + }] + ); +} + #[test] fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { let runtime = ApplyPatchRuntime::new(); @@ -88,6 +118,7 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { let expected_patch = action.patch.clone(); let req = ApplyPatchRequest { action, + environment_id: "local".to_string(), file_paths: vec![path], changes: HashMap::new(), exec_approval_requirement: ExecApprovalRequirement::NeedsApproval { @@ -109,7 +140,10 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { ); assert_eq!( payload.tool_input, - serde_json::json!({ "command": expected_patch }) + serde_json::json!({ + "command": expected_patch, + "environment_id": "local", + }) ); } @@ -127,6 +161,7 @@ fn file_system_sandbox_context_uses_active_attempt() { }; let req = ApplyPatchRequest { action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), + environment_id: "local".to_string(), file_paths: vec![path.clone()], changes: HashMap::new(), exec_approval_requirement: ExecApprovalRequirement::Skip { @@ -168,7 +203,7 @@ fn file_system_sandbox_context_uses_active_attempt() { let expected_permissions = PermissionProfile::from_runtime_permissions(&file_system_policy, network_policy); assert_eq!(sandbox.permissions, expected_permissions); - assert_eq!(sandbox.cwd, Some(path.clone())); + assert_eq!(sandbox.cwd, Some(req.action.cwd)); assert_eq!( sandbox.windows_sandbox_level, WindowsSandboxLevel::RestrictedToken @@ -184,6 +219,7 @@ fn no_sandbox_attempt_has_no_file_system_context() { .abs(); let req = ApplyPatchRequest { action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), + environment_id: "local".to_string(), file_paths: vec![path.clone()], changes: HashMap::new(), exec_approval_requirement: ExecApprovalRequirement::Skip { diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index ddad2cbf4e..9094e71217 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -260,8 +260,10 @@ pub fn build_tool_registry_builder( if config.environment_mode.has_environment() && let Some(apply_patch_tool_type) = &config.apply_patch_tool_type { + let multi_environment = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); builder.register_handler(Arc::new(ApplyPatchHandler::new( apply_patch_tool_type.clone(), + multi_environment, ))); } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 3fa8797c50..abb850f59f 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -118,7 +118,10 @@ 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(matches!( + config.environment_mode, + ToolEnvironmentMode::Multiple + )), ToolSpec::WebSearch { external_web_access: Some(true), filters: None, diff --git a/codex-rs/core/tests/suite/tool_harness.rs b/codex-rs/core/tests/suite/tool_harness.rs index a69ec3f7f6..5fc6c40577 100644 --- a/codex-rs/core/tests/suite/tool_harness.rs +++ b/codex-rs/core/tests/suite/tool_harness.rs @@ -3,6 +3,8 @@ use std::fs; use assert_matches::assert_matches; +use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_features::Feature; use codex_protocol::items::TurnItem; use codex_protocol::models::PermissionProfile; @@ -10,10 +12,13 @@ use codex_protocol::plan_tool::StepStatus; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; +use core_test_support::PathBufExt; use core_test_support::assert_regex_match; use core_test_support::responses; use core_test_support::responses::ResponsesRequest; +use core_test_support::responses::ev_apply_patch_custom_tool_call; use core_test_support::responses::ev_apply_patch_function_call; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -440,6 +445,157 @@ A {file_name} Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn apply_patch_multi_environment_uses_remote_executor() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::ApplyPatchFreeform) + .expect("test config should allow feature update"); + }); + let test = builder.build_remote_aware(&server).await?; + + let local_cwd = test.cwd.path().to_path_buf().abs(); + let remote_cwd = test.config.cwd.clone(); + let turn_environments = vec![ + TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: local_cwd.clone(), + }, + TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: remote_cwd.clone(), + }, + ]; + + let function_file_name = "remote-function-target.txt"; + let function_local_path = test.cwd.path().join(function_file_name); + let function_remote_path = remote_cwd.join(function_file_name); + let function_call_id = "apply-patch-remote-function"; + let function_patch = format!( + r#"*** Begin Patch +*** Add File: {function_file_name} ++patched via remote function +*** End Patch"# + ); + let function_arguments = serde_json::to_string(&json!({ + "input": function_patch, + "environment_id": REMOTE_ENVIRONMENT_ID, + }))?; + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(function_call_id, "apply_patch", &function_arguments), + ev_completed("resp-1"), + ]), + ) + .await; + let function_output_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-1", "function patch complete"), + ev_completed("resp-2"), + ]), + ) + .await; + + test.submit_turn_with_environments( + "apply the patch to the remote environment via function arguments", + Some(turn_environments.clone()), + ) + .await?; + + assert!( + !function_local_path.exists(), + "function-call targeted patch should not write locally" + ); + let function_remote_contents = test.fs().read_file(&function_remote_path, None).await?; + assert_eq!( + String::from_utf8(function_remote_contents)?, + "patched via remote function\n" + ); + + let (function_output_text, function_success_flag) = + call_output(&function_output_mock.single_request(), function_call_id); + assert!( + function_output_text.contains("Success."), + "expected success output for remote function apply_patch, got {function_output_text:?}" + ); + if let Some(function_success_flag) = function_success_flag { + assert!(function_success_flag, "expected success=true when provided"); + } + + let freeform_file_name = "remote-freeform-target.txt"; + let freeform_local_path = test.cwd.path().join(freeform_file_name); + let freeform_remote_path = remote_cwd.join(freeform_file_name); + let freeform_call_id = "apply-patch-remote-freeform"; + let freeform_patch = format!( + r#"*** Begin Patch +*** Environment ID: {REMOTE_ENVIRONMENT_ID} +*** Add File: {freeform_file_name} ++patched via remote freeform +*** End Patch"# + ); + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-3"), + ev_apply_patch_custom_tool_call(freeform_call_id, &freeform_patch), + ev_completed("resp-3"), + ]), + ) + .await; + let freeform_output_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-2", "freeform patch complete"), + ev_completed("resp-4"), + ]), + ) + .await; + + test.submit_turn_with_environments( + "apply the patch to the remote environment via freeform header", + Some(turn_environments), + ) + .await?; + + assert!( + !freeform_local_path.exists(), + "freeform targeted patch should not write locally" + ); + let freeform_remote_contents = test.fs().read_file(&freeform_remote_path, None).await?; + assert_eq!( + String::from_utf8(freeform_remote_contents)?, + "patched via remote freeform\n" + ); + + let freeform_output = freeform_output_mock + .single_request() + .custom_tool_call_output(freeform_call_id); + assert_eq!( + freeform_output.get("call_id").and_then(Value::as_str), + Some(freeform_call_id) + ); + let freeform_output_text = freeform_output + .get("output") + .and_then(Value::as_str) + .unwrap_or_default(); + assert!( + freeform_output_text.contains("Success."), + "expected success output for remote freeform apply_patch, got {freeform_output_text:?}" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn apply_patch_reports_parse_diagnostics() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index 4e2d5c08e5..d2972457f0 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -280,6 +280,8 @@ async fn run_codex_tool_session_inner( call_id, turn_id: _, started_at_ms: _, + environment_id: _, + cwd: _, reason, grant_root, changes, diff --git a/codex-rs/protocol/src/approvals.rs b/codex-rs/protocol/src/approvals.rs index ace096359c..dfb9579a73 100644 --- a/codex-rs/protocol/src/approvals.rs +++ b/codex-rs/protocol/src/approvals.rs @@ -380,6 +380,10 @@ pub struct ApplyPatchApprovalRequestEvent { pub turn_id: String, #[ts(type = "number")] pub started_at_ms: i64, + /// Stable selected environment id for this patch apply. + #[serde(default)] + pub environment_id: String, + pub cwd: AbsolutePathBuf, pub changes: HashMap, /// Optional explanatory reason (e.g. request for extra write access). #[serde(skip_serializing_if = "Option::is_none")] diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 1a7986ad31..73925b667b 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -249,6 +249,7 @@ impl App { thread_id, thread_label, id: params.item_id.clone(), + environment_id: None, reason: params.reason.clone(), cwd: self .thread_cwd(thread_id) diff --git a/codex-rs/tui/src/approval_events.rs b/codex-rs/tui/src/approval_events.rs index 57e0959d56..1e95b26109 100644 --- a/codex-rs/tui/src/approval_events.rs +++ b/codex-rs/tui/src/approval_events.rs @@ -111,6 +111,9 @@ pub(crate) struct ApplyPatchApprovalRequestEvent { pub(crate) call_id: String, #[serde(default)] pub(crate) turn_id: String, + #[serde(default)] + pub(crate) environment_id: String, + pub(crate) cwd: AbsolutePathBuf, pub(crate) changes: HashMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub(crate) reason: Option, diff --git a/codex-rs/tui/src/bottom_pane/approval_overlay.rs b/codex-rs/tui/src/bottom_pane/approval_overlay.rs index 1089d3668e..06fbcc9c87 100644 --- a/codex-rs/tui/src/bottom_pane/approval_overlay.rs +++ b/codex-rs/tui/src/bottom_pane/approval_overlay.rs @@ -90,6 +90,7 @@ pub(crate) enum ApprovalRequest { thread_id: ThreadId, thread_label: Option, id: String, + environment_id: Option, reason: Option, cwd: AbsolutePathBuf, changes: HashMap, @@ -692,6 +693,7 @@ fn build_header(request: &ApprovalRequest) -> Box { } ApprovalRequest::ApplyPatch { thread_label, + environment_id, reason, .. } => { @@ -716,6 +718,15 @@ fn build_header(request: &ApprovalRequest) -> Box { .wrap(Wrap { trim: false }), )); } + if let Some(environment_id) = environment_id { + if !header.is_empty() { + header.push(Box::new(Line::from(""))); + } + header.push(Box::new(Line::from(vec![ + "Environment: ".into(), + environment_id.clone().bold(), + ]))); + } Box::new(ColumnRenderable::with(header)) } ApprovalRequest::McpElicitation { @@ -1979,6 +1990,7 @@ mod tests { thread_id: ThreadId::new(), thread_label: Some("Banach [worker]".to_string()), id: "test".to_string(), + environment_id: None, reason: None, cwd: absolute_path("/tmp"), changes, diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 35d523aa5b..94cd1d375a 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -1621,10 +1621,13 @@ fn exec_approval_request_from_params( fn patch_approval_request_from_params( params: FileChangeRequestApprovalParams, + fallback_cwd: &AbsolutePathBuf, ) -> ApplyPatchApprovalRequestEvent { ApplyPatchApprovalRequestEvent { call_id: params.item_id, turn_id: params.turn_id, + environment_id: String::new(), + cwd: fallback_cwd.clone(), changes: HashMap::new(), reason: params.reason, grant_root: params.grant_root, @@ -4545,15 +4548,16 @@ impl ChatWidget { thread_id: self.thread_id.unwrap_or_default(), thread_label: None, id: ev.call_id, + environment_id: (!ev.environment_id.is_empty()).then_some(ev.environment_id.clone()), reason: ev.reason, changes: ev.changes.clone(), - cwd: self.config.cwd.clone(), + cwd: ev.cwd.clone(), }; self.bottom_pane .push_approval_request(request, &self.config.features); self.request_redraw(); self.notify(Notification::EditApprovalRequested { - cwd: self.config.cwd.to_path_buf(), + cwd: ev.cwd.to_path_buf(), changes: ev.changes.keys().cloned().collect(), }); } @@ -6190,9 +6194,10 @@ impl ChatWidget { ); } ServerRequest::FileChangeRequestApproval { params, .. } => { + let fallback_cwd = self.config.cwd.clone(); self.on_apply_patch_approval_request( id, - patch_approval_request_from_params(params), + patch_approval_request_from_params(params, &fallback_cwd), ); } ServerRequest::McpServerElicitationRequest { request_id, params } => { diff --git a/codex-rs/tui/src/chatwidget/slash_dispatch.rs b/codex-rs/tui/src/chatwidget/slash_dispatch.rs index 6d1278ea2d..409dbb4115 100644 --- a/codex-rs/tui/src/chatwidget/slash_dispatch.rs +++ b/codex-rs/tui/src/chatwidget/slash_dispatch.rs @@ -430,11 +430,17 @@ impl ChatWidget { use crate::approval_events::ApplyPatchApprovalRequestEvent; use crate::diff_model::FileChange; + let Ok(cwd) = std::env::temp_dir().try_into() else { + self.add_error_message("Temp dir is not absolute.".to_string()); + return; + }; self.on_apply_patch_approval_request( "1".to_string(), ApplyPatchApprovalRequestEvent { call_id: "1".to_string(), turn_id: "turn-1".to_string(), + environment_id: "local".to_string(), + cwd, changes: HashMap::from([ ( PathBuf::from("/tmp/test.txt"), diff --git a/codex-rs/tui/src/chatwidget/tests/exec_flow.rs b/codex-rs/tui/src/chatwidget/tests/exec_flow.rs index 0045e7d126..9851baa7da 100644 --- a/codex-rs/tui/src/chatwidget/tests/exec_flow.rs +++ b/codex-rs/tui/src/chatwidget/tests/exec_flow.rs @@ -1215,6 +1215,10 @@ async fn approval_modal_patch_snapshot() -> anyhow::Result<()> { let ev = ApplyPatchApprovalRequestEvent { call_id: "call-approve-patch".into(), turn_id: "turn-approve-patch".into(), + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes, reason: Some("The model wants to apply changes".into()), grant_root: Some(PathBuf::from("/tmp")), @@ -1335,6 +1339,10 @@ async fn apply_patch_events_emit_history_cells() { let ev = ApplyPatchApprovalRequestEvent { call_id: "c1".into(), turn_id: "turn-c1".into(), + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes, reason: None, grant_root: None, @@ -1401,6 +1409,10 @@ async fn apply_patch_manual_approval_adjusts_header() { ApplyPatchApprovalRequestEvent { call_id: "c1".into(), turn_id: "turn-c1".into(), + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes: proposed_changes, reason: None, grant_root: None, @@ -1443,6 +1455,10 @@ async fn apply_patch_manual_flow_snapshot() { ApplyPatchApprovalRequestEvent { call_id: "c1".into(), turn_id: "turn-c1".into(), + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes: proposed_changes, reason: Some("Manual review required".into()), grant_root: None, @@ -1486,6 +1502,10 @@ async fn apply_patch_approval_sends_op_with_call_id() { let ev = ApplyPatchApprovalRequestEvent { call_id: "call-999".into(), turn_id: "turn-999".into(), + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes, reason: None, grant_root: None, @@ -1531,6 +1551,10 @@ async fn apply_patch_full_flow_integration_like() { ApplyPatchApprovalRequestEvent { call_id: "call-1".into(), turn_id: "turn-call-1".into(), + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes, reason: None, grant_root: None, @@ -1606,6 +1630,10 @@ async fn apply_patch_untrusted_shows_approval_modal() -> anyhow::Result<()> { ApplyPatchApprovalRequestEvent { call_id: "call-1".into(), turn_id: "turn-call-1".into(), + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes, reason: None, grant_root: None, @@ -1661,6 +1689,10 @@ async fn apply_patch_request_omits_diff_summary_from_modal() -> anyhow::Result<( ApplyPatchApprovalRequestEvent { call_id: "call-apply".into(), turn_id: "turn-apply".into(), + environment_id: "local".into(), + cwd: std::env::temp_dir() + .try_into() + .expect("temp dir should be absolute"), changes, reason: None, grant_root: None,