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