When the model reads file, we'll track the version it read, and let it know if the user makes edits to the buffer. This helps prevent edit failures because it'll know to re-read the file before. Release Notes: - N/A
Unified · Split
| @@ -692,6 +692,7 @@ name = "assistant_tool" · 691 unmodified lines | ||
| 692 | 692 | version = "0.1.0" |
| 693 | 693 | dependencies = [ |
| 694 | 694 | "anyhow", |
| 695 | + "clock", | |
| 695 | 696 | "collections", |
| 696 | 697 | "derive_more", |
| 697 | 698 | "gpui", |
| @@ -361,7 +361,7 @@ impl ActiveThread { · 360 unmodified lines | ||
| 361 | 361 | if self.thread.read(cx).all_tools_finished() { |
| 362 | 362 | let pending_refresh_buffers = self.thread.update(cx, |thread, cx| { |
| 363 | 363 | thread.action_log().update(cx, |action_log, _cx| { |
| 364 | − action_log.take_pending_refresh_buffers() | |
| 364 | + action_log.take_stale_buffers_in_context() | |
| 365 | 365 | }) |
| 366 | 366 | }); |
| 367 | 367 | |
| @@ -1,3 +1,4 @@ | ||
| 1 | +use std::fmt::Write as _; | |
| 1 | 2 | use std::io::Write; |
| 2 | 3 | use std::sync::Arc; |
| 3 | 4 | |
| @@ -560,9 +561,39 @@ impl Thread { · 556 unmodified lines | ||
| 560 | 561 | request.messages.push(context_message); |
| 561 | 562 | } |
| 562 | 563 | |
| 564 | + self.attach_stale_files(&mut request.messages, cx); | |
| 565 | + | |
| 563 | 566 | request |
| 564 | 567 | } |
| 565 | 568 | |
| 569 | + fn attach_stale_files(&self, messages: &mut Vec<LanguageModelRequestMessage>, cx: &App) { | |
| 570 | + const STALE_FILES_HEADER: &str = "These files changed since last read:"; | |
| 571 | + | |
| 572 | + let mut stale_message = String::new(); | |
| 573 | + | |
| 574 | + for stale_file in self.action_log.read(cx).stale_buffers(cx) { | |
| 575 | + let Some(file) = stale_file.read(cx).file() else { | |
| 576 | + continue; | |
| 577 | + }; | |
| 578 | + | |
| 579 | + if stale_message.is_empty() { | |
| 580 | + write!(&mut stale_message, "{}", STALE_FILES_HEADER).ok(); | |
| 581 | + } | |
| 582 | + | |
| 583 | + writeln!(&mut stale_message, "- {}", file.path().display()).ok(); | |
| 584 | + } | |
| 585 | + | |
| 586 | + if !stale_message.is_empty() { | |
| 587 | + let context_message = LanguageModelRequestMessage { | |
| 588 | + role: Role::User, | |
| 589 | + content: vec![stale_message.into()], | |
| 590 | + cache: false, | |
| 591 | + }; | |
| 592 | + | |
| 593 | + messages.push(context_message); | |
| 594 | + } | |
| 595 | + } | |
| 596 | + | |
| 566 | 597 | pub fn stream_completion( |
| 567 | 598 | &mut self, |
| 568 | 599 | request: LanguageModelRequest, |
| @@ -14,6 +14,7 @@ path = "src/assistant_tool.rs" · 13 unmodified lines | ||
| 14 | 14 | [dependencies] |
| 15 | 15 | anyhow.workspace = true |
| 16 | 16 | collections.workspace = true |
| 17 | +clock.workspace = true | |
| 17 | 18 | derive_more.workspace = true |
| 18 | 19 | gpui.workspace = true |
| 19 | 20 | language.workspace = true |
| @@ -4,7 +4,7 @@ mod tool_working_set; · 3 unmodified lines | ||
| 4 | 4 | use std::sync::Arc; |
| 5 | 5 | |
| 6 | 6 | use anyhow::Result; |
| 7 | −use collections::HashSet; | |
| 7 | +use collections::{HashMap, HashSet}; | |
| 8 | 8 | use gpui::Context; |
| 9 | 9 | use gpui::{App, Entity, SharedString, Task}; |
| 10 | 10 | use language::Buffer; |
| @@ -58,31 +58,53 @@ pub trait Tool: 'static + Send + Sync { · 47 unmodified lines | ||
| 58 | 58 | /// Tracks actions performed by tools in a thread |
| 59 | 59 | #[derive(Debug)] |
| 60 | 60 | pub struct ActionLog { |
| 61 | − changed_buffers: HashSet<Entity<Buffer>>, | |
| 62 | − pending_refresh: HashSet<Entity<Buffer>>, | |
| 61 | + /// Buffers that user manually added to the context, and whose content has | |
| 62 | + /// changed since the model last saw them. | |
| 63 | + stale_buffers_in_context: HashSet<Entity<Buffer>>, | |
| 64 | + /// Buffers that we want to notify the model about when they change. | |
| 65 | + tracked_buffers: HashMap<Entity<Buffer>, TrackedBuffer>, | |
| 66 | +} | |
| 67 | + | |
| 68 | +#[derive(Debug, Default)] | |
| 69 | +struct TrackedBuffer { | |
| 70 | + version: clock::Global, | |
| 63 | 71 | } |
| 64 | 72 | |
| 65 | 73 | impl ActionLog { |
| 66 | 74 | /// Creates a new, empty action log. |
| 67 | 75 | pub fn new() -> Self { |
| 68 | 76 | Self { |
| 69 | − changed_buffers: HashSet::default(), | |
| 70 | − pending_refresh: HashSet::default(), | |
| 77 | + stale_buffers_in_context: HashSet::default(), | |
| 78 | + tracked_buffers: HashMap::default(), | |
| 79 | + } | |
| 80 | + } | |
| 81 | + | |
| 82 | + /// Track a buffer as read, so we can notify the model about user edits. | |
| 83 | + pub fn buffer_read(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) { | |
| 84 | + let tracked_buffer = self.tracked_buffers.entry(buffer.clone()).or_default(); | |
| 85 | + tracked_buffer.version = buffer.read(cx).version(); | |
| 86 | + } | |
| 87 | + | |
| 88 | + /// Mark a buffer as edited, so we can refresh it in the context | |
| 89 | + pub fn buffer_edited(&mut self, buffers: HashSet<Entity<Buffer>>, cx: &mut Context<Self>) { | |
| 90 | + for buffer in &buffers { | |
| 91 | + let tracked_buffer = self.tracked_buffers.entry(buffer.clone()).or_default(); | |
| 92 | + tracked_buffer.version = buffer.read(cx).version(); | |
| 71 | 93 | } |
| 94 | + | |
| 95 | + self.stale_buffers_in_context.extend(buffers); | |
| 72 | 96 | } |
| 73 | 97 | |
| 74 | − /// Registers buffers that have changed and need refreshing. | |
| 75 | − pub fn notify_buffers_changed( | |
| 76 | − &mut self, | |
| 77 | − buffers: HashSet<Entity<Buffer>>, | |
| 78 | − _cx: &mut Context<Self>, | |
| 79 | − ) { | |
| 80 | − self.changed_buffers.extend(buffers.clone()); | |
| 81 | − self.pending_refresh.extend(buffers); | |
| 98 | + /// Iterate over buffers changed since last read or edited by the model | |
| 99 | + pub fn stale_buffers<'a>(&'a self, cx: &'a App) -> impl Iterator<Item = &'a Entity<Buffer>> { | |
| 100 | + self.tracked_buffers | |
| 101 | + .iter() | |
| 102 | + .filter(|(buffer, tracked)| tracked.version != buffer.read(cx).version) | |
| 103 | + .map(|(buffer, _)| buffer) | |
| 82 | 104 | } |
| 83 | 105 | |
| 84 | 106 | /// Takes and returns the set of buffers pending refresh, clearing internal state. |
| 85 | − pub fn take_pending_refresh_buffers(&mut self) -> HashSet<Entity<Buffer>> { | |
| 86 | − std::mem::take(&mut self.pending_refresh) | |
| 107 | + pub fn take_stale_buffers_in_context(&mut self) -> HashSet<Entity<Buffer>> { | |
| 108 | + std::mem::take(&mut self.stale_buffers_in_context) | |
| 87 | 109 | } |
| 88 | 110 | } |
| @@ -309,9 +309,7 @@ impl EditToolRequest { · 308 unmodified lines | ||
| 309 | 309 | } |
| 310 | 310 | |
| 311 | 311 | self.action_log |
| 312 | − .update(cx, |log, cx| { | |
| 313 | − log.notify_buffers_changed(self.changed_buffers, cx) | |
| 314 | − }) | |
| 312 | + .update(cx, |log, cx| log.buffer_edited(self.changed_buffers, cx)) | |
| 315 | 313 | .log_err(); |
| 316 | 314 | |
| 317 | 315 | let errors = self.parser.errors(); |
| @@ -49,7 +49,7 @@ impl Tool for ReadFileTool { · 48 unmodified lines | ||
| 49 | 49 | input: serde_json::Value, |
| 50 | 50 | _messages: &[LanguageModelRequestMessage], |
| 51 | 51 | project: Entity<Project>, |
| 52 | − _action_log: Entity<ActionLog>, | |
| 52 | + action_log: Entity<ActionLog>, | |
| 53 | 53 | cx: &mut App, |
| 54 | 54 | ) -> Task<Result<String>> { |
| 55 | 55 | let input = match serde_json::from_value::<ReadFileToolInput>(input) { |
| @@ -60,14 +60,15 @@ impl Tool for ReadFileTool { · 4 unmodified lines | ||
| 60 | 60 | let Some(project_path) = project.read(cx).find_project_path(&input.path, cx) else { |
| 61 | 61 | return Task::ready(Err(anyhow!("Path not found in project"))); |
| 62 | 62 | }; |
| 63 | − cx.spawn(|cx| async move { | |
| 63 | + | |
| 64 | + cx.spawn(|mut cx| async move { | |
| 64 | 65 | let buffer = cx |
| 65 | 66 | .update(|cx| { |
| 66 | 67 | project.update(cx, |project, cx| project.open_buffer(project_path, cx)) |
| 67 | 68 | })? |
| 68 | 69 | .await?; |
| 69 | 70 | |
| 70 | − buffer.read_with(&cx, |buffer, _cx| { | |
| 71 | + let result = buffer.read_with(&cx, |buffer, _cx| { | |
| 71 | 72 | if buffer |
| 72 | 73 | .file() |
| 73 | 74 | .map_or(false, |file| file.disk_state().exists()) |
| @@ -76,7 +77,13 @@ impl Tool for ReadFileTool { · 2 unmodified lines | ||
| 76 | 77 | } else { |
| 77 | 78 | Err(anyhow!("File does not exist")) |
| 78 | 79 | } |
| 79 | − })? | |
| 80 | + })??; | |
| 81 | + | |
| 82 | + action_log.update(&mut cx, |log, cx| { | |
| 83 | + log.buffer_read(buffer, cx); | |
| 84 | + })?; | |
| 85 | + | |
| 86 | + anyhow::Ok(result) | |
| 80 | 87 | }) |
| 81 | 88 | } |
| 82 | 89 | } |