Address apply_patch review feedback

Co-authored-by: Codex <noreply@openai.com>
2026-05-28 06:55:01 +00:00 · 2026-05-11 10:51:10 -07:00
parent 7eb107f797
commit aea830453c
9 changed files with 295 additions and 61 deletions
--- a/codex-rs/apply-patch/src/streaming_parser.rs
+++ b/codex-rs/apply-patch/src/streaming_parser.rs
@@ -48,7 +48,7 @@ impl StreamingPatchParser {
    // The live streaming parser only needs to keep the patch preview flowing.
    // Environment selection and validation happen on the final tool invocation,
    // so here we just tolerate and skip the optional preamble line.
-    fn maybe_skip_environment_id_preamble_line(&self, line: &str) -> bool {
+    fn is_environment_id_preamble_line(&self, line: &str) -> bool {
        line.starts_with(ENVIRONMENT_ID_MARKER)
    }

@@ -170,10 +170,9 @@ impl StreamingPatchParser {
                ))
            }
            StreamingParserMode::StartedPatch => {
-                if self.maybe_skip_environment_id_preamble_line(line) {
+                if self.is_environment_id_preamble_line(line) {
                    return Ok(());
                }
-                let trimmed = line.trim();
                if self.handle_hunk_headers_and_end_patch(trimmed)? {
                    return Ok(());
                }
--- a/codex-rs/core/src/tools/handlers/apply_patch.rs
+++ b/codex-rs/core/src/tools/handlers/apply_patch.rs
@@ -37,6 +37,7 @@ use codex_apply_patch::ApplyPatchAction;
 use codex_apply_patch::ApplyPatchFileChange;
 use codex_apply_patch::Hunk;
 use codex_apply_patch::StreamingPatchParser;
+use codex_exec_server::Environment;
 use codex_exec_server::ExecutorFileSystem;
 use codex_features::Feature;
 use codex_protocol::models::AdditionalPermissionProfile;
@@ -413,6 +414,7 @@ impl ToolHandler for ApplyPatchHandler {

                        let req = ApplyPatchRequest {
                            environment_id: turn_environment.environment_id.clone(),
+                            environment: Arc::clone(&turn_environment.environment),
                            action: apply.action,
                            file_paths,
                            changes,
@@ -482,6 +484,7 @@ pub(crate) async fn intercept_apply_patch(
    cwd: &AbsolutePathBuf,
    fs: &dyn ExecutorFileSystem,
    environment_id: &str,
+    environment: Arc<Environment>,
    session: Arc<Session>,
    turn: Arc<TurnContext>,
    tracker: Option<&SharedTurnDiffTracker>,
@@ -524,6 +527,7 @@ pub(crate) async fn intercept_apply_patch(

                    let req = ApplyPatchRequest {
                        environment_id: environment_id.to_string(),
+                        environment,
                        action: apply.action,
                        file_paths: approval_keys,
                        changes,
--- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs
+++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs
@@ -260,12 +260,14 @@ async fn intercepted_apply_patch_verification_uses_local_sandbox() {
    ];
    let (session, mut turn) = make_session_and_context().await;
    turn.permission_profile = PermissionProfile::read_only();
+    let environment = session.services.environment_manager.local_environment();

    let result = intercept_apply_patch(
        &command,
        &cwd,
        LOCAL_FS.as_ref(),
        codex_exec_server::LOCAL_ENVIRONMENT_ID,
+        environment,
        Arc::new(session),
        Arc::new(turn),
        /*tracker*/ None,
--- a/codex-rs/core/src/tools/handlers/shell.rs
+++ b/codex-rs/core/src/tools/handlers/shell.rs
@@ -199,6 +199,7 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result<FunctionToolOutput, Func
        &exec_params.cwd,
        fs.as_ref(),
        &turn_environment.environment_id,
+        Arc::clone(&turn_environment.environment),
        session.clone(),
        turn.clone(),
        Some(&tracker),
--- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs
+++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs
@@ -274,6 +274,7 @@ impl ToolHandler for ExecCommandHandler {
            &cwd,
            fs.as_ref(),
            &turn_environment.environment_id,
+            Arc::clone(&environment),
            context.session.clone(),
            context.turn.clone(),
            Some(&tracker),
--- a/codex-rs/core/src/tools/runtimes/apply_patch.rs
+++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs
@@ -19,6 +19,7 @@ use crate::tools::sandboxing::ToolRuntime;
 use crate::tools::sandboxing::with_cached_approval;
 use codex_apply_patch::AppliedPatchDelta;
 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;
@@ -45,6 +46,7 @@ pub(crate) struct ApplyPatchApprovalKey {
 #[derive(Debug)]
 pub struct ApplyPatchRequest {
    pub environment_id: String,
+    pub environment: std::sync::Arc<Environment>,
    pub action: ApplyPatchAction,
    pub file_paths: Vec<AbsolutePathBuf>,
    pub changes: std::collections::HashMap<PathBuf, FileChange>,
@@ -85,18 +87,6 @@ impl ApplyPatchRuntime {
        }
    }

-    fn approval_reason(req: &ApplyPatchRequest, reason: Option<String>) -> Option<String> {
-        let context = format!(
-            "Environment `{}`, cwd `{}`.",
-            req.environment_id,
-            req.action.cwd.display()
-        );
-        Some(match reason {
-            Some(reason) => format!("{reason}\n{context}"),
-            None => context,
-        })
-    }
-
    fn file_system_sandbox_context_for_attempt(
        req: &ApplyPatchRequest,
        attempt: &SandboxAttempt<'_>,
@@ -162,13 +152,12 @@ impl Approvable<ApplyPatchRequest> for ApplyPatchRuntime {
                return ReviewDecision::Approved;
            }
            if let Some(reason) = retry_reason {
-                let reason = ApplyPatchRuntime::approval_reason(req, Some(reason));
                let rx_approve = session
                    .request_patch_approval(
                        turn,
                        call_id,
                        changes.clone(),
-                        reason,
+                        Some(reason),
                        /*grant_root*/ None,
                    )
                    .await;
@@ -180,10 +169,9 @@ impl Approvable<ApplyPatchRequest> for ApplyPatchRuntime {
                "apply_patch",
                approval_keys,
                || async move {
-                    let reason = ApplyPatchRuntime::approval_reason(req, /*reason*/ None);
                    let rx_approve = session
                        .request_patch_approval(
-                            turn, call_id, changes, reason, /*grant_root*/ None,
+                            turn, call_id, changes, None, /*grant_root*/ None,
                        )
                        .await;
                    rx_approve.await.unwrap_or_default()
@@ -234,19 +222,10 @@ impl ToolRuntime<ApplyPatchRequest, ApplyPatchRuntimeOutput> for ApplyPatchRunti
        &mut self,
        req: &ApplyPatchRequest,
        attempt: &SandboxAttempt<'_>,
-        ctx: &ToolCtx,
+        _ctx: &ToolCtx,
    ) -> Result<ApplyPatchRuntimeOutput, ToolError> {
-        let turn_environment = ctx
-            .turn
-            .environments
-            .turn_environments
-            .iter()
-            .find(|environment| environment.environment_id == req.environment_id)
-            .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();
--- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs
+++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs
@@ -15,6 +15,11 @@ 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<codex_exec_server::Environment> {
+    Arc::new(codex_exec_server::Environment::default_for_tests())
+}

 #[test]
 fn wants_no_sandbox_approval_granular_respects_sandbox_flag() {
@@ -50,6 +55,7 @@ fn guardian_review_request_includes_patch_context() {
    let expected_patch = action.patch.clone();
    let request = ApplyPatchRequest {
        environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(),
+        environment: test_environment(),
        action,
        file_paths: vec![path.clone()],
        changes: HashMap::from([(
@@ -89,6 +95,7 @@ fn permission_request_payload_uses_apply_patch_hook_name_and_aliases() {
    let expected_patch = action.patch.clone();
    let req = ApplyPatchRequest {
        environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(),
+        environment: test_environment(),
        action,
        file_paths: vec![path],
        changes: HashMap::new(),
@@ -123,6 +130,7 @@ fn approval_keys_include_environment_id() {
        .abs();
    let req = ApplyPatchRequest {
        environment_id: "remote".to_string(),
+        environment: test_environment(),
        action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()),
        file_paths: vec![path.clone()],
        changes: HashMap::new(),
@@ -155,6 +163,7 @@ fn sandbox_cwd_uses_patch_action_cwd() {
        .abs();
    let req = ApplyPatchRequest {
        environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(),
+        environment: test_environment(),
        action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()),
        file_paths: vec![path.clone()],
        changes: HashMap::new(),
@@ -169,33 +178,6 @@ fn sandbox_cwd_uses_patch_action_cwd() {
    assert_eq!(runtime.sandbox_cwd(&req), Some(&req.action.cwd));
 }

-#[test]
-fn approval_reason_includes_environment_and_cwd() {
-    let path = std::env::temp_dir()
-        .join("apply-patch-approval-reason.txt")
-        .abs();
-    let req = ApplyPatchRequest {
-        environment_id: "remote".to_string(),
-        action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()),
-        file_paths: vec![path],
-        changes: HashMap::new(),
-        exec_approval_requirement: ExecApprovalRequirement::NeedsApproval {
-            reason: None,
-            proposed_execpolicy_amendment: None,
-        },
-        additional_permissions: None,
-        permissions_preapproved: false,
-    };
-
-    assert_eq!(
-        ApplyPatchRuntime::approval_reason(&req, Some("retry".to_string())),
-        Some(format!(
-            "retry\nEnvironment `remote`, cwd `{}`.",
-            req.action.cwd.display()
-        ))
-    );
-}
-
 #[test]
 fn file_system_sandbox_context_uses_active_attempt() {
    let path = std::env::temp_dir()
@@ -210,6 +192,7 @@ fn file_system_sandbox_context_uses_active_attempt() {
    };
    let req = ApplyPatchRequest {
        environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(),
+        environment: test_environment(),
        action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()),
        file_paths: vec![path.clone()],
        changes: HashMap::new(),
@@ -268,6 +251,7 @@ fn no_sandbox_attempt_has_no_file_system_context() {
        .abs();
    let req = ApplyPatchRequest {
        environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(),
+        environment: test_environment(),
        action: ApplyPatchAction::new_add_for_test(&path, "hello".to_string()),
        file_paths: vec![path.clone()],
        changes: HashMap::new(),
--- a/codex-rs/core/tests/suite/remote_env.rs
+++ b/codex-rs/core/tests/suite/remote_env.rs
@@ -1,5 +1,7 @@
 use anyhow::Context;
 use anyhow::Result;
+use codex_config::types::ApprovalsReviewer;
+use codex_core::config::Constrained;
 use codex_exec_server::CopyOptions;
 use codex_exec_server::CreateDirectoryOptions;
 use codex_exec_server::FileSystemSandboxContext;
@@ -13,7 +15,14 @@ use codex_protocol::permissions::FileSystemPath;
 use codex_protocol::permissions::FileSystemSandboxEntry;
 use codex_protocol::permissions::FileSystemSandboxPolicy;
 use codex_protocol::permissions::NetworkSandboxPolicy;
+use codex_protocol::protocol::ApplyPatchApprovalRequestEvent;
+use codex_protocol::protocol::AskForApproval;
+use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::Op;
+use codex_protocol::protocol::ReviewDecision;
+use codex_protocol::protocol::SandboxPolicy;
 use codex_protocol::protocol::TurnEnvironmentSelection;
+use codex_protocol::user_input::UserInput;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use core_test_support::PathBufExt;
 use core_test_support::PathExt;
@@ -30,6 +39,7 @@ use core_test_support::skip_if_no_network;
 use core_test_support::test_codex::TestCodex;
 use core_test_support::test_codex::test_codex;
 use core_test_support::test_codex::test_env;
+use core_test_support::wait_for_event;
 use pretty_assertions::assert_eq;
 use serde_json::Value;
 use serde_json::json;
@@ -51,6 +61,76 @@ async fn unified_exec_test(server: &wiremock::MockServer) -> Result<TestCodex> {
    builder.build_remote_aware(server).await
 }

+async fn submit_turn_with_approval_and_environments(
+    test: &TestCodex,
+    prompt: &str,
+    environments: Vec<TurnEnvironmentSelection>,
+) -> Result<()> {
+    test.codex
+        .submit(Op::UserTurn {
+            environments: Some(environments),
+            items: vec![UserInput::Text {
+                text: prompt.into(),
+                text_elements: Vec::new(),
+            }],
+            final_output_json_schema: None,
+            cwd: test.cwd.path().to_path_buf(),
+            approval_policy: AskForApproval::OnRequest,
+            approvals_reviewer: Some(ApprovalsReviewer::User),
+            sandbox_policy: SandboxPolicy::new_workspace_write_policy(),
+            permission_profile: None,
+            model: test.session_configured.model.clone(),
+            effort: None,
+            summary: None,
+            service_tier: None,
+            collaboration_mode: None,
+            personality: None,
+        })
+        .await?;
+
+    Ok(())
+}
+
+async fn expect_patch_approval(
+    test: &TestCodex,
+    expected_call_id: &str,
+) -> ApplyPatchApprovalRequestEvent {
+    let event = wait_for_event(&test.codex, |event| {
+        matches!(
+            event,
+            EventMsg::ApplyPatchApprovalRequest(_) | EventMsg::TurnComplete(_)
+        )
+    })
+    .await;
+
+    match event {
+        EventMsg::ApplyPatchApprovalRequest(approval) => {
+            assert_eq!(approval.call_id, expected_call_id);
+            approval
+        }
+        EventMsg::TurnComplete(_) => panic!("expected patch approval request before completion"),
+        other => panic!("unexpected event: {other:?}"),
+    }
+}
+
+async fn wait_for_completion_without_patch_approval(test: &TestCodex) {
+    let event = wait_for_event(&test.codex, |event| {
+        matches!(
+            event,
+            EventMsg::ApplyPatchApprovalRequest(_) | EventMsg::TurnComplete(_)
+        )
+    })
+    .await;
+
+    match event {
+        EventMsg::TurnComplete(_) => {}
+        EventMsg::ApplyPatchApprovalRequest(event) => {
+            panic!("unexpected patch approval request: {:?}", event.call_id)
+        }
+        other => panic!("unexpected event: {other:?}"),
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn remote_test_env_can_connect_and_use_filesystem() -> Result<()> {
    let Some(_remote_env) = get_remote_test_env() else {
@@ -346,6 +426,193 @@ async fn apply_patch_freeform_routes_to_selected_remote_environment() -> Result<
    Ok(())
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn apply_patch_approvals_are_remembered_per_environment() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    let Some(_remote_env) = get_remote_test_env() else {
+        return Ok(());
+    };
+
+    let server = start_mock_server().await;
+    let mut builder = test_codex().with_config(|config| {
+        config.include_apply_patch_tool = true;
+        config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
+        config.approvals_reviewer = ApprovalsReviewer::User;
+    });
+    let test = builder.build_remote_aware(&server).await?;
+    let local_cwd = TempDir::new()?;
+    let remote_cwd = PathBuf::from(format!(
+        "/tmp/codex-remote-apply-patch-approval-cwd-{}",
+        SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis()
+    ))
+    .abs();
+    test.fs()
+        .create_directory(
+            &remote_cwd,
+            CreateDirectoryOptions { recursive: true },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    let target_path = PathBuf::from(format!(
+        "/tmp/codex-apply-patch-approval-scope-{}.txt",
+        SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis()
+    ))
+    .abs();
+    let _ = fs::remove_file(&target_path);
+    test.fs()
+        .remove(
+            &target_path,
+            RemoveOptions {
+                recursive: false,
+                force: true,
+            },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    let environments = vec![
+        TurnEnvironmentSelection {
+            environment_id: LOCAL_ENVIRONMENT_ID.to_string(),
+            cwd: local_cwd.path().abs(),
+        },
+        TurnEnvironmentSelection {
+            environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
+            cwd: remote_cwd.clone(),
+        },
+    ];
+    let local_patch = format!(
+        "*** Begin Patch\n*** Environment ID: {LOCAL_ENVIRONMENT_ID}\n*** Add File: {}\n+local\n*** End Patch",
+        target_path.display()
+    );
+    let remote_patch = format!(
+        "*** Begin Patch\n*** Environment ID: {REMOTE_ENVIRONMENT_ID}\n*** Add File: {}\n+remote\n*** End Patch",
+        target_path.display()
+    );
+    let remote_update_patch = format!(
+        "*** Begin Patch\n*** Environment ID: {REMOTE_ENVIRONMENT_ID}\n*** Update File: {}\n@@\n-remote\n+remote updated\n*** End Patch",
+        target_path.display()
+    );
+
+    mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-local-1"),
+                ev_apply_patch_custom_tool_call("call-local", &local_patch),
+                ev_completed("resp-local-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-local-2"),
+                ev_assistant_message("msg-local", "done"),
+                ev_completed("resp-local-2"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-remote-1"),
+                ev_apply_patch_custom_tool_call("call-remote", &remote_patch),
+                ev_completed("resp-remote-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-remote-2"),
+                ev_assistant_message("msg-remote", "done"),
+                ev_completed("resp-remote-2"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-remote-3"),
+                ev_apply_patch_custom_tool_call("call-remote-followup", &remote_update_patch),
+                ev_completed("resp-remote-3"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-remote-4"),
+                ev_assistant_message("msg-remote-followup", "done"),
+                ev_completed("resp-remote-4"),
+            ]),
+        ],
+    )
+    .await;
+
+    submit_turn_with_approval_and_environments(
+        &test,
+        "apply patch in local environment",
+        environments.clone(),
+    )
+    .await?;
+    let approval = expect_patch_approval(&test, "call-local").await;
+    test.codex
+        .submit(Op::PatchApproval {
+            id: approval.call_id,
+            decision: ReviewDecision::ApprovedForSession,
+        })
+        .await?;
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    assert_eq!(fs::read_to_string(&target_path)?, "local\n");
+
+    submit_turn_with_approval_and_environments(
+        &test,
+        "apply patch in remote environment",
+        environments.clone(),
+    )
+    .await?;
+    let approval = expect_patch_approval(&test, "call-remote").await;
+    test.codex
+        .submit(Op::PatchApproval {
+            id: approval.call_id,
+            decision: ReviewDecision::ApprovedForSession,
+        })
+        .await?;
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    assert_eq!(
+        test.fs()
+            .read_file_text(&target_path, /*sandbox*/ None)
+            .await?,
+        "remote\n"
+    );
+
+    submit_turn_with_approval_and_environments(
+        &test,
+        "apply patch again in remote environment",
+        environments,
+    )
+    .await?;
+    wait_for_completion_without_patch_approval(&test).await;
+    assert_eq!(
+        test.fs()
+            .read_file_text(&target_path, /*sandbox*/ None)
+            .await?,
+        "remote updated\n"
+    );
+
+    let _ = fs::remove_file(&target_path);
+    test.fs()
+        .remove(
+            &target_path,
+            RemoveOptions {
+                recursive: false,
+                force: true,
+            },
+            /*sandbox*/ None,
+        )
+        .await?;
+    test.fs()
+        .remove(
+            &remote_cwd,
+            RemoveOptions {
+                recursive: true,
+                force: true,
+            },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn apply_patch_intercepted_exec_command_routes_to_selected_remote_environment() -> Result<()>
 {
--- a/codex-rs/mcp-server/tests/suite/codex_tool.rs
+++ b/codex-rs/mcp-server/tests/suite/codex_tool.rs
@@ -296,10 +296,7 @@ async fn patch_approval_triggers_elicitation() -> anyhow::Result<()> {
        Some(create_expected_patch_approval_elicitation_request_params(
            expected_changes,
            /*grant_root*/ None, // No grant_root expected
-            Some(format!(
-                "Environment `local`, cwd `{}`.",
-                cwd.path().display()
-            )),
+            /*reason*/ None,
            codex_request_id.to_string(),
            params.codex_event_id.clone(),
            params.thread_id,