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::StreamExt;
10use gpui::{App, 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 stream = model.stream_completion_text(llm_request, &cx);
229 let mut chunks = stream.await?;
230
231 let mut request = Self {
232 parser: EditActionParser::new(),
233 // we start with the success header so we don't need to shift the output in the common case
234 output: Self::SUCCESS_OUTPUT_HEADER.to_string(),
235 changed_buffers: HashSet::default(),
236 bad_searches: Vec::new(),
237 action_log,
238 project,
239 tool_log,
240 };
241
242 while let Some(chunk) = chunks.stream.next().await {
243 request.process_response_chunk(&chunk?, cx).await?;
244 }
245
246 request.finalize(cx).await
247 })
248 }
249
250 async fn process_response_chunk(&mut self, chunk: &str, cx: &mut AsyncApp) -> Result<()> {
251 let new_actions = self.parser.parse_chunk(chunk);
252
253 if let Some((ref log, req_id)) = self.tool_log {
254 log.update(cx, |log, cx| {
255 log.push_editor_response_chunk(req_id, chunk, &new_actions, cx)
256 })
257 .log_err();
258 }
259
260 for action in new_actions {
261 self.apply_action(action, cx).await?;
262 }
263
264 Ok(())
265 }
266
267 async fn apply_action(
268 &mut self,
269 (action, source): (EditAction, String),
270 cx: &mut AsyncApp,
271 ) -> Result<()> {
272 let project_path = self.project.read_with(cx, |project, cx| {
273 project
274 .find_project_path(action.file_path(), cx)
275 .context("Path not found in project")
276 })??;
277
278 let buffer = self
279 .project
280 .update(cx, |project, cx| project.open_buffer(project_path, cx))?
281 .await?;
282
283 let result = match action {
284 EditAction::Replace {
285 old,
286 new,
287 file_path,
288 } => {
289 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
290
291 cx.background_executor()
292 .spawn(Self::replace_diff(old, new, file_path, snapshot))
293 .await
294 }
295 EditAction::Write { content, .. } => Ok(DiffResult::Diff(
296 buffer
297 .read_with(cx, |buffer, cx| buffer.diff(content, cx))?
298 .await,
299 )),
300 }?;
301
302 match result {
303 DiffResult::BadSearch(invalid_replace) => {
304 self.bad_searches.push(invalid_replace);
305 }
306 DiffResult::Diff(diff) => {
307 let _clock = buffer.update(cx, |buffer, cx| buffer.apply_diff(diff, cx))?;
308
309 write!(&mut self.output, "\n\n{}", source)?;
310 self.changed_buffers.insert(buffer);
311 }
312 }
313
314 Ok(())
315 }
316
317 async fn replace_diff(
318 old: String,
319 new: String,
320 file_path: std::path::PathBuf,
321 snapshot: language::BufferSnapshot,
322 ) -> Result<DiffResult> {
323 if snapshot.is_empty() {
324 let exists = snapshot
325 .file()
326 .map_or(false, |file| file.disk_state().exists());
327
328 return Ok(DiffResult::BadSearch(BadSearch::EmptyBuffer {
329 file_path: file_path.display().to_string(),
330 exists,
331 search: old,
332 }));
333 }
334
335 let result =
336 // Try to match exactly
337 replace_exact(&old, &new, &snapshot)
338 .await
339 // If that fails, try being flexible about indentation
340 .or_else(|| replace_with_flexible_indent(&old, &new, &snapshot));
341
342 let Some(diff) = result else {
343 return anyhow::Ok(DiffResult::BadSearch(BadSearch::NoMatch {
344 search: old,
345 file_path: file_path.display().to_string(),
346 }));
347 };
348
349 anyhow::Ok(DiffResult::Diff(diff))
350 }
351
352 const SUCCESS_OUTPUT_HEADER: &str = "Successfully applied. Here's a list of changes:";
353 const ERROR_OUTPUT_HEADER_NO_EDITS: &str = "I couldn't apply any edits!";
354 const ERROR_OUTPUT_HEADER_WITH_EDITS: &str =
355 "Errors occurred. First, here's a list of the edits we managed to apply:";
356
357 async fn finalize(self, cx: &mut AsyncApp) -> Result<String> {
358 let changed_buffer_count = self.changed_buffers.len();
359
360 // Save each buffer once at the end
361 for buffer in &self.changed_buffers {
362 self.project
363 .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
364 .await?;
365 }
366
367 self.action_log
368 .update(cx, |log, cx| log.buffer_edited(self.changed_buffers, cx))
369 .log_err();
370
371 let errors = self.parser.errors();
372
373 if errors.is_empty() && self.bad_searches.is_empty() {
374 if changed_buffer_count == 0 {
375 return Err(anyhow!(
376 "The instructions didn't lead to any changes. You might need to consult the file contents first."
377 ));
378 }
379
380 Ok(self.output)
381 } else {
382 let mut output = self.output;
383
384 if output.is_empty() {
385 output.replace_range(
386 0..Self::SUCCESS_OUTPUT_HEADER.len(),
387 Self::ERROR_OUTPUT_HEADER_NO_EDITS,
388 );
389 } else {
390 output.replace_range(
391 0..Self::SUCCESS_OUTPUT_HEADER.len(),
392 Self::ERROR_OUTPUT_HEADER_WITH_EDITS,
393 );
394 }
395
396 if !self.bad_searches.is_empty() {
397 writeln!(
398 &mut output,
399 "\n\n# {} SEARCH/REPLACE block(s) failed to match:\n",
400 self.bad_searches.len()
401 )?;
402
403 for bad_search in self.bad_searches {
404 match bad_search {
405 BadSearch::NoMatch { file_path, search } => {
406 writeln!(
407 &mut output,
408 "## No exact match in: `{}`\n```\n{}\n```\n",
409 file_path, search,
410 )?;
411 }
412 BadSearch::EmptyBuffer {
413 file_path,
414 exists: true,
415 search,
416 } => {
417 writeln!(
418 &mut output,
419 "## No match because `{}` is empty:\n```\n{}\n```\n",
420 file_path, search,
421 )?;
422 }
423 BadSearch::EmptyBuffer {
424 file_path,
425 exists: false,
426 search,
427 } => {
428 writeln!(
429 &mut output,
430 "## No match because `{}` does not exist:\n```\n{}\n```\n",
431 file_path, search,
432 )?;
433 }
434 }
435 }
436
437 write!(&mut output,
438 "The SEARCH section must exactly match an existing block of lines including all white \
439 space, comments, indentation, docstrings, etc."
440 )?;
441 }
442
443 if !errors.is_empty() {
444 writeln!(
445 &mut output,
446 "\n\n# {} SEARCH/REPLACE blocks failed to parse:",
447 errors.len()
448 )?;
449
450 for error in errors {
451 writeln!(&mut output, "- {}", error)?;
452 }
453 }
454
455 if changed_buffer_count > 0 {
456 writeln!(
457 &mut output,
458 "\n\nThe other SEARCH/REPLACE blocks were applied successfully. Do not re-send them!",
459 )?;
460 }
461
462 writeln!(
463 &mut output,
464 "{}You can fix errors by running the tool again. You can include instructions, \
465 but errors are part of the conversation so you don't need to repeat them.",
466 if changed_buffer_count == 0 {
467 "\n\n"
468 } else {
469 ""
470 }
471 )?;
472
473 Err(anyhow!(output))
474 }
475 }
476}