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

Add locations to native agent tool calls, and wire them up to UI (#36058)

Browse source 
Release Notes:

- N/A

---------

Co-authored-by: Conrad <conrad@zed.dev>
This commit is contained in:
Cole Miller 2025-08-12 21:48:28 -04:00 • committed by GitHub
parent d78bd8f1d7
commit 1957e1f642
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 257 additions and 114 deletions
Download patch file Download diff file Expand all files Collapse all files

1
Cargo.lock generated
View file

@ -231,6 +231,7 @@ dependencies = [
"task",
"tempfile",
"terminal",
"text",
"theme",
"tree-sitter-rust",
"ui",

155
crates/acp_thread/src/acp_thread.rs
View file

@ -13,9 +13,9 @@ use agent_client_protocol::{self as acp};
use anyhow::{Context as _, Result};
use editor::Bias;
use futures::{FutureExt, channel::oneshot, future::BoxFuture};
use gpui::{AppContext, Context, Entity, EventEmitter, SharedString, Task};
use gpui::{AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
use itertools::Itertools;
use language::{Anchor, Buffer, BufferSnapshot, LanguageRegistry, Point, text_diff};
use language::{Anchor, Buffer, BufferSnapshot, LanguageRegistry, Point, ToPoint, text_diff};
use markdown::Markdown;
use project::{AgentLocation, Project};
use std::collections::HashMap;
@ -122,9 +122,17 @@ impl AgentThreadEntry {
}
}
pub fn locations(&self) -> Option<&[acp::ToolCallLocation]> {
if let AgentThreadEntry::ToolCall(ToolCall { locations, .. }) = self {
Some(locations)
pub fn location(&self, ix: usize) -> Option<(acp::ToolCallLocation, AgentLocation)> {
if let AgentThreadEntry::ToolCall(ToolCall {
locations,
resolved_locations,
..
}) = self
{
Some((
locations.get(ix)?.clone(),
resolved_locations.get(ix)?.clone()?,
))
} else {
None
}
@ -139,6 +147,7 @@ pub struct ToolCall {
pub content: Vec<ToolCallContent>,
pub status: ToolCallStatus,
pub locations: Vec<acp::ToolCallLocation>,
pub resolved_locations: Vec<Option<AgentLocation>>,
pub raw_input: Option<serde_json::Value>,
pub raw_output: Option<serde_json::Value>,
}
@ -167,6 +176,7 @@ impl ToolCall {
.map(|content| ToolCallContent::from_acp(content, language_registry.clone(), cx))
.collect(),
locations: tool_call.locations,
resolved_locations: Vec::default(),
status,
raw_input: tool_call.raw_input,
raw_output: tool_call.raw_output,
@ -260,6 +270,57 @@ impl ToolCall {
}
markdown
}
async fn resolve_location(
location: acp::ToolCallLocation,
project: WeakEntity<Project>,
cx: &mut AsyncApp,
) -> Option<AgentLocation> {
let buffer = project
.update(cx, |project, cx| {
if let Some(path) = project.project_path_for_absolute_path(&location.path, cx) {
Some(project.open_buffer(path, cx))
} else {
None
}
})
.ok()??;
let buffer = buffer.await.log_err()?;
let position = buffer
.update(cx, |buffer, _| {
if let Some(row) = location.line {
let snapshot = buffer.snapshot();
let column = snapshot.indent_size_for_line(row).len;
let point = snapshot.clip_point(Point::new(row, column), Bias::Left);
snapshot.anchor_before(point)
} else {
Anchor::MIN
}
})
.ok()?;
Some(AgentLocation {
buffer: buffer.downgrade(),
position,
})
}
fn resolve_locations(
&self,
project: Entity<Project>,
cx: &mut App,
) -> Task<Vec<Option<AgentLocation>>> {
let locations = self.locations.clone();
project.update(cx, |_, cx| {
cx.spawn(async move |project, cx| {
let mut new_locations = Vec::new();
for location in locations {
new_locations.push(Self::resolve_location(location, project.clone(), cx).await);
}
new_locations
})
})
}
}
#[derive(Debug)]
@ -804,7 +865,11 @@ impl AcpThread {
.context("Tool call not found")?;
match update {
ToolCallUpdate::UpdateFields(update) => {
let location_updated = update.fields.locations.is_some();
current_call.update_fields(update.fields, languages, cx);
if location_updated {
self.resolve_locations(update.id.clone(), cx);
}
}
ToolCallUpdate::UpdateDiff(update) => {
current_call.content.clear();
@ -841,8 +906,7 @@ impl AcpThread {
) {
let language_registry = self.project.read(cx).languages().clone();
let call = ToolCall::from_acp(tool_call, status, language_registry, cx);
let location = call.locations.last().cloned();
let id = call.id.clone();
if let Some((ix, current_call)) = self.tool_call_mut(&call.id) {
*current_call = call;
@ -850,11 +914,9 @@ impl AcpThread {
cx.emit(AcpThreadEvent::EntryUpdated(ix));
} else {
self.push_entry(AgentThreadEntry::ToolCall(call), cx);
}
};
if let Some(location) = location {
self.set_project_location(location, cx)
}
self.resolve_locations(id, cx);
}
fn tool_call_mut(&mut self, id: &acp::ToolCallId) -> Option<(usize, &mut ToolCall)> {
@ -875,35 +937,50 @@ impl AcpThread {
})
}
pub fn set_project_location(&self, location: acp::ToolCallLocation, cx: &mut Context<Self>) {
self.project.update(cx, |project, cx| {
let Some(path) = project.project_path_for_absolute_path(&location.path, cx) else {
return;
};
let buffer = project.open_buffer(path, cx);
cx.spawn(async move |project, cx| {
let buffer = buffer.await?;
project.update(cx, |project, cx| {
let position = if let Some(line) = location.line {
let snapshot = buffer.read(cx).snapshot();
let point = snapshot.clip_point(Point::new(line, 0), Bias::Left);
snapshot.anchor_before(point)
} else {
Anchor::MIN
};
project.set_agent_location(
Some(AgentLocation {
buffer: buffer.downgrade(),
position,
}),
cx,
);
})
pub fn resolve_locations(&mut self, id: acp::ToolCallId, cx: &mut Context<Self>) {
let project = self.project.clone();
let Some((_, tool_call)) = self.tool_call_mut(&id) else {
return;
};
let task = tool_call.resolve_locations(project, cx);
cx.spawn(async move |this, cx| {
let resolved_locations = task.await;
this.update(cx, |this, cx| {
let project = this.project.clone();
let Some((ix, tool_call)) = this.tool_call_mut(&id) else {
return;
};
if let Some(Some(location)) = resolved_locations.last() {
project.update(cx, |project, cx| {
if let Some(agent_location) = project.agent_location() {
let should_ignore = agent_location.buffer == location.buffer
&& location
.buffer
.update(cx, |buffer, _| {
let snapshot = buffer.snapshot();
let old_position =
agent_location.position.to_point(&snapshot);
let new_position = location.position.to_point(&snapshot);
// ignore this so that when we get updates from the edit tool
// the position doesn't reset to the startof line
old_position.row == new_position.row
&& old_position.column > new_position.column
})
.ok()
.unwrap_or_default();
if !should_ignore {
project.set_agent_location(Some(location.clone()), cx);
}
}
});
}
if tool_call.resolved_locations != resolved_locations {
tool_call.resolved_locations = resolved_locations;
cx.emit(AcpThreadEvent::EntryUpdated(ix));
}
})
.detach_and_log_err(cx);
});
})
.detach();
}
pub fn request_tool_call_authorization(

1
crates/agent2/Cargo.toml
View file

@ -49,6 +49,7 @@ settings.workspace = true
smol.workspace = true
task.workspace = true
terminal.workspace = true
text.workspace = true
ui.workspace = true
util.workspace = true
uuid.workspace = true

42
crates/agent2/src/tools/edit_file_tool.rs
View file

@ -1,12 +1,13 @@
use crate::{AgentTool, Thread, ToolCallEventStream};
use acp_thread::Diff;
use agent_client_protocol as acp;
use agent_client_protocol::{self as acp, ToolCallLocation, ToolCallUpdateFields};
use anyhow::{Context as _, Result, anyhow};
use assistant_tools::edit_agent::{EditAgent, EditAgentOutput, EditAgentOutputEvent, EditFormat};
use cloud_llm_client::CompletionIntent;
use collections::HashSet;
use gpui::{App, AppContext, AsyncApp, Entity, Task};
use indoc::formatdoc;
use language::ToPoint;
use language::language_settings::{self, FormatOnSave};
use language_model::LanguageModelToolResultContent;
use paths;
@ -225,6 +226,16 @@ impl AgentTool for EditFileTool {
Ok(path) => path,
Err(err) => return Task::ready(Err(anyhow!(err))),
};
let abs_path = project.read(cx).absolute_path(&project_path, cx);
if let Some(abs_path) = abs_path.clone() {
event_stream.update_fields(ToolCallUpdateFields {
locations: Some(vec![acp::ToolCallLocation {
path: abs_path,
line: None,
}]),
..Default::default()
});
}
let request = self.thread.update(cx, |thread, cx| {
thread.build_completion_request(CompletionIntent::ToolResults, cx)
@ -283,13 +294,38 @@ impl AgentTool for EditFileTool {
let mut hallucinated_old_text = false;
let mut ambiguous_ranges = Vec::new();
let mut emitted_location = false;
while let Some(event) = events.next().await {
match event {
EditAgentOutputEvent::Edited => {},
EditAgentOutputEvent::Edited(range) => {
if !emitted_location {
let line = buffer.update(cx, |buffer, _cx| {
range.start.to_point(&buffer.snapshot()).row
}).ok();
if let Some(abs_path) = abs_path.clone() {
event_stream.update_fields(ToolCallUpdateFields {
locations: Some(vec![ToolCallLocation { path: abs_path, line }]),
..Default::default()
});
}
emitted_location = true;
}
},
EditAgentOutputEvent::UnresolvedEditRange => hallucinated_old_text = true,
EditAgentOutputEvent::AmbiguousEditRange(ranges) => ambiguous_ranges = ranges,
EditAgentOutputEvent::ResolvingEditRange(range) => {
diff.update(cx, |card, cx| card.reveal_range(range, cx))?;
diff.update(cx, |card, cx| card.reveal_range(range.clone(), cx))?;
// if !emitted_location {
// let line = buffer.update(cx, |buffer, _cx| {
// range.start.to_point(&buffer.snapshot()).row
// }).ok();
// if let Some(abs_path) = abs_path.clone() {
// event_stream.update_fields(ToolCallUpdateFields {
// locations: Some(vec![ToolCallLocation { path: abs_path, line }]),
// ..Default::default()
// });
// }
// }
}
}
}

63
crates/agent2/src/tools/read_file_tool.rs
View file

@ -1,10 +1,10 @@
use action_log::ActionLog;
use agent_client_protocol::{self as acp};
use agent_client_protocol::{self as acp, ToolCallUpdateFields};
use anyhow::{Context as _, Result, anyhow};
use assistant_tool::outline;
use gpui::{App, Entity, SharedString, Task};
use indoc::formatdoc;
use language::{Anchor, Point};
use language::Point;
use language_model::{LanguageModelImage, LanguageModelToolResultContent};
use project::{AgentLocation, ImageItem, Project, WorktreeSettings, image_store};
use schemars::JsonSchema;
@ -97,7 +97,7 @@ impl AgentTool for ReadFileTool {
fn run(
self: Arc<Self>,
input: Self::Input,
_event_stream: ToolCallEventStream,
event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<LanguageModelToolResultContent>> {
let Some(project_path) = self.project.read(cx).find_project_path(&input.path, cx) else {
@ -166,7 +166,9 @@ impl AgentTool for ReadFileTool {
cx.spawn(async move |cx| {
let buffer = cx
.update(|cx| {
project.update(cx, |project, cx| project.open_buffer(project_path, cx))
project.update(cx, |project, cx| {
project.open_buffer(project_path.clone(), cx)
})
})?
.await?;
if buffer.read_with(cx, |buffer, _| {
@ -178,19 +180,10 @@ impl AgentTool for ReadFileTool {
anyhow::bail!("{file_path} not found");
}
project.update(cx, |project, cx| {
project.set_agent_location(
Some(AgentLocation {
buffer: buffer.downgrade(),
position: Anchor::MIN,
}),
cx,
);
})?;
let mut anchor = None;
// Check if specific line ranges are provided
if input.start_line.is_some() || input.end_line.is_some() {
let mut anchor = None;
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.
@ -214,18 +207,6 @@ impl AgentTool for ReadFileTool {
log.buffer_read(buffer.clone(), cx);
})?;
if let Some(anchor) = anchor {
project.update(cx, |project, cx| {
project.set_agent_location(
Some(AgentLocation {
buffer: buffer.downgrade(),
position: anchor,
}),
cx,
);
})?;
}
Ok(result.into())
} else {
// No line ranges specified, so check file size to see if it's too big.
@ -236,7 +217,7 @@ impl AgentTool for ReadFileTool {
let result = buffer.read_with(cx, |buffer, _cx| buffer.text())?;
action_log.update(cx, |log, cx| {
log.buffer_read(buffer, cx);
log.buffer_read(buffer.clone(), cx);
})?;
Ok(result.into())
@ -244,7 +225,8 @@ impl AgentTool for ReadFileTool {
// File is too big, so return the outline
// and a suggestion to read again with line numbers.
let outline =
outline::file_outline(project, file_path, action_log, None, cx).await?;
outline::file_outline(project.clone(), file_path, action_log, None, cx)
.await?;
Ok(formatdoc! {"
This file was too big to read all at once.
@ -261,7 +243,28 @@ impl AgentTool for ReadFileTool {
}
.into())
}
}
};
project.update(cx, |project, cx| {
if let Some(abs_path) = project.absolute_path(&project_path, cx) {
project.set_agent_location(
Some(AgentLocation {
buffer: buffer.downgrade(),
position: anchor.unwrap_or(text::Anchor::MIN),
}),
cx,
);
event_stream.update_fields(ToolCallUpdateFields {
locations: Some(vec![acp::ToolCallLocation {
path: abs_path,
line: input.start_line.map(|line| line.saturating_sub(1)),
}]),
..Default::default()
});
}
})?;
result
})
}
}

38
crates/agent_ui/src/acp/thread_view.rs
View file

@ -27,6 +27,7 @@ use language::{Buffer, Language};
use markdown::{HeadingLevelStyles, Markdown, MarkdownElement, MarkdownStyle};
use parking_lot::Mutex;
use project::{CompletionIntent, Project};
use rope::Point;
use settings::{Settings as _, SettingsStore};
use std::path::PathBuf;
use std::{
@ -2679,26 +2680,24 @@ impl AcpThreadView {
window: &mut Window,
cx: &mut Context<Self>,
) -> Option<()> {
let location = self
let (tool_call_location, agent_location) = self
.thread()?
.read(cx)
.entries()
.get(entry_ix)?
.locations()?
.get(location_ix)?;
.location(location_ix)?;
let project_path = self
.project
.read(cx)
.find_project_path(&location.path, cx)?;
.find_project_path(&tool_call_location.path, cx)?;
let open_task = self
.workspace
.update(cx, |worskpace, cx| {
worskpace.open_path(project_path, None, true, window, cx)
.update(cx, |workspace, cx| {
workspace.open_path(project_path, None, true, window, cx)
})
.log_err()?;
window
.spawn(cx, async move |cx| {
let item = open_task.await?;
@ -2708,17 +2707,22 @@ impl AcpThreadView {
};
active_editor.update_in(cx, |editor, window, cx| {
let snapshot = editor.buffer().read(cx).snapshot(cx);
let first_hunk = editor
.diff_hunks_in_ranges(
&[editor::Anchor::min()..editor::Anchor::max()],
&snapshot,
)
.next();
if let Some(first_hunk) = first_hunk {
let first_hunk_start = first_hunk.multi_buffer_range().start;
let multibuffer = editor.buffer().read(cx);
let buffer = multibuffer.as_singleton();
if agent_location.buffer.upgrade() == buffer {
let excerpt_id = multibuffer.excerpt_ids().first().cloned();
let anchor = editor::Anchor::in_buffer(
excerpt_id.unwrap(),
buffer.unwrap().read(cx).remote_id(),
agent_location.position,
);
editor.change_selections(Default::default(), window, cx, |selections| {
selections.select_anchor_ranges([first_hunk_start..first_hunk_start]);
selections.select_anchor_ranges([anchor..anchor]);
})
} else {
let row = tool_call_location.line.unwrap_or_default();
editor.change_selections(Default::default(), window, cx, |selections| {
selections.select_ranges([Point::new(row, 0)..Point::new(row, 0)]);
})
}
})?;

69
crates/assistant_tools/src/edit_agent.rs
View file

@ -65,7 +65,7 @@ pub enum EditAgentOutputEvent {
ResolvingEditRange(Range<Anchor>),
UnresolvedEditRange,
AmbiguousEditRange(Vec<Range<usize>>),
Edited,
Edited(Range<Anchor>),
}
#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
@ -178,7 +178,9 @@ impl EditAgent {
)
});
output_events_tx
.unbounded_send(EditAgentOutputEvent::Edited)
.unbounded_send(EditAgentOutputEvent::Edited(
language::Anchor::MIN..language::Anchor::MAX,
))
.ok();
})?;
@ -200,7 +202,9 @@ impl EditAgent {
});
})?;
output_events_tx
.unbounded_send(EditAgentOutputEvent::Edited)
.unbounded_send(EditAgentOutputEvent::Edited(
language::Anchor::MIN..language::Anchor::MAX,
))
.ok();
}
}
@ -336,8 +340,8 @@ impl EditAgent {
// Edit the buffer and report edits to the action log as part of the
// same effect cycle, otherwise the edit will be reported as if the
// user made it.
cx.update(|cx| {
let max_edit_end = buffer.update(cx, |buffer, cx| {
let (min_edit_start, max_edit_end) = cx.update(|cx| {
let (min_edit_start, max_edit_end) = buffer.update(cx, |buffer, cx| {
buffer.edit(edits.iter().cloned(), None, cx);
let max_edit_end = buffer
.summaries_for_anchors::<Point, _>(
@ -345,7 +349,16 @@ impl EditAgent {
)
.max()
.unwrap();
buffer.anchor_before(max_edit_end)
let min_edit_start = buffer
.summaries_for_anchors::<Point, _>(
edits.iter().map(|(range, _)| &range.start),
)
.min()
.unwrap();
(
buffer.anchor_after(min_edit_start),
buffer.anchor_before(max_edit_end),
)
});
self.action_log
.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
@ -358,9 +371,10 @@ impl EditAgent {
cx,
);
});
(min_edit_start, max_edit_end)
})?;
output_events
.unbounded_send(EditAgentOutputEvent::Edited)
.unbounded_send(EditAgentOutputEvent::Edited(min_edit_start..max_edit_end))
.ok();
}
@ -755,6 +769,7 @@ mod tests {
use gpui::{AppContext, TestAppContext};
use indoc::indoc;
use language_model::fake_provider::FakeLanguageModel;
use pretty_assertions::assert_matches;
use project::{AgentLocation, Project};
use rand::prelude::*;
use rand::rngs::StdRng;
@ -992,7 +1007,10 @@ mod tests {
model.send_last_completion_stream_text_chunk("<new_text>abX");
cx.run_until_parked();
assert_eq!(drain_events(&mut events), [EditAgentOutputEvent::Edited]);
assert_matches!(
drain_events(&mut events).as_slice(),
[EditAgentOutputEvent::Edited(_)]
);
assert_eq!(
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
"abXc\ndef\nghi\njkl"
@ -1007,7 +1025,10 @@ mod tests {
model.send_last_completion_stream_text_chunk("cY");
cx.run_until_parked();
assert_eq!(drain_events(&mut events), [EditAgentOutputEvent::Edited]);
assert_matches!(
drain_events(&mut events).as_slice(),
[EditAgentOutputEvent::Edited { .. }]
);
assert_eq!(
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
"abXcY\ndef\nghi\njkl"
@ -1118,9 +1139,9 @@ mod tests {
model.send_last_completion_stream_text_chunk("GHI</new_text>");
cx.run_until_parked();
assert_eq!(
drain_events(&mut events),
vec![EditAgentOutputEvent::Edited]
assert_matches!(
drain_events(&mut events).as_slice(),
[EditAgentOutputEvent::Edited { .. }]
);
assert_eq!(
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
@ -1165,9 +1186,9 @@ mod tests {
);
cx.run_until_parked();
assert_eq!(
drain_events(&mut events),
vec![EditAgentOutputEvent::Edited]
assert_matches!(
drain_events(&mut events).as_slice(),
[EditAgentOutputEvent::Edited(_)]
);
assert_eq!(
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
@ -1183,9 +1204,9 @@ mod tests {
chunks_tx.unbounded_send("```\njkl\n").unwrap();
cx.run_until_parked();
assert_eq!(
drain_events(&mut events),
vec![EditAgentOutputEvent::Edited]
assert_matches!(
drain_events(&mut events).as_slice(),
[EditAgentOutputEvent::Edited { .. }]
);
assert_eq!(
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
@ -1201,9 +1222,9 @@ mod tests {
chunks_tx.unbounded_send("mno\n").unwrap();
cx.run_until_parked();
assert_eq!(
drain_events(&mut events),
vec![EditAgentOutputEvent::Edited]
assert_matches!(
drain_events(&mut events).as_slice(),
[EditAgentOutputEvent::Edited { .. }]
);
assert_eq!(
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
@ -1219,9 +1240,9 @@ mod tests {
chunks_tx.unbounded_send("pqr\n```").unwrap();
cx.run_until_parked();
assert_eq!(
drain_events(&mut events),
vec![EditAgentOutputEvent::Edited]
assert_matches!(
drain_events(&mut events).as_slice(),
[EditAgentOutputEvent::Edited(_)],
);
assert_eq!(
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),

2
crates/assistant_tools/src/edit_file_tool.rs
View file

@ -307,7 +307,7 @@ impl Tool for EditFileTool {
let mut ambiguous_ranges = Vec::new();
while let Some(event) = events.next().await {
match event {
EditAgentOutputEvent::Edited => {
EditAgentOutputEvent::Edited { .. } => {
if let Some(card) = card_clone.as_ref() {
card.update(cx, |card, cx| card.update_diff(cx))?;
}