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