Support multi-environment apply_patch selection (#21617)

## Summary - add multi-environment apply_patch routing for both freeform and function-call tool flows - parse and reconcile the optional environment selector in the main apply_patch parser, then verify against the selected environment in the handler - carry environment_id through runtime and approval surfaces so remote-targeted patches stay explicit end to end ## Testing - just fmt - remote exec-server e2e: `cargo test -p codex-core --test all apply_patch_multi_environment_uses_remote_executor -- --nocapture` on dev via `scripts/test-remote-env.sh` --------- Co-authored-by: Codex <noreply@openai.com>
2026-07-01 00:31:56 +08:00 · 2026-05-11 16:33:44 -07:00
parent bb6134c028
commit 22e84c49d0
18 changed files with 991 additions and 123 deletions
@@ -139,6 +139,28 @@ fn workspace_write_with_read_only_root(read_only_root: AbsolutePathBuf) -> Permi
    )
 }

+#[cfg(unix)]
+fn workspace_write_with_unreadable_path(unreadable_path: AbsolutePathBuf) -> PermissionProfile {
+    let file_system_sandbox_policy = FileSystemSandboxPolicy::restricted(vec![
+        FileSystemSandboxEntry {
+            path: FileSystemPath::Path {
+                path: unreadable_path,
+            },
+            access: FileSystemAccessMode::None,
+        },
+        FileSystemSandboxEntry {
+            path: FileSystemPath::Special {
+                value: FileSystemSpecialPath::project_roots(/*subpath*/ None),
+            },
+            access: FileSystemAccessMode::Write,
+        },
+    ]);
+    PermissionProfile::from_runtime_permissions(
+        &file_system_sandbox_policy,
+        NetworkSandboxPolicy::Restricted,
+    )
+}
+
 #[cfg(unix)]
 fn create_file_symlink(source: &std::path::Path, link: &std::path::Path) -> std::io::Result<()> {
    std::os::unix::fs::symlink(source, link)
@@ -720,6 +742,59 @@ async fn apply_patch_cli_rejects_path_traversal_outside_workspace(
    Ok(())
 }

+#[cfg(unix)]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[test_case(ApplyPatchModelOutput::Shell ; "shell")]
+#[test_case(ApplyPatchModelOutput::ShellViaHeredoc ; "shell_heredoc")]
+#[test_case(ApplyPatchModelOutput::ShellCommandViaHeredoc ; "shell_command_heredoc")]
+async fn intercepted_apply_patch_verification_uses_local_sandbox(
+    model_output: ApplyPatchModelOutput,
+) -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    skip_if_remote!(Ok(()), "symlink setup needs local filesystem link creation");
+
+    let harness = apply_patch_harness().await?;
+    let denied_target = harness.path("denied-target.txt");
+    std::fs::write(&denied_target, "outside content\n")?;
+
+    let link_rel = "soft-link.txt";
+    create_file_symlink(&denied_target, &harness.path(link_rel))?;
+
+    let patch = format!(
+        r#"*** Begin Patch
+*** Update File: {link_rel}
+@@
+-outside content
+pwned
+*** End Patch"#
+    );
+    let call_id = "apply-sandboxed-read";
+    mount_apply_patch(&harness, call_id, &patch, "fail", model_output).await;
+
+    harness
+        .submit_with_permission_profile(
+            "attempt to read denied target via intercepted apply_patch",
+            workspace_write_with_unreadable_path(AbsolutePathBuf::try_from(denied_target.clone())?),
+        )
+        .await?;
+
+    let out = harness.apply_patch_output(call_id, model_output).await;
+    assert!(
+        out.contains("apply_patch verification failed"),
+        "expected sandboxed verification failure: {out}"
+    );
+    assert!(
+        out.contains("Failed to read"),
+        "expected read failure: {out}"
+    );
+    assert_eq!(
+        std::fs::read_to_string(&denied_target)?,
+        "outside content\n",
+        "verification failure should leave the denied target unchanged"
+    );
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[test_case(ApplyPatchModelOutput::Freeform ; "freeform")]
 #[test_case(ApplyPatchModelOutput::Shell ; "shell")]
@@ -1,5 +1,7 @@
 use anyhow::Context;
 use anyhow::Result;
+use codex_config::types::ApprovalsReviewer;
+use codex_core::config::Constrained;
 use codex_exec_server::CopyOptions;
 use codex_exec_server::CreateDirectoryOptions;
 use codex_exec_server::FileSystemSandboxContext;
@@ -13,11 +15,19 @@ use codex_protocol::permissions::FileSystemPath;
 use codex_protocol::permissions::FileSystemSandboxEntry;
 use codex_protocol::permissions::FileSystemSandboxPolicy;
 use codex_protocol::permissions::NetworkSandboxPolicy;
+use codex_protocol::protocol::ApplyPatchApprovalRequestEvent;
+use codex_protocol::protocol::AskForApproval;
+use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::Op;
+use codex_protocol::protocol::ReviewDecision;
+use codex_protocol::protocol::SandboxPolicy;
 use codex_protocol::protocol::TurnEnvironmentSelection;
+use codex_protocol::user_input::UserInput;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use core_test_support::PathBufExt;
 use core_test_support::PathExt;
 use core_test_support::get_remote_test_env;
+use core_test_support::responses::ev_apply_patch_custom_tool_call;
 use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_function_call;
@@ -29,6 +39,7 @@ use core_test_support::skip_if_no_network;
 use core_test_support::test_codex::TestCodex;
 use core_test_support::test_codex::test_codex;
 use core_test_support::test_codex::test_env;
+use core_test_support::wait_for_event;
 use pretty_assertions::assert_eq;
 use serde_json::Value;
 use serde_json::json;
@@ -50,6 +61,76 @@ async fn unified_exec_test(server: &wiremock::MockServer) -> Result<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 {
@@ -258,6 +339,375 @@ async fn exec_command_routes_to_selected_remote_environment() -> Result<()> {
    Ok(())
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn apply_patch_freeform_routes_to_selected_remote_environment() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    let Some(_remote_env) = get_remote_test_env() else {
+        return Ok(());
+    };
+
+    let server = start_mock_server().await;
+    let mut builder = test_codex().with_config(|config| {
+        config.include_apply_patch_tool = true;
+    });
+    let test = builder.build_remote_aware(&server).await?;
+    let local_cwd = TempDir::new()?;
+    let file_name = "apply_patch_remote_freeform.txt";
+    let remote_cwd = PathBuf::from(format!(
+        "/tmp/codex-remote-apply-patch-freeform-{}",
+        SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis()
+    ))
+    .abs();
+    test.fs()
+        .create_directory(
+            &remote_cwd,
+            CreateDirectoryOptions { recursive: true },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    let patch = format!(
+        "*** Begin Patch\n*** Environment ID: {REMOTE_ENVIRONMENT_ID}\n*** Add File: {file_name}\n+patched remote freeform\n*** End Patch"
+    );
+    let call_id = "apply-patch-remote-freeform";
+    mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_apply_patch_custom_tool_call(call_id, &patch),
+                ev_completed("resp-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_assistant_message("msg-1", "done"),
+                ev_completed("resp-2"),
+            ]),
+        ],
+    )
+    .await;
+
+    test.submit_turn_with_environments(
+        "apply patch to remote environment",
+        Some(vec![
+            TurnEnvironmentSelection {
+                environment_id: LOCAL_ENVIRONMENT_ID.to_string(),
+                cwd: local_cwd.path().abs(),
+            },
+            TurnEnvironmentSelection {
+                environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
+                cwd: remote_cwd.clone(),
+            },
+        ]),
+    )
+    .await?;
+
+    let remote_contents = test
+        .fs()
+        .read_file_text(&remote_cwd.join(file_name), /*sandbox*/ None)
+        .await?;
+    assert_eq!(remote_contents, "patched remote freeform\n");
+    assert!(
+        !local_cwd.path().join(file_name).exists(),
+        "freeform apply_patch should not create the file in the local environment"
+    );
+
+    test.fs()
+        .remove(
+            &remote_cwd,
+            RemoveOptions {
+                recursive: true,
+                force: true,
+            },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn apply_patch_approvals_are_remembered_per_environment() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+    let Some(_remote_env) = get_remote_test_env() else {
+        return Ok(());
+    };
+
+    let server = start_mock_server().await;
+    let mut builder = test_codex().with_config(|config| {
+        config.include_apply_patch_tool = true;
+        config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
+        config.approvals_reviewer = ApprovalsReviewer::User;
+    });
+    let test = builder.build_remote_aware(&server).await?;
+    let local_cwd = TempDir::new()?;
+    let remote_cwd = PathBuf::from(format!(
+        "/tmp/codex-remote-apply-patch-approval-cwd-{}",
+        SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis()
+    ))
+    .abs();
+    test.fs()
+        .create_directory(
+            &remote_cwd,
+            CreateDirectoryOptions { recursive: true },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    let target_path = PathBuf::from(format!(
+        "/tmp/codex-apply-patch-approval-scope-{}.txt",
+        SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis()
+    ))
+    .abs();
+    let _ = fs::remove_file(&target_path);
+    test.fs()
+        .remove(
+            &target_path,
+            RemoveOptions {
+                recursive: false,
+                force: true,
+            },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    let environments = vec![
+        TurnEnvironmentSelection {
+            environment_id: LOCAL_ENVIRONMENT_ID.to_string(),
+            cwd: local_cwd.path().abs(),
+        },
+        TurnEnvironmentSelection {
+            environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
+            cwd: remote_cwd.clone(),
+        },
+    ];
+    let local_patch = format!(
+        "*** Begin Patch\n*** Environment ID: {LOCAL_ENVIRONMENT_ID}\n*** Add File: {}\n+local\n*** End Patch",
+        target_path.display()
+    );
+    let remote_patch = format!(
+        "*** Begin Patch\n*** Environment ID: {REMOTE_ENVIRONMENT_ID}\n*** Add File: {}\n+remote\n*** End Patch",
+        target_path.display()
+    );
+    let remote_update_patch = format!(
+        "*** Begin Patch\n*** Environment ID: {REMOTE_ENVIRONMENT_ID}\n*** Update File: {}\n@@\n-remote\n+remote updated\n*** End Patch",
+        target_path.display()
+    );
+
+    mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-local-1"),
+                ev_apply_patch_custom_tool_call("call-local", &local_patch),
+                ev_completed("resp-local-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-local-2"),
+                ev_assistant_message("msg-local", "done"),
+                ev_completed("resp-local-2"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-remote-1"),
+                ev_apply_patch_custom_tool_call("call-remote", &remote_patch),
+                ev_completed("resp-remote-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-remote-2"),
+                ev_assistant_message("msg-remote", "done"),
+                ev_completed("resp-remote-2"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-remote-3"),
+                ev_apply_patch_custom_tool_call("call-remote-followup", &remote_update_patch),
+                ev_completed("resp-remote-3"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-remote-4"),
+                ev_assistant_message("msg-remote-followup", "done"),
+                ev_completed("resp-remote-4"),
+            ]),
+        ],
+    )
+    .await;
+
+    submit_turn_with_approval_and_environments(
+        &test,
+        "apply patch in local environment",
+        environments.clone(),
+    )
+    .await?;
+    let approval = expect_patch_approval(&test, "call-local").await;
+    test.codex
+        .submit(Op::PatchApproval {
+            id: approval.call_id,
+            decision: ReviewDecision::ApprovedForSession,
+        })
+        .await?;
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    assert_eq!(fs::read_to_string(&target_path)?, "local\n");
+
+    submit_turn_with_approval_and_environments(
+        &test,
+        "apply patch in remote environment",
+        environments.clone(),
+    )
+    .await?;
+    let approval = expect_patch_approval(&test, "call-remote").await;
+    test.codex
+        .submit(Op::PatchApproval {
+            id: approval.call_id,
+            decision: ReviewDecision::ApprovedForSession,
+        })
+        .await?;
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    assert_eq!(
+        test.fs()
+            .read_file_text(&target_path, /*sandbox*/ None)
+            .await?,
+        "remote\n"
+    );
+
+    submit_turn_with_approval_and_environments(
+        &test,
+        "apply patch again in remote environment",
+        environments,
+    )
+    .await?;
+    wait_for_completion_without_patch_approval(&test).await;
+    assert_eq!(
+        test.fs()
+            .read_file_text(&target_path, /*sandbox*/ None)
+            .await?,
+        "remote updated\n"
+    );
+
+    let _ = fs::remove_file(&target_path);
+    test.fs()
+        .remove(
+            &target_path,
+            RemoveOptions {
+                recursive: false,
+                force: true,
+            },
+            /*sandbox*/ None,
+        )
+        .await?;
+    test.fs()
+        .remove(
+            &remote_cwd,
+            RemoveOptions {
+                recursive: true,
+                force: true,
+            },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn apply_patch_intercepted_exec_command_routes_to_selected_remote_environment() -> Result<()>
+{
+    skip_if_no_network!(Ok(()));
+    let Some(_remote_env) = get_remote_test_env() else {
+        return Ok(());
+    };
+
+    let server = start_mock_server().await;
+    let test = unified_exec_test(&server).await?;
+    let local_cwd = TempDir::new()?;
+    let file_name = "apply_patch_remote_exec.txt";
+    let remote_cwd = PathBuf::from(format!(
+        "/tmp/codex-remote-apply-patch-exec-{}",
+        SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis()
+    ))
+    .abs();
+    test.fs()
+        .create_directory(
+            &remote_cwd,
+            CreateDirectoryOptions { recursive: true },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    let patch =
+        format!("*** Begin Patch\n*** Add File: {file_name}\n+patched remote exec\n*** End Patch");
+    let command = format!("apply_patch <<'EOF'\n{patch}\nEOF\n");
+    let call_id = "apply-patch-remote-exec";
+    mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_function_call(
+                    call_id,
+                    "exec_command",
+                    &serde_json::to_string(&json!({
+                        "shell": "/bin/sh",
+                        "cmd": command,
+                        "login": false,
+                        "yield_time_ms": 5_000,
+                        "environment_id": REMOTE_ENVIRONMENT_ID,
+                    }))?,
+                ),
+                ev_completed("resp-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_assistant_message("msg-1", "done"),
+                ev_completed("resp-2"),
+            ]),
+        ],
+    )
+    .await;
+
+    test.submit_turn_with_environments(
+        "apply patch through exec command to remote environment",
+        Some(vec![
+            TurnEnvironmentSelection {
+                environment_id: LOCAL_ENVIRONMENT_ID.to_string(),
+                cwd: local_cwd.path().abs(),
+            },
+            TurnEnvironmentSelection {
+                environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
+                cwd: remote_cwd.clone(),
+            },
+        ]),
+    )
+    .await?;
+
+    let remote_contents = test
+        .fs()
+        .read_file_text(&remote_cwd.join(file_name), /*sandbox*/ None)
+        .await?;
+    assert_eq!(remote_contents, "patched remote exec\n");
+    assert!(
+        !local_cwd.path().join(file_name).exists(),
+        "intercepted apply_patch should not create the file in the local environment"
+    );
+
+    test.fs()
+        .remove(
+            &remote_cwd,
+            RemoveOptions {
+                recursive: true,
+                force: true,
+            },
+            /*sandbox*/ None,
+        )
+        .await?;
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn remote_test_env_sandboxed_read_allows_readable_root() -> Result<()> {
    skip_if_no_network!(Ok(()));