Gate streaming apply_patch parser

This commit is contained in:
Akshay Nathan
2026-05-04 11:04:13 -07:00
parent 66e5d99eb2
commit b459ba0f4d
10 changed files with 230 additions and 17 deletions

View File

@@ -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 <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 <path> &&`) 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(),
})

View File

@@ -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<PathBuf, ApplyPatchFileChange>,
hunks: Vec<Hunk>,
/// 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,
}

View File

@@ -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<ApplyPatchArgs, ParseError> {
parse_patch_with_mode(patch, ParsePatchMode::Legacy)
}
pub fn parse_patch_with_mode(
patch: &str,
mode: ParsePatchMode,
) -> Result<ApplyPatchArgs, ParseError> {
if mode == ParsePatchMode::Streaming {
return crate::streaming_parser::parse_patch(patch);
}
let mode = if PARSE_IN_STRICT_MODE {
ParseMode::Strict
} else {

View File

@@ -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<ApplyPatchArgs, ParseError> {
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<String, ParseError> {
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, "<<EOF" | "<<'EOF'" | "<<\"EOF\"")
&& last.ends_with("EOF")
&& lines.len() >= 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() {

View File

@@ -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"
},

View File

@@ -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<String> {
}
}
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

View File

@@ -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();

View File

@@ -199,8 +199,8 @@ impl ToolRuntime<ApplyPatchRequest, ExecToolCallOutput> 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,

View File

@@ -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",

View File

@@ -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!(