diff --git a/codex-rs/apply-patch/src/invocation.rs b/codex-rs/apply-patch/src/invocation.rs index 075c94c60c..427aa4e876 100644 --- a/codex-rs/apply-patch/src/invocation.rs +++ b/codex-rs/apply-patch/src/invocation.rs @@ -19,7 +19,8 @@ use crate::IoError; use crate::MaybeApplyPatchVerified; use crate::parser::Hunk; use crate::parser::ParseError; -use crate::parser::parse_patch; +use crate::parser::ParsePatchMode; +use crate::parser::parse_patch_with_mode; use crate::unified_diff_from_chunks; use std::str::Utf8Error; use tree_sitter::LanguageError; @@ -102,17 +103,24 @@ fn extract_apply_patch_from_shell( } // TODO: make private once we remove tests in lib.rs +#[cfg_attr(not(test), allow(dead_code))] pub fn maybe_parse_apply_patch(argv: &[String]) -> MaybeApplyPatch { + maybe_parse_apply_patch_with_mode(argv, ParsePatchMode::Legacy) +} + +pub fn maybe_parse_apply_patch_with_mode(argv: &[String], mode: ParsePatchMode) -> MaybeApplyPatch { match argv { // Direct invocation: apply_patch - [cmd, body] if APPLY_PATCH_COMMANDS.contains(&cmd.as_str()) => match parse_patch(body) { - Ok(source) => MaybeApplyPatch::Body(source), - Err(e) => MaybeApplyPatch::PatchParseError(e), - }, + [cmd, body] if APPLY_PATCH_COMMANDS.contains(&cmd.as_str()) => { + match parse_patch_with_mode(body, mode) { + Ok(source) => MaybeApplyPatch::Body(source), + Err(e) => MaybeApplyPatch::PatchParseError(e), + } + } // Shell heredoc form: (optional `cd &&`) apply_patch <<'EOF' ... _ => match parse_shell_script(argv) { Some((shell, script)) => match extract_apply_patch_from_shell(shell, script) { - Ok((body, workdir)) => match parse_patch(&body) { + Ok((body, workdir)) => match parse_patch_with_mode(&body, mode) { Ok(mut source) => { source.workdir = workdir; MaybeApplyPatch::Body(source) @@ -136,21 +144,31 @@ pub async fn maybe_parse_apply_patch_verified( cwd: &AbsolutePathBuf, fs: &dyn ExecutorFileSystem, sandbox: Option<&codex_exec_server::FileSystemSandboxContext>, +) -> MaybeApplyPatchVerified { + maybe_parse_apply_patch_verified_with_mode(argv, ParsePatchMode::Legacy, cwd, fs, sandbox).await +} + +pub async fn maybe_parse_apply_patch_verified_with_mode( + argv: &[String], + mode: ParsePatchMode, + cwd: &AbsolutePathBuf, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&codex_exec_server::FileSystemSandboxContext>, ) -> MaybeApplyPatchVerified { // Detect a raw patch body passed directly as the command or as the body of a shell // script. In these cases, report an explicit error rather than applying the patch. if let [body] = argv - && parse_patch(body).is_ok() + && parse_patch_with_mode(body, mode).is_ok() { return MaybeApplyPatchVerified::CorrectnessError(ApplyPatchError::ImplicitInvocation); } if let Some((_, script)) = parse_shell_script(argv) - && parse_patch(script).is_ok() + && parse_patch_with_mode(script, mode).is_ok() { return MaybeApplyPatchVerified::CorrectnessError(ApplyPatchError::ImplicitInvocation); } - match maybe_parse_apply_patch(argv) { + match maybe_parse_apply_patch_with_mode(argv, mode) { MaybeApplyPatch::Body(ApplyPatchArgs { patch, hunks, @@ -161,13 +179,15 @@ pub async fn maybe_parse_apply_patch_verified( .map(|dir| cwd.join(Path::new(dir))) .unwrap_or_else(|| cwd.clone()); let mut changes = HashMap::new(); - for hunk in hunks { + for hunk in &hunks { let path = hunk.resolve_path(&effective_cwd); match hunk { Hunk::AddFile { contents, .. } => { changes.insert( path.into_path_buf(), - ApplyPatchFileChange::Add { content: contents }, + ApplyPatchFileChange::Add { + content: contents.clone(), + }, ); } Hunk::DeleteFile { .. } => { @@ -203,7 +223,9 @@ pub async fn maybe_parse_apply_patch_verified( path.into_path_buf(), ApplyPatchFileChange::Update { unified_diff, - move_path: move_path.map(|p| effective_cwd.join(p).into_path_buf()), + move_path: move_path + .as_ref() + .map(|p| effective_cwd.join(p).into_path_buf()), new_content: contents, }, ); @@ -212,6 +234,7 @@ pub async fn maybe_parse_apply_patch_verified( } MaybeApplyPatchVerified::Body(ApplyPatchAction { changes, + hunks, patch, cwd: effective_cwd, }) @@ -377,6 +400,8 @@ fn extract_apply_patch_from_bash( #[cfg(test)] mod tests { use super::*; + use crate::UpdateFileChunk; + use crate::parse_patch; use crate::unified_diff_from_chunks; use assert_matches::assert_matches; use codex_exec_server::LOCAL_FS; @@ -796,6 +821,16 @@ PATCH"#, new_content: "updated session directory content\n".to_string(), }, )]), + hunks: vec![Hunk::UpdateFile { + path: PathBuf::from(relative_path), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["session directory content".to_string()], + new_lines: vec!["updated session directory content".to_string()], + is_end_of_file: false, + }], + }], patch: argv[1].clone(), cwd: AbsolutePathBuf::from_absolute_path(session_dir.path()).unwrap(), }) diff --git a/codex-rs/apply-patch/src/lib.rs b/codex-rs/apply-patch/src/lib.rs index 7a47b1ea48..1af69782bc 100644 --- a/codex-rs/apply-patch/src/lib.rs +++ b/codex-rs/apply-patch/src/lib.rs @@ -19,6 +19,7 @@ use codex_utils_absolute_path::AbsolutePathBuf; pub use parser::Hunk; pub use parser::ParseError; use parser::ParseError::*; +pub use parser::ParsePatchMode; pub use parser::UpdateFileChunk; pub use parser::parse_patch; use similar::TextDiff; @@ -26,6 +27,7 @@ pub use streaming_parser::StreamingPatchParser; use thiserror::Error; pub use invocation::maybe_parse_apply_patch_verified; +pub use invocation::maybe_parse_apply_patch_verified_with_mode; pub use standalone_executable::main; use crate::invocation::ExtractHeredocError; @@ -135,6 +137,7 @@ pub enum MaybeApplyPatchVerified { #[derive(Debug, PartialEq)] pub struct ApplyPatchAction { changes: HashMap, + hunks: Vec, /// The raw patch argument that can be used to apply the patch. i.e., if the /// original arg was parsed in "lenient" mode with a @@ -155,6 +158,10 @@ impl ApplyPatchAction { &self.changes } + pub fn hunks(&self) -> &[Hunk] { + &self.hunks + } + /// Should be used exclusively for testing. (Not worth the overhead of /// creating a feature flag for this.) pub fn new_add_for_test(path: &AbsolutePathBuf, content: String) -> Self { @@ -170,10 +177,15 @@ impl ApplyPatchAction { + {content} *** End Patch"#, ); + let hunk = Hunk::AddFile { + path: path.to_path_buf(), + contents: content.clone(), + }; let changes = HashMap::from([(path.to_path_buf(), ApplyPatchFileChange::Add { content })]); #[expect(clippy::expect_used)] Self { changes, + hunks: vec![hunk], cwd: path.parent().expect("path should have parent"), patch, } diff --git a/codex-rs/apply-patch/src/parser.rs b/codex-rs/apply-patch/src/parser.rs index 5403055272..a3b8f21dfe 100644 --- a/codex-rs/apply-patch/src/parser.rs +++ b/codex-rs/apply-patch/src/parser.rs @@ -123,7 +123,24 @@ pub struct UpdateFileChunk { pub is_end_of_file: bool, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ParsePatchMode { + Legacy, + Streaming, +} + pub fn parse_patch(patch: &str) -> Result { + parse_patch_with_mode(patch, ParsePatchMode::Legacy) +} + +pub fn parse_patch_with_mode( + patch: &str, + mode: ParsePatchMode, +) -> Result { + if mode == ParsePatchMode::Streaming { + return crate::streaming_parser::parse_patch(patch); + } + let mode = if PARSE_IN_STRICT_MODE { ParseMode::Strict } else { diff --git a/codex-rs/apply-patch/src/streaming_parser.rs b/codex-rs/apply-patch/src/streaming_parser.rs index 4acfad6720..941cb51084 100644 --- a/codex-rs/apply-patch/src/streaming_parser.rs +++ b/codex-rs/apply-patch/src/streaming_parser.rs @@ -1,5 +1,6 @@ use std::path::PathBuf; +use crate::ApplyPatchArgs; use crate::parser::ADD_FILE_MARKER; use crate::parser::BEGIN_PATCH_MARKER; use crate::parser::CHANGE_CONTEXT_MARKER; @@ -16,6 +17,67 @@ use crate::parser::UpdateFileChunk; use Hunk::*; use ParseError::*; +pub fn parse_patch(patch: &str) -> Result { + let patch = normalize_patch_text(patch)?; + let mut parser = StreamingPatchParser::default(); + parser.push_delta(&patch)?; + let hunks = parser.finish()?; + Ok(ApplyPatchArgs { + patch, + hunks, + workdir: None, + }) +} + +fn normalize_patch_text(patch: &str) -> Result { + let lines: Vec<&str> = patch.trim().lines().collect(); + let patch_lines = match check_patch_boundaries(&lines) { + Ok(lines) => lines, + Err(original_parse_error) => match lines.as_slice() { + [first, .., last] + if matches!(*first, "<= 4 => + { + let inner_lines = &lines[1..lines.len() - 1]; + check_patch_boundaries(inner_lines)? + } + _ => return Err(original_parse_error), + }, + }; + Ok(patch_lines.join("\n")) +} + +fn check_patch_boundaries<'a>(lines: &'a [&'a str]) -> Result<&'a [&'a str], ParseError> { + let (first_line, last_line) = match lines { + [] => (None, None), + [first] => (Some(first), Some(first)), + [first, .., last] => (Some(first), Some(last)), + }; + check_start_and_end_lines(first_line, last_line)?; + Ok(lines) +} + +fn check_start_and_end_lines( + first_line: Option<&&str>, + last_line: Option<&&str>, +) -> Result<(), ParseError> { + let first_line = first_line.map(|line| line.trim()); + let last_line = last_line.map(|line| line.trim()); + + match (first_line, last_line) { + (Some(first), Some(last)) if first == BEGIN_PATCH_MARKER && last == END_PATCH_MARKER => { + Ok(()) + } + (Some(first), _) if first != BEGIN_PATCH_MARKER => Err(InvalidPatchError(String::from( + "The first line of the patch must be '*** Begin Patch'", + ))), + _ => Err(InvalidPatchError(String::from( + "The last line of the patch must be '*** End Patch'", + ))), + } +} + #[derive(Debug, Default, Clone)] pub struct StreamingPatchParser { line_buffer: String, @@ -366,6 +428,36 @@ mod tests { use std::path::PathBuf; use super::*; + use crate::parser; + + #[test] + fn test_parse_patch_matches_legacy_parser_for_complete_input() { + let patch = "\ +*** Begin Patch +*** Add File: added.txt ++hello +*** Update File: changed.txt +@@ +-old ++new +*** Delete File: removed.txt +*** End Patch"; + + assert_eq!(parse_patch(patch), parser::parse_patch(patch)); + } + + #[test] + fn test_parse_patch_supports_lenient_heredoc_wrapper() { + let patch = "\ +<<'EOF' +*** Begin Patch +*** Add File: added.txt ++hello +*** End Patch +EOF"; + + assert_eq!(parse_patch(patch), parser::parse_patch(patch)); + } #[test] fn test_streaming_patch_parser_streams_complete_lines_before_end_patch() { diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index a75e287134..36754e8772 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -364,6 +364,9 @@ "apply_patch_streaming_events": { "type": "boolean" }, + "apply_patch_streaming_parser": { + "type": "boolean" + }, "apps": { "type": "boolean" }, @@ -3833,6 +3836,9 @@ "apply_patch_streaming_events": { "type": "boolean" }, + "apply_patch_streaming_parser": { + "type": "boolean" + }, "apps": { "type": "boolean" }, diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 294e161483..f5df92989d 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -35,6 +35,7 @@ use crate::tools::sandboxing::ToolCtx; use codex_apply_patch::ApplyPatchAction; use codex_apply_patch::ApplyPatchFileChange; use codex_apply_patch::Hunk; +use codex_apply_patch::ParsePatchMode; use codex_apply_patch::StreamingPatchParser; use codex_exec_server::ExecutorFileSystem; use codex_features::Feature; @@ -255,6 +256,14 @@ fn apply_patch_payload_command(payload: &ToolPayload) -> Option { } } +fn apply_patch_mode(turn: &TurnContext) -> ParsePatchMode { + if turn.features.enabled(Feature::ApplyPatchStreamingParser) { + ParsePatchMode::Streaming + } else { + ParsePatchMode::Legacy + } +} + async fn effective_patch_permissions( session: &Session, turn: &TurnContext, @@ -373,8 +382,10 @@ impl ToolHandler for ApplyPatchHandler { .environment .is_remote() .then(|| turn.file_system_sandbox_context(/*additional_permissions*/ None)); - match codex_apply_patch::maybe_parse_apply_patch_verified( + let parse_mode = apply_patch_mode(turn.as_ref()); + match codex_apply_patch::maybe_parse_apply_patch_verified_with_mode( &command, + parse_mode, &cwd, fs.as_ref(), sandbox.as_ref(), @@ -478,8 +489,15 @@ pub(crate) async fn intercept_apply_patch( .primary_environment() .filter(|env| env.environment.is_remote()) .map(|_| turn.file_system_sandbox_context(/*additional_permissions*/ None)); - match codex_apply_patch::maybe_parse_apply_patch_verified(command, cwd, fs, sandbox.as_ref()) - .await + let parse_mode = apply_patch_mode(turn.as_ref()); + match codex_apply_patch::maybe_parse_apply_patch_verified_with_mode( + command, + parse_mode, + cwd, + fs, + sandbox.as_ref(), + ) + .await { codex_apply_patch::MaybeApplyPatchVerified::Body(changes) => { session 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 04472e4623..98931f2263 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs @@ -99,6 +99,18 @@ async fn post_tool_use_payload_uses_patch_input_and_tool_output() { ); } +#[tokio::test] +async fn apply_patch_mode_follows_streaming_parser_feature() { + let (_, mut turn) = make_session_and_context().await; + assert_eq!(apply_patch_mode(&turn), ParsePatchMode::Legacy); + + turn.features + .enable(Feature::ApplyPatchStreamingParser) + .expect("enable feature"); + + assert_eq!(apply_patch_mode(&turn), ParsePatchMode::Streaming); +} + #[test] fn diff_consumer_does_not_stream_json_tool_call_arguments() { let mut consumer = ApplyPatchArgumentDiffConsumer::default(); diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index e720243f2b..1e7cb6ccbe 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -199,8 +199,8 @@ impl ToolRuntime for ApplyPatchRuntime { let sandbox = Self::file_system_sandbox_context_for_attempt(req, attempt); let mut stdout = Vec::new(); let mut stderr = Vec::new(); - let result = codex_apply_patch::apply_patch( - &req.action.patch, + let result = codex_apply_patch::apply_hunks( + req.action.hunks(), &req.action.cwd, &mut stdout, &mut stderr, diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index acb6fa195e..017db015b3 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -99,6 +99,8 @@ pub enum Feature { ApplyPatchFreeform, /// Stream structured progress while apply_patch input is being generated. ApplyPatchStreamingEvents, + /// Use the streaming apply_patch parser for completed patch verification and application. + ApplyPatchStreamingParser, /// Allow exec tools to request additional permissions while staying sandboxed. ExecPermissionApprovals, /// Expose the built-in request_permissions tool. @@ -815,6 +817,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::ApplyPatchStreamingParser, + key: "apply_patch_streaming_parser", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::ExecPermissionApprovals, key: "exec_permission_approvals", diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index da0e36442c..09b8f45386 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -110,6 +110,19 @@ fn request_permissions_is_under_development() { assert_eq!(Feature::ExecPermissionApprovals.default_enabled(), false); } +#[test] +fn apply_patch_streaming_parser_is_under_development() { + assert_eq!( + feature_for_key("apply_patch_streaming_parser"), + Some(Feature::ApplyPatchStreamingParser) + ); + assert_eq!( + Feature::ApplyPatchStreamingParser.stage(), + Stage::UnderDevelopment + ); + assert_eq!(Feature::ApplyPatchStreamingParser.default_enabled(), false); +} + #[test] fn request_permissions_tool_is_under_development() { assert_eq!(