agent2: Port edit_file tool (#35844)
TODO: - [x] Authorization - [x] Restore tests Release Notes: - N/A --------- Co-authored-by: Antonio Scandurra <me@as-cii.com> Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>
This commit is contained in:
Agus Zubiaga
GitHub
parent
d705585a2e
commit
2526dcb5a5
20 changed files with 2075 additions and 414 deletions
|
|
@ -1,10 +1,11 @@
|
|||
use agent_client_protocol::{self as acp};
|
||||
use anyhow::{anyhow, Result};
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use assistant_tool::{outline, ActionLog};
|
||||
use gpui::{Entity, Task};
|
||||
use indoc::formatdoc;
|
||||
use language::{Anchor, Point};
|
||||
use project::{AgentLocation, Project, WorktreeSettings};
|
||||
use language_model::{LanguageModelImage, LanguageModelToolResultContent};
|
||||
use project::{image_store, AgentLocation, ImageItem, Project, WorktreeSettings};
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use settings::Settings;
|
||||
|
|
@ -59,6 +60,7 @@ impl ReadFileTool {
|
|||
|
||||
impl AgentTool for ReadFileTool {
|
||||
type Input = ReadFileToolInput;
|
||||
type Output = LanguageModelToolResultContent;
|
||||
|
||||
fn name(&self) -> SharedString {
|
||||
"read_file".into()
|
||||
|
|
@ -91,9 +93,9 @@ impl AgentTool for ReadFileTool {
|
|||
fn run(
|
||||
self: Arc<Self>,
|
||||
input: Self::Input,
|
||||
event_stream: ToolCallEventStream,
|
||||
_event_stream: ToolCallEventStream,
|
||||
cx: &mut App,
|
||||
) -> Task<Result<String>> {
|
||||
) -> Task<Result<LanguageModelToolResultContent>> {
|
||||
let Some(project_path) = self.project.read(cx).find_project_path(&input.path, cx) else {
|
||||
return Task::ready(Err(anyhow!("Path {} not found in project", &input.path)));
|
||||
};
|
||||
|
|
@ -132,51 +134,27 @@ impl AgentTool for ReadFileTool {
|
|||
|
||||
let file_path = input.path.clone();
|
||||
|
||||
event_stream.send_update(acp::ToolCallUpdateFields {
|
||||
locations: Some(vec![acp::ToolCallLocation {
|
||||
path: project_path.path.to_path_buf(),
|
||||
line: input.start_line,
|
||||
// TODO (tracked): use full range
|
||||
}]),
|
||||
..Default::default()
|
||||
});
|
||||
if image_store::is_image_file(&self.project, &project_path, cx) {
|
||||
return cx.spawn(async move |cx| {
|
||||
let image_entity: Entity<ImageItem> = cx
|
||||
.update(|cx| {
|
||||
self.project.update(cx, |project, cx| {
|
||||
project.open_image(project_path.clone(), cx)
|
||||
})
|
||||
})?
|
||||
.await?;
|
||||
|
||||
// TODO (tracked): images
|
||||
// if image_store::is_image_file(&self.project, &project_path, cx) {
|
||||
// let model = &self.thread.read(cx).selected_model;
|
||||
let image =
|
||||
image_entity.read_with(cx, |image_item, _| Arc::clone(&image_item.image))?;
|
||||
|
||||
// if !model.supports_images() {
|
||||
// return Task::ready(Err(anyhow!(
|
||||
// "Attempted to read an image, but Zed doesn't currently support sending images to {}.",
|
||||
// model.name().0
|
||||
// )))
|
||||
// .into();
|
||||
// }
|
||||
let language_model_image = cx
|
||||
.update(|cx| LanguageModelImage::from_image(image, cx))?
|
||||
.await
|
||||
.context("processing image")?;
|
||||
|
||||
// return cx.spawn(async move |cx| -> Result<ToolResultOutput> {
|
||||
// let image_entity: Entity<ImageItem> = cx
|
||||
// .update(|cx| {
|
||||
// self.project.update(cx, |project, cx| {
|
||||
// project.open_image(project_path.clone(), cx)
|
||||
// })
|
||||
// })?
|
||||
// .await?;
|
||||
|
||||
// let image =
|
||||
// image_entity.read_with(cx, |image_item, _| Arc::clone(&image_item.image))?;
|
||||
|
||||
// let language_model_image = cx
|
||||
// .update(|cx| LanguageModelImage::from_image(image, cx))?
|
||||
// .await
|
||||
// .context("processing image")?;
|
||||
|
||||
// Ok(ToolResultOutput {
|
||||
// content: ToolResultContent::Image(language_model_image),
|
||||
// output: None,
|
||||
// })
|
||||
// });
|
||||
// }
|
||||
//
|
||||
Ok(language_model_image.into())
|
||||
});
|
||||
}
|
||||
|
||||
let project = self.project.clone();
|
||||
let action_log = self.action_log.clone();
|
||||
|
|
@ -244,7 +222,7 @@ impl AgentTool for ReadFileTool {
|
|||
})?;
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
Ok(result.into())
|
||||
} else {
|
||||
// No line ranges specified, so check file size to see if it's too big.
|
||||
let file_size = buffer.read_with(cx, |buffer, _cx| buffer.text().len())?;
|
||||
|
|
@ -257,7 +235,7 @@ impl AgentTool for ReadFileTool {
|
|||
log.buffer_read(buffer, cx);
|
||||
})?;
|
||||
|
||||
Ok(result)
|
||||
Ok(result.into())
|
||||
} else {
|
||||
// File is too big, so return the outline
|
||||
// and a suggestion to read again with line numbers.
|
||||
|
|
@ -276,7 +254,8 @@ impl AgentTool for ReadFileTool {
|
|||
|
||||
Alternatively, you can fall back to the `grep` tool (if available)
|
||||
to search the file for specific content."
|
||||
})
|
||||
}
|
||||
.into())
|
||||
}
|
||||
}
|
||||
})
|
||||
|
|
@ -285,8 +264,6 @@ impl AgentTool for ReadFileTool {
|
|||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use crate::TestToolCallEventStream;
|
||||
|
||||
use super::*;
|
||||
use gpui::{AppContext, TestAppContext, UpdateGlobal as _};
|
||||
use language::{tree_sitter_rust, Language, LanguageConfig, LanguageMatcher};
|
||||
|
|
@ -304,7 +281,7 @@ mod test {
|
|||
let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;
|
||||
let action_log = cx.new(|_| ActionLog::new(project.clone()));
|
||||
let tool = Arc::new(ReadFileTool::new(project, action_log));
|
||||
let event_stream = TestToolCallEventStream::new();
|
||||
let (event_stream, _) = ToolCallEventStream::test();
|
||||
|
||||
let result = cx
|
||||
.update(|cx| {
|
||||
|
|
@ -313,7 +290,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.run(input, event_stream.stream(), cx)
|
||||
tool.run(input, event_stream, cx)
|
||||
})
|
||||
.await;
|
||||
assert_eq!(
|
||||
|
|
@ -321,6 +298,7 @@ mod test {
|
|||
"root/nonexistent_file.txt not found"
|
||||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_read_small_file(cx: &mut TestAppContext) {
|
||||
init_test(cx);
|
||||
|
|
@ -336,7 +314,6 @@ mod test {
|
|||
let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;
|
||||
let action_log = cx.new(|_| ActionLog::new(project.clone()));
|
||||
let tool = Arc::new(ReadFileTool::new(project, action_log));
|
||||
let event_stream = TestToolCallEventStream::new();
|
||||
let result = cx
|
||||
.update(|cx| {
|
||||
let input = ReadFileToolInput {
|
||||
|
|
@ -344,10 +321,10 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.run(input, event_stream.stream(), cx)
|
||||
tool.run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert_eq!(result.unwrap(), "This is a small file content");
|
||||
assert_eq!(result.unwrap(), "This is a small file content".into());
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
|
|
@ -367,18 +344,18 @@ mod test {
|
|||
language_registry.add(Arc::new(rust_lang()));
|
||||
let action_log = cx.new(|_| ActionLog::new(project.clone()));
|
||||
let tool = Arc::new(ReadFileTool::new(project, action_log));
|
||||
let event_stream = TestToolCallEventStream::new();
|
||||
let content = cx
|
||||
let result = cx
|
||||
.update(|cx| {
|
||||
let input = ReadFileToolInput {
|
||||
path: "root/large_file.rs".into(),
|
||||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let content = result.to_str().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
content.lines().skip(4).take(6).collect::<Vec<_>>(),
|
||||
|
|
@ -399,10 +376,11 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.run(input, event_stream.stream(), cx)
|
||||
tool.run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
let content = result.unwrap();
|
||||
.await
|
||||
.unwrap();
|
||||
let content = result.to_str().unwrap();
|
||||
let expected_content = (0..1000)
|
||||
.flat_map(|i| {
|
||||
vec![
|
||||
|
|
@ -438,7 +416,6 @@ mod test {
|
|||
|
||||
let action_log = cx.new(|_| ActionLog::new(project.clone()));
|
||||
let tool = Arc::new(ReadFileTool::new(project, action_log));
|
||||
let event_stream = TestToolCallEventStream::new();
|
||||
let result = cx
|
||||
.update(|cx| {
|
||||
let input = ReadFileToolInput {
|
||||
|
|
@ -446,10 +423,10 @@ mod test {
|
|||
start_line: Some(2),
|
||||
end_line: Some(4),
|
||||
};
|
||||
tool.run(input, event_stream.stream(), cx)
|
||||
tool.run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert_eq!(result.unwrap(), "Line 2\nLine 3\nLine 4");
|
||||
assert_eq!(result.unwrap(), "Line 2\nLine 3\nLine 4".into());
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
|
|
@ -467,7 +444,6 @@ mod test {
|
|||
let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;
|
||||
let action_log = cx.new(|_| ActionLog::new(project.clone()));
|
||||
let tool = Arc::new(ReadFileTool::new(project, action_log));
|
||||
let event_stream = TestToolCallEventStream::new();
|
||||
|
||||
// start_line of 0 should be treated as 1
|
||||
let result = cx
|
||||
|
|
@ -477,10 +453,10 @@ mod test {
|
|||
start_line: Some(0),
|
||||
end_line: Some(2),
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert_eq!(result.unwrap(), "Line 1\nLine 2");
|
||||
assert_eq!(result.unwrap(), "Line 1\nLine 2".into());
|
||||
|
||||
// end_line of 0 should result in at least 1 line
|
||||
let result = cx
|
||||
|
|
@ -490,10 +466,10 @@ mod test {
|
|||
start_line: Some(1),
|
||||
end_line: Some(0),
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert_eq!(result.unwrap(), "Line 1");
|
||||
assert_eq!(result.unwrap(), "Line 1".into());
|
||||
|
||||
// when start_line > end_line, should still return at least 1 line
|
||||
let result = cx
|
||||
|
|
@ -503,10 +479,10 @@ mod test {
|
|||
start_line: Some(3),
|
||||
end_line: Some(2),
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert_eq!(result.unwrap(), "Line 3");
|
||||
assert_eq!(result.unwrap(), "Line 3".into());
|
||||
}
|
||||
|
||||
fn init_test(cx: &mut TestAppContext) {
|
||||
|
|
@ -612,7 +588,6 @@ mod test {
|
|||
let project = Project::test(fs.clone(), [path!("/project_root").as_ref()], cx).await;
|
||||
let action_log = cx.new(|_| ActionLog::new(project.clone()));
|
||||
let tool = Arc::new(ReadFileTool::new(project, action_log));
|
||||
let event_stream = TestToolCallEventStream::new();
|
||||
|
||||
// Reading a file outside the project worktree should fail
|
||||
let result = cx
|
||||
|
|
@ -622,7 +597,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
|
|
@ -638,7 +613,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
|
|
@ -654,7 +629,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
|
|
@ -669,7 +644,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
|
|
@ -685,7 +660,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
|
|
@ -700,7 +675,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
|
|
@ -715,7 +690,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
|
|
@ -731,11 +706,11 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(result.is_ok(), "Should be able to read normal files");
|
||||
assert_eq!(result.unwrap(), "Normal file content");
|
||||
assert_eq!(result.unwrap(), "Normal file content".into());
|
||||
|
||||
// Path traversal attempts with .. should fail
|
||||
let result = cx
|
||||
|
|
@ -745,7 +720,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.run(input, event_stream.stream(), cx)
|
||||
tool.run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
|
|
@ -826,7 +801,6 @@ mod test {
|
|||
|
||||
let action_log = cx.new(|_| ActionLog::new(project.clone()));
|
||||
let tool = Arc::new(ReadFileTool::new(project.clone(), action_log.clone()));
|
||||
let event_stream = TestToolCallEventStream::new();
|
||||
|
||||
// Test reading allowed files in worktree1
|
||||
let result = cx
|
||||
|
|
@ -836,12 +810,15 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result, "fn main() { println!(\"Hello from worktree1\"); }");
|
||||
assert_eq!(
|
||||
result,
|
||||
"fn main() { println!(\"Hello from worktree1\"); }".into()
|
||||
);
|
||||
|
||||
// Test reading private file in worktree1 should fail
|
||||
let result = cx
|
||||
|
|
@ -851,7 +828,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
|
||||
|
|
@ -872,7 +849,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
|
||||
|
|
@ -893,14 +870,14 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
result,
|
||||
"export function greet() { return 'Hello from worktree2'; }"
|
||||
"export function greet() { return 'Hello from worktree2'; }".into()
|
||||
);
|
||||
|
||||
// Test reading private file in worktree2 should fail
|
||||
|
|
@ -911,7 +888,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
|
||||
|
|
@ -932,7 +909,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
|
||||
|
|
@ -954,7 +931,7 @@ mod test {
|
|||
start_line: None,
|
||||
end_line: None,
|
||||
};
|
||||
tool.clone().run(input, event_stream.stream(), cx)
|
||||
tool.clone().run(input, ToolCallEventStream::test().0, cx)
|
||||
})
|
||||
.await;
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue