agent2: More efficent read file tool (#38407)

Ben Brandt created 3 months ago

Before we were always reading the entire file into memory as a string.
Now we only read the range that is actually requested.

Release Notes:

- N/A

Change summary

crates/agent2/src/tools/read_file_tool.rs | 21 ++++++++++-----------
1 file changed, 10 insertions(+), 11 deletions(-)

Detailed changes

crates/agent2/src/tools/read_file_tool.rs

@@ -201,7 +201,6 @@ impl AgentTool for ReadFileTool {
             // Check if specific line ranges are provided
             let result = if input.start_line.is_some() || input.end_line.is_some() {
                 let result = buffer.read_with(cx, |buffer, _cx| {
-                    let text = buffer.text();
                     // .max(1) because despite instructions to be 1-indexed, sometimes the model passes 0.
                     let start = input.start_line.unwrap_or(1).max(1);
                     let start_row = start - 1;
@@ -210,13 +209,13 @@ impl AgentTool for ReadFileTool {
                         anchor = Some(buffer.anchor_before(Point::new(start_row, column)));
                     }
 
-                    let lines = text.split('\n').skip(start_row as usize);
-                    if let Some(end) = input.end_line {
-                        let count = end.saturating_sub(start).saturating_add(1); // Ensure at least 1 line
-                        itertools::intersperse(lines.take(count as usize), "\n").collect::<String>()
-                    } else {
-                        itertools::intersperse(lines, "\n").collect::<String>()
+                    let mut end_row = input.end_line.unwrap_or(u32::MAX);
+                    if end_row <= start_row {
+                        end_row = start_row + 1; // read at least one lines
                     }
+                    let start = buffer.anchor_before(Point::new(start_row, 0));
+                    let end = buffer.anchor_before(Point::new(end_row, 0));
+                    buffer.text_for_range(start..end).collect::<String>()
                 })?;
 
                 action_log.update(cx, |log, cx| {
@@ -445,7 +444,7 @@ mod test {
                 tool.run(input, ToolCallEventStream::test().0, cx)
             })
             .await;
-        assert_eq!(result.unwrap(), "Line 2\nLine 3\nLine 4".into());
+        assert_eq!(result.unwrap(), "Line 2\nLine 3\nLine 4\n".into());
     }
 
     #[gpui::test]
@@ -475,7 +474,7 @@ mod test {
                 tool.clone().run(input, ToolCallEventStream::test().0, cx)
             })
             .await;
-        assert_eq!(result.unwrap(), "Line 1\nLine 2".into());
+        assert_eq!(result.unwrap(), "Line 1\nLine 2\n".into());
 
         // end_line of 0 should result in at least 1 line
         let result = cx
@@ -488,7 +487,7 @@ mod test {
                 tool.clone().run(input, ToolCallEventStream::test().0, cx)
             })
             .await;
-        assert_eq!(result.unwrap(), "Line 1".into());
+        assert_eq!(result.unwrap(), "Line 1\n".into());
 
         // when start_line > end_line, should still return at least 1 line
         let result = cx
@@ -501,7 +500,7 @@ mod test {
                 tool.clone().run(input, ToolCallEventStream::test().0, cx)
             })
             .await;
-        assert_eq!(result.unwrap(), "Line 3".into());
+        assert_eq!(result.unwrap(), "Line 3\n".into());
     }
 
     fn init_test(cx: &mut TestAppContext) {