From 32ae8e03c87080cd0786c278525699b00fa2d1c5 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 15:51:09 -0700 Subject: [PATCH] codex: route apply_patch environment metadata Co-authored-by: Codex --- codex-rs/apply-patch/src/invocation.rs | 1 + codex-rs/apply-patch/src/lib.rs | 1 + codex-rs/apply-patch/src/parser.rs | 53 +++++++++++- codex-rs/core/src/apply_patch.rs | 2 +- .../core/src/tools/handlers/apply_patch.rs | 84 +++++++++++++++---- .../core/src/tools/runtimes/apply_patch.rs | 17 ++-- .../src/tools/runtimes/apply_patch_tests.rs | 31 +++++-- codex-rs/tools/src/apply_patch_tool.rs | 2 + codex-rs/tools/src/tool_apply_patch.lark | 4 +- 9 files changed, 160 insertions(+), 35 deletions(-) diff --git a/codex-rs/apply-patch/src/invocation.rs b/codex-rs/apply-patch/src/invocation.rs index 075c94c60c..660221f950 100644 --- a/codex-rs/apply-patch/src/invocation.rs +++ b/codex-rs/apply-patch/src/invocation.rs @@ -155,6 +155,7 @@ pub async fn maybe_parse_apply_patch_verified( patch, hunks, workdir, + environment_id: _, }) => { let effective_cwd = workdir .as_ref() diff --git a/codex-rs/apply-patch/src/lib.rs b/codex-rs/apply-patch/src/lib.rs index 7a47b1ea48..62e7366c31 100644 --- a/codex-rs/apply-patch/src/lib.rs +++ b/codex-rs/apply-patch/src/lib.rs @@ -97,6 +97,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..2e11a92dc0 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? hunk+ end_patch //! begin_patch: "*** Begin Patch" LF +//! environment: "*** Environment ID: " environment_id LF //! end_patch: "*** End Patch" LF? //! //! hunk: add_hunk | delete_hunk | update_hunk @@ -40,6 +41,7 @@ pub(crate) const MOVE_TO_MARKER: &str = "*** Move to: "; pub(crate) const EOF_MARKER: &str = "*** End of File"; pub(crate) const CHANGE_CONTEXT_MARKER: &str = "@@ "; pub(crate) const EMPTY_CHANGE_CONTEXT_MARKER: &str = "@@"; +const ENVIRONMENT_ID_MARKER: &str = "*** Environment ID: "; /// Currently, the only OpenAI model that knowingly requires lenient parsing is /// gpt-4.1. While we could try to require everyone to pass in a strictness @@ -177,10 +179,29 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result check_patch_boundaries_strict(&lines)?, ParseMode::Lenient => check_patch_boundaries_lenient(&lines)?, }; + let mut patch_lines = patch_lines.to_vec(); + let mut hunk_lines = hunk_lines; + let mut line_number = 2; + let environment_id = match hunk_lines.first() { + Some(line) => match line.strip_prefix(ENVIRONMENT_ID_MARKER) { + Some("") => { + return Err(InvalidPatchError( + "environment_id cannot be empty".to_string(), + )); + } + Some(environment_id) => { + patch_lines.remove(1); + hunk_lines = &hunk_lines[1..]; + line_number += 1; + Some(environment_id.to_string()) + } + None => None, + }, + None => None, + }; 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,6 +213,7 @@ fn parse_patch_text(patch: &str, mode: ParseMode) -> Result Result { + codex_apply_patch::parse_patch(&input).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "apply_patch verification failed: {}", + codex_apply_patch::ApplyPatchError::from(err) + )) + }) +} + fn convert_apply_patch_hunks_to_protocol(hunks: &[Hunk]) -> HashMap { hunks .iter() @@ -259,6 +270,7 @@ fn apply_patch_payload_command(payload: &ToolPayload) -> Option { async fn effective_patch_permissions( session: &Session, turn: &TurnContext, + cwd: &AbsolutePathBuf, action: &ApplyPatchAction, ) -> ( Vec, @@ -277,9 +289,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; @@ -351,12 +363,28 @@ impl ToolHandler for ApplyPatchHandler { .. } = invocation; - let patch_input = match payload { + let (patch_input, environment_id) = match payload { ToolPayload::Function { arguments } => { let args: ApplyPatchToolArgs = parse_arguments(&arguments)?; - args.input + let parsed_input = parse_apply_patch_input(args.input)?; + let environment_id = match (args.environment_id, parsed_input.environment_id) { + (Some(argument_environment_id), Some(header_environment_id)) + if argument_environment_id != header_environment_id => + { + return Err(FunctionCallError::RespondToModel( + "apply_patch environment_id argument conflicts with patch header metadata" + .to_string(), + )); + } + (Some(argument_environment_id), _) => Some(argument_environment_id), + (None, header_environment_id) => header_environment_id, + }; + (parsed_input.patch, environment_id) + } + ToolPayload::Custom { input } => { + let parsed_input = parse_apply_patch_input(input)?; + (parsed_input.patch, parsed_input.environment_id) } - ToolPayload::Custom { input } => input, _ => { return Err(FunctionCallError::RespondToModel( "apply_patch handler received unsupported payload".to_string(), @@ -366,18 +394,25 @@ impl ToolHandler for ApplyPatchHandler { // 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 { + let Some(turn_environment) = turn + .environments + .get_or_primary(environment_id.as_deref()) + .cloned() + 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)); + let cwd = turn_environment.cwd.clone(); + let command = vec!["apply_patch".to_string(), patch_input.clone()]; + let environment = Arc::clone(&turn_environment.environment); + let fs = environment.get_filesystem(); + let sandbox = turn_environment.environment.is_remote().then(|| { + let mut context = + turn.file_system_sandbox_context(/*additional_permissions*/ None); + context.cwd = Some(cwd.clone()); + context + }); match codex_apply_patch::maybe_parse_apply_patch_verified( &command, &cwd, @@ -388,7 +423,13 @@ impl ToolHandler for ApplyPatchHandler { { 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, + &changes, + ) + .await; match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes) .await { @@ -410,6 +451,7 @@ impl ToolHandler for ApplyPatchHandler { let req = ApplyPatchRequest { action: apply.action, + environment, file_paths, changes, exec_approval_requirement: apply.exec_approval_requirement, @@ -497,7 +539,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.cwd, + &changes, + ) + .await; match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes) .await { @@ -506,6 +554,11 @@ pub(crate) async fn intercept_apply_patch( Ok(Some(FunctionToolOutput::from_text(content, Some(true)))) } InternalApplyPatchInvocation::DelegateToRuntime(apply) => { + let Some(turn_environment) = turn.environments.primary() else { + return Err(FunctionCallError::RespondToModel( + "apply_patch is unavailable in this session".to_string(), + )); + }; let changes = convert_apply_patch_to_protocol(&apply.action); let emitter = ToolEmitter::apply_patch(changes.clone(), apply.auto_approved); let event_ctx = ToolEventCtx::new( @@ -518,6 +571,7 @@ pub(crate) async fn intercept_apply_patch( let req = ApplyPatchRequest { action: apply.action, + environment: turn_environment.environment.clone(), file_paths: approval_keys, changes, exec_approval_requirement: apply.exec_approval_requirement, diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index fb330fa913..abf9bfad4f 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -18,6 +18,7 @@ use crate::tools::sandboxing::ToolError; use crate::tools::sandboxing::ToolRuntime; use crate::tools::sandboxing::with_cached_approval; use codex_apply_patch::ApplyPatchAction; +use codex_exec_server::Environment; use codex_exec_server::FileSystemSandboxContext; use codex_protocol::error::CodexErr; use codex_protocol::error::SandboxErr; @@ -33,11 +34,12 @@ use codex_sandboxing::policy_transforms::effective_permission_profile; use codex_utils_absolute_path::AbsolutePathBuf; use futures::future::BoxFuture; use std::path::PathBuf; +use std::sync::Arc; use std::time::Instant; -#[derive(Debug)] pub struct ApplyPatchRequest { pub action: ApplyPatchAction, + pub environment: Arc, pub file_paths: Vec, pub changes: std::collections::HashMap, pub exec_approval_requirement: ExecApprovalRequirement, @@ -77,7 +79,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, @@ -185,17 +187,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.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..186ab6f110 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::tools::sandboxing::SandboxAttempt; +use codex_exec_server::Environment; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::models::FileSystemPermissions; @@ -15,6 +16,14 @@ use codex_sandboxing::policy_transforms::effective_network_sandbox_policy; use core_test_support::PathBufExt; use pretty_assertions::assert_eq; use std::collections::HashMap; +use std::sync::Arc; + +fn test_environment() -> Arc { + Arc::new( + Environment::create_for_tests(Some("ws://127.0.0.1:0".to_string())) + .expect("test environment"), + ) +} #[test] fn wants_no_sandbox_approval_granular_respects_sandbox_flag() { @@ -40,8 +49,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(); @@ -50,6 +59,7 @@ fn guardian_review_request_includes_patch_context() { let expected_patch = action.patch.clone(); let request = ApplyPatchRequest { action, + environment: test_environment(), file_paths: vec![path.clone()], changes: HashMap::from([( path.to_path_buf(), @@ -78,8 +88,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") @@ -88,6 +98,7 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { let expected_patch = action.patch.clone(); let req = ApplyPatchRequest { action, + environment: test_environment(), file_paths: vec![path], changes: HashMap::new(), exec_approval_requirement: ExecApprovalRequirement::NeedsApproval { @@ -113,8 +124,8 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() { ); } -#[test] -fn file_system_sandbox_context_uses_active_attempt() { +#[tokio::test] +async fn file_system_sandbox_context_uses_active_attempt() { let path = std::env::temp_dir() .join("apply-patch-runtime-attempt.txt") .abs(); @@ -127,6 +138,7 @@ fn file_system_sandbox_context_uses_active_attempt() { }; let req = ApplyPatchRequest { action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), + environment: test_environment(), file_paths: vec![path.clone()], changes: HashMap::new(), exec_approval_requirement: ExecApprovalRequirement::Skip { @@ -168,7 +180,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.clone())); assert_eq!( sandbox.windows_sandbox_level, WindowsSandboxLevel::RestrictedToken @@ -177,13 +189,14 @@ 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 { action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()), + environment: test_environment(), file_paths: vec![path.clone()], changes: HashMap::new(), exec_approval_requirement: ExecApprovalRequirement::Skip { diff --git a/codex-rs/tools/src/apply_patch_tool.rs b/codex-rs/tools/src/apply_patch_tool.rs index 469bb52367..8ba8eaad25 100644 --- a/codex-rs/tools/src/apply_patch_tool.rs +++ b/codex-rs/tools/src/apply_patch_tool.rs @@ -82,6 +82,8 @@ It is important to remember: #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ApplyPatchToolArgs { pub input: String, + #[serde(default)] + pub environment_id: Option, } /// Returns a custom tool that can be used to edit files. Well-suited for GPT-5 models diff --git a/codex-rs/tools/src/tool_apply_patch.lark b/codex-rs/tools/src/tool_apply_patch.lark index 5aa41b0af7..b2d82c7edd 100644 --- a/codex-rs/tools/src/tool_apply_patch.lark +++ b/codex-rs/tools/src/tool_apply_patch.lark @@ -1,5 +1,6 @@ -start: begin_patch hunk+ end_patch +start: begin_patch environment? hunk+ end_patch begin_patch: "*** Begin Patch" LF +environment: "*** Environment ID: " environment_id LF end_patch: "*** End Patch" LF? hunk: add_hunk | delete_hunk | update_hunk @@ -8,6 +9,7 @@ delete_hunk: "*** Delete File: " filename LF update_hunk: "*** Update File: " filename LF change_move? change? filename: /(.+)/ +environment_id: /(.+)/ add_line: "+" /(.*)/ LF -> line change_move: "*** Move to: " filename LF