diff --git a/codex-rs/apply-patch/src/streaming_parser.rs b/codex-rs/apply-patch/src/streaming_parser.rs index b7efdf1f14..86084af0b8 100644 --- 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(()); } diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 54dcffaf46..a66d3f8204 100644 --- 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, session: Arc, turn: Arc, 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, diff --git a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs index ef710c282b..07e9a7516a 100644 --- 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, diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index bf686ac9f8..5cf0f24bfa 100644 --- 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, pub action: ApplyPatchAction, pub file_paths: Vec, pub changes: std::collections::HashMap, @@ -85,18 +87,6 @@ impl ApplyPatchRuntime { } } - fn approval_reason(req: &ApplyPatchRequest, reason: Option) -> Option { - 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 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 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 for ApplyPatchRunti &mut self, req: &ApplyPatchRequest, attempt: &SandboxAttempt<'_>, - ctx: &ToolCtx, + _ctx: &ToolCtx, ) -> Result { - 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(); 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 8bf445bc71..a5d59a46cf 100644 --- 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 { + 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(), diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 1d50f5fa57..dc6d5cefc5 100644 --- 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 { builder.build_remote_aware(server).await } +async fn submit_turn_with_approval_and_environments( + test: &TestCodex, + prompt: &str, + environments: Vec, +) -> 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<()> { diff --git a/codex-rs/mcp-server/tests/suite/codex_tool.rs b/codex-rs/mcp-server/tests/suite/codex_tool.rs index 78a840e1bc..ebe77c147c 100644 --- 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,