file_update.rs (18477B)
1 use serde_json::Value; 2 use similar::{ChangeTag, TextDiff}; 3 use std::path::Path; 4 5 /// Represents a proposed file modification from an AI tool call 6 #[derive(Debug, Clone)] 7 pub struct FileUpdate { 8 pub file_path: String, 9 pub update_type: FileUpdateType, 10 /// Cached diff lines (computed eagerly at construction) 11 diff_lines: Vec<DiffLine>, 12 } 13 14 #[derive(Debug, Clone)] 15 pub enum FileUpdateType { 16 /// Edit: replace old_string with new_string 17 Edit { 18 old_string: String, 19 new_string: String, 20 }, 21 /// Write: create/overwrite entire file 22 Write { content: String }, 23 /// Unified diff from an external tool (e.g. Codex) 24 UnifiedDiff { diff: String }, 25 } 26 27 /// A single line in a diff 28 #[derive(Debug, Clone)] 29 pub struct DiffLine { 30 pub tag: DiffTag, 31 pub content: String, 32 } 33 34 #[derive(Debug, Clone, Copy, PartialEq, Eq)] 35 pub enum DiffTag { 36 Equal, 37 Delete, 38 Insert, 39 } 40 41 impl From<ChangeTag> for DiffTag { 42 fn from(tag: ChangeTag) -> Self { 43 match tag { 44 ChangeTag::Equal => DiffTag::Equal, 45 ChangeTag::Delete => DiffTag::Delete, 46 ChangeTag::Insert => DiffTag::Insert, 47 } 48 } 49 } 50 51 /// Result of expanding diff context by reading the actual file from disk. 52 pub struct ExpandedDiffContext { 53 /// Extra Equal lines loaded above the diff 54 pub above: Vec<DiffLine>, 55 /// Extra Equal lines loaded below the diff 56 pub below: Vec<DiffLine>, 57 /// 1-based line number in the file where the first displayed line starts 58 pub start_line: usize, 59 /// Whether there are more lines above that could be loaded 60 pub has_more_above: bool, 61 /// Whether there are more lines below that could be loaded 62 pub has_more_below: bool, 63 } 64 65 impl FileUpdate { 66 /// Create a new FileUpdate, computing the diff eagerly 67 pub fn new(file_path: String, update_type: FileUpdateType) -> Self { 68 let diff_lines = Self::compute_diff_for(&update_type); 69 Self { 70 file_path, 71 update_type, 72 diff_lines, 73 } 74 } 75 76 /// Get the cached diff lines 77 pub fn diff_lines(&self) -> &[DiffLine] { 78 &self.diff_lines 79 } 80 81 /// Try to parse a FileUpdate from a tool name and tool input JSON 82 pub fn from_tool_call(tool_name: &str, tool_input: &Value) -> Option<Self> { 83 let obj = tool_input.as_object()?; 84 85 match tool_name { 86 "Edit" => { 87 let file_path = obj.get("file_path")?.as_str()?.to_string(); 88 let old_string = obj.get("old_string")?.as_str()?.to_string(); 89 let new_string = obj.get("new_string")?.as_str()?.to_string(); 90 91 Some(FileUpdate::new( 92 file_path, 93 FileUpdateType::Edit { 94 old_string, 95 new_string, 96 }, 97 )) 98 } 99 "Write" => { 100 let file_path = obj.get("file_path")?.as_str()?.to_string(); 101 let content = obj.get("content")?.as_str()?.to_string(); 102 103 Some(FileUpdate::new( 104 file_path, 105 FileUpdateType::Write { content }, 106 )) 107 } 108 _ => None, 109 } 110 } 111 112 /// Returns true if this is an Edit that changes at most `max_lines` lines 113 /// (deleted + inserted lines). Never returns true for Write operations. 114 /// 115 /// This counts actual changed lines using a diff, not total lines in the 116 /// strings. This is important because Claude Code typically includes 117 /// surrounding context lines for matching, so even a 1-line change may 118 /// have multi-line old_string/new_string. 119 pub fn is_small_edit(&self, max_lines: usize) -> bool { 120 match &self.update_type { 121 FileUpdateType::Edit { 122 old_string, 123 new_string, 124 } => { 125 let diff = TextDiff::from_lines(old_string.as_str(), new_string.as_str()); 126 let mut deleted_lines = 0; 127 let mut inserted_lines = 0; 128 for change in diff.iter_all_changes() { 129 match change.tag() { 130 ChangeTag::Delete => deleted_lines += 1, 131 ChangeTag::Insert => inserted_lines += 1, 132 ChangeTag::Equal => {} 133 } 134 } 135 deleted_lines <= max_lines && inserted_lines <= max_lines 136 } 137 FileUpdateType::Write { .. } | FileUpdateType::UnifiedDiff { .. } => false, 138 } 139 } 140 141 /// Read the file from disk and expand context around the edit. 142 /// 143 /// Returns `None` if this is not an Edit, the file can't be read, 144 /// or `old_string` can't be found in the file. 145 pub fn expanded_context( 146 &self, 147 extra_above: usize, 148 extra_below: usize, 149 ) -> Option<ExpandedDiffContext> { 150 let FileUpdateType::Edit { old_string, .. } = &self.update_type else { 151 return None; 152 }; 153 154 let file_content = std::fs::read_to_string(Path::new(&self.file_path)).ok()?; 155 156 // Find where old_string appears in the file 157 let byte_offset = file_content.find(old_string.as_str())?; 158 159 // Count newlines before the match to get 0-based start line index 160 let start_idx = file_content[..byte_offset] 161 .chars() 162 .filter(|&c| c == '\n') 163 .count(); 164 165 let file_lines: Vec<&str> = file_content.lines().collect(); 166 let total_lines = file_lines.len(); 167 168 let old_line_count = old_string.lines().count(); 169 let end_idx = start_idx + old_line_count; // exclusive, 0-based 170 171 // Extra lines above 172 let above_start = start_idx.saturating_sub(extra_above); 173 let above: Vec<DiffLine> = file_lines[above_start..start_idx] 174 .iter() 175 .map(|line| DiffLine { 176 tag: DiffTag::Equal, 177 content: format!("{}\n", line), 178 }) 179 .collect(); 180 181 // Extra lines below 182 let below_end = (end_idx + extra_below).min(total_lines); 183 let below: Vec<DiffLine> = file_lines[end_idx..below_end] 184 .iter() 185 .map(|line| DiffLine { 186 tag: DiffTag::Equal, 187 content: format!("{}\n", line), 188 }) 189 .collect(); 190 191 Some(ExpandedDiffContext { 192 start_line: above_start + 1, // 1-based 193 has_more_above: above_start > 0, 194 has_more_below: below_end < total_lines, 195 above, 196 below, 197 }) 198 } 199 200 /// Compute the diff lines for an update type (internal helper) 201 fn compute_diff_for(update_type: &FileUpdateType) -> Vec<DiffLine> { 202 match update_type { 203 FileUpdateType::Edit { 204 old_string, 205 new_string, 206 } => { 207 let diff = TextDiff::from_lines(old_string.as_str(), new_string.as_str()); 208 diff.iter_all_changes() 209 .map(|change| DiffLine { 210 tag: change.tag().into(), 211 content: change.value().to_string(), 212 }) 213 .collect() 214 } 215 FileUpdateType::Write { content } => { 216 // For writes, everything is an insertion 217 content 218 .lines() 219 .map(|line| DiffLine { 220 tag: DiffTag::Insert, 221 content: format!("{}\n", line), 222 }) 223 .collect() 224 } 225 FileUpdateType::UnifiedDiff { diff } => { 226 // Parse unified diff format: lines starting with '+'/'-'/' ' 227 // Skip header lines (---/+++/@@ lines) 228 diff.lines() 229 .filter(|line| { 230 !line.starts_with("---") 231 && !line.starts_with("+++") 232 && !line.starts_with("@@") 233 }) 234 .map(|line| { 235 if let Some(rest) = line.strip_prefix('+') { 236 DiffLine { 237 tag: DiffTag::Insert, 238 content: format!("{}\n", rest), 239 } 240 } else if let Some(rest) = line.strip_prefix('-') { 241 DiffLine { 242 tag: DiffTag::Delete, 243 content: format!("{}\n", rest), 244 } 245 } else { 246 // Context line (starts with ' ' or is bare) 247 let content = line.strip_prefix(' ').unwrap_or(line); 248 DiffLine { 249 tag: DiffTag::Equal, 250 content: format!("{}\n", content), 251 } 252 } 253 }) 254 .collect() 255 } 256 } 257 } 258 } 259 260 #[cfg(test)] 261 mod tests { 262 use super::*; 263 use serde_json::json; 264 265 #[test] 266 fn test_is_small_edit_single_line() { 267 let update = FileUpdate::new( 268 "test.rs".to_string(), 269 FileUpdateType::Edit { 270 old_string: "foo".to_string(), 271 new_string: "bar".to_string(), 272 }, 273 ); 274 assert!( 275 update.is_small_edit(2), 276 "Single line without newline should be small" 277 ); 278 } 279 280 #[test] 281 fn test_is_small_edit_single_line_with_newline() { 282 let update = FileUpdate::new( 283 "test.rs".to_string(), 284 FileUpdateType::Edit { 285 old_string: "foo\n".to_string(), 286 new_string: "bar\n".to_string(), 287 }, 288 ); 289 assert!( 290 update.is_small_edit(2), 291 "Single line with trailing newline should be small" 292 ); 293 } 294 295 #[test] 296 fn test_is_small_edit_two_lines() { 297 let update = FileUpdate::new( 298 "test.rs".to_string(), 299 FileUpdateType::Edit { 300 old_string: "foo\nbar".to_string(), 301 new_string: "baz\nqux".to_string(), 302 }, 303 ); 304 assert!( 305 update.is_small_edit(2), 306 "Two lines without trailing newline should be small" 307 ); 308 } 309 310 #[test] 311 fn test_is_small_edit_two_lines_with_newline() { 312 let update = FileUpdate::new( 313 "test.rs".to_string(), 314 FileUpdateType::Edit { 315 old_string: "foo\nbar\n".to_string(), 316 new_string: "baz\nqux\n".to_string(), 317 }, 318 ); 319 assert!( 320 update.is_small_edit(2), 321 "Two lines with trailing newline should be small" 322 ); 323 } 324 325 #[test] 326 fn test_is_small_edit_three_lines_not_small() { 327 let update = FileUpdate::new( 328 "test.rs".to_string(), 329 FileUpdateType::Edit { 330 old_string: "foo\nbar\nbaz".to_string(), 331 new_string: "a\nb\nc".to_string(), 332 }, 333 ); 334 assert!(!update.is_small_edit(2), "Three lines should NOT be small"); 335 } 336 337 #[test] 338 fn test_is_small_edit_write_never_small() { 339 let update = FileUpdate::new( 340 "test.rs".to_string(), 341 FileUpdateType::Write { 342 content: "x".to_string(), 343 }, 344 ); 345 assert!( 346 !update.is_small_edit(2), 347 "Write operations should never be small" 348 ); 349 } 350 351 #[test] 352 fn test_is_small_edit_old_small_new_large() { 353 let update = FileUpdate::new( 354 "test.rs".to_string(), 355 FileUpdateType::Edit { 356 old_string: "foo".to_string(), 357 new_string: "a\nb\nc\nd".to_string(), 358 }, 359 ); 360 assert!( 361 !update.is_small_edit(2), 362 "Large new_string should NOT be small" 363 ); 364 } 365 366 #[test] 367 fn test_is_small_edit_old_large_new_small() { 368 let update = FileUpdate::new( 369 "test.rs".to_string(), 370 FileUpdateType::Edit { 371 old_string: "a\nb\nc\nd".to_string(), 372 new_string: "foo".to_string(), 373 }, 374 ); 375 assert!( 376 !update.is_small_edit(2), 377 "Large old_string should NOT be small" 378 ); 379 } 380 381 #[test] 382 fn test_from_tool_call_edit() { 383 let input = json!({ 384 "file_path": "/path/to/file.rs", 385 "old_string": "old", 386 "new_string": "new" 387 }); 388 let update = FileUpdate::from_tool_call("Edit", &input).unwrap(); 389 assert_eq!(update.file_path, "/path/to/file.rs"); 390 assert!(update.is_small_edit(2)); 391 } 392 393 #[test] 394 fn test_from_tool_call_write() { 395 let input = json!({ 396 "file_path": "/path/to/file.rs", 397 "content": "content" 398 }); 399 let update = FileUpdate::from_tool_call("Write", &input).unwrap(); 400 assert_eq!(update.file_path, "/path/to/file.rs"); 401 assert!(!update.is_small_edit(2)); 402 } 403 404 #[test] 405 fn test_from_tool_call_unknown_tool() { 406 let input = json!({ 407 "file_path": "/path/to/file.rs" 408 }); 409 assert!(FileUpdate::from_tool_call("Bash", &input).is_none()); 410 } 411 412 #[test] 413 fn test_is_small_edit_realistic_claude_edit_with_context() { 414 // Claude Code typically sends old_string/new_string with context lines 415 // for matching. Even a "small" 1-line change includes context. 416 // This tests what an actual Edit tool call might look like. 417 let input = json!({ 418 "file_path": "/path/to/file.rs", 419 "old_string": " fn foo() {\n let x = 1;\n }", 420 "new_string": " fn foo() {\n let x = 2;\n }" 421 }); 422 let update = FileUpdate::from_tool_call("Edit", &input).unwrap(); 423 // Only 1 line actually changed (let x = 1 -> let x = 2) 424 // The context lines (fn foo() and }) are the same 425 assert!( 426 update.is_small_edit(2), 427 "1-line actual change should be small, even with 3 lines of context" 428 ); 429 } 430 431 #[test] 432 fn test_is_small_edit_minimal_change() { 433 // A truly minimal single-line change 434 let input = json!({ 435 "file_path": "/path/to/file.rs", 436 "old_string": "let x = 1;", 437 "new_string": "let x = 2;" 438 }); 439 let update = FileUpdate::from_tool_call("Edit", &input).unwrap(); 440 assert!( 441 update.is_small_edit(2), 442 "Single-line change should be small" 443 ); 444 } 445 446 // ----------------------------------------------------------------------- 447 // UnifiedDiff tests 448 // ----------------------------------------------------------------------- 449 450 #[test] 451 fn test_unified_diff_basic() { 452 let update = FileUpdate::new( 453 "test.rs".to_string(), 454 FileUpdateType::UnifiedDiff { 455 diff: "--- a/test.rs\n+++ b/test.rs\n@@ -1,3 +1,3 @@\n context\n-old line\n+new line\n more context\n" 456 .to_string(), 457 }, 458 ); 459 let lines = FileUpdate::compute_diff_for(&update.update_type); 460 assert_eq!(lines.len(), 4); 461 assert_eq!(lines[0].tag, DiffTag::Equal); 462 assert_eq!(lines[0].content, "context\n"); 463 assert_eq!(lines[1].tag, DiffTag::Delete); 464 assert_eq!(lines[1].content, "old line\n"); 465 assert_eq!(lines[2].tag, DiffTag::Insert); 466 assert_eq!(lines[2].content, "new line\n"); 467 assert_eq!(lines[3].tag, DiffTag::Equal); 468 assert_eq!(lines[3].content, "more context\n"); 469 } 470 471 #[test] 472 fn test_unified_diff_skips_headers() { 473 let update = FileUpdate::new( 474 "test.rs".to_string(), 475 FileUpdateType::UnifiedDiff { 476 diff: "--- a/old.rs\n+++ b/new.rs\n@@ -10,4 +10,4 @@\n+added\n".to_string(), 477 }, 478 ); 479 let lines = FileUpdate::compute_diff_for(&update.update_type); 480 assert_eq!(lines.len(), 1); 481 assert_eq!(lines[0].tag, DiffTag::Insert); 482 assert_eq!(lines[0].content, "added\n"); 483 } 484 485 #[test] 486 fn test_unified_diff_delete_only() { 487 let update = FileUpdate::new( 488 "test.rs".to_string(), 489 FileUpdateType::UnifiedDiff { 490 diff: "-removed line 1\n-removed line 2\n".to_string(), 491 }, 492 ); 493 let lines = FileUpdate::compute_diff_for(&update.update_type); 494 assert_eq!(lines.len(), 2); 495 assert!(lines.iter().all(|l| l.tag == DiffTag::Delete)); 496 } 497 498 #[test] 499 fn test_unified_diff_insert_only() { 500 let update = FileUpdate::new( 501 "test.rs".to_string(), 502 FileUpdateType::UnifiedDiff { 503 diff: "+new line 1\n+new line 2\n+new line 3\n".to_string(), 504 }, 505 ); 506 let lines = FileUpdate::compute_diff_for(&update.update_type); 507 assert_eq!(lines.len(), 3); 508 assert!(lines.iter().all(|l| l.tag == DiffTag::Insert)); 509 } 510 511 #[test] 512 fn test_unified_diff_empty() { 513 let update = FileUpdate::new( 514 "test.rs".to_string(), 515 FileUpdateType::UnifiedDiff { 516 diff: String::new(), 517 }, 518 ); 519 let lines = FileUpdate::compute_diff_for(&update.update_type); 520 assert!(lines.is_empty()); 521 } 522 523 #[test] 524 fn test_unified_diff_is_never_small_edit() { 525 let update = FileUpdate::new( 526 "test.rs".to_string(), 527 FileUpdateType::UnifiedDiff { 528 diff: "+x\n".to_string(), 529 }, 530 ); 531 assert!(!update.is_small_edit(100)); 532 } 533 534 #[test] 535 fn test_line_count_behavior() { 536 // Document how lines().count() behaves 537 assert_eq!("foo".lines().count(), 1); 538 assert_eq!("foo\n".lines().count(), 1); // trailing newline doesn't add line 539 assert_eq!("foo\nbar".lines().count(), 2); 540 assert_eq!("foo\nbar\n".lines().count(), 2); 541 assert_eq!("foo\nbar\nbaz".lines().count(), 3); 542 assert_eq!("foo\nbar\nbaz\n".lines().count(), 3); 543 // Empty strings 544 assert_eq!("".lines().count(), 0); 545 assert_eq!("\n".lines().count(), 1); // just a newline counts as 1 empty line 546 } 547 }