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 ui::IconName;
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 icon(&self) -> IconName {
91 IconName::Pencil
92 }
93
94 fn input_schema(&self) -> serde_json::Value {
95 let schema = schemars::schema_for!(EditFilesToolInput);
96 serde_json::to_value(&schema).unwrap()
97 }
98
99 fn ui_text(&self, input: &serde_json::Value) -> String {
100 match serde_json::from_value::<EditFilesToolInput>(input.clone()) {
101 Ok(input) => input.display_description,
102 Err(_) => "Edit files".to_string(),
103 }
104 }
105
106 fn run(
107 self: Arc<Self>,
108 input: serde_json::Value,
109 messages: &[LanguageModelRequestMessage],
110 project: Entity<Project>,
111 action_log: Entity<ActionLog>,
112 cx: &mut App,
113 ) -> Task<Result<String>> {
114 let input = match serde_json::from_value::<EditFilesToolInput>(input) {
115 Ok(input) => input,
116 Err(err) => return Task::ready(Err(anyhow!(err))),
117 };
118
119 match EditToolLog::try_global(cx) {
120 Some(log) => {
121 let req_id = log.update(cx, |log, cx| {
122 log.new_request(input.edit_instructions.clone(), cx)
123 });
124
125 let task = EditToolRequest::new(
126 input,
127 messages,
128 project,
129 action_log,
130 Some((log.clone(), req_id)),
131 cx,
132 );
133
134 cx.spawn(async move |cx| {
135 let result = task.await;
136
137 let str_result = match &result {
138 Ok(out) => Ok(out.clone()),
139 Err(err) => Err(err.to_string()),
140 };
141
142 log.update(cx, |log, cx| log.set_tool_output(req_id, str_result, cx))
143 .log_err();
144
145 result
146 })
147 }
148
149 None => EditToolRequest::new(input, messages, project, action_log, None, cx),
150 }
151 }
152}
153
154struct EditToolRequest {
155 parser: EditActionParser,
156 editor_response: EditorResponse,
157 project: Entity<Project>,
158 action_log: Entity<ActionLog>,
159 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
160}
161
162enum EditorResponse {
163 /// The editor model hasn't produced any actions yet.
164 /// If we don't have any by the end, we'll return its message to the architect model.
165 Message(String),
166 /// The editor model produced at least one action.
167 Actions {
168 applied: Vec<AppliedAction>,
169 search_errors: Vec<SearchError>,
170 },
171}
172
173struct AppliedAction {
174 source: String,
175 buffer: Entity<language::Buffer>,
176}
177
178#[derive(Debug)]
179enum DiffResult {
180 Diff(language::Diff),
181 SearchError(SearchError),
182}
183
184#[derive(Debug)]
185enum SearchError {
186 NoMatch {
187 file_path: String,
188 search: String,
189 },
190 EmptyBuffer {
191 file_path: String,
192 search: String,
193 exists: bool,
194 },
195}
196
197impl EditToolRequest {
198 fn new(
199 input: EditFilesToolInput,
200 messages: &[LanguageModelRequestMessage],
201 project: Entity<Project>,
202 action_log: Entity<ActionLog>,
203 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
204 cx: &mut App,
205 ) -> Task<Result<String>> {
206 let model_registry = LanguageModelRegistry::read_global(cx);
207 let Some(model) = model_registry.editor_model() else {
208 return Task::ready(Err(anyhow!("No editor model configured")));
209 };
210
211 let mut messages = messages.to_vec();
212 // Remove the last tool use (this run) to prevent an invalid request
213 'outer: for message in messages.iter_mut().rev() {
214 for (index, content) in message.content.iter().enumerate().rev() {
215 match content {
216 MessageContent::ToolUse(_) => {
217 message.content.remove(index);
218 break 'outer;
219 }
220 MessageContent::ToolResult(_) => {
221 // If we find any tool results before a tool use, the request is already valid
222 break 'outer;
223 }
224 MessageContent::Text(_) | MessageContent::Image(_) => {}
225 }
226 }
227 }
228
229 messages.push(LanguageModelRequestMessage {
230 role: Role::User,
231 content: vec![
232 include_str!("./edit_files_tool/edit_prompt.md").into(),
233 input.edit_instructions.into(),
234 ],
235 cache: false,
236 });
237
238 cx.spawn(async move |cx| {
239 let llm_request = LanguageModelRequest {
240 messages,
241 tools: vec![],
242 stop: vec![],
243 temperature: Some(0.0),
244 };
245
246 let (mut tx, mut rx) = mpsc::channel::<String>(32);
247 let stream = model.stream_completion_text(llm_request, &cx);
248 let reader_task = cx.background_spawn(async move {
249 let mut chunks = stream.await?;
250
251 while let Some(chunk) = chunks.stream.next().await {
252 if let Some(chunk) = chunk.log_err() {
253 // we don't process here because the API fails
254 // if we take too long between reads
255 tx.send(chunk).await?
256 }
257 }
258 tx.close().await?;
259 anyhow::Ok(())
260 });
261
262 let mut request = Self {
263 parser: EditActionParser::new(),
264 editor_response: EditorResponse::Message(String::with_capacity(256)),
265 action_log,
266 project,
267 tool_log,
268 };
269
270 while let Some(chunk) = rx.next().await {
271 request.process_response_chunk(&chunk, cx).await?;
272 }
273
274 reader_task.await?;
275
276 request.finalize(cx).await
277 })
278 }
279
280 async fn process_response_chunk(&mut self, chunk: &str, cx: &mut AsyncApp) -> Result<()> {
281 let new_actions = self.parser.parse_chunk(chunk);
282
283 if let EditorResponse::Message(ref mut message) = self.editor_response {
284 if new_actions.is_empty() {
285 message.push_str(chunk);
286 }
287 }
288
289 if let Some((ref log, req_id)) = self.tool_log {
290 log.update(cx, |log, cx| {
291 log.push_editor_response_chunk(req_id, chunk, &new_actions, cx)
292 })
293 .log_err();
294 }
295
296 for action in new_actions {
297 self.apply_action(action, cx).await?;
298 }
299
300 Ok(())
301 }
302
303 async fn apply_action(
304 &mut self,
305 (action, source): (EditAction, String),
306 cx: &mut AsyncApp,
307 ) -> Result<()> {
308 let project_path = self.project.read_with(cx, |project, cx| {
309 project
310 .find_project_path(action.file_path(), cx)
311 .context("Path not found in project")
312 })??;
313
314 let buffer = self
315 .project
316 .update(cx, |project, cx| project.open_buffer(project_path, cx))?
317 .await?;
318
319 let result = match action {
320 EditAction::Replace {
321 old,
322 new,
323 file_path,
324 } => {
325 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
326
327 cx.background_executor()
328 .spawn(Self::replace_diff(old, new, file_path, snapshot))
329 .await
330 }
331 EditAction::Write { content, .. } => Ok(DiffResult::Diff(
332 buffer
333 .read_with(cx, |buffer, cx| buffer.diff(content, cx))?
334 .await,
335 )),
336 }?;
337
338 match result {
339 DiffResult::SearchError(error) => {
340 self.push_search_error(error);
341 }
342 DiffResult::Diff(diff) => {
343 let _clock = buffer.update(cx, |buffer, cx| buffer.apply_diff(diff, cx))?;
344
345 self.push_applied_action(AppliedAction { source, buffer });
346 }
347 }
348
349 anyhow::Ok(())
350 }
351
352 fn push_search_error(&mut self, error: SearchError) {
353 match &mut self.editor_response {
354 EditorResponse::Message(_) => {
355 self.editor_response = EditorResponse::Actions {
356 applied: Vec::new(),
357 search_errors: vec![error],
358 };
359 }
360 EditorResponse::Actions { search_errors, .. } => {
361 search_errors.push(error);
362 }
363 }
364 }
365
366 fn push_applied_action(&mut self, action: AppliedAction) {
367 match &mut self.editor_response {
368 EditorResponse::Message(_) => {
369 self.editor_response = EditorResponse::Actions {
370 applied: vec![action],
371 search_errors: Vec::new(),
372 };
373 }
374 EditorResponse::Actions { applied, .. } => {
375 applied.push(action);
376 }
377 }
378 }
379
380 async fn replace_diff(
381 old: String,
382 new: String,
383 file_path: std::path::PathBuf,
384 snapshot: language::BufferSnapshot,
385 ) -> Result<DiffResult> {
386 if snapshot.is_empty() {
387 let exists = snapshot
388 .file()
389 .map_or(false, |file| file.disk_state().exists());
390
391 let error = SearchError::EmptyBuffer {
392 file_path: file_path.display().to_string(),
393 exists,
394 search: old,
395 };
396
397 return Ok(DiffResult::SearchError(error));
398 }
399
400 let replace_result =
401 // Try to match exactly
402 replace_exact(&old, &new, &snapshot)
403 .await
404 // If that fails, try being flexible about indentation
405 .or_else(|| replace_with_flexible_indent(&old, &new, &snapshot));
406
407 let Some(diff) = replace_result else {
408 let error = SearchError::NoMatch {
409 search: old,
410 file_path: file_path.display().to_string(),
411 };
412
413 return Ok(DiffResult::SearchError(error));
414 };
415
416 Ok(DiffResult::Diff(diff))
417 }
418
419 async fn finalize(self, cx: &mut AsyncApp) -> Result<String> {
420 match self.editor_response {
421 EditorResponse::Message(message) => Err(anyhow!(
422 "No edits were applied! You might need to provide more context.\n\n{}",
423 message
424 )),
425 EditorResponse::Actions {
426 applied,
427 search_errors,
428 } => {
429 let mut output = String::with_capacity(1024);
430
431 let parse_errors = self.parser.errors();
432 let has_errors = !search_errors.is_empty() || !parse_errors.is_empty();
433
434 if has_errors {
435 let error_count = search_errors.len() + parse_errors.len();
436
437 if applied.is_empty() {
438 writeln!(
439 &mut output,
440 "{} errors occurred! No edits were applied.",
441 error_count,
442 )?;
443 } else {
444 writeln!(
445 &mut output,
446 "{} errors occurred, but {} edits were correctly applied.",
447 error_count,
448 applied.len(),
449 )?;
450
451 writeln!(
452 &mut output,
453 "# {} SEARCH/REPLACE block(s) applied:\n\nDo not re-send these since they are already applied!\n",
454 applied.len()
455 )?;
456 }
457 } else {
458 write!(
459 &mut output,
460 "Successfully applied! Here's a list of applied edits:"
461 )?;
462 }
463
464 let mut changed_buffers = HashSet::default();
465
466 for action in applied {
467 changed_buffers.insert(action.buffer);
468 write!(&mut output, "\n\n{}", action.source)?;
469 }
470
471 for buffer in &changed_buffers {
472 self.project
473 .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
474 .await?;
475 }
476
477 self.action_log
478 .update(cx, |log, cx| log.buffer_edited(changed_buffers.clone(), cx))
479 .log_err();
480
481 if !search_errors.is_empty() {
482 writeln!(
483 &mut output,
484 "\n\n## {} SEARCH/REPLACE block(s) failed to match:\n",
485 search_errors.len()
486 )?;
487
488 for error in search_errors {
489 match error {
490 SearchError::NoMatch { file_path, search } => {
491 writeln!(
492 &mut output,
493 "### No exact match in: `{}`\n```\n{}\n```\n",
494 file_path, search,
495 )?;
496 }
497 SearchError::EmptyBuffer {
498 file_path,
499 exists: true,
500 search,
501 } => {
502 writeln!(
503 &mut output,
504 "### No match because `{}` is empty:\n```\n{}\n```\n",
505 file_path, search,
506 )?;
507 }
508 SearchError::EmptyBuffer {
509 file_path,
510 exists: false,
511 search,
512 } => {
513 writeln!(
514 &mut output,
515 "### No match because `{}` does not exist:\n```\n{}\n```\n",
516 file_path, search,
517 )?;
518 }
519 }
520 }
521
522 write!(&mut output,
523 "The SEARCH section must exactly match an existing block of lines including all white \
524 space, comments, indentation, docstrings, etc."
525 )?;
526 }
527
528 if !parse_errors.is_empty() {
529 writeln!(
530 &mut output,
531 "\n\n## {} SEARCH/REPLACE blocks failed to parse:",
532 parse_errors.len()
533 )?;
534
535 for error in parse_errors {
536 writeln!(&mut output, "- {}", error)?;
537 }
538 }
539
540 if has_errors {
541 writeln!(&mut output,
542 "\n\nYou can fix errors by running the tool again. You can include instructions, \
543 but errors are part of the conversation so you don't need to repeat them.",
544 )?;
545
546 Err(anyhow!(output))
547 } else {
548 Ok(output)
549 }
550 }
551 }
552 }
553}