codex: address apply_patch target propagation feedback

Co-authored-by: Codex <noreply@openai.com>

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 {

codex-rs/core/src/codex_delegate.rs

@@ -528,6 +528,7 @@ async fn handle_patch_approval(
         grant_root,
         ..
     } = event;
+    let cwd = cwd.unwrap_or_else(|| parent_ctx.cwd.clone());
     let approval_id = call_id.clone();
     let guardian_decision = if routes_approval_to_guardian(parent_ctx) {
         let files = changes

codex-rs/core/src/session/mod.rs

@@ -2019,7 +2019,7 @@ impl Session {
             turn_id: turn_context.sub_id.clone(),
             started_at_ms: now_unix_timestamp_ms(),
             environment_id,
-            cwd,
+            cwd: Some(cwd),
             changes,
             reason,
             grant_root,

codex-rs/core/src/tools/handlers/apply_patch.rs

@@ -25,6 +25,7 @@ use crate::tools::handlers::apply_patch_spec::ApplyPatchToolArgs;
 use crate::tools::handlers::apply_patch_spec::create_apply_patch_freeform_tool;
 use crate::tools::handlers::apply_patch_spec::create_apply_patch_json_tool;
 use crate::tools::handlers::parse_arguments;
+use crate::tools::handlers::resolve_tool_environment;
 use crate::tools::hook_names::HookToolName;
 use crate::tools::orchestrator::ToolOrchestrator;
 use crate::tools::registry::PostToolUsePayload;
@@ -278,14 +279,52 @@ fn write_permissions_for_paths(
 /// The apply_patch tool can arrive as the older JSON/function shape or as a
 /// freeform custom tool call. Both represent the same file edit operation, so
 /// hooks see the raw patch body in `tool_input.command` either way.
-fn apply_patch_payload_command(payload: &ToolPayload) -> Option<String> {
-    match payload {
-        ToolPayload::Function { arguments } => parse_arguments::<ApplyPatchToolArgs>(arguments)
-            .ok()
-            .map(|args| args.input),
-        ToolPayload::Custom { input } => Some(input.clone()),
+fn apply_patch_hook_tool_input(invocation: &ToolInvocation) -> Option<serde_json::Value> {
+    let (command, function_environment_id) = match &invocation.payload {
+        ToolPayload::Function { arguments } => {
+            let args = parse_arguments::<ApplyPatchToolArgs>(arguments).ok()?;
+            Some((args.input, args.environment_id))
+        }
+        ToolPayload::Custom { input } => Some((input.clone(), None)),
         _ => None,
+    }?;
+
+    let parsed_environment_id = codex_apply_patch::parse_patch_with_options(
+        &command,
+        ParseOptions {
+            allow_environment_id: true,
+        },
+    )
+    .ok()
+    .and_then(|args| args.environment_id);
+    let requested_environment_id = function_environment_id
+        .as_deref()
+        .or(parsed_environment_id.as_deref())
+        .filter(|environment_id| !environment_id.trim().is_empty());
+    let selected_environment =
+        resolve_tool_environment(invocation.turn.as_ref(), requested_environment_id)
+            .ok()
+            .flatten();
+
+    let mut tool_input = serde_json::json!({ "command": command });
+    if let Some(object) = tool_input.as_object_mut() {
+        if let Some(environment) = selected_environment {
+            object.insert(
+                "environment_id".to_string(),
+                serde_json::Value::String(environment.environment_id.clone()),
+            );
+            object.insert(
+                "cwd".to_string(),
+                serde_json::json!(environment.cwd.clone()),
+            );
+        } else if let Some(environment_id) = requested_environment_id {
+            object.insert(
+                "environment_id".to_string(),
+                serde_json::Value::String(environment_id.to_string()),
+            );
+        }
     }
+    Some(tool_input)
 }
 
 fn reconcile_environment_id(
@@ -415,9 +454,9 @@ impl ToolHandler for ApplyPatchHandler {
     }
 
     fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
-        apply_patch_payload_command(&invocation.payload).map(|command| PreToolUsePayload {
+        apply_patch_hook_tool_input(invocation).map(|tool_input| PreToolUsePayload {
             tool_name: HookToolName::apply_patch(),
-            tool_input: serde_json::json!({ "command": command }),
+            tool_input,
         })
     }
 
@@ -431,9 +470,7 @@ impl ToolHandler for ApplyPatchHandler {
         Some(PostToolUsePayload {
             tool_name: HookToolName::apply_patch(),
             tool_use_id: invocation.call_id.clone(),
-            tool_input: serde_json::json!({
-                "command": apply_patch_payload_command(&invocation.payload)?,
-            }),
+            tool_input: apply_patch_hook_tool_input(invocation)?,
             tool_response,
         })
     }
@@ -607,17 +644,16 @@ pub(crate) async fn intercept_apply_patch(
     command: &[String],
     cwd: &AbsolutePathBuf,
     fs: &dyn ExecutorFileSystem,
+    environment_id: &str,
+    environment_is_remote: bool,
     session: Arc<Session>,
     turn: Arc<TurnContext>,
     tracker: Option<&SharedTurnDiffTracker>,
     call_id: &str,
     tool_name: &str,
 ) -> Result<Option<FunctionToolOutput>, FunctionCallError> {
-    let sandbox = turn
-        .environments
-        .primary()
-        .filter(|env| env.environment.is_remote())
-        .map(|_| turn.file_system_sandbox_context(/*additional_permissions*/ None));
+    let sandbox = environment_is_remote
+        .then(|| turn.file_system_sandbox_context(/*additional_permissions*/ None));
     match codex_apply_patch::maybe_parse_apply_patch_verified(command, cwd, fs, sandbox.as_ref())
         .await
     {
@@ -658,11 +694,7 @@ 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(),
+                        environment_id: environment_id.to_string(),
                         file_paths: approval_keys,
                         changes,
                         exec_approval_requirement: apply.exec_approval_requirement,

codex-rs/core/src/tools/handlers/apply_patch_spec.rs

@@ -85,6 +85,8 @@ 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.
 "#;
+const APPLY_PATCH_FREEFORM_TOOL_DESCRIPTION: &str = "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, so do not wrap the patch in JSON.";
+const APPLY_PATCH_FREEFORM_TOOL_MULTI_ENVIRONMENT_DESCRIPTION: &str = "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, so do not wrap the patch in JSON. When multiple environments are available, you may target a non-default environment by adding `*** Environment ID: <id>` immediately after `*** Begin Patch`, using an id from the turn environment context.";
 
 /// TODO(dylan): deprecate once we get rid of json tool
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -98,7 +100,11 @@ pub struct ApplyPatchToolArgs {
 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(),
+        description: if multi_environment {
+            APPLY_PATCH_FREEFORM_TOOL_MULTI_ENVIRONMENT_DESCRIPTION.to_string()
+        } else {
+            APPLY_PATCH_FREEFORM_TOOL_DESCRIPTION.to_string()
+        },
         format: FreeformToolFormat {
             r#type: "grammar".to_string(),
             syntax: "lark".to_string(),

codex-rs/core/src/tools/handlers/apply_patch_spec_tests.rs

@@ -9,9 +9,7 @@ fn create_apply_patch_freeform_tool_matches_expected_spec() {
         create_apply_patch_freeform_tool(/*multi_environment*/ true),
         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(),
+            description: APPLY_PATCH_FREEFORM_TOOL_MULTI_ENVIRONMENT_DESCRIPTION.to_string(),
             format: FreeformToolFormat {
                 r#type: "grammar".to_string(),
                 syntax: "lark".to_string(),
@@ -65,6 +63,7 @@ fn strict_apply_patch_tools_do_not_advertise_multi_environment() {
         .replace(APPLY_PATCH_WITH_ENV_START, APPLY_PATCH_STRICT_START)
         .replace(APPLY_PATCH_ENVIRONMENT_RULE, "");
     assert_eq!(freeform.format.definition, expected_strict);
+    assert_eq!(freeform.description, APPLY_PATCH_FREEFORM_TOOL_DESCRIPTION);
 
     let json_tool = create_apply_patch_json_tool(/*multi_environment*/ false);
     let ToolSpec::Function(json_tool) = json_tool else {

codex-rs/core/src/tools/handlers/apply_patch_tests.rs

@@ -49,13 +49,18 @@ async fn pre_tool_use_payload_uses_json_patch_input() {
         arguments: json!({ "input": patch }).to_string(),
     };
     let invocation = invocation_for_payload(payload).await;
+    let environment = invocation.turn.environments.primary().expect("primary env");
     let handler = ApplyPatchHandler::default();
 
     assert_eq!(
         handler.pre_tool_use_payload(&invocation),
         Some(PreToolUsePayload {
             tool_name: HookToolName::apply_patch(),
-            tool_input: json!({ "command": patch }),
+            tool_input: json!({
+                "command": patch,
+                "environment_id": environment.environment_id.clone(),
+                "cwd": environment.cwd.clone(),
+            }),
         })
     );
 }
@@ -67,13 +72,18 @@ async fn pre_tool_use_payload_uses_freeform_patch_input() {
         input: patch.to_string(),
     };
     let invocation = invocation_for_payload(payload).await;
+    let environment = invocation.turn.environments.primary().expect("primary env");
     let handler = ApplyPatchHandler::default();
 
     assert_eq!(
         handler.pre_tool_use_payload(&invocation),
         Some(PreToolUsePayload {
             tool_name: HookToolName::apply_patch(),
-            tool_input: json!({ "command": patch }),
+            tool_input: json!({
+                "command": patch,
+                "environment_id": environment.environment_id.clone(),
+                "cwd": environment.cwd.clone(),
+            }),
         })
     );
 }
@@ -85,6 +95,7 @@ async fn post_tool_use_payload_uses_patch_input_and_tool_output() {
         input: patch.to_string(),
     };
     let invocation = invocation_for_payload(payload).await;
+    let environment = invocation.turn.environments.primary().expect("primary env");
     let output = ApplyPatchToolOutput::from_text("Success. Updated files.".to_string());
     let handler = ApplyPatchHandler::default();
 
@@ -93,7 +104,11 @@ async fn post_tool_use_payload_uses_patch_input_and_tool_output() {
         Some(PostToolUsePayload {
             tool_name: HookToolName::apply_patch(),
             tool_use_id: "call-apply-patch".to_string(),
-            tool_input: json!({ "command": patch }),
+            tool_input: json!({
+                "command": patch,
+                "environment_id": environment.environment_id.clone(),
+                "cwd": environment.cwd.clone(),
+            }),
             tool_response: json!("Success. Updated files."),
         })
     );

codex-rs/core/src/tools/handlers/shell.rs

@@ -197,6 +197,8 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, Func
         &exec_params.command,
         &exec_params.cwd,
         fs.as_ref(),
+        &turn_environment.environment_id,
+        turn_environment.environment.is_remote(),
         session.clone(),
         turn.clone(),
         Some(&tracker),

codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs

@@ -273,6 +273,8 @@ impl ToolHandler for ExecCommandHandler {
             &command,
             &cwd,
             fs.as_ref(),
+            &turn_environment.environment_id,
+            turn_environment.environment.is_remote(),
             context.session.clone(),
             context.turn.clone(),
             Some(&tracker),