1mod connection;
2mod diff;
3mod mention;
4mod terminal;
5
6use collections::HashSet;
7pub use connection::*;
8pub use diff::*;
9use language::language_settings::FormatOnSave;
10pub use mention::*;
11use project::lsp_store::{FormatTrigger, LspFormatTarget};
12use serde::{Deserialize, Serialize};
13pub use terminal::*;
14
15use action_log::ActionLog;
16use agent_client_protocol as acp;
17use anyhow::{Context as _, Result, anyhow};
18use editor::Bias;
19use futures::{FutureExt, channel::oneshot, future::BoxFuture};
20use gpui::{AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
21use itertools::Itertools;
22use language::{Anchor, Buffer, BufferSnapshot, LanguageRegistry, Point, ToPoint, text_diff};
23use markdown::Markdown;
24use project::{AgentLocation, Project, git_store::GitStoreCheckpoint};
25use std::collections::HashMap;
26use std::error::Error;
27use std::fmt::{Formatter, Write};
28use std::ops::Range;
29use std::process::ExitStatus;
30use std::rc::Rc;
31use std::time::{Duration, Instant};
32use std::{fmt::Display, mem, path::PathBuf, sync::Arc};
33use ui::App;
34use util::ResultExt;
35
36#[derive(Debug)]
37pub struct UserMessage {
38 pub id: Option<UserMessageId>,
39 pub content: ContentBlock,
40 pub chunks: Vec<acp::ContentBlock>,
41 pub checkpoint: Option<Checkpoint>,
42}
43
44#[derive(Debug)]
45pub struct Checkpoint {
46 git_checkpoint: GitStoreCheckpoint,
47 pub show: bool,
48}
49
50impl UserMessage {
51 fn to_markdown(&self, cx: &App) -> String {
52 let mut markdown = String::new();
53 if self
54 .checkpoint
55 .as_ref()
56 .is_some_and(|checkpoint| checkpoint.show)
57 {
58 writeln!(markdown, "## User (checkpoint)").unwrap();
59 } else {
60 writeln!(markdown, "## User").unwrap();
61 }
62 writeln!(markdown).unwrap();
63 writeln!(markdown, "{}", self.content.to_markdown(cx)).unwrap();
64 writeln!(markdown).unwrap();
65 markdown
66 }
67}
68
69#[derive(Debug, PartialEq)]
70pub struct AssistantMessage {
71 pub chunks: Vec<AssistantMessageChunk>,
72}
73
74impl AssistantMessage {
75 pub fn to_markdown(&self, cx: &App) -> String {
76 format!(
77 "## Assistant\n\n{}\n\n",
78 self.chunks
79 .iter()
80 .map(|chunk| chunk.to_markdown(cx))
81 .join("\n\n")
82 )
83 }
84}
85
86#[derive(Debug, PartialEq)]
87pub enum AssistantMessageChunk {
88 Message { block: ContentBlock },
89 Thought { block: ContentBlock },
90}
91
92impl AssistantMessageChunk {
93 pub fn from_str(chunk: &str, language_registry: &Arc<LanguageRegistry>, cx: &mut App) -> Self {
94 Self::Message {
95 block: ContentBlock::new(chunk.into(), language_registry, cx),
96 }
97 }
98
99 fn to_markdown(&self, cx: &App) -> String {
100 match self {
101 Self::Message { block } => block.to_markdown(cx).to_string(),
102 Self::Thought { block } => {
103 format!("<thinking>\n{}\n</thinking>", block.to_markdown(cx))
104 }
105 }
106 }
107}
108
109#[derive(Debug)]
110pub enum AgentThreadEntry {
111 UserMessage(UserMessage),
112 AssistantMessage(AssistantMessage),
113 ToolCall(ToolCall),
114}
115
116impl AgentThreadEntry {
117 pub fn to_markdown(&self, cx: &App) -> String {
118 match self {
119 Self::UserMessage(message) => message.to_markdown(cx),
120 Self::AssistantMessage(message) => message.to_markdown(cx),
121 Self::ToolCall(tool_call) => tool_call.to_markdown(cx),
122 }
123 }
124
125 pub fn user_message(&self) -> Option<&UserMessage> {
126 if let AgentThreadEntry::UserMessage(message) = self {
127 Some(message)
128 } else {
129 None
130 }
131 }
132
133 pub fn diffs(&self) -> impl Iterator<Item = &Entity<Diff>> {
134 if let AgentThreadEntry::ToolCall(call) = self {
135 itertools::Either::Left(call.diffs())
136 } else {
137 itertools::Either::Right(std::iter::empty())
138 }
139 }
140
141 pub fn terminals(&self) -> impl Iterator<Item = &Entity<Terminal>> {
142 if let AgentThreadEntry::ToolCall(call) = self {
143 itertools::Either::Left(call.terminals())
144 } else {
145 itertools::Either::Right(std::iter::empty())
146 }
147 }
148
149 pub fn location(&self, ix: usize) -> Option<(acp::ToolCallLocation, AgentLocation)> {
150 if let AgentThreadEntry::ToolCall(ToolCall {
151 locations,
152 resolved_locations,
153 ..
154 }) = self
155 {
156 Some((
157 locations.get(ix)?.clone(),
158 resolved_locations.get(ix)?.clone()?,
159 ))
160 } else {
161 None
162 }
163 }
164}
165
166#[derive(Debug)]
167pub struct ToolCall {
168 pub id: acp::ToolCallId,
169 pub label: Entity<Markdown>,
170 pub kind: acp::ToolKind,
171 pub content: Vec<ToolCallContent>,
172 pub status: ToolCallStatus,
173 pub locations: Vec<acp::ToolCallLocation>,
174 pub resolved_locations: Vec<Option<AgentLocation>>,
175 pub raw_input: Option<serde_json::Value>,
176 pub raw_output: Option<serde_json::Value>,
177}
178
179impl ToolCall {
180 fn from_acp(
181 tool_call: acp::ToolCall,
182 status: ToolCallStatus,
183 language_registry: Arc<LanguageRegistry>,
184 cx: &mut App,
185 ) -> Self {
186 Self {
187 id: tool_call.id,
188 label: cx.new(|cx| {
189 Markdown::new(
190 tool_call.title.into(),
191 Some(language_registry.clone()),
192 None,
193 cx,
194 )
195 }),
196 kind: tool_call.kind,
197 content: tool_call
198 .content
199 .into_iter()
200 .map(|content| ToolCallContent::from_acp(content, language_registry.clone(), cx))
201 .collect(),
202 locations: tool_call.locations,
203 resolved_locations: Vec::default(),
204 status,
205 raw_input: tool_call.raw_input,
206 raw_output: tool_call.raw_output,
207 }
208 }
209
210 fn update_fields(
211 &mut self,
212 fields: acp::ToolCallUpdateFields,
213 language_registry: Arc<LanguageRegistry>,
214 cx: &mut App,
215 ) {
216 let acp::ToolCallUpdateFields {
217 kind,
218 status,
219 title,
220 content,
221 locations,
222 raw_input,
223 raw_output,
224 } = fields;
225
226 if let Some(kind) = kind {
227 self.kind = kind;
228 }
229
230 if let Some(status) = status {
231 self.status = status.into();
232 }
233
234 if let Some(title) = title {
235 self.label.update(cx, |label, cx| {
236 label.replace(title, cx);
237 });
238 }
239
240 if let Some(content) = content {
241 let new_content_len = content.len();
242 let mut content = content.into_iter();
243
244 // Reuse existing content if we can
245 for (old, new) in self.content.iter_mut().zip(content.by_ref()) {
246 old.update_from_acp(new, language_registry.clone(), cx);
247 }
248 for new in content {
249 self.content.push(ToolCallContent::from_acp(
250 new,
251 language_registry.clone(),
252 cx,
253 ))
254 }
255 self.content.truncate(new_content_len);
256 }
257
258 if let Some(locations) = locations {
259 self.locations = locations;
260 }
261
262 if let Some(raw_input) = raw_input {
263 self.raw_input = Some(raw_input);
264 }
265
266 if let Some(raw_output) = raw_output {
267 if self.content.is_empty()
268 && let Some(markdown) = markdown_for_raw_output(&raw_output, &language_registry, cx)
269 {
270 self.content
271 .push(ToolCallContent::ContentBlock(ContentBlock::Markdown {
272 markdown,
273 }));
274 }
275 self.raw_output = Some(raw_output);
276 }
277 }
278
279 pub fn diffs(&self) -> impl Iterator<Item = &Entity<Diff>> {
280 self.content.iter().filter_map(|content| match content {
281 ToolCallContent::Diff(diff) => Some(diff),
282 ToolCallContent::ContentBlock(_) => None,
283 ToolCallContent::Terminal(_) => None,
284 })
285 }
286
287 pub fn terminals(&self) -> impl Iterator<Item = &Entity<Terminal>> {
288 self.content.iter().filter_map(|content| match content {
289 ToolCallContent::Terminal(terminal) => Some(terminal),
290 ToolCallContent::ContentBlock(_) => None,
291 ToolCallContent::Diff(_) => None,
292 })
293 }
294
295 fn to_markdown(&self, cx: &App) -> String {
296 let mut markdown = format!(
297 "**Tool Call: {}**\nStatus: {}\n\n",
298 self.label.read(cx).source(),
299 self.status
300 );
301 for content in &self.content {
302 markdown.push_str(content.to_markdown(cx).as_str());
303 markdown.push_str("\n\n");
304 }
305 markdown
306 }
307
308 async fn resolve_location(
309 location: acp::ToolCallLocation,
310 project: WeakEntity<Project>,
311 cx: &mut AsyncApp,
312 ) -> Option<AgentLocation> {
313 let buffer = project
314 .update(cx, |project, cx| {
315 project
316 .project_path_for_absolute_path(&location.path, cx)
317 .map(|path| project.open_buffer(path, cx))
318 })
319 .ok()??;
320 let buffer = buffer.await.log_err()?;
321 let position = buffer
322 .update(cx, |buffer, _| {
323 if let Some(row) = location.line {
324 let snapshot = buffer.snapshot();
325 let column = snapshot.indent_size_for_line(row).len;
326 let point = snapshot.clip_point(Point::new(row, column), Bias::Left);
327 snapshot.anchor_before(point)
328 } else {
329 Anchor::MIN
330 }
331 })
332 .ok()?;
333
334 Some(AgentLocation {
335 buffer: buffer.downgrade(),
336 position,
337 })
338 }
339
340 fn resolve_locations(
341 &self,
342 project: Entity<Project>,
343 cx: &mut App,
344 ) -> Task<Vec<Option<AgentLocation>>> {
345 let locations = self.locations.clone();
346 project.update(cx, |_, cx| {
347 cx.spawn(async move |project, cx| {
348 let mut new_locations = Vec::new();
349 for location in locations {
350 new_locations.push(Self::resolve_location(location, project.clone(), cx).await);
351 }
352 new_locations
353 })
354 })
355 }
356}
357
358#[derive(Debug)]
359pub enum ToolCallStatus {
360 /// The tool call hasn't started running yet, but we start showing it to
361 /// the user.
362 Pending,
363 /// The tool call is waiting for confirmation from the user.
364 WaitingForConfirmation {
365 options: Vec<acp::PermissionOption>,
366 respond_tx: oneshot::Sender<acp::PermissionOptionId>,
367 },
368 /// The tool call is currently running.
369 InProgress,
370 /// The tool call completed successfully.
371 Completed,
372 /// The tool call failed.
373 Failed,
374 /// The user rejected the tool call.
375 Rejected,
376 /// The user canceled generation so the tool call was canceled.
377 Canceled,
378}
379
380impl From<acp::ToolCallStatus> for ToolCallStatus {
381 fn from(status: acp::ToolCallStatus) -> Self {
382 match status {
383 acp::ToolCallStatus::Pending => Self::Pending,
384 acp::ToolCallStatus::InProgress => Self::InProgress,
385 acp::ToolCallStatus::Completed => Self::Completed,
386 acp::ToolCallStatus::Failed => Self::Failed,
387 }
388 }
389}
390
391impl Display for ToolCallStatus {
392 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
393 write!(
394 f,
395 "{}",
396 match self {
397 ToolCallStatus::Pending => "Pending",
398 ToolCallStatus::WaitingForConfirmation { .. } => "Waiting for confirmation",
399 ToolCallStatus::InProgress => "In Progress",
400 ToolCallStatus::Completed => "Completed",
401 ToolCallStatus::Failed => "Failed",
402 ToolCallStatus::Rejected => "Rejected",
403 ToolCallStatus::Canceled => "Canceled",
404 }
405 )
406 }
407}
408
409#[derive(Debug, PartialEq, Clone)]
410pub enum ContentBlock {
411 Empty,
412 Markdown { markdown: Entity<Markdown> },
413 ResourceLink { resource_link: acp::ResourceLink },
414}
415
416impl ContentBlock {
417 pub fn new(
418 block: acp::ContentBlock,
419 language_registry: &Arc<LanguageRegistry>,
420 cx: &mut App,
421 ) -> Self {
422 let mut this = Self::Empty;
423 this.append(block, language_registry, cx);
424 this
425 }
426
427 pub fn new_combined(
428 blocks: impl IntoIterator<Item = acp::ContentBlock>,
429 language_registry: Arc<LanguageRegistry>,
430 cx: &mut App,
431 ) -> Self {
432 let mut this = Self::Empty;
433 for block in blocks {
434 this.append(block, &language_registry, cx);
435 }
436 this
437 }
438
439 pub fn append(
440 &mut self,
441 block: acp::ContentBlock,
442 language_registry: &Arc<LanguageRegistry>,
443 cx: &mut App,
444 ) {
445 if matches!(self, ContentBlock::Empty)
446 && let acp::ContentBlock::ResourceLink(resource_link) = block
447 {
448 *self = ContentBlock::ResourceLink { resource_link };
449 return;
450 }
451
452 let new_content = self.block_string_contents(block);
453
454 match self {
455 ContentBlock::Empty => {
456 *self = Self::create_markdown_block(new_content, language_registry, cx);
457 }
458 ContentBlock::Markdown { markdown } => {
459 markdown.update(cx, |markdown, cx| markdown.append(&new_content, cx));
460 }
461 ContentBlock::ResourceLink { resource_link } => {
462 let existing_content = Self::resource_link_md(&resource_link.uri);
463 let combined = format!("{}\n{}", existing_content, new_content);
464
465 *self = Self::create_markdown_block(combined, language_registry, cx);
466 }
467 }
468 }
469
470 fn create_markdown_block(
471 content: String,
472 language_registry: &Arc<LanguageRegistry>,
473 cx: &mut App,
474 ) -> ContentBlock {
475 ContentBlock::Markdown {
476 markdown: cx
477 .new(|cx| Markdown::new(content.into(), Some(language_registry.clone()), None, cx)),
478 }
479 }
480
481 fn block_string_contents(&self, block: acp::ContentBlock) -> String {
482 match block {
483 acp::ContentBlock::Text(text_content) => text_content.text,
484 acp::ContentBlock::ResourceLink(resource_link) => {
485 Self::resource_link_md(&resource_link.uri)
486 }
487 acp::ContentBlock::Resource(acp::EmbeddedResource {
488 resource:
489 acp::EmbeddedResourceResource::TextResourceContents(acp::TextResourceContents {
490 uri,
491 ..
492 }),
493 ..
494 }) => Self::resource_link_md(&uri),
495 acp::ContentBlock::Image(image) => Self::image_md(&image),
496 acp::ContentBlock::Audio(_) | acp::ContentBlock::Resource(_) => String::new(),
497 }
498 }
499
500 fn resource_link_md(uri: &str) -> String {
501 if let Some(uri) = MentionUri::parse(uri).log_err() {
502 uri.as_link().to_string()
503 } else {
504 uri.to_string()
505 }
506 }
507
508 fn image_md(_image: &acp::ImageContent) -> String {
509 "`Image`".into()
510 }
511
512 fn to_markdown<'a>(&'a self, cx: &'a App) -> &'a str {
513 match self {
514 ContentBlock::Empty => "",
515 ContentBlock::Markdown { markdown } => markdown.read(cx).source(),
516 ContentBlock::ResourceLink { resource_link } => &resource_link.uri,
517 }
518 }
519
520 pub fn markdown(&self) -> Option<&Entity<Markdown>> {
521 match self {
522 ContentBlock::Empty => None,
523 ContentBlock::Markdown { markdown } => Some(markdown),
524 ContentBlock::ResourceLink { .. } => None,
525 }
526 }
527
528 pub fn resource_link(&self) -> Option<&acp::ResourceLink> {
529 match self {
530 ContentBlock::ResourceLink { resource_link } => Some(resource_link),
531 _ => None,
532 }
533 }
534}
535
536#[derive(Debug)]
537pub enum ToolCallContent {
538 ContentBlock(ContentBlock),
539 Diff(Entity<Diff>),
540 Terminal(Entity<Terminal>),
541}
542
543impl ToolCallContent {
544 pub fn from_acp(
545 content: acp::ToolCallContent,
546 language_registry: Arc<LanguageRegistry>,
547 cx: &mut App,
548 ) -> Self {
549 match content {
550 acp::ToolCallContent::Content { content } => {
551 Self::ContentBlock(ContentBlock::new(content, &language_registry, cx))
552 }
553 acp::ToolCallContent::Diff { diff } => Self::Diff(cx.new(|cx| {
554 Diff::finalized(
555 diff.path,
556 diff.old_text,
557 diff.new_text,
558 language_registry,
559 cx,
560 )
561 })),
562 }
563 }
564
565 pub fn update_from_acp(
566 &mut self,
567 new: acp::ToolCallContent,
568 language_registry: Arc<LanguageRegistry>,
569 cx: &mut App,
570 ) {
571 let needs_update = match (&self, &new) {
572 (Self::Diff(old_diff), acp::ToolCallContent::Diff { diff: new_diff }) => {
573 old_diff.read(cx).needs_update(
574 new_diff.old_text.as_deref().unwrap_or(""),
575 &new_diff.new_text,
576 cx,
577 )
578 }
579 _ => true,
580 };
581
582 if needs_update {
583 *self = Self::from_acp(new, language_registry, cx);
584 }
585 }
586
587 pub fn to_markdown(&self, cx: &App) -> String {
588 match self {
589 Self::ContentBlock(content) => content.to_markdown(cx).to_string(),
590 Self::Diff(diff) => diff.read(cx).to_markdown(cx),
591 Self::Terminal(terminal) => terminal.read(cx).to_markdown(cx),
592 }
593 }
594}
595
596#[derive(Debug, PartialEq)]
597pub enum ToolCallUpdate {
598 UpdateFields(acp::ToolCallUpdate),
599 UpdateDiff(ToolCallUpdateDiff),
600 UpdateTerminal(ToolCallUpdateTerminal),
601}
602
603impl ToolCallUpdate {
604 fn id(&self) -> &acp::ToolCallId {
605 match self {
606 Self::UpdateFields(update) => &update.id,
607 Self::UpdateDiff(diff) => &diff.id,
608 Self::UpdateTerminal(terminal) => &terminal.id,
609 }
610 }
611}
612
613impl From<acp::ToolCallUpdate> for ToolCallUpdate {
614 fn from(update: acp::ToolCallUpdate) -> Self {
615 Self::UpdateFields(update)
616 }
617}
618
619impl From<ToolCallUpdateDiff> for ToolCallUpdate {
620 fn from(diff: ToolCallUpdateDiff) -> Self {
621 Self::UpdateDiff(diff)
622 }
623}
624
625#[derive(Debug, PartialEq)]
626pub struct ToolCallUpdateDiff {
627 pub id: acp::ToolCallId,
628 pub diff: Entity<Diff>,
629}
630
631impl From<ToolCallUpdateTerminal> for ToolCallUpdate {
632 fn from(terminal: ToolCallUpdateTerminal) -> Self {
633 Self::UpdateTerminal(terminal)
634 }
635}
636
637#[derive(Debug, PartialEq)]
638pub struct ToolCallUpdateTerminal {
639 pub id: acp::ToolCallId,
640 pub terminal: Entity<Terminal>,
641}
642
643#[derive(Debug, Default)]
644pub struct Plan {
645 pub entries: Vec<PlanEntry>,
646}
647
648#[derive(Debug)]
649pub struct PlanStats<'a> {
650 pub in_progress_entry: Option<&'a PlanEntry>,
651 pub pending: u32,
652 pub completed: u32,
653}
654
655impl Plan {
656 pub fn is_empty(&self) -> bool {
657 self.entries.is_empty()
658 }
659
660 pub fn stats(&self) -> PlanStats<'_> {
661 let mut stats = PlanStats {
662 in_progress_entry: None,
663 pending: 0,
664 completed: 0,
665 };
666
667 for entry in &self.entries {
668 match &entry.status {
669 acp::PlanEntryStatus::Pending => {
670 stats.pending += 1;
671 }
672 acp::PlanEntryStatus::InProgress => {
673 stats.in_progress_entry = stats.in_progress_entry.or(Some(entry));
674 }
675 acp::PlanEntryStatus::Completed => {
676 stats.completed += 1;
677 }
678 }
679 }
680
681 stats
682 }
683}
684
685#[derive(Debug)]
686pub struct PlanEntry {
687 pub content: Entity<Markdown>,
688 pub priority: acp::PlanEntryPriority,
689 pub status: acp::PlanEntryStatus,
690}
691
692impl PlanEntry {
693 pub fn from_acp(entry: acp::PlanEntry, cx: &mut App) -> Self {
694 Self {
695 content: cx.new(|cx| Markdown::new(entry.content.into(), None, None, cx)),
696 priority: entry.priority,
697 status: entry.status,
698 }
699 }
700}
701
702#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
703pub struct TokenUsage {
704 pub max_tokens: u64,
705 pub used_tokens: u64,
706}
707
708impl TokenUsage {
709 pub fn ratio(&self) -> TokenUsageRatio {
710 #[cfg(debug_assertions)]
711 let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
712 .unwrap_or("0.8".to_string())
713 .parse()
714 .unwrap();
715 #[cfg(not(debug_assertions))]
716 let warning_threshold: f32 = 0.8;
717
718 // When the maximum is unknown because there is no selected model,
719 // avoid showing the token limit warning.
720 if self.max_tokens == 0 {
721 TokenUsageRatio::Normal
722 } else if self.used_tokens >= self.max_tokens {
723 TokenUsageRatio::Exceeded
724 } else if self.used_tokens as f32 / self.max_tokens as f32 >= warning_threshold {
725 TokenUsageRatio::Warning
726 } else {
727 TokenUsageRatio::Normal
728 }
729 }
730}
731
732#[derive(Debug, Clone, PartialEq, Eq)]
733pub enum TokenUsageRatio {
734 Normal,
735 Warning,
736 Exceeded,
737}
738
739#[derive(Debug, Clone)]
740pub struct RetryStatus {
741 pub last_error: SharedString,
742 pub attempt: usize,
743 pub max_attempts: usize,
744 pub started_at: Instant,
745 pub duration: Duration,
746}
747
748pub struct AcpThread {
749 title: SharedString,
750 entries: Vec<AgentThreadEntry>,
751 plan: Plan,
752 project: Entity<Project>,
753 action_log: Entity<ActionLog>,
754 shared_buffers: HashMap<Entity<Buffer>, BufferSnapshot>,
755 send_task: Option<Task<()>>,
756 connection: Rc<dyn AgentConnection>,
757 session_id: acp::SessionId,
758 token_usage: Option<TokenUsage>,
759}
760
761#[derive(Debug)]
762pub enum AcpThreadEvent {
763 NewEntry,
764 TitleUpdated,
765 TokenUsageUpdated,
766 EntryUpdated(usize),
767 EntriesRemoved(Range<usize>),
768 ToolAuthorizationRequired,
769 Retry(RetryStatus),
770 Stopped,
771 Error,
772 LoadError(LoadError),
773}
774
775impl EventEmitter<AcpThreadEvent> for AcpThread {}
776
777#[derive(PartialEq, Eq)]
778pub enum ThreadStatus {
779 Idle,
780 WaitingForToolConfirmation,
781 Generating,
782}
783
784#[derive(Debug, Clone)]
785pub enum LoadError {
786 NotInstalled {
787 error_message: SharedString,
788 install_message: SharedString,
789 install_command: String,
790 },
791 Unsupported {
792 error_message: SharedString,
793 upgrade_message: SharedString,
794 upgrade_command: String,
795 },
796 Exited {
797 status: ExitStatus,
798 },
799 Other(SharedString),
800}
801
802impl Display for LoadError {
803 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
804 match self {
805 LoadError::NotInstalled { error_message, .. }
806 | LoadError::Unsupported { error_message, .. } => {
807 write!(f, "{error_message}")
808 }
809 LoadError::Exited { status } => write!(f, "Server exited with status {status}"),
810 LoadError::Other(msg) => write!(f, "{}", msg),
811 }
812 }
813}
814
815impl Error for LoadError {}
816
817impl AcpThread {
818 pub fn new(
819 title: impl Into<SharedString>,
820 connection: Rc<dyn AgentConnection>,
821 project: Entity<Project>,
822 action_log: Entity<ActionLog>,
823 session_id: acp::SessionId,
824 ) -> Self {
825 Self {
826 action_log,
827 shared_buffers: Default::default(),
828 entries: Default::default(),
829 plan: Default::default(),
830 title: title.into(),
831 project,
832 send_task: None,
833 connection,
834 session_id,
835 token_usage: None,
836 }
837 }
838
839 pub fn connection(&self) -> &Rc<dyn AgentConnection> {
840 &self.connection
841 }
842
843 pub fn action_log(&self) -> &Entity<ActionLog> {
844 &self.action_log
845 }
846
847 pub fn project(&self) -> &Entity<Project> {
848 &self.project
849 }
850
851 pub fn title(&self) -> SharedString {
852 self.title.clone()
853 }
854
855 pub fn entries(&self) -> &[AgentThreadEntry] {
856 &self.entries
857 }
858
859 pub fn session_id(&self) -> &acp::SessionId {
860 &self.session_id
861 }
862
863 pub fn status(&self) -> ThreadStatus {
864 if self.send_task.is_some() {
865 if self.waiting_for_tool_confirmation() {
866 ThreadStatus::WaitingForToolConfirmation
867 } else {
868 ThreadStatus::Generating
869 }
870 } else {
871 ThreadStatus::Idle
872 }
873 }
874
875 pub fn token_usage(&self) -> Option<&TokenUsage> {
876 self.token_usage.as_ref()
877 }
878
879 pub fn has_pending_edit_tool_calls(&self) -> bool {
880 for entry in self.entries.iter().rev() {
881 match entry {
882 AgentThreadEntry::UserMessage(_) => return false,
883 AgentThreadEntry::ToolCall(
884 call @ ToolCall {
885 status: ToolCallStatus::InProgress | ToolCallStatus::Pending,
886 ..
887 },
888 ) if call.diffs().next().is_some() => {
889 return true;
890 }
891 AgentThreadEntry::ToolCall(_) | AgentThreadEntry::AssistantMessage(_) => {}
892 }
893 }
894
895 false
896 }
897
898 pub fn used_tools_since_last_user_message(&self) -> bool {
899 for entry in self.entries.iter().rev() {
900 match entry {
901 AgentThreadEntry::UserMessage(..) => return false,
902 AgentThreadEntry::AssistantMessage(..) => continue,
903 AgentThreadEntry::ToolCall(..) => return true,
904 }
905 }
906
907 false
908 }
909
910 pub fn handle_session_update(
911 &mut self,
912 update: acp::SessionUpdate,
913 cx: &mut Context<Self>,
914 ) -> Result<(), acp::Error> {
915 match update {
916 acp::SessionUpdate::UserMessageChunk { content } => {
917 self.push_user_content_block(None, content, cx);
918 }
919 acp::SessionUpdate::AgentMessageChunk { content } => {
920 self.push_assistant_content_block(content, false, cx);
921 }
922 acp::SessionUpdate::AgentThoughtChunk { content } => {
923 self.push_assistant_content_block(content, true, cx);
924 }
925 acp::SessionUpdate::ToolCall(tool_call) => {
926 self.upsert_tool_call(tool_call, cx)?;
927 }
928 acp::SessionUpdate::ToolCallUpdate(tool_call_update) => {
929 self.update_tool_call(tool_call_update, cx)?;
930 }
931 acp::SessionUpdate::Plan(plan) => {
932 self.update_plan(plan, cx);
933 }
934 }
935 Ok(())
936 }
937
938 pub fn push_user_content_block(
939 &mut self,
940 message_id: Option<UserMessageId>,
941 chunk: acp::ContentBlock,
942 cx: &mut Context<Self>,
943 ) {
944 let language_registry = self.project.read(cx).languages().clone();
945 let entries_len = self.entries.len();
946
947 if let Some(last_entry) = self.entries.last_mut()
948 && let AgentThreadEntry::UserMessage(UserMessage {
949 id,
950 content,
951 chunks,
952 ..
953 }) = last_entry
954 {
955 *id = message_id.or(id.take());
956 content.append(chunk.clone(), &language_registry, cx);
957 chunks.push(chunk);
958 let idx = entries_len - 1;
959 cx.emit(AcpThreadEvent::EntryUpdated(idx));
960 } else {
961 let content = ContentBlock::new(chunk.clone(), &language_registry, cx);
962 self.push_entry(
963 AgentThreadEntry::UserMessage(UserMessage {
964 id: message_id,
965 content,
966 chunks: vec![chunk],
967 checkpoint: None,
968 }),
969 cx,
970 );
971 }
972 }
973
974 pub fn push_assistant_content_block(
975 &mut self,
976 chunk: acp::ContentBlock,
977 is_thought: bool,
978 cx: &mut Context<Self>,
979 ) {
980 let language_registry = self.project.read(cx).languages().clone();
981 let entries_len = self.entries.len();
982 if let Some(last_entry) = self.entries.last_mut()
983 && let AgentThreadEntry::AssistantMessage(AssistantMessage { chunks }) = last_entry
984 {
985 let idx = entries_len - 1;
986 cx.emit(AcpThreadEvent::EntryUpdated(idx));
987 match (chunks.last_mut(), is_thought) {
988 (Some(AssistantMessageChunk::Message { block }), false)
989 | (Some(AssistantMessageChunk::Thought { block }), true) => {
990 block.append(chunk, &language_registry, cx)
991 }
992 _ => {
993 let block = ContentBlock::new(chunk, &language_registry, cx);
994 if is_thought {
995 chunks.push(AssistantMessageChunk::Thought { block })
996 } else {
997 chunks.push(AssistantMessageChunk::Message { block })
998 }
999 }
1000 }
1001 } else {
1002 let block = ContentBlock::new(chunk, &language_registry, cx);
1003 let chunk = if is_thought {
1004 AssistantMessageChunk::Thought { block }
1005 } else {
1006 AssistantMessageChunk::Message { block }
1007 };
1008
1009 self.push_entry(
1010 AgentThreadEntry::AssistantMessage(AssistantMessage {
1011 chunks: vec![chunk],
1012 }),
1013 cx,
1014 );
1015 }
1016 }
1017
1018 fn push_entry(&mut self, entry: AgentThreadEntry, cx: &mut Context<Self>) {
1019 self.entries.push(entry);
1020 cx.emit(AcpThreadEvent::NewEntry);
1021 }
1022
1023 pub fn can_set_title(&mut self, cx: &mut Context<Self>) -> bool {
1024 self.connection.set_title(&self.session_id, cx).is_some()
1025 }
1026
1027 pub fn set_title(&mut self, title: SharedString, cx: &mut Context<Self>) -> Task<Result<()>> {
1028 if title != self.title {
1029 self.title = title.clone();
1030 cx.emit(AcpThreadEvent::TitleUpdated);
1031 if let Some(set_title) = self.connection.set_title(&self.session_id, cx) {
1032 return set_title.run(title, cx);
1033 }
1034 }
1035 Task::ready(Ok(()))
1036 }
1037
1038 pub fn update_token_usage(&mut self, usage: Option<TokenUsage>, cx: &mut Context<Self>) {
1039 self.token_usage = usage;
1040 cx.emit(AcpThreadEvent::TokenUsageUpdated);
1041 }
1042
1043 pub fn update_retry_status(&mut self, status: RetryStatus, cx: &mut Context<Self>) {
1044 cx.emit(AcpThreadEvent::Retry(status));
1045 }
1046
1047 pub fn update_tool_call(
1048 &mut self,
1049 update: impl Into<ToolCallUpdate>,
1050 cx: &mut Context<Self>,
1051 ) -> Result<()> {
1052 let update = update.into();
1053 let languages = self.project.read(cx).languages().clone();
1054
1055 let (ix, current_call) = self
1056 .tool_call_mut(update.id())
1057 .context("Tool call not found")?;
1058 match update {
1059 ToolCallUpdate::UpdateFields(update) => {
1060 let location_updated = update.fields.locations.is_some();
1061 current_call.update_fields(update.fields, languages, cx);
1062 if location_updated {
1063 self.resolve_locations(update.id, cx);
1064 }
1065 }
1066 ToolCallUpdate::UpdateDiff(update) => {
1067 current_call.content.clear();
1068 current_call
1069 .content
1070 .push(ToolCallContent::Diff(update.diff));
1071 }
1072 ToolCallUpdate::UpdateTerminal(update) => {
1073 current_call.content.clear();
1074 current_call
1075 .content
1076 .push(ToolCallContent::Terminal(update.terminal));
1077 }
1078 }
1079
1080 cx.emit(AcpThreadEvent::EntryUpdated(ix));
1081
1082 Ok(())
1083 }
1084
1085 /// Updates a tool call if id matches an existing entry, otherwise inserts a new one.
1086 pub fn upsert_tool_call(
1087 &mut self,
1088 tool_call: acp::ToolCall,
1089 cx: &mut Context<Self>,
1090 ) -> Result<(), acp::Error> {
1091 let status = tool_call.status.into();
1092 self.upsert_tool_call_inner(tool_call.into(), status, cx)
1093 }
1094
1095 /// Fails if id does not match an existing entry.
1096 pub fn upsert_tool_call_inner(
1097 &mut self,
1098 tool_call_update: acp::ToolCallUpdate,
1099 status: ToolCallStatus,
1100 cx: &mut Context<Self>,
1101 ) -> Result<(), acp::Error> {
1102 let language_registry = self.project.read(cx).languages().clone();
1103 let id = tool_call_update.id.clone();
1104
1105 if let Some((ix, current_call)) = self.tool_call_mut(&id) {
1106 current_call.update_fields(tool_call_update.fields, language_registry, cx);
1107 current_call.status = status;
1108
1109 cx.emit(AcpThreadEvent::EntryUpdated(ix));
1110 } else {
1111 let call =
1112 ToolCall::from_acp(tool_call_update.try_into()?, status, language_registry, cx);
1113 self.push_entry(AgentThreadEntry::ToolCall(call), cx);
1114 };
1115
1116 self.resolve_locations(id, cx);
1117 Ok(())
1118 }
1119
1120 fn tool_call_mut(&mut self, id: &acp::ToolCallId) -> Option<(usize, &mut ToolCall)> {
1121 // The tool call we are looking for is typically the last one, or very close to the end.
1122 // At the moment, it doesn't seem like a hashmap would be a good fit for this use case.
1123 self.entries
1124 .iter_mut()
1125 .enumerate()
1126 .rev()
1127 .find_map(|(index, tool_call)| {
1128 if let AgentThreadEntry::ToolCall(tool_call) = tool_call
1129 && &tool_call.id == id
1130 {
1131 Some((index, tool_call))
1132 } else {
1133 None
1134 }
1135 })
1136 }
1137
1138 pub fn tool_call(&mut self, id: &acp::ToolCallId) -> Option<(usize, &ToolCall)> {
1139 self.entries
1140 .iter()
1141 .enumerate()
1142 .rev()
1143 .find_map(|(index, tool_call)| {
1144 if let AgentThreadEntry::ToolCall(tool_call) = tool_call
1145 && &tool_call.id == id
1146 {
1147 Some((index, tool_call))
1148 } else {
1149 None
1150 }
1151 })
1152 }
1153
1154 pub fn resolve_locations(&mut self, id: acp::ToolCallId, cx: &mut Context<Self>) {
1155 let project = self.project.clone();
1156 let Some((_, tool_call)) = self.tool_call_mut(&id) else {
1157 return;
1158 };
1159 let task = tool_call.resolve_locations(project, cx);
1160 cx.spawn(async move |this, cx| {
1161 let resolved_locations = task.await;
1162 this.update(cx, |this, cx| {
1163 let project = this.project.clone();
1164 let Some((ix, tool_call)) = this.tool_call_mut(&id) else {
1165 return;
1166 };
1167 if let Some(Some(location)) = resolved_locations.last() {
1168 project.update(cx, |project, cx| {
1169 if let Some(agent_location) = project.agent_location() {
1170 let should_ignore = agent_location.buffer == location.buffer
1171 && location
1172 .buffer
1173 .update(cx, |buffer, _| {
1174 let snapshot = buffer.snapshot();
1175 let old_position =
1176 agent_location.position.to_point(&snapshot);
1177 let new_position = location.position.to_point(&snapshot);
1178 // ignore this so that when we get updates from the edit tool
1179 // the position doesn't reset to the startof line
1180 old_position.row == new_position.row
1181 && old_position.column > new_position.column
1182 })
1183 .ok()
1184 .unwrap_or_default();
1185 if !should_ignore {
1186 project.set_agent_location(Some(location.clone()), cx);
1187 }
1188 }
1189 });
1190 }
1191 if tool_call.resolved_locations != resolved_locations {
1192 tool_call.resolved_locations = resolved_locations;
1193 cx.emit(AcpThreadEvent::EntryUpdated(ix));
1194 }
1195 })
1196 })
1197 .detach();
1198 }
1199
1200 pub fn request_tool_call_authorization(
1201 &mut self,
1202 tool_call: acp::ToolCallUpdate,
1203 options: Vec<acp::PermissionOption>,
1204 cx: &mut Context<Self>,
1205 ) -> Result<oneshot::Receiver<acp::PermissionOptionId>, acp::Error> {
1206 let (tx, rx) = oneshot::channel();
1207
1208 let status = ToolCallStatus::WaitingForConfirmation {
1209 options,
1210 respond_tx: tx,
1211 };
1212
1213 self.upsert_tool_call_inner(tool_call, status, cx)?;
1214 cx.emit(AcpThreadEvent::ToolAuthorizationRequired);
1215 Ok(rx)
1216 }
1217
1218 pub fn authorize_tool_call(
1219 &mut self,
1220 id: acp::ToolCallId,
1221 option_id: acp::PermissionOptionId,
1222 option_kind: acp::PermissionOptionKind,
1223 cx: &mut Context<Self>,
1224 ) {
1225 let Some((ix, call)) = self.tool_call_mut(&id) else {
1226 return;
1227 };
1228
1229 let new_status = match option_kind {
1230 acp::PermissionOptionKind::RejectOnce | acp::PermissionOptionKind::RejectAlways => {
1231 ToolCallStatus::Rejected
1232 }
1233 acp::PermissionOptionKind::AllowOnce | acp::PermissionOptionKind::AllowAlways => {
1234 ToolCallStatus::InProgress
1235 }
1236 };
1237
1238 let curr_status = mem::replace(&mut call.status, new_status);
1239
1240 if let ToolCallStatus::WaitingForConfirmation { respond_tx, .. } = curr_status {
1241 respond_tx.send(option_id).log_err();
1242 } else if cfg!(debug_assertions) {
1243 panic!("tried to authorize an already authorized tool call");
1244 }
1245
1246 cx.emit(AcpThreadEvent::EntryUpdated(ix));
1247 }
1248
1249 /// Returns true if the last turn is awaiting tool authorization
1250 pub fn waiting_for_tool_confirmation(&self) -> bool {
1251 for entry in self.entries.iter().rev() {
1252 match &entry {
1253 AgentThreadEntry::ToolCall(call) => match call.status {
1254 ToolCallStatus::WaitingForConfirmation { .. } => return true,
1255 ToolCallStatus::Pending
1256 | ToolCallStatus::InProgress
1257 | ToolCallStatus::Completed
1258 | ToolCallStatus::Failed
1259 | ToolCallStatus::Rejected
1260 | ToolCallStatus::Canceled => continue,
1261 },
1262 AgentThreadEntry::UserMessage(_) | AgentThreadEntry::AssistantMessage(_) => {
1263 // Reached the beginning of the turn
1264 return false;
1265 }
1266 }
1267 }
1268 false
1269 }
1270
1271 pub fn plan(&self) -> &Plan {
1272 &self.plan
1273 }
1274
1275 pub fn update_plan(&mut self, request: acp::Plan, cx: &mut Context<Self>) {
1276 let new_entries_len = request.entries.len();
1277 let mut new_entries = request.entries.into_iter();
1278
1279 // Reuse existing markdown to prevent flickering
1280 for (old, new) in self.plan.entries.iter_mut().zip(new_entries.by_ref()) {
1281 let PlanEntry {
1282 content,
1283 priority,
1284 status,
1285 } = old;
1286 content.update(cx, |old, cx| {
1287 old.replace(new.content, cx);
1288 });
1289 *priority = new.priority;
1290 *status = new.status;
1291 }
1292 for new in new_entries {
1293 self.plan.entries.push(PlanEntry::from_acp(new, cx))
1294 }
1295 self.plan.entries.truncate(new_entries_len);
1296
1297 cx.notify();
1298 }
1299
1300 fn clear_completed_plan_entries(&mut self, cx: &mut Context<Self>) {
1301 self.plan
1302 .entries
1303 .retain(|entry| !matches!(entry.status, acp::PlanEntryStatus::Completed));
1304 cx.notify();
1305 }
1306
1307 #[cfg(any(test, feature = "test-support"))]
1308 pub fn send_raw(
1309 &mut self,
1310 message: &str,
1311 cx: &mut Context<Self>,
1312 ) -> BoxFuture<'static, Result<()>> {
1313 self.send(
1314 vec![acp::ContentBlock::Text(acp::TextContent {
1315 text: message.to_string(),
1316 annotations: None,
1317 })],
1318 cx,
1319 )
1320 }
1321
1322 pub fn send(
1323 &mut self,
1324 message: Vec<acp::ContentBlock>,
1325 cx: &mut Context<Self>,
1326 ) -> BoxFuture<'static, Result<()>> {
1327 let block = ContentBlock::new_combined(
1328 message.clone(),
1329 self.project.read(cx).languages().clone(),
1330 cx,
1331 );
1332 let request = acp::PromptRequest {
1333 prompt: message.clone(),
1334 session_id: self.session_id.clone(),
1335 };
1336 let git_store = self.project.read(cx).git_store().clone();
1337
1338 let message_id = if self.connection.truncate(&self.session_id, cx).is_some() {
1339 Some(UserMessageId::new())
1340 } else {
1341 None
1342 };
1343
1344 self.run_turn(cx, async move |this, cx| {
1345 this.update(cx, |this, cx| {
1346 this.push_entry(
1347 AgentThreadEntry::UserMessage(UserMessage {
1348 id: message_id.clone(),
1349 content: block,
1350 chunks: message,
1351 checkpoint: None,
1352 }),
1353 cx,
1354 );
1355 })
1356 .ok();
1357
1358 let old_checkpoint = git_store
1359 .update(cx, |git, cx| git.checkpoint(cx))?
1360 .await
1361 .context("failed to get old checkpoint")
1362 .log_err();
1363 this.update(cx, |this, cx| {
1364 if let Some((_ix, message)) = this.last_user_message() {
1365 message.checkpoint = old_checkpoint.map(|git_checkpoint| Checkpoint {
1366 git_checkpoint,
1367 show: false,
1368 });
1369 }
1370 this.connection.prompt(message_id, request, cx)
1371 })?
1372 .await
1373 })
1374 }
1375
1376 pub fn resume(&mut self, cx: &mut Context<Self>) -> BoxFuture<'static, Result<()>> {
1377 self.run_turn(cx, async move |this, cx| {
1378 this.update(cx, |this, cx| {
1379 this.connection
1380 .resume(&this.session_id, cx)
1381 .map(|resume| resume.run(cx))
1382 })?
1383 .context("resuming a session is not supported")?
1384 .await
1385 })
1386 }
1387
1388 fn run_turn(
1389 &mut self,
1390 cx: &mut Context<Self>,
1391 f: impl 'static + AsyncFnOnce(WeakEntity<Self>, &mut AsyncApp) -> Result<acp::PromptResponse>,
1392 ) -> BoxFuture<'static, Result<()>> {
1393 self.clear_completed_plan_entries(cx);
1394
1395 let (tx, rx) = oneshot::channel();
1396 let cancel_task = self.cancel(cx);
1397
1398 self.send_task = Some(cx.spawn(async move |this, cx| {
1399 cancel_task.await;
1400 tx.send(f(this, cx).await).ok();
1401 }));
1402
1403 cx.spawn(async move |this, cx| {
1404 let response = rx.await;
1405
1406 this.update(cx, |this, cx| this.update_last_checkpoint(cx))?
1407 .await?;
1408
1409 this.update(cx, |this, cx| {
1410 this.project
1411 .update(cx, |project, cx| project.set_agent_location(None, cx));
1412 match response {
1413 Ok(Err(e)) => {
1414 this.send_task.take();
1415 cx.emit(AcpThreadEvent::Error);
1416 Err(e)
1417 }
1418 result => {
1419 let canceled = matches!(
1420 result,
1421 Ok(Ok(acp::PromptResponse {
1422 stop_reason: acp::StopReason::Cancelled
1423 }))
1424 );
1425
1426 // We only take the task if the current prompt wasn't canceled.
1427 //
1428 // This prompt may have been canceled because another one was sent
1429 // while it was still generating. In these cases, dropping `send_task`
1430 // would cause the next generation to be canceled.
1431 if !canceled {
1432 this.send_task.take();
1433 }
1434
1435 // Truncate entries if the last prompt was refused.
1436 if let Ok(Ok(acp::PromptResponse {
1437 stop_reason: acp::StopReason::Refusal,
1438 })) = result
1439 && let Some((ix, _)) = this.last_user_message()
1440 {
1441 let range = ix..this.entries.len();
1442 this.entries.truncate(ix);
1443 cx.emit(AcpThreadEvent::EntriesRemoved(range));
1444 }
1445
1446 cx.emit(AcpThreadEvent::Stopped);
1447 Ok(())
1448 }
1449 }
1450 })?
1451 })
1452 .boxed()
1453 }
1454
1455 pub fn cancel(&mut self, cx: &mut Context<Self>) -> Task<()> {
1456 let Some(send_task) = self.send_task.take() else {
1457 return Task::ready(());
1458 };
1459
1460 for entry in self.entries.iter_mut() {
1461 if let AgentThreadEntry::ToolCall(call) = entry {
1462 let cancel = matches!(
1463 call.status,
1464 ToolCallStatus::Pending
1465 | ToolCallStatus::WaitingForConfirmation { .. }
1466 | ToolCallStatus::InProgress
1467 );
1468
1469 if cancel {
1470 call.status = ToolCallStatus::Canceled;
1471 }
1472 }
1473 }
1474
1475 self.connection.cancel(&self.session_id, cx);
1476
1477 // Wait for the send task to complete
1478 cx.foreground_executor().spawn(send_task)
1479 }
1480
1481 /// Rewinds this thread to before the entry at `index`, removing it and all
1482 /// subsequent entries while reverting any changes made from that point.
1483 pub fn rewind(&mut self, id: UserMessageId, cx: &mut Context<Self>) -> Task<Result<()>> {
1484 let Some(truncate) = self.connection.truncate(&self.session_id, cx) else {
1485 return Task::ready(Err(anyhow!("not supported")));
1486 };
1487 let Some(message) = self.user_message(&id) else {
1488 return Task::ready(Err(anyhow!("message not found")));
1489 };
1490
1491 let checkpoint = message
1492 .checkpoint
1493 .as_ref()
1494 .map(|c| c.git_checkpoint.clone());
1495
1496 let git_store = self.project.read(cx).git_store().clone();
1497 cx.spawn(async move |this, cx| {
1498 if let Some(checkpoint) = checkpoint {
1499 git_store
1500 .update(cx, |git, cx| git.restore_checkpoint(checkpoint, cx))?
1501 .await?;
1502 }
1503
1504 cx.update(|cx| truncate.run(id.clone(), cx))?.await?;
1505 this.update(cx, |this, cx| {
1506 if let Some((ix, _)) = this.user_message_mut(&id) {
1507 let range = ix..this.entries.len();
1508 this.entries.truncate(ix);
1509 cx.emit(AcpThreadEvent::EntriesRemoved(range));
1510 }
1511 })
1512 })
1513 }
1514
1515 fn update_last_checkpoint(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
1516 let git_store = self.project.read(cx).git_store().clone();
1517
1518 let old_checkpoint = if let Some((_, message)) = self.last_user_message() {
1519 if let Some(checkpoint) = message.checkpoint.as_ref() {
1520 checkpoint.git_checkpoint.clone()
1521 } else {
1522 return Task::ready(Ok(()));
1523 }
1524 } else {
1525 return Task::ready(Ok(()));
1526 };
1527
1528 let new_checkpoint = git_store.update(cx, |git, cx| git.checkpoint(cx));
1529 cx.spawn(async move |this, cx| {
1530 let new_checkpoint = new_checkpoint
1531 .await
1532 .context("failed to get new checkpoint")
1533 .log_err();
1534 if let Some(new_checkpoint) = new_checkpoint {
1535 let equal = git_store
1536 .update(cx, |git, cx| {
1537 git.compare_checkpoints(old_checkpoint.clone(), new_checkpoint, cx)
1538 })?
1539 .await
1540 .unwrap_or(true);
1541 this.update(cx, |this, cx| {
1542 let (ix, message) = this.last_user_message().context("no user message")?;
1543 let checkpoint = message.checkpoint.as_mut().context("no checkpoint")?;
1544 checkpoint.show = !equal;
1545 cx.emit(AcpThreadEvent::EntryUpdated(ix));
1546 anyhow::Ok(())
1547 })??;
1548 }
1549
1550 Ok(())
1551 })
1552 }
1553
1554 fn last_user_message(&mut self) -> Option<(usize, &mut UserMessage)> {
1555 self.entries
1556 .iter_mut()
1557 .enumerate()
1558 .rev()
1559 .find_map(|(ix, entry)| {
1560 if let AgentThreadEntry::UserMessage(message) = entry {
1561 Some((ix, message))
1562 } else {
1563 None
1564 }
1565 })
1566 }
1567
1568 fn user_message(&self, id: &UserMessageId) -> Option<&UserMessage> {
1569 self.entries.iter().find_map(|entry| {
1570 if let AgentThreadEntry::UserMessage(message) = entry {
1571 if message.id.as_ref() == Some(id) {
1572 Some(message)
1573 } else {
1574 None
1575 }
1576 } else {
1577 None
1578 }
1579 })
1580 }
1581
1582 fn user_message_mut(&mut self, id: &UserMessageId) -> Option<(usize, &mut UserMessage)> {
1583 self.entries.iter_mut().enumerate().find_map(|(ix, entry)| {
1584 if let AgentThreadEntry::UserMessage(message) = entry {
1585 if message.id.as_ref() == Some(id) {
1586 Some((ix, message))
1587 } else {
1588 None
1589 }
1590 } else {
1591 None
1592 }
1593 })
1594 }
1595
1596 pub fn read_text_file(
1597 &self,
1598 path: PathBuf,
1599 line: Option<u32>,
1600 limit: Option<u32>,
1601 reuse_shared_snapshot: bool,
1602 cx: &mut Context<Self>,
1603 ) -> Task<Result<String>> {
1604 let project = self.project.clone();
1605 let action_log = self.action_log.clone();
1606 cx.spawn(async move |this, cx| {
1607 let load = project.update(cx, |project, cx| {
1608 let path = project
1609 .project_path_for_absolute_path(&path, cx)
1610 .context("invalid path")?;
1611 anyhow::Ok(project.open_buffer(path, cx))
1612 });
1613 let buffer = load??.await?;
1614
1615 let snapshot = if reuse_shared_snapshot {
1616 this.read_with(cx, |this, _| {
1617 this.shared_buffers.get(&buffer.clone()).cloned()
1618 })
1619 .log_err()
1620 .flatten()
1621 } else {
1622 None
1623 };
1624
1625 let snapshot = if let Some(snapshot) = snapshot {
1626 snapshot
1627 } else {
1628 action_log.update(cx, |action_log, cx| {
1629 action_log.buffer_read(buffer.clone(), cx);
1630 })?;
1631 project.update(cx, |project, cx| {
1632 let position = buffer
1633 .read(cx)
1634 .snapshot()
1635 .anchor_before(Point::new(line.unwrap_or_default(), 0));
1636 project.set_agent_location(
1637 Some(AgentLocation {
1638 buffer: buffer.downgrade(),
1639 position,
1640 }),
1641 cx,
1642 );
1643 })?;
1644
1645 buffer.update(cx, |buffer, _| buffer.snapshot())?
1646 };
1647
1648 this.update(cx, |this, _| {
1649 let text = snapshot.text();
1650 this.shared_buffers.insert(buffer.clone(), snapshot);
1651 if line.is_none() && limit.is_none() {
1652 return Ok(text);
1653 }
1654 let limit = limit.unwrap_or(u32::MAX) as usize;
1655 let Some(line) = line else {
1656 return Ok(text.lines().take(limit).collect::<String>());
1657 };
1658
1659 let count = text.lines().count();
1660 if count < line as usize {
1661 anyhow::bail!("There are only {} lines", count);
1662 }
1663 Ok(text
1664 .lines()
1665 .skip(line as usize + 1)
1666 .take(limit)
1667 .collect::<String>())
1668 })?
1669 })
1670 }
1671
1672 pub fn write_text_file(
1673 &self,
1674 path: PathBuf,
1675 content: String,
1676 cx: &mut Context<Self>,
1677 ) -> Task<Result<()>> {
1678 let project = self.project.clone();
1679 let action_log = self.action_log.clone();
1680 cx.spawn(async move |this, cx| {
1681 let load = project.update(cx, |project, cx| {
1682 let path = project
1683 .project_path_for_absolute_path(&path, cx)
1684 .context("invalid path")?;
1685 anyhow::Ok(project.open_buffer(path, cx))
1686 });
1687 let buffer = load??.await?;
1688 let snapshot = this.update(cx, |this, cx| {
1689 this.shared_buffers
1690 .get(&buffer)
1691 .cloned()
1692 .unwrap_or_else(|| buffer.read(cx).snapshot())
1693 })?;
1694 let edits = cx
1695 .background_executor()
1696 .spawn(async move {
1697 let old_text = snapshot.text();
1698 text_diff(old_text.as_str(), &content)
1699 .into_iter()
1700 .map(|(range, replacement)| {
1701 (
1702 snapshot.anchor_after(range.start)
1703 ..snapshot.anchor_before(range.end),
1704 replacement,
1705 )
1706 })
1707 .collect::<Vec<_>>()
1708 })
1709 .await;
1710
1711 project.update(cx, |project, cx| {
1712 project.set_agent_location(
1713 Some(AgentLocation {
1714 buffer: buffer.downgrade(),
1715 position: edits
1716 .last()
1717 .map(|(range, _)| range.end)
1718 .unwrap_or(Anchor::MIN),
1719 }),
1720 cx,
1721 );
1722 })?;
1723
1724 let format_on_save = cx.update(|cx| {
1725 action_log.update(cx, |action_log, cx| {
1726 action_log.buffer_read(buffer.clone(), cx);
1727 });
1728
1729 let format_on_save = buffer.update(cx, |buffer, cx| {
1730 buffer.edit(edits, None, cx);
1731
1732 let settings = language::language_settings::language_settings(
1733 buffer.language().map(|l| l.name()),
1734 buffer.file(),
1735 cx,
1736 );
1737
1738 settings.format_on_save != FormatOnSave::Off
1739 });
1740 action_log.update(cx, |action_log, cx| {
1741 action_log.buffer_edited(buffer.clone(), cx);
1742 });
1743 format_on_save
1744 })?;
1745
1746 if format_on_save {
1747 let format_task = project.update(cx, |project, cx| {
1748 project.format(
1749 HashSet::from_iter([buffer.clone()]),
1750 LspFormatTarget::Buffers,
1751 false,
1752 FormatTrigger::Save,
1753 cx,
1754 )
1755 })?;
1756 format_task.await.log_err();
1757
1758 action_log.update(cx, |action_log, cx| {
1759 action_log.buffer_edited(buffer.clone(), cx);
1760 })?;
1761 }
1762
1763 project
1764 .update(cx, |project, cx| project.save_buffer(buffer, cx))?
1765 .await
1766 })
1767 }
1768
1769 pub fn to_markdown(&self, cx: &App) -> String {
1770 self.entries.iter().map(|e| e.to_markdown(cx)).collect()
1771 }
1772
1773 pub fn emit_load_error(&mut self, error: LoadError, cx: &mut Context<Self>) {
1774 cx.emit(AcpThreadEvent::LoadError(error));
1775 }
1776}
1777
1778fn markdown_for_raw_output(
1779 raw_output: &serde_json::Value,
1780 language_registry: &Arc<LanguageRegistry>,
1781 cx: &mut App,
1782) -> Option<Entity<Markdown>> {
1783 match raw_output {
1784 serde_json::Value::Null => None,
1785 serde_json::Value::Bool(value) => Some(cx.new(|cx| {
1786 Markdown::new(
1787 value.to_string().into(),
1788 Some(language_registry.clone()),
1789 None,
1790 cx,
1791 )
1792 })),
1793 serde_json::Value::Number(value) => Some(cx.new(|cx| {
1794 Markdown::new(
1795 value.to_string().into(),
1796 Some(language_registry.clone()),
1797 None,
1798 cx,
1799 )
1800 })),
1801 serde_json::Value::String(value) => Some(cx.new(|cx| {
1802 Markdown::new(
1803 value.clone().into(),
1804 Some(language_registry.clone()),
1805 None,
1806 cx,
1807 )
1808 })),
1809 value => Some(cx.new(|cx| {
1810 Markdown::new(
1811 format!("```json\n{}\n```", value).into(),
1812 Some(language_registry.clone()),
1813 None,
1814 cx,
1815 )
1816 })),
1817 }
1818}
1819
1820#[cfg(test)]
1821mod tests {
1822 use super::*;
1823 use anyhow::anyhow;
1824 use futures::{channel::mpsc, future::LocalBoxFuture, select};
1825 use gpui::{App, AsyncApp, TestAppContext, WeakEntity};
1826 use indoc::indoc;
1827 use project::{FakeFs, Fs};
1828 use rand::Rng as _;
1829 use serde_json::json;
1830 use settings::SettingsStore;
1831 use smol::stream::StreamExt as _;
1832 use std::{
1833 any::Any,
1834 cell::RefCell,
1835 path::Path,
1836 rc::Rc,
1837 sync::atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
1838 time::Duration,
1839 };
1840 use util::path;
1841
1842 fn init_test(cx: &mut TestAppContext) {
1843 env_logger::try_init().ok();
1844 cx.update(|cx| {
1845 let settings_store = SettingsStore::test(cx);
1846 cx.set_global(settings_store);
1847 Project::init_settings(cx);
1848 language::init(cx);
1849 });
1850 }
1851
1852 #[gpui::test]
1853 async fn test_push_user_content_block(cx: &mut gpui::TestAppContext) {
1854 init_test(cx);
1855
1856 let fs = FakeFs::new(cx.executor());
1857 let project = Project::test(fs, [], cx).await;
1858 let connection = Rc::new(FakeAgentConnection::new());
1859 let thread = cx
1860 .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
1861 .await
1862 .unwrap();
1863
1864 // Test creating a new user message
1865 thread.update(cx, |thread, cx| {
1866 thread.push_user_content_block(
1867 None,
1868 acp::ContentBlock::Text(acp::TextContent {
1869 annotations: None,
1870 text: "Hello, ".to_string(),
1871 }),
1872 cx,
1873 );
1874 });
1875
1876 thread.update(cx, |thread, cx| {
1877 assert_eq!(thread.entries.len(), 1);
1878 if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[0] {
1879 assert_eq!(user_msg.id, None);
1880 assert_eq!(user_msg.content.to_markdown(cx), "Hello, ");
1881 } else {
1882 panic!("Expected UserMessage");
1883 }
1884 });
1885
1886 // Test appending to existing user message
1887 let message_1_id = UserMessageId::new();
1888 thread.update(cx, |thread, cx| {
1889 thread.push_user_content_block(
1890 Some(message_1_id.clone()),
1891 acp::ContentBlock::Text(acp::TextContent {
1892 annotations: None,
1893 text: "world!".to_string(),
1894 }),
1895 cx,
1896 );
1897 });
1898
1899 thread.update(cx, |thread, cx| {
1900 assert_eq!(thread.entries.len(), 1);
1901 if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[0] {
1902 assert_eq!(user_msg.id, Some(message_1_id));
1903 assert_eq!(user_msg.content.to_markdown(cx), "Hello, world!");
1904 } else {
1905 panic!("Expected UserMessage");
1906 }
1907 });
1908
1909 // Test creating new user message after assistant message
1910 thread.update(cx, |thread, cx| {
1911 thread.push_assistant_content_block(
1912 acp::ContentBlock::Text(acp::TextContent {
1913 annotations: None,
1914 text: "Assistant response".to_string(),
1915 }),
1916 false,
1917 cx,
1918 );
1919 });
1920
1921 let message_2_id = UserMessageId::new();
1922 thread.update(cx, |thread, cx| {
1923 thread.push_user_content_block(
1924 Some(message_2_id.clone()),
1925 acp::ContentBlock::Text(acp::TextContent {
1926 annotations: None,
1927 text: "New user message".to_string(),
1928 }),
1929 cx,
1930 );
1931 });
1932
1933 thread.update(cx, |thread, cx| {
1934 assert_eq!(thread.entries.len(), 3);
1935 if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[2] {
1936 assert_eq!(user_msg.id, Some(message_2_id));
1937 assert_eq!(user_msg.content.to_markdown(cx), "New user message");
1938 } else {
1939 panic!("Expected UserMessage at index 2");
1940 }
1941 });
1942 }
1943
1944 #[gpui::test]
1945 async fn test_thinking_concatenation(cx: &mut gpui::TestAppContext) {
1946 init_test(cx);
1947
1948 let fs = FakeFs::new(cx.executor());
1949 let project = Project::test(fs, [], cx).await;
1950 let connection = Rc::new(FakeAgentConnection::new().on_user_message(
1951 |_, thread, mut cx| {
1952 async move {
1953 thread.update(&mut cx, |thread, cx| {
1954 thread
1955 .handle_session_update(
1956 acp::SessionUpdate::AgentThoughtChunk {
1957 content: "Thinking ".into(),
1958 },
1959 cx,
1960 )
1961 .unwrap();
1962 thread
1963 .handle_session_update(
1964 acp::SessionUpdate::AgentThoughtChunk {
1965 content: "hard!".into(),
1966 },
1967 cx,
1968 )
1969 .unwrap();
1970 })?;
1971 Ok(acp::PromptResponse {
1972 stop_reason: acp::StopReason::EndTurn,
1973 })
1974 }
1975 .boxed_local()
1976 },
1977 ));
1978
1979 let thread = cx
1980 .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
1981 .await
1982 .unwrap();
1983
1984 thread
1985 .update(cx, |thread, cx| thread.send_raw("Hello from Zed!", cx))
1986 .await
1987 .unwrap();
1988
1989 let output = thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
1990 assert_eq!(
1991 output,
1992 indoc! {r#"
1993 ## User
1994
1995 Hello from Zed!
1996
1997 ## Assistant
1998
1999 <thinking>
2000 Thinking hard!
2001 </thinking>
2002
2003 "#}
2004 );
2005 }
2006
2007 #[gpui::test]
2008 async fn test_edits_concurrently_to_user(cx: &mut TestAppContext) {
2009 init_test(cx);
2010
2011 let fs = FakeFs::new(cx.executor());
2012 fs.insert_tree(path!("/tmp"), json!({"foo": "one\ntwo\nthree\n"}))
2013 .await;
2014 let project = Project::test(fs.clone(), [], cx).await;
2015 let (read_file_tx, read_file_rx) = oneshot::channel::<()>();
2016 let read_file_tx = Rc::new(RefCell::new(Some(read_file_tx)));
2017 let connection = Rc::new(FakeAgentConnection::new().on_user_message(
2018 move |_, thread, mut cx| {
2019 let read_file_tx = read_file_tx.clone();
2020 async move {
2021 let content = thread
2022 .update(&mut cx, |thread, cx| {
2023 thread.read_text_file(path!("/tmp/foo").into(), None, None, false, cx)
2024 })
2025 .unwrap()
2026 .await
2027 .unwrap();
2028 assert_eq!(content, "one\ntwo\nthree\n");
2029 read_file_tx.take().unwrap().send(()).unwrap();
2030 thread
2031 .update(&mut cx, |thread, cx| {
2032 thread.write_text_file(
2033 path!("/tmp/foo").into(),
2034 "one\ntwo\nthree\nfour\nfive\n".to_string(),
2035 cx,
2036 )
2037 })
2038 .unwrap()
2039 .await
2040 .unwrap();
2041 Ok(acp::PromptResponse {
2042 stop_reason: acp::StopReason::EndTurn,
2043 })
2044 }
2045 .boxed_local()
2046 },
2047 ));
2048
2049 let (worktree, pathbuf) = project
2050 .update(cx, |project, cx| {
2051 project.find_or_create_worktree(path!("/tmp/foo"), true, cx)
2052 })
2053 .await
2054 .unwrap();
2055 let buffer = project
2056 .update(cx, |project, cx| {
2057 project.open_buffer((worktree.read(cx).id(), pathbuf), cx)
2058 })
2059 .await
2060 .unwrap();
2061
2062 let thread = cx
2063 .update(|cx| connection.new_thread(project, Path::new(path!("/tmp")), cx))
2064 .await
2065 .unwrap();
2066
2067 let request = thread.update(cx, |thread, cx| {
2068 thread.send_raw("Extend the count in /tmp/foo", cx)
2069 });
2070 read_file_rx.await.ok();
2071 buffer.update(cx, |buffer, cx| {
2072 buffer.edit([(0..0, "zero\n".to_string())], None, cx);
2073 });
2074 cx.run_until_parked();
2075 assert_eq!(
2076 buffer.read_with(cx, |buffer, _| buffer.text()),
2077 "zero\none\ntwo\nthree\nfour\nfive\n"
2078 );
2079 assert_eq!(
2080 String::from_utf8(fs.read_file_sync(path!("/tmp/foo")).unwrap()).unwrap(),
2081 "zero\none\ntwo\nthree\nfour\nfive\n"
2082 );
2083 request.await.unwrap();
2084 }
2085
2086 #[gpui::test]
2087 async fn test_succeeding_canceled_toolcall(cx: &mut TestAppContext) {
2088 init_test(cx);
2089
2090 let fs = FakeFs::new(cx.executor());
2091 let project = Project::test(fs, [], cx).await;
2092 let id = acp::ToolCallId("test".into());
2093
2094 let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2095 let id = id.clone();
2096 move |_, thread, mut cx| {
2097 let id = id.clone();
2098 async move {
2099 thread
2100 .update(&mut cx, |thread, cx| {
2101 thread.handle_session_update(
2102 acp::SessionUpdate::ToolCall(acp::ToolCall {
2103 id: id.clone(),
2104 title: "Label".into(),
2105 kind: acp::ToolKind::Fetch,
2106 status: acp::ToolCallStatus::InProgress,
2107 content: vec![],
2108 locations: vec![],
2109 raw_input: None,
2110 raw_output: None,
2111 }),
2112 cx,
2113 )
2114 })
2115 .unwrap()
2116 .unwrap();
2117 Ok(acp::PromptResponse {
2118 stop_reason: acp::StopReason::EndTurn,
2119 })
2120 }
2121 .boxed_local()
2122 }
2123 }));
2124
2125 let thread = cx
2126 .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2127 .await
2128 .unwrap();
2129
2130 let request = thread.update(cx, |thread, cx| {
2131 thread.send_raw("Fetch https://example.com", cx)
2132 });
2133
2134 run_until_first_tool_call(&thread, cx).await;
2135
2136 thread.read_with(cx, |thread, _| {
2137 assert!(matches!(
2138 thread.entries[1],
2139 AgentThreadEntry::ToolCall(ToolCall {
2140 status: ToolCallStatus::InProgress,
2141 ..
2142 })
2143 ));
2144 });
2145
2146 thread.update(cx, |thread, cx| thread.cancel(cx)).await;
2147
2148 thread.read_with(cx, |thread, _| {
2149 assert!(matches!(
2150 &thread.entries[1],
2151 AgentThreadEntry::ToolCall(ToolCall {
2152 status: ToolCallStatus::Canceled,
2153 ..
2154 })
2155 ));
2156 });
2157
2158 thread
2159 .update(cx, |thread, cx| {
2160 thread.handle_session_update(
2161 acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate {
2162 id,
2163 fields: acp::ToolCallUpdateFields {
2164 status: Some(acp::ToolCallStatus::Completed),
2165 ..Default::default()
2166 },
2167 }),
2168 cx,
2169 )
2170 })
2171 .unwrap();
2172
2173 request.await.unwrap();
2174
2175 thread.read_with(cx, |thread, _| {
2176 assert!(matches!(
2177 thread.entries[1],
2178 AgentThreadEntry::ToolCall(ToolCall {
2179 status: ToolCallStatus::Completed,
2180 ..
2181 })
2182 ));
2183 });
2184 }
2185
2186 #[gpui::test]
2187 async fn test_no_pending_edits_if_tool_calls_are_completed(cx: &mut TestAppContext) {
2188 init_test(cx);
2189 let fs = FakeFs::new(cx.background_executor.clone());
2190 fs.insert_tree(path!("/test"), json!({})).await;
2191 let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
2192
2193 let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2194 move |_, thread, mut cx| {
2195 async move {
2196 thread
2197 .update(&mut cx, |thread, cx| {
2198 thread.handle_session_update(
2199 acp::SessionUpdate::ToolCall(acp::ToolCall {
2200 id: acp::ToolCallId("test".into()),
2201 title: "Label".into(),
2202 kind: acp::ToolKind::Edit,
2203 status: acp::ToolCallStatus::Completed,
2204 content: vec![acp::ToolCallContent::Diff {
2205 diff: acp::Diff {
2206 path: "/test/test.txt".into(),
2207 old_text: None,
2208 new_text: "foo".into(),
2209 },
2210 }],
2211 locations: vec![],
2212 raw_input: None,
2213 raw_output: None,
2214 }),
2215 cx,
2216 )
2217 })
2218 .unwrap()
2219 .unwrap();
2220 Ok(acp::PromptResponse {
2221 stop_reason: acp::StopReason::EndTurn,
2222 })
2223 }
2224 .boxed_local()
2225 }
2226 }));
2227
2228 let thread = cx
2229 .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2230 .await
2231 .unwrap();
2232
2233 cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["Hi".into()], cx)))
2234 .await
2235 .unwrap();
2236
2237 assert!(cx.read(|cx| !thread.read(cx).has_pending_edit_tool_calls()));
2238 }
2239
2240 #[gpui::test(iterations = 10)]
2241 async fn test_checkpoints(cx: &mut TestAppContext) {
2242 init_test(cx);
2243 let fs = FakeFs::new(cx.background_executor.clone());
2244 fs.insert_tree(
2245 path!("/test"),
2246 json!({
2247 ".git": {}
2248 }),
2249 )
2250 .await;
2251 let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
2252
2253 let simulate_changes = Arc::new(AtomicBool::new(true));
2254 let next_filename = Arc::new(AtomicUsize::new(0));
2255 let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2256 let simulate_changes = simulate_changes.clone();
2257 let next_filename = next_filename.clone();
2258 let fs = fs.clone();
2259 move |request, thread, mut cx| {
2260 let fs = fs.clone();
2261 let simulate_changes = simulate_changes.clone();
2262 let next_filename = next_filename.clone();
2263 async move {
2264 if simulate_changes.load(SeqCst) {
2265 let filename = format!("/test/file-{}", next_filename.fetch_add(1, SeqCst));
2266 fs.write(Path::new(&filename), b"").await?;
2267 }
2268
2269 let acp::ContentBlock::Text(content) = &request.prompt[0] else {
2270 panic!("expected text content block");
2271 };
2272 thread.update(&mut cx, |thread, cx| {
2273 thread
2274 .handle_session_update(
2275 acp::SessionUpdate::AgentMessageChunk {
2276 content: content.text.to_uppercase().into(),
2277 },
2278 cx,
2279 )
2280 .unwrap();
2281 })?;
2282 Ok(acp::PromptResponse {
2283 stop_reason: acp::StopReason::EndTurn,
2284 })
2285 }
2286 .boxed_local()
2287 }
2288 }));
2289 let thread = cx
2290 .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2291 .await
2292 .unwrap();
2293
2294 cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["Lorem".into()], cx)))
2295 .await
2296 .unwrap();
2297 thread.read_with(cx, |thread, cx| {
2298 assert_eq!(
2299 thread.to_markdown(cx),
2300 indoc! {"
2301 ## User (checkpoint)
2302
2303 Lorem
2304
2305 ## Assistant
2306
2307 LOREM
2308
2309 "}
2310 );
2311 });
2312 assert_eq!(fs.files(), vec![Path::new(path!("/test/file-0"))]);
2313
2314 cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["ipsum".into()], cx)))
2315 .await
2316 .unwrap();
2317 thread.read_with(cx, |thread, cx| {
2318 assert_eq!(
2319 thread.to_markdown(cx),
2320 indoc! {"
2321 ## User (checkpoint)
2322
2323 Lorem
2324
2325 ## Assistant
2326
2327 LOREM
2328
2329 ## User (checkpoint)
2330
2331 ipsum
2332
2333 ## Assistant
2334
2335 IPSUM
2336
2337 "}
2338 );
2339 });
2340 assert_eq!(
2341 fs.files(),
2342 vec![
2343 Path::new(path!("/test/file-0")),
2344 Path::new(path!("/test/file-1"))
2345 ]
2346 );
2347
2348 // Checkpoint isn't stored when there are no changes.
2349 simulate_changes.store(false, SeqCst);
2350 cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["dolor".into()], cx)))
2351 .await
2352 .unwrap();
2353 thread.read_with(cx, |thread, cx| {
2354 assert_eq!(
2355 thread.to_markdown(cx),
2356 indoc! {"
2357 ## User (checkpoint)
2358
2359 Lorem
2360
2361 ## Assistant
2362
2363 LOREM
2364
2365 ## User (checkpoint)
2366
2367 ipsum
2368
2369 ## Assistant
2370
2371 IPSUM
2372
2373 ## User
2374
2375 dolor
2376
2377 ## Assistant
2378
2379 DOLOR
2380
2381 "}
2382 );
2383 });
2384 assert_eq!(
2385 fs.files(),
2386 vec![
2387 Path::new(path!("/test/file-0")),
2388 Path::new(path!("/test/file-1"))
2389 ]
2390 );
2391
2392 // Rewinding the conversation truncates the history and restores the checkpoint.
2393 thread
2394 .update(cx, |thread, cx| {
2395 let AgentThreadEntry::UserMessage(message) = &thread.entries[2] else {
2396 panic!("unexpected entries {:?}", thread.entries)
2397 };
2398 thread.rewind(message.id.clone().unwrap(), cx)
2399 })
2400 .await
2401 .unwrap();
2402 thread.read_with(cx, |thread, cx| {
2403 assert_eq!(
2404 thread.to_markdown(cx),
2405 indoc! {"
2406 ## User (checkpoint)
2407
2408 Lorem
2409
2410 ## Assistant
2411
2412 LOREM
2413
2414 "}
2415 );
2416 });
2417 assert_eq!(fs.files(), vec![Path::new(path!("/test/file-0"))]);
2418 }
2419
2420 #[gpui::test]
2421 async fn test_refusal(cx: &mut TestAppContext) {
2422 init_test(cx);
2423 let fs = FakeFs::new(cx.background_executor.clone());
2424 fs.insert_tree(path!("/"), json!({})).await;
2425 let project = Project::test(fs.clone(), [path!("/").as_ref()], cx).await;
2426
2427 let refuse_next = Arc::new(AtomicBool::new(false));
2428 let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2429 let refuse_next = refuse_next.clone();
2430 move |request, thread, mut cx| {
2431 let refuse_next = refuse_next.clone();
2432 async move {
2433 if refuse_next.load(SeqCst) {
2434 return Ok(acp::PromptResponse {
2435 stop_reason: acp::StopReason::Refusal,
2436 });
2437 }
2438
2439 let acp::ContentBlock::Text(content) = &request.prompt[0] else {
2440 panic!("expected text content block");
2441 };
2442 thread.update(&mut cx, |thread, cx| {
2443 thread
2444 .handle_session_update(
2445 acp::SessionUpdate::AgentMessageChunk {
2446 content: content.text.to_uppercase().into(),
2447 },
2448 cx,
2449 )
2450 .unwrap();
2451 })?;
2452 Ok(acp::PromptResponse {
2453 stop_reason: acp::StopReason::EndTurn,
2454 })
2455 }
2456 .boxed_local()
2457 }
2458 }));
2459 let thread = cx
2460 .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2461 .await
2462 .unwrap();
2463
2464 cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx)))
2465 .await
2466 .unwrap();
2467 thread.read_with(cx, |thread, cx| {
2468 assert_eq!(
2469 thread.to_markdown(cx),
2470 indoc! {"
2471 ## User
2472
2473 hello
2474
2475 ## Assistant
2476
2477 HELLO
2478
2479 "}
2480 );
2481 });
2482
2483 // Simulate refusing the second message, ensuring the conversation gets
2484 // truncated to before sending it.
2485 refuse_next.store(true, SeqCst);
2486 cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["world".into()], cx)))
2487 .await
2488 .unwrap();
2489 thread.read_with(cx, |thread, cx| {
2490 assert_eq!(
2491 thread.to_markdown(cx),
2492 indoc! {"
2493 ## User
2494
2495 hello
2496
2497 ## Assistant
2498
2499 HELLO
2500
2501 "}
2502 );
2503 });
2504 }
2505
2506 async fn run_until_first_tool_call(
2507 thread: &Entity<AcpThread>,
2508 cx: &mut TestAppContext,
2509 ) -> usize {
2510 let (mut tx, mut rx) = mpsc::channel::<usize>(1);
2511
2512 let subscription = cx.update(|cx| {
2513 cx.subscribe(thread, move |thread, _, cx| {
2514 for (ix, entry) in thread.read(cx).entries.iter().enumerate() {
2515 if matches!(entry, AgentThreadEntry::ToolCall(_)) {
2516 return tx.try_send(ix).unwrap();
2517 }
2518 }
2519 })
2520 });
2521
2522 select! {
2523 _ = futures::FutureExt::fuse(smol::Timer::after(Duration::from_secs(10))) => {
2524 panic!("Timeout waiting for tool call")
2525 }
2526 ix = rx.next().fuse() => {
2527 drop(subscription);
2528 ix.unwrap()
2529 }
2530 }
2531 }
2532
2533 #[derive(Clone, Default)]
2534 struct FakeAgentConnection {
2535 auth_methods: Vec<acp::AuthMethod>,
2536 sessions: Arc<parking_lot::Mutex<HashMap<acp::SessionId, WeakEntity<AcpThread>>>>,
2537 on_user_message: Option<
2538 Rc<
2539 dyn Fn(
2540 acp::PromptRequest,
2541 WeakEntity<AcpThread>,
2542 AsyncApp,
2543 ) -> LocalBoxFuture<'static, Result<acp::PromptResponse>>
2544 + 'static,
2545 >,
2546 >,
2547 }
2548
2549 impl FakeAgentConnection {
2550 fn new() -> Self {
2551 Self {
2552 auth_methods: Vec::new(),
2553 on_user_message: None,
2554 sessions: Arc::default(),
2555 }
2556 }
2557
2558 #[expect(unused)]
2559 fn with_auth_methods(mut self, auth_methods: Vec<acp::AuthMethod>) -> Self {
2560 self.auth_methods = auth_methods;
2561 self
2562 }
2563
2564 fn on_user_message(
2565 mut self,
2566 handler: impl Fn(
2567 acp::PromptRequest,
2568 WeakEntity<AcpThread>,
2569 AsyncApp,
2570 ) -> LocalBoxFuture<'static, Result<acp::PromptResponse>>
2571 + 'static,
2572 ) -> Self {
2573 self.on_user_message.replace(Rc::new(handler));
2574 self
2575 }
2576 }
2577
2578 impl AgentConnection for FakeAgentConnection {
2579 fn auth_methods(&self) -> &[acp::AuthMethod] {
2580 &self.auth_methods
2581 }
2582
2583 fn new_thread(
2584 self: Rc<Self>,
2585 project: Entity<Project>,
2586 _cwd: &Path,
2587 cx: &mut App,
2588 ) -> Task<gpui::Result<Entity<AcpThread>>> {
2589 let session_id = acp::SessionId(
2590 rand::thread_rng()
2591 .sample_iter(&rand::distributions::Alphanumeric)
2592 .take(7)
2593 .map(char::from)
2594 .collect::<String>()
2595 .into(),
2596 );
2597 let action_log = cx.new(|_| ActionLog::new(project.clone()));
2598 let thread = cx.new(|_cx| {
2599 AcpThread::new(
2600 "Test",
2601 self.clone(),
2602 project,
2603 action_log,
2604 session_id.clone(),
2605 )
2606 });
2607 self.sessions.lock().insert(session_id, thread.downgrade());
2608 Task::ready(Ok(thread))
2609 }
2610
2611 fn authenticate(&self, method: acp::AuthMethodId, _cx: &mut App) -> Task<gpui::Result<()>> {
2612 if self.auth_methods().iter().any(|m| m.id == method) {
2613 Task::ready(Ok(()))
2614 } else {
2615 Task::ready(Err(anyhow!("Invalid Auth Method")))
2616 }
2617 }
2618
2619 fn prompt(
2620 &self,
2621 _id: Option<UserMessageId>,
2622 params: acp::PromptRequest,
2623 cx: &mut App,
2624 ) -> Task<gpui::Result<acp::PromptResponse>> {
2625 let sessions = self.sessions.lock();
2626 let thread = sessions.get(¶ms.session_id).unwrap();
2627 if let Some(handler) = &self.on_user_message {
2628 let handler = handler.clone();
2629 let thread = thread.clone();
2630 cx.spawn(async move |cx| handler(params, thread, cx.clone()).await)
2631 } else {
2632 Task::ready(Ok(acp::PromptResponse {
2633 stop_reason: acp::StopReason::EndTurn,
2634 }))
2635 }
2636 }
2637
2638 fn prompt_capabilities(&self) -> acp::PromptCapabilities {
2639 acp::PromptCapabilities {
2640 image: true,
2641 audio: true,
2642 embedded_context: true,
2643 }
2644 }
2645
2646 fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
2647 let sessions = self.sessions.lock();
2648 let thread = sessions.get(session_id).unwrap().clone();
2649
2650 cx.spawn(async move |cx| {
2651 thread
2652 .update(cx, |thread, cx| thread.cancel(cx))
2653 .unwrap()
2654 .await
2655 })
2656 .detach();
2657 }
2658
2659 fn truncate(
2660 &self,
2661 session_id: &acp::SessionId,
2662 _cx: &mut App,
2663 ) -> Option<Rc<dyn AgentSessionTruncate>> {
2664 Some(Rc::new(FakeAgentSessionEditor {
2665 _session_id: session_id.clone(),
2666 }))
2667 }
2668
2669 fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
2670 self
2671 }
2672 }
2673
2674 struct FakeAgentSessionEditor {
2675 _session_id: acp::SessionId,
2676 }
2677
2678 impl AgentSessionTruncate for FakeAgentSessionEditor {
2679 fn run(&self, _message_id: UserMessageId, _cx: &mut App) -> Task<Result<()>> {
2680 Task::ready(Ok(()))
2681 }
2682 }
2683}