Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

assistant edit tool: Support \r\n around markers (#26538)

Browse source 
This should fix the tests on Windows

Release Notes:

- N/A
This commit is contained in:
Agus Zubiaga 2025-03-12 11:00:16 -03:00 • committed by GitHub
parent d94001f445
commit 6b56fee6b0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 60 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

87
crates/assistant_tools/src/edit_files_tool/edit_action.rs
View file

@ -78,10 +78,10 @@ impl EditActionParser {
pub fn parse_chunk(&mut self, input: &str) -> Vec<EditAction> { pub fn parse_chunk(&mut self, input: &str) -> Vec<EditAction> {
use State::*; use State::*;
const FENCE: &[u8] = b"\n```"; const FENCE: &[u8] = b"```";
const SEARCH_MARKER: &[u8] = b"<<<<<<< SEARCH\n"; const SEARCH_MARKER: &[u8] = b"<<<<<<< SEARCH";
const DIVIDER: &[u8] = b"=======\n"; const DIVIDER: &[u8] = b"=======";
const NL_DIVIDER: &[u8] = b"\n=======\n"; const NL_DIVIDER: &[u8] = b"\n=======";
const REPLACE_MARKER: &[u8] = b">>>>>>> REPLACE"; const REPLACE_MARKER: &[u8] = b">>>>>>> REPLACE";
const NL_REPLACE_MARKER: &[u8] = b"\n>>>>>>> REPLACE"; const NL_REPLACE_MARKER: &[u8] = b"\n>>>>>>> REPLACE";
@ -96,8 +96,8 @@ impl EditActionParser {
self.column += 1; self.column += 1;
} }
match self.state { match &self.state {
Default => match match_marker(byte, FENCE, &mut self.marker_ix) { Default => match match_marker(byte, FENCE, false, &mut self.marker_ix) {
MarkerMatch::Complete => { MarkerMatch::Complete => {
self.to_state(OpenFence); self.to_state(OpenFence);
} }
@ -105,12 +105,11 @@ impl EditActionParser {
MarkerMatch::None => { MarkerMatch::None => {
if self.marker_ix > 0 { if self.marker_ix > 0 {
self.marker_ix = 0; self.marker_ix = 0;
} else if self.pre_fence_line.ends_with(b"\n") {
self.pre_fence_line.clear(); self.pre_fence_line.clear();
} }
if byte != b'\n' { self.pre_fence_line.push(byte);
self.pre_fence_line.push(byte);
}
} }
}, },
OpenFence => { OpenFence => {
@ -120,7 +119,7 @@ impl EditActionParser {
} }
} }
SearchMarker => { SearchMarker => {
if self.expect_marker(byte, SEARCH_MARKER) { if self.expect_marker(byte, SEARCH_MARKER, true) {
self.to_state(SearchBlock); self.to_state(SearchBlock);
} }
} }
@ -129,6 +128,7 @@ impl EditActionParser {
byte, byte,
DIVIDER, DIVIDER,
NL_DIVIDER, NL_DIVIDER,
true,
&mut self.marker_ix, &mut self.marker_ix,
&mut self.old_bytes, &mut self.old_bytes,
) { ) {
@ -140,6 +140,7 @@ impl EditActionParser {
byte, byte,
REPLACE_MARKER, REPLACE_MARKER,
NL_REPLACE_MARKER, NL_REPLACE_MARKER,
true,
&mut self.marker_ix, &mut self.marker_ix,
&mut self.new_bytes, &mut self.new_bytes,
) { ) {
@ -147,10 +148,11 @@ impl EditActionParser {
} }
} }
CloseFence => { CloseFence => {
if self.expect_marker(byte, FENCE) { if self.expect_marker(byte, FENCE, false) {
if let Some(action) = self.action() { if let Some(action) = self.action() {
actions.push(action); actions.push(action);
} }
self.errors();
self.reset(); self.reset();
} }
} }
@ -171,7 +173,17 @@ impl EditActionParser {
return None; return None;
} }
let file_path = String::from_utf8(std::mem::take(&mut self.pre_fence_line)).log_err()?; let mut pre_fence_line = std::mem::take(&mut self.pre_fence_line);
if pre_fence_line.ends_with(b"\n") {
pre_fence_line.pop();
if pre_fence_line.ends_with(b"\r") {
pre_fence_line.pop();
}
}
let file_path = String::from_utf8(pre_fence_line).log_err()?;
let content = String::from_utf8(std::mem::take(&mut self.new_bytes)).log_err()?; let content = String::from_utf8(std::mem::take(&mut self.new_bytes)).log_err()?;
if self.old_bytes.is_empty() { if self.old_bytes.is_empty() {
@ -187,8 +199,8 @@ impl EditActionParser {
} }
} }
fn expect_marker(&mut self, byte: u8, marker: &'static [u8]) -> bool { fn expect_marker(&mut self, byte: u8, marker: &'static [u8], trailing_newline: bool) -> bool {
match match_marker(byte, marker, &mut self.marker_ix) { match match_marker(byte, marker, trailing_newline, &mut self.marker_ix) {
MarkerMatch::Complete => true, MarkerMatch::Complete => true,
MarkerMatch::Partial => false, MarkerMatch::Partial => false,
MarkerMatch::None => { MarkerMatch::None => {
@ -230,14 +242,27 @@ enum MarkerMatch {
Complete, Complete,
} }
fn match_marker(byte: u8, marker: &[u8], marker_ix: &mut usize) -> MarkerMatch { fn match_marker(
if byte == marker[*marker_ix] { byte: u8,
*marker_ix += 1; marker: &[u8],
trailing_newline: bool,
if *marker_ix >= marker.len() { marker_ix: &mut usize,
) -> MarkerMatch {
if trailing_newline && *marker_ix >= marker.len() {
if byte == b'\n' {
MarkerMatch::Complete
} else if byte == b'\r' {
MarkerMatch::Complete MarkerMatch::Complete
} else { } else {
MarkerMatch::None
}
} else if byte == marker[*marker_ix] {
*marker_ix += 1;
if *marker_ix < marker.len() || trailing_newline {
MarkerMatch::Partial MarkerMatch::Partial
} else {
MarkerMatch::Complete
} }
} else { } else {
MarkerMatch::None MarkerMatch::None
@ -248,6 +273,7 @@ fn collect_until_marker(
byte: u8, byte: u8,
marker: &[u8], marker: &[u8],
nl_marker: &[u8], nl_marker: &[u8],
trailing_newline: bool,
marker_ix: &mut usize, marker_ix: &mut usize,
buf: &mut Vec<u8>, buf: &mut Vec<u8>,
) -> bool { ) -> bool {
@ -258,7 +284,7 @@ fn collect_until_marker(
nl_marker nl_marker
}; };
match match_marker(byte, marker, marker_ix) { match match_marker(byte, marker, trailing_newline, marker_ix) {
MarkerMatch::Complete => true, MarkerMatch::Complete => true,
MarkerMatch::Partial => false, MarkerMatch::Partial => false,
MarkerMatch::None => { MarkerMatch::None => {
@ -267,7 +293,7 @@ fn collect_until_marker(
*marker_ix = 0; *marker_ix = 0;
// The beginning of marker might match current byte // The beginning of marker might match current byte
match match_marker(byte, marker, marker_ix) { match match_marker(byte, marker, trailing_newline, marker_ix) {
MarkerMatch::Complete => return true, MarkerMatch::Complete => return true,
MarkerMatch::Partial => return false, MarkerMatch::Partial => return false,
MarkerMatch::None => { /* no match, keep collecting */ } MarkerMatch::None => { /* no match, keep collecting */ }
@ -615,7 +641,7 @@ fn replacement() {}"#;
// Check parser is in the correct state // Check parser is in the correct state
assert_eq!(parser.state, State::SearchBlock); assert_eq!(parser.state, State::SearchBlock);
assert_eq!(parser.pre_fence_line, b"src/main.rs"); assert_eq!(parser.pre_fence_line, b"src/main.rs\n");
// Continue parsing // Continue parsing
let actions2 = parser.parse_chunk("original code\n=======\n"); let actions2 = parser.parse_chunk("original code\n=======\n");
@ -626,7 +652,7 @@ fn replacement() {}"#;
// After complete parsing, state should reset // After complete parsing, state should reset
assert_eq!(parser.state, State::Default); assert_eq!(parser.state, State::Default);
assert!(parser.pre_fence_line.is_empty()); assert_eq!(parser.pre_fence_line, b"\n");
assert!(parser.old_bytes.is_empty()); assert!(parser.old_bytes.is_empty());
assert!(parser.new_bytes.is_empty()); assert!(parser.new_bytes.is_empty());
@ -659,7 +685,7 @@ fn replacement() {}
assert_eq!( assert_eq!(
error.kind, error.kind,
ParseErrorKind::ExpectedMarker { ParseErrorKind::ExpectedMarker {
expected: b"<<<<<<< SEARCH\n", expected: b"<<<<<<< SEARCH",
found: b'W' found: b'W'
} }
); );
@ -779,8 +805,15 @@ fn new_utils_func() {}
} }
); );
// Ensure we have no parsing errors // The system prompt includes some text that would produce errors
assert!(errors.is_empty(), "Parsing errors found: {:?}", errors); assert_eq!(
errors[0].to_string(),
"input:102:1: Expected marker \"<<<<<<< SEARCH\", found '3'"
);
assert_eq!(
errors[1].to_string(),
"input:109:0: Expected marker \"<<<<<<< SEARCH\", found '\\n'"
);
} }
#[test] #[test]
@ -800,7 +833,7 @@ fn replacement() {}
assert_eq!(parser.errors().len(), 1); assert_eq!(parser.errors().len(), 1);
let error = &parser.errors()[0]; let error = &parser.errors()[0];
let expected_error = r#"input:3:9: Expected marker "<<<<<<< SEARCH\n", found 'W'"#; let expected_error = r#"input:3:9: Expected marker "<<<<<<< SEARCH", found 'W'"#;
assert_eq!(format!("{}", error), expected_error); assert_eq!(format!("{}", error), expected_error);
} }