1mod edit_action;
2pub mod log;
3
4use crate::replace::{replace_exact, replace_with_flexible_indent};
5use anyhow::{anyhow, Context, Result};
6use assistant_tool::{ActionLog, Tool};
7use collections::HashSet;
8use edit_action::{EditAction, EditActionParser};
9use futures::{channel::mpsc, SinkExt, StreamExt};
10use gpui::{App, AppContext, AsyncApp, Entity, Task};
11use language_model::{
12 LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,
13};
14use log::{EditToolLog, EditToolRequestId};
15use project::Project;
16use schemars::JsonSchema;
17use serde::{Deserialize, Serialize};
18use std::fmt::Write;
19use std::sync::Arc;
20use util::ResultExt;
21
22#[derive(Debug, Serialize, Deserialize, JsonSchema)]
23pub struct EditFilesToolInput {
24 /// High-level edit instructions. These will be interpreted by a smaller
25 /// model, so explain the changes you want that model to make and which
26 /// file paths need changing. The description should be concise and clear.
27 ///
28 /// WARNING: When specifying which file paths need changing, you MUST
29 /// start each path with one of the project's root directories.
30 ///
31 /// WARNING: NEVER include code blocks or snippets in edit instructions.
32 /// Only provide natural language descriptions of the changes needed! The tool will
33 /// reject any instructions that contain code blocks or snippets.
34 ///
35 /// The following examples assume we have two root directories in the project:
36 /// - root-1
37 /// - root-2
38 ///
39 /// <example>
40 /// If you want to introduce a new quit function to kill the process, your
41 /// instructions should be: "Add a new `quit` function to
42 /// `root-1/src/main.rs` to kill the process".
43 ///
44 /// Notice how the file path starts with root-1. Without that, the path
45 /// would be ambiguous and the call would fail!
46 /// </example>
47 ///
48 /// <example>
49 /// If you want to change documentation to always start with a capital
50 /// letter, your instructions should be: "In `root-2/db.js`,
51 /// `root-2/inMemory.js` and `root-2/sql.js`, change all the documentation
52 /// to start with a capital letter".
53 ///
54 /// Notice how we never specify code snippets in the instructions!
55 /// </example>
56 pub edit_instructions: String,
57
58 /// A user-friendly description of what changes are being made.
59 /// This will be shown to the user in the UI to describe the edit operation. The screen real estate for this UI will be extremely
60 /// constrained, so make the description extremely terse.
61 ///
62 /// <example>
63 /// For fixing a broken authentication system:
64 /// "Fix auth bug in login flow"
65 /// </example>
66 ///
67 /// <example>
68 /// For adding unit tests to a module:
69 /// "Add tests for user profile logic"
70 /// </example>
71 pub display_description: String,
72}
73
74pub struct EditFilesTool;
75
76impl Tool for EditFilesTool {
77 fn name(&self) -> String {
78 "edit-files".into()
79 }
80
81 fn needs_confirmation(&self) -> bool {
82 true
83 }
84
85 fn description(&self) -> String {
86 include_str!("./edit_files_tool/description.md").into()
87 }
88
89 fn input_schema(&self) -> serde_json::Value {
90 let schema = schemars::schema_for!(EditFilesToolInput);
91 serde_json::to_value(&schema).unwrap()
92 }
93
94 fn ui_text(&self, input: &serde_json::Value) -> String {
95 match serde_json::from_value::<EditFilesToolInput>(input.clone()) {
96 Ok(input) => input.display_description,
97 Err(_) => "Edit files".to_string(),
98 }
99 }
100
101 fn run(
102 self: Arc<Self>,
103 input: serde_json::Value,
104 messages: &[LanguageModelRequestMessage],
105 project: Entity<Project>,
106 action_log: Entity<ActionLog>,
107 cx: &mut App,
108 ) -> Task<Result<String>> {
109 let input = match serde_json::from_value::<EditFilesToolInput>(input) {
110 Ok(input) => input,
111 Err(err) => return Task::ready(Err(anyhow!(err))),
112 };
113
114 match EditToolLog::try_global(cx) {
115 Some(log) => {
116 let req_id = log.update(cx, |log, cx| {
117 log.new_request(input.edit_instructions.clone(), cx)
118 });
119
120 let task = EditToolRequest::new(
121 input,
122 messages,
123 project,
124 action_log,
125 Some((log.clone(), req_id)),
126 cx,
127 );
128
129 cx.spawn(async move |cx| {
130 let result = task.await;
131
132 let str_result = match &result {
133 Ok(out) => Ok(out.clone()),
134 Err(err) => Err(err.to_string()),
135 };
136
137 log.update(cx, |log, cx| log.set_tool_output(req_id, str_result, cx))
138 .log_err();
139
140 result
141 })
142 }
143
144 None => EditToolRequest::new(input, messages, project, action_log, None, cx),
145 }
146 }
147}
148
149struct EditToolRequest {
150 parser: EditActionParser,
151 output: String,
152 changed_buffers: HashSet<Entity<language::Buffer>>,
153 bad_searches: Vec<BadSearch>,
154 project: Entity<Project>,
155 action_log: Entity<ActionLog>,
156 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
157}
158
159#[derive(Debug)]
160enum DiffResult {
161 BadSearch(BadSearch),
162 Diff(language::Diff),
163}
164
165#[derive(Debug)]
166enum BadSearch {
167 NoMatch {
168 file_path: String,
169 search: String,
170 },
171 EmptyBuffer {
172 file_path: String,
173 search: String,
174 exists: bool,
175 },
176}
177
178impl EditToolRequest {
179 fn new(
180 input: EditFilesToolInput,
181 messages: &[LanguageModelRequestMessage],
182 project: Entity<Project>,
183 action_log: Entity<ActionLog>,
184 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
185 cx: &mut App,
186 ) -> Task<Result<String>> {
187 let model_registry = LanguageModelRegistry::read_global(cx);
188 let Some(model) = model_registry.editor_model() else {
189 return Task::ready(Err(anyhow!("No editor model configured")));
190 };
191
192 let mut messages = messages.to_vec();
193 // Remove the last tool use (this run) to prevent an invalid request
194 'outer: for message in messages.iter_mut().rev() {
195 for (index, content) in message.content.iter().enumerate().rev() {
196 match content {
197 MessageContent::ToolUse(_) => {
198 message.content.remove(index);
199 break 'outer;
200 }
201 MessageContent::ToolResult(_) => {
202 // If we find any tool results before a tool use, the request is already valid
203 break 'outer;
204 }
205 MessageContent::Text(_) | MessageContent::Image(_) => {}
206 }
207 }
208 }
209
210 messages.push(LanguageModelRequestMessage {
211 role: Role::User,
212 content: vec![
213 include_str!("./edit_files_tool/edit_prompt.md").into(),
214 input.edit_instructions.into(),
215 ],
216 cache: false,
217 });
218
219 cx.spawn(async move |cx| {
220 let llm_request = LanguageModelRequest {
221 messages,
222 tools: vec![],
223 stop: vec![],
224 temperature: Some(0.0),
225 };
226
227 let (mut tx, mut rx) = mpsc::channel::<String>(32);
228 let stream = model.stream_completion_text(llm_request, &cx);
229 let reader_task = cx.background_spawn(async move {
230 let mut chunks = stream.await?;
231
232 while let Some(chunk) = chunks.stream.next().await {
233 if let Some(chunk) = chunk.log_err() {
234 // we don't process here because the API fails
235 // if we take too long between reads
236 tx.send(chunk).await?
237 }
238 }
239 tx.close().await?;
240 anyhow::Ok(())
241 });
242
243 let mut request = Self {
244 parser: EditActionParser::new(),
245 output: Self::SUCCESS_OUTPUT_HEADER.to_string(),
246 changed_buffers: HashSet::default(),
247 bad_searches: Vec::new(),
248 action_log,
249 project,
250 tool_log,
251 };
252
253 while let Some(chunk) = rx.next().await {
254 request.process_response_chunk(&chunk, cx).await?;
255 }
256
257 reader_task.await?;
258
259 request.finalize(cx).await
260 })
261 }
262
263 async fn process_response_chunk(&mut self, chunk: &str, cx: &mut AsyncApp) -> Result<()> {
264 let new_actions = self.parser.parse_chunk(chunk);
265
266 if let Some((ref log, req_id)) = self.tool_log {
267 log.update(cx, |log, cx| {
268 log.push_editor_response_chunk(req_id, chunk, &new_actions, cx)
269 })
270 .log_err();
271 }
272
273 for action in new_actions {
274 self.apply_action(action, cx).await?;
275 }
276
277 Ok(())
278 }
279
280 async fn apply_action(
281 &mut self,
282 (action, source): (EditAction, String),
283 cx: &mut AsyncApp,
284 ) -> Result<()> {
285 let project_path = self.project.read_with(cx, |project, cx| {
286 project
287 .find_project_path(action.file_path(), cx)
288 .context("Path not found in project")
289 })??;
290
291 let buffer = self
292 .project
293 .update(cx, |project, cx| project.open_buffer(project_path, cx))?
294 .await?;
295
296 let result = match action {
297 EditAction::Replace {
298 old,
299 new,
300 file_path,
301 } => {
302 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
303
304 cx.background_executor()
305 .spawn(Self::replace_diff(old, new, file_path, snapshot))
306 .await
307 }
308 EditAction::Write { content, .. } => Ok(DiffResult::Diff(
309 buffer
310 .read_with(cx, |buffer, cx| buffer.diff(content, cx))?
311 .await,
312 )),
313 }?;
314
315 match result {
316 DiffResult::BadSearch(invalid_replace) => {
317 self.bad_searches.push(invalid_replace);
318 }
319 DiffResult::Diff(diff) => {
320 let _clock = buffer.update(cx, |buffer, cx| buffer.apply_diff(diff, cx))?;
321
322 write!(&mut self.output, "\n\n{}", source)?;
323 self.changed_buffers.insert(buffer);
324 }
325 }
326
327 Ok(())
328 }
329
330 async fn replace_diff(
331 old: String,
332 new: String,
333 file_path: std::path::PathBuf,
334 snapshot: language::BufferSnapshot,
335 ) -> Result<DiffResult> {
336 if snapshot.is_empty() {
337 let exists = snapshot
338 .file()
339 .map_or(false, |file| file.disk_state().exists());
340
341 return Ok(DiffResult::BadSearch(BadSearch::EmptyBuffer {
342 file_path: file_path.display().to_string(),
343 exists,
344 search: old,
345 }));
346 }
347
348 let result =
349 // Try to match exactly
350 replace_exact(&old, &new, &snapshot)
351 .await
352 // If that fails, try being flexible about indentation
353 .or_else(|| replace_with_flexible_indent(&old, &new, &snapshot));
354
355 let Some(diff) = result else {
356 return anyhow::Ok(DiffResult::BadSearch(BadSearch::NoMatch {
357 search: old,
358 file_path: file_path.display().to_string(),
359 }));
360 };
361
362 anyhow::Ok(DiffResult::Diff(diff))
363 }
364
365 const SUCCESS_OUTPUT_HEADER: &str = "Successfully applied. Here's a list of changes:";
366 const ERROR_OUTPUT_HEADER_NO_EDITS: &str = "I couldn't apply any edits!";
367 const ERROR_OUTPUT_HEADER_WITH_EDITS: &str =
368 "Errors occurred. First, here's a list of the edits we managed to apply:";
369
370 async fn finalize(self, cx: &mut AsyncApp) -> Result<String> {
371 let changed_buffer_count = self.changed_buffers.len();
372
373 // Save each buffer once at the end
374 for buffer in &self.changed_buffers {
375 self.project
376 .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
377 .await?;
378 }
379
380 self.action_log
381 .update(cx, |log, cx| log.buffer_edited(self.changed_buffers, cx))
382 .log_err();
383
384 let errors = self.parser.errors();
385
386 if errors.is_empty() && self.bad_searches.is_empty() {
387 if changed_buffer_count == 0 {
388 return Err(anyhow!(
389 "The instructions didn't lead to any changes. You might need to consult the file contents first."
390 ));
391 }
392
393 Ok(self.output)
394 } else {
395 let mut output = self.output;
396
397 if output.is_empty() {
398 output.replace_range(
399 0..Self::SUCCESS_OUTPUT_HEADER.len(),
400 Self::ERROR_OUTPUT_HEADER_NO_EDITS,
401 );
402 } else {
403 output.replace_range(
404 0..Self::SUCCESS_OUTPUT_HEADER.len(),
405 Self::ERROR_OUTPUT_HEADER_WITH_EDITS,
406 );
407 }
408
409 if !self.bad_searches.is_empty() {
410 writeln!(
411 &mut output,
412 "\n\n# {} SEARCH/REPLACE block(s) failed to match:\n",
413 self.bad_searches.len()
414 )?;
415
416 for bad_search in self.bad_searches {
417 match bad_search {
418 BadSearch::NoMatch { file_path, search } => {
419 writeln!(
420 &mut output,
421 "## No exact match in: `{}`\n```\n{}\n```\n",
422 file_path, search,
423 )?;
424 }
425 BadSearch::EmptyBuffer {
426 file_path,
427 exists: true,
428 search,
429 } => {
430 writeln!(
431 &mut output,
432 "## No match because `{}` is empty:\n```\n{}\n```\n",
433 file_path, search,
434 )?;
435 }
436 BadSearch::EmptyBuffer {
437 file_path,
438 exists: false,
439 search,
440 } => {
441 writeln!(
442 &mut output,
443 "## No match because `{}` does not exist:\n```\n{}\n```\n",
444 file_path, search,
445 )?;
446 }
447 }
448 }
449
450 write!(&mut output,
451 "The SEARCH section must exactly match an existing block of lines including all white \
452 space, comments, indentation, docstrings, etc."
453 )?;
454 }
455
456 if !errors.is_empty() {
457 writeln!(
458 &mut output,
459 "\n\n# {} SEARCH/REPLACE blocks failed to parse:",
460 errors.len()
461 )?;
462
463 for error in errors {
464 writeln!(&mut output, "- {}", error)?;
465 }
466 }
467
468 if changed_buffer_count > 0 {
469 writeln!(
470 &mut output,
471 "\n\nThe other SEARCH/REPLACE blocks were applied successfully. Do not re-send them!",
472 )?;
473 }
474
475 writeln!(
476 &mut output,
477 "{}You can fix errors by running the tool again. You can include instructions, \
478 but errors are part of the conversation so you don't need to repeat them.",
479 if changed_buffer_count == 0 {
480 "\n\n"
481 } else {
482 ""
483 }
484 )?;
485
486 Err(anyhow!(output))
487 }
488 }
489}