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