1use std::fmt::Write as _;
2use std::io::Write;
3use std::ops::Range;
4use std::sync::Arc;
5
6use agent_rules::load_worktree_rules_file;
7use anyhow::{Context as _, Result, anyhow};
8use assistant_settings::AssistantSettings;
9use assistant_tool::{ActionLog, Tool, ToolWorkingSet};
10use chrono::{DateTime, Utc};
11use collections::{BTreeMap, HashMap};
12use fs::Fs;
13use futures::future::Shared;
14use futures::{FutureExt, StreamExt as _};
15use git::repository::DiffType;
16use gpui::{App, AppContext, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
17use language_model::{
18 ConfiguredModel, LanguageModel, LanguageModelCompletionEvent, LanguageModelRegistry,
19 LanguageModelRequest, LanguageModelRequestMessage, LanguageModelRequestTool,
20 LanguageModelToolResult, LanguageModelToolUseId, MaxMonthlySpendReachedError, MessageContent,
21 PaymentRequiredError, Role, StopReason, TokenUsage,
22};
23use project::git_store::{GitStore, GitStoreCheckpoint, RepositoryState};
24use project::{Project, Worktree};
25use prompt_store::{AssistantSystemPromptContext, PromptBuilder, WorktreeInfoForSystemPrompt};
26use schemars::JsonSchema;
27use serde::{Deserialize, Serialize};
28use settings::Settings;
29use util::{ResultExt as _, TryFutureExt as _, post_inc};
30use uuid::Uuid;
31
32use crate::context::{AssistantContext, ContextId, format_context_as_string};
33use crate::thread_store::{
34 SerializedMessage, SerializedMessageSegment, SerializedThread, SerializedToolResult,
35 SerializedToolUse,
36};
37use crate::tool_use::{PendingToolUse, ToolUse, ToolUseState, USING_TOOL_MARKER};
38
39#[derive(Debug, Clone, Copy)]
40pub enum RequestKind {
41 Chat,
42 /// Used when summarizing a thread.
43 Summarize,
44}
45
46#[derive(
47 Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,
48)]
49pub struct ThreadId(Arc<str>);
50
51impl ThreadId {
52 pub fn new() -> Self {
53 Self(Uuid::new_v4().to_string().into())
54 }
55}
56
57impl std::fmt::Display for ThreadId {
58 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59 write!(f, "{}", self.0)
60 }
61}
62
63impl From<&str> for ThreadId {
64 fn from(value: &str) -> Self {
65 Self(value.into())
66 }
67}
68
69#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
70pub struct MessageId(pub(crate) usize);
71
72impl MessageId {
73 fn post_inc(&mut self) -> Self {
74 Self(post_inc(&mut self.0))
75 }
76}
77
78/// A message in a [`Thread`].
79#[derive(Debug, Clone)]
80pub struct Message {
81 pub id: MessageId,
82 pub role: Role,
83 pub segments: Vec<MessageSegment>,
84 pub context: String,
85}
86
87impl Message {
88 /// Returns whether the message contains any meaningful text that should be displayed
89 /// The model sometimes runs tool without producing any text or just a marker ([`USING_TOOL_MARKER`])
90 pub fn should_display_content(&self) -> bool {
91 self.segments.iter().all(|segment| segment.should_display())
92 }
93
94 pub fn push_thinking(&mut self, text: &str) {
95 if let Some(MessageSegment::Thinking(segment)) = self.segments.last_mut() {
96 segment.push_str(text);
97 } else {
98 self.segments
99 .push(MessageSegment::Thinking(text.to_string()));
100 }
101 }
102
103 pub fn push_text(&mut self, text: &str) {
104 if let Some(MessageSegment::Text(segment)) = self.segments.last_mut() {
105 segment.push_str(text);
106 } else {
107 self.segments.push(MessageSegment::Text(text.to_string()));
108 }
109 }
110
111 pub fn to_string(&self) -> String {
112 let mut result = String::new();
113
114 if !self.context.is_empty() {
115 result.push_str(&self.context);
116 }
117
118 for segment in &self.segments {
119 match segment {
120 MessageSegment::Text(text) => result.push_str(text),
121 MessageSegment::Thinking(text) => {
122 result.push_str("<think>");
123 result.push_str(text);
124 result.push_str("</think>");
125 }
126 }
127 }
128
129 result
130 }
131}
132
133#[derive(Debug, Clone, PartialEq, Eq)]
134pub enum MessageSegment {
135 Text(String),
136 Thinking(String),
137}
138
139impl MessageSegment {
140 pub fn text_mut(&mut self) -> &mut String {
141 match self {
142 Self::Text(text) => text,
143 Self::Thinking(text) => text,
144 }
145 }
146
147 pub fn should_display(&self) -> bool {
148 // We add USING_TOOL_MARKER when making a request that includes tool uses
149 // without non-whitespace text around them, and this can cause the model
150 // to mimic the pattern, so we consider those segments not displayable.
151 match self {
152 Self::Text(text) => text.is_empty() || text.trim() == USING_TOOL_MARKER,
153 Self::Thinking(text) => text.is_empty() || text.trim() == USING_TOOL_MARKER,
154 }
155 }
156}
157
158#[derive(Debug, Clone, Serialize, Deserialize)]
159pub struct ProjectSnapshot {
160 pub worktree_snapshots: Vec<WorktreeSnapshot>,
161 pub unsaved_buffer_paths: Vec<String>,
162 pub timestamp: DateTime<Utc>,
163}
164
165#[derive(Debug, Clone, Serialize, Deserialize)]
166pub struct WorktreeSnapshot {
167 pub worktree_path: String,
168 pub git_state: Option<GitState>,
169}
170
171#[derive(Debug, Clone, Serialize, Deserialize)]
172pub struct GitState {
173 pub remote_url: Option<String>,
174 pub head_sha: Option<String>,
175 pub current_branch: Option<String>,
176 pub diff: Option<String>,
177}
178
179#[derive(Clone)]
180pub struct ThreadCheckpoint {
181 message_id: MessageId,
182 git_checkpoint: GitStoreCheckpoint,
183}
184
185#[derive(Copy, Clone, Debug)]
186pub enum ThreadFeedback {
187 Positive,
188 Negative,
189}
190
191pub enum LastRestoreCheckpoint {
192 Pending {
193 message_id: MessageId,
194 },
195 Error {
196 message_id: MessageId,
197 error: String,
198 },
199}
200
201impl LastRestoreCheckpoint {
202 pub fn message_id(&self) -> MessageId {
203 match self {
204 LastRestoreCheckpoint::Pending { message_id } => *message_id,
205 LastRestoreCheckpoint::Error { message_id, .. } => *message_id,
206 }
207 }
208}
209
210#[derive(Clone, Debug, Default, Serialize, Deserialize)]
211pub enum DetailedSummaryState {
212 #[default]
213 NotGenerated,
214 Generating {
215 message_id: MessageId,
216 },
217 Generated {
218 text: SharedString,
219 message_id: MessageId,
220 },
221}
222
223#[derive(Default)]
224pub struct TotalTokenUsage {
225 pub total: usize,
226 pub max: usize,
227 pub ratio: TokenUsageRatio,
228}
229
230#[derive(Default, PartialEq, Eq)]
231pub enum TokenUsageRatio {
232 #[default]
233 Normal,
234 Warning,
235 Exceeded,
236}
237
238/// A thread of conversation with the LLM.
239pub struct Thread {
240 id: ThreadId,
241 updated_at: DateTime<Utc>,
242 summary: Option<SharedString>,
243 pending_summary: Task<Option<()>>,
244 detailed_summary_state: DetailedSummaryState,
245 messages: Vec<Message>,
246 next_message_id: MessageId,
247 context: BTreeMap<ContextId, AssistantContext>,
248 context_by_message: HashMap<MessageId, Vec<ContextId>>,
249 system_prompt_context: Option<AssistantSystemPromptContext>,
250 checkpoints_by_message: HashMap<MessageId, ThreadCheckpoint>,
251 completion_count: usize,
252 pending_completions: Vec<PendingCompletion>,
253 project: Entity<Project>,
254 prompt_builder: Arc<PromptBuilder>,
255 tools: Arc<ToolWorkingSet>,
256 tool_use: ToolUseState,
257 action_log: Entity<ActionLog>,
258 last_restore_checkpoint: Option<LastRestoreCheckpoint>,
259 pending_checkpoint: Option<ThreadCheckpoint>,
260 initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
261 cumulative_token_usage: TokenUsage,
262 feedback: Option<ThreadFeedback>,
263}
264
265impl Thread {
266 pub fn new(
267 project: Entity<Project>,
268 tools: Arc<ToolWorkingSet>,
269 prompt_builder: Arc<PromptBuilder>,
270 cx: &mut Context<Self>,
271 ) -> Self {
272 Self {
273 id: ThreadId::new(),
274 updated_at: Utc::now(),
275 summary: None,
276 pending_summary: Task::ready(None),
277 detailed_summary_state: DetailedSummaryState::NotGenerated,
278 messages: Vec::new(),
279 next_message_id: MessageId(0),
280 context: BTreeMap::default(),
281 context_by_message: HashMap::default(),
282 system_prompt_context: None,
283 checkpoints_by_message: HashMap::default(),
284 completion_count: 0,
285 pending_completions: Vec::new(),
286 project: project.clone(),
287 prompt_builder,
288 tools: tools.clone(),
289 last_restore_checkpoint: None,
290 pending_checkpoint: None,
291 tool_use: ToolUseState::new(tools.clone()),
292 action_log: cx.new(|_| ActionLog::new(project.clone())),
293 initial_project_snapshot: {
294 let project_snapshot = Self::project_snapshot(project, cx);
295 cx.foreground_executor()
296 .spawn(async move { Some(project_snapshot.await) })
297 .shared()
298 },
299 cumulative_token_usage: TokenUsage::default(),
300 feedback: None,
301 }
302 }
303
304 pub fn deserialize(
305 id: ThreadId,
306 serialized: SerializedThread,
307 project: Entity<Project>,
308 tools: Arc<ToolWorkingSet>,
309 prompt_builder: Arc<PromptBuilder>,
310 cx: &mut Context<Self>,
311 ) -> Self {
312 let next_message_id = MessageId(
313 serialized
314 .messages
315 .last()
316 .map(|message| message.id.0 + 1)
317 .unwrap_or(0),
318 );
319 let tool_use =
320 ToolUseState::from_serialized_messages(tools.clone(), &serialized.messages, |_| true);
321
322 Self {
323 id,
324 updated_at: serialized.updated_at,
325 summary: Some(serialized.summary),
326 pending_summary: Task::ready(None),
327 detailed_summary_state: serialized.detailed_summary_state,
328 messages: serialized
329 .messages
330 .into_iter()
331 .map(|message| Message {
332 id: message.id,
333 role: message.role,
334 segments: message
335 .segments
336 .into_iter()
337 .map(|segment| match segment {
338 SerializedMessageSegment::Text { text } => MessageSegment::Text(text),
339 SerializedMessageSegment::Thinking { text } => {
340 MessageSegment::Thinking(text)
341 }
342 })
343 .collect(),
344 context: message.context,
345 })
346 .collect(),
347 next_message_id,
348 context: BTreeMap::default(),
349 context_by_message: HashMap::default(),
350 system_prompt_context: None,
351 checkpoints_by_message: HashMap::default(),
352 completion_count: 0,
353 pending_completions: Vec::new(),
354 last_restore_checkpoint: None,
355 pending_checkpoint: None,
356 project: project.clone(),
357 prompt_builder,
358 tools,
359 tool_use,
360 action_log: cx.new(|_| ActionLog::new(project)),
361 initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
362 cumulative_token_usage: serialized.cumulative_token_usage,
363 feedback: None,
364 }
365 }
366
367 pub fn id(&self) -> &ThreadId {
368 &self.id
369 }
370
371 pub fn is_empty(&self) -> bool {
372 self.messages.is_empty()
373 }
374
375 pub fn updated_at(&self) -> DateTime<Utc> {
376 self.updated_at
377 }
378
379 pub fn touch_updated_at(&mut self) {
380 self.updated_at = Utc::now();
381 }
382
383 pub fn summary(&self) -> Option<SharedString> {
384 self.summary.clone()
385 }
386
387 pub const DEFAULT_SUMMARY: SharedString = SharedString::new_static("New Thread");
388
389 pub fn summary_or_default(&self) -> SharedString {
390 self.summary.clone().unwrap_or(Self::DEFAULT_SUMMARY)
391 }
392
393 pub fn set_summary(&mut self, new_summary: impl Into<SharedString>, cx: &mut Context<Self>) {
394 let Some(current_summary) = &self.summary else {
395 // Don't allow setting summary until generated
396 return;
397 };
398
399 let mut new_summary = new_summary.into();
400
401 if new_summary.is_empty() {
402 new_summary = Self::DEFAULT_SUMMARY;
403 }
404
405 if current_summary != &new_summary {
406 self.summary = Some(new_summary);
407 cx.emit(ThreadEvent::SummaryChanged);
408 }
409 }
410
411 pub fn latest_detailed_summary_or_text(&self) -> SharedString {
412 self.latest_detailed_summary()
413 .unwrap_or_else(|| self.text().into())
414 }
415
416 fn latest_detailed_summary(&self) -> Option<SharedString> {
417 if let DetailedSummaryState::Generated { text, .. } = &self.detailed_summary_state {
418 Some(text.clone())
419 } else {
420 None
421 }
422 }
423
424 pub fn message(&self, id: MessageId) -> Option<&Message> {
425 self.messages.iter().find(|message| message.id == id)
426 }
427
428 pub fn messages(&self) -> impl Iterator<Item = &Message> {
429 self.messages.iter()
430 }
431
432 pub fn is_generating(&self) -> bool {
433 !self.pending_completions.is_empty() || !self.all_tools_finished()
434 }
435
436 pub fn tools(&self) -> &Arc<ToolWorkingSet> {
437 &self.tools
438 }
439
440 pub fn pending_tool(&self, id: &LanguageModelToolUseId) -> Option<&PendingToolUse> {
441 self.tool_use
442 .pending_tool_uses()
443 .into_iter()
444 .find(|tool_use| &tool_use.id == id)
445 }
446
447 pub fn tools_needing_confirmation(&self) -> impl Iterator<Item = &PendingToolUse> {
448 self.tool_use
449 .pending_tool_uses()
450 .into_iter()
451 .filter(|tool_use| tool_use.status.needs_confirmation())
452 }
453
454 pub fn has_pending_tool_uses(&self) -> bool {
455 !self.tool_use.pending_tool_uses().is_empty()
456 }
457
458 pub fn checkpoint_for_message(&self, id: MessageId) -> Option<ThreadCheckpoint> {
459 self.checkpoints_by_message.get(&id).cloned()
460 }
461
462 pub fn restore_checkpoint(
463 &mut self,
464 checkpoint: ThreadCheckpoint,
465 cx: &mut Context<Self>,
466 ) -> Task<Result<()>> {
467 self.last_restore_checkpoint = Some(LastRestoreCheckpoint::Pending {
468 message_id: checkpoint.message_id,
469 });
470 cx.emit(ThreadEvent::CheckpointChanged);
471 cx.notify();
472
473 let git_store = self.project().read(cx).git_store().clone();
474 let restore = git_store.update(cx, |git_store, cx| {
475 git_store.restore_checkpoint(checkpoint.git_checkpoint.clone(), cx)
476 });
477
478 cx.spawn(async move |this, cx| {
479 let result = restore.await;
480 this.update(cx, |this, cx| {
481 if let Err(err) = result.as_ref() {
482 this.last_restore_checkpoint = Some(LastRestoreCheckpoint::Error {
483 message_id: checkpoint.message_id,
484 error: err.to_string(),
485 });
486 } else {
487 this.truncate(checkpoint.message_id, cx);
488 this.last_restore_checkpoint = None;
489 }
490 this.pending_checkpoint = None;
491 cx.emit(ThreadEvent::CheckpointChanged);
492 cx.notify();
493 })?;
494 result
495 })
496 }
497
498 fn finalize_pending_checkpoint(&mut self, cx: &mut Context<Self>) {
499 let pending_checkpoint = if self.is_generating() {
500 return;
501 } else if let Some(checkpoint) = self.pending_checkpoint.take() {
502 checkpoint
503 } else {
504 return;
505 };
506
507 let git_store = self.project.read(cx).git_store().clone();
508 let final_checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
509 cx.spawn(async move |this, cx| match final_checkpoint.await {
510 Ok(final_checkpoint) => {
511 let equal = git_store
512 .update(cx, |store, cx| {
513 store.compare_checkpoints(
514 pending_checkpoint.git_checkpoint.clone(),
515 final_checkpoint.clone(),
516 cx,
517 )
518 })?
519 .await
520 .unwrap_or(false);
521
522 if equal {
523 git_store
524 .update(cx, |store, cx| {
525 store.delete_checkpoint(pending_checkpoint.git_checkpoint, cx)
526 })?
527 .detach();
528 } else {
529 this.update(cx, |this, cx| {
530 this.insert_checkpoint(pending_checkpoint, cx)
531 })?;
532 }
533
534 git_store
535 .update(cx, |store, cx| {
536 store.delete_checkpoint(final_checkpoint, cx)
537 })?
538 .detach();
539
540 Ok(())
541 }
542 Err(_) => this.update(cx, |this, cx| {
543 this.insert_checkpoint(pending_checkpoint, cx)
544 }),
545 })
546 .detach();
547 }
548
549 fn insert_checkpoint(&mut self, checkpoint: ThreadCheckpoint, cx: &mut Context<Self>) {
550 self.checkpoints_by_message
551 .insert(checkpoint.message_id, checkpoint);
552 cx.emit(ThreadEvent::CheckpointChanged);
553 cx.notify();
554 }
555
556 pub fn last_restore_checkpoint(&self) -> Option<&LastRestoreCheckpoint> {
557 self.last_restore_checkpoint.as_ref()
558 }
559
560 pub fn truncate(&mut self, message_id: MessageId, cx: &mut Context<Self>) {
561 let Some(message_ix) = self
562 .messages
563 .iter()
564 .rposition(|message| message.id == message_id)
565 else {
566 return;
567 };
568 for deleted_message in self.messages.drain(message_ix..) {
569 self.context_by_message.remove(&deleted_message.id);
570 self.checkpoints_by_message.remove(&deleted_message.id);
571 }
572 cx.notify();
573 }
574
575 pub fn context_for_message(&self, id: MessageId) -> impl Iterator<Item = &AssistantContext> {
576 self.context_by_message
577 .get(&id)
578 .into_iter()
579 .flat_map(|context| {
580 context
581 .iter()
582 .filter_map(|context_id| self.context.get(&context_id))
583 })
584 }
585
586 /// Returns whether all of the tool uses have finished running.
587 pub fn all_tools_finished(&self) -> bool {
588 // If the only pending tool uses left are the ones with errors, then
589 // that means that we've finished running all of the pending tools.
590 self.tool_use
591 .pending_tool_uses()
592 .iter()
593 .all(|tool_use| tool_use.status.is_error())
594 }
595
596 pub fn tool_uses_for_message(&self, id: MessageId, cx: &App) -> Vec<ToolUse> {
597 self.tool_use.tool_uses_for_message(id, cx)
598 }
599
600 pub fn tool_results_for_message(&self, id: MessageId) -> Vec<&LanguageModelToolResult> {
601 self.tool_use.tool_results_for_message(id)
602 }
603
604 pub fn tool_result(&self, id: &LanguageModelToolUseId) -> Option<&LanguageModelToolResult> {
605 self.tool_use.tool_result(id)
606 }
607
608 pub fn message_has_tool_results(&self, message_id: MessageId) -> bool {
609 self.tool_use.message_has_tool_results(message_id)
610 }
611
612 pub fn insert_user_message(
613 &mut self,
614 text: impl Into<String>,
615 context: Vec<AssistantContext>,
616 git_checkpoint: Option<GitStoreCheckpoint>,
617 cx: &mut Context<Self>,
618 ) -> MessageId {
619 let text = text.into();
620
621 let message_id = self.insert_message(Role::User, vec![MessageSegment::Text(text)], cx);
622
623 // Filter out contexts that have already been included in previous messages
624 let new_context: Vec<_> = context
625 .into_iter()
626 .filter(|ctx| !self.context.contains_key(&ctx.id()))
627 .collect();
628
629 if !new_context.is_empty() {
630 if let Some(context_string) = format_context_as_string(new_context.iter(), cx) {
631 if let Some(message) = self.messages.iter_mut().find(|m| m.id == message_id) {
632 message.context = context_string;
633 }
634 }
635
636 self.action_log.update(cx, |log, cx| {
637 // Track all buffers added as context
638 for ctx in &new_context {
639 match ctx {
640 AssistantContext::File(file_ctx) => {
641 log.buffer_added_as_context(file_ctx.context_buffer.buffer.clone(), cx);
642 }
643 AssistantContext::Directory(dir_ctx) => {
644 for context_buffer in &dir_ctx.context_buffers {
645 log.buffer_added_as_context(context_buffer.buffer.clone(), cx);
646 }
647 }
648 AssistantContext::Symbol(symbol_ctx) => {
649 log.buffer_added_as_context(
650 symbol_ctx.context_symbol.buffer.clone(),
651 cx,
652 );
653 }
654 AssistantContext::FetchedUrl(_) | AssistantContext::Thread(_) => {}
655 }
656 }
657 });
658 }
659
660 let context_ids = new_context
661 .iter()
662 .map(|context| context.id())
663 .collect::<Vec<_>>();
664 self.context.extend(
665 new_context
666 .into_iter()
667 .map(|context| (context.id(), context)),
668 );
669 self.context_by_message.insert(message_id, context_ids);
670
671 if let Some(git_checkpoint) = git_checkpoint {
672 self.pending_checkpoint = Some(ThreadCheckpoint {
673 message_id,
674 git_checkpoint,
675 });
676 }
677 message_id
678 }
679
680 pub fn insert_message(
681 &mut self,
682 role: Role,
683 segments: Vec<MessageSegment>,
684 cx: &mut Context<Self>,
685 ) -> MessageId {
686 let id = self.next_message_id.post_inc();
687 self.messages.push(Message {
688 id,
689 role,
690 segments,
691 context: String::new(),
692 });
693 self.touch_updated_at();
694 cx.emit(ThreadEvent::MessageAdded(id));
695 id
696 }
697
698 pub fn edit_message(
699 &mut self,
700 id: MessageId,
701 new_role: Role,
702 new_segments: Vec<MessageSegment>,
703 cx: &mut Context<Self>,
704 ) -> bool {
705 let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
706 return false;
707 };
708 message.role = new_role;
709 message.segments = new_segments;
710 self.touch_updated_at();
711 cx.emit(ThreadEvent::MessageEdited(id));
712 true
713 }
714
715 pub fn delete_message(&mut self, id: MessageId, cx: &mut Context<Self>) -> bool {
716 let Some(index) = self.messages.iter().position(|message| message.id == id) else {
717 return false;
718 };
719 self.messages.remove(index);
720 self.context_by_message.remove(&id);
721 self.touch_updated_at();
722 cx.emit(ThreadEvent::MessageDeleted(id));
723 true
724 }
725
726 /// Returns the representation of this [`Thread`] in a textual form.
727 ///
728 /// This is the representation we use when attaching a thread as context to another thread.
729 pub fn text(&self) -> String {
730 let mut text = String::new();
731
732 for message in &self.messages {
733 text.push_str(match message.role {
734 language_model::Role::User => "User:",
735 language_model::Role::Assistant => "Assistant:",
736 language_model::Role::System => "System:",
737 });
738 text.push('\n');
739
740 for segment in &message.segments {
741 match segment {
742 MessageSegment::Text(content) => text.push_str(content),
743 MessageSegment::Thinking(content) => {
744 text.push_str(&format!("<think>{}</think>", content))
745 }
746 }
747 }
748 text.push('\n');
749 }
750
751 text
752 }
753
754 /// Serializes this thread into a format for storage or telemetry.
755 pub fn serialize(&self, cx: &mut Context<Self>) -> Task<Result<SerializedThread>> {
756 let initial_project_snapshot = self.initial_project_snapshot.clone();
757 cx.spawn(async move |this, cx| {
758 let initial_project_snapshot = initial_project_snapshot.await;
759 this.read_with(cx, |this, cx| SerializedThread {
760 version: SerializedThread::VERSION.to_string(),
761 summary: this.summary_or_default(),
762 updated_at: this.updated_at(),
763 messages: this
764 .messages()
765 .map(|message| SerializedMessage {
766 id: message.id,
767 role: message.role,
768 segments: message
769 .segments
770 .iter()
771 .map(|segment| match segment {
772 MessageSegment::Text(text) => {
773 SerializedMessageSegment::Text { text: text.clone() }
774 }
775 MessageSegment::Thinking(text) => {
776 SerializedMessageSegment::Thinking { text: text.clone() }
777 }
778 })
779 .collect(),
780 tool_uses: this
781 .tool_uses_for_message(message.id, cx)
782 .into_iter()
783 .map(|tool_use| SerializedToolUse {
784 id: tool_use.id,
785 name: tool_use.name,
786 input: tool_use.input,
787 })
788 .collect(),
789 tool_results: this
790 .tool_results_for_message(message.id)
791 .into_iter()
792 .map(|tool_result| SerializedToolResult {
793 tool_use_id: tool_result.tool_use_id.clone(),
794 is_error: tool_result.is_error,
795 content: tool_result.content.clone(),
796 })
797 .collect(),
798 context: message.context.clone(),
799 })
800 .collect(),
801 initial_project_snapshot,
802 cumulative_token_usage: this.cumulative_token_usage.clone(),
803 detailed_summary_state: this.detailed_summary_state.clone(),
804 })
805 })
806 }
807
808 pub fn set_system_prompt_context(&mut self, context: AssistantSystemPromptContext) {
809 self.system_prompt_context = Some(context);
810 }
811
812 pub fn system_prompt_context(&self) -> &Option<AssistantSystemPromptContext> {
813 &self.system_prompt_context
814 }
815
816 pub fn load_system_prompt_context(
817 &self,
818 cx: &App,
819 ) -> Task<(AssistantSystemPromptContext, Option<ThreadError>)> {
820 let project = self.project.read(cx);
821 let tasks = project
822 .visible_worktrees(cx)
823 .map(|worktree| {
824 Self::load_worktree_info_for_system_prompt(
825 project.fs().clone(),
826 worktree.read(cx),
827 cx,
828 )
829 })
830 .collect::<Vec<_>>();
831
832 cx.spawn(async |_cx| {
833 let results = futures::future::join_all(tasks).await;
834 let mut first_err = None;
835 let worktrees = results
836 .into_iter()
837 .map(|(worktree, err)| {
838 if first_err.is_none() && err.is_some() {
839 first_err = err;
840 }
841 worktree
842 })
843 .collect::<Vec<_>>();
844 (AssistantSystemPromptContext::new(worktrees), first_err)
845 })
846 }
847
848 fn load_worktree_info_for_system_prompt(
849 fs: Arc<dyn Fs>,
850 worktree: &Worktree,
851 cx: &App,
852 ) -> Task<(WorktreeInfoForSystemPrompt, Option<ThreadError>)> {
853 let root_name = worktree.root_name().into();
854 let abs_path = worktree.abs_path();
855
856 let rules_task = load_worktree_rules_file(fs, worktree, cx);
857 let Some(rules_task) = rules_task else {
858 return Task::ready((
859 WorktreeInfoForSystemPrompt {
860 root_name,
861 abs_path,
862 rules_file: None,
863 },
864 None,
865 ));
866 };
867
868 cx.spawn(async move |_| {
869 let (rules_file, rules_file_error) = match rules_task.await {
870 Ok(rules_file) => (Some(rules_file), None),
871 Err(err) => (
872 None,
873 Some(ThreadError::Message {
874 header: "Error loading rules file".into(),
875 message: format!("{err}").into(),
876 }),
877 ),
878 };
879 let worktree_info = WorktreeInfoForSystemPrompt {
880 root_name,
881 abs_path,
882 rules_file,
883 };
884 (worktree_info, rules_file_error)
885 })
886 }
887
888 pub fn send_to_model(
889 &mut self,
890 model: Arc<dyn LanguageModel>,
891 request_kind: RequestKind,
892 cx: &mut Context<Self>,
893 ) {
894 let mut request = self.to_completion_request(request_kind, cx);
895 if model.supports_tools() {
896 request.tools = {
897 let mut tools = Vec::new();
898 tools.extend(self.tools().enabled_tools(cx).into_iter().map(|tool| {
899 LanguageModelRequestTool {
900 name: tool.name(),
901 description: tool.description(),
902 input_schema: tool.input_schema(model.tool_input_format()),
903 }
904 }));
905
906 tools
907 };
908 }
909
910 self.stream_completion(request, model, cx);
911 }
912
913 pub fn used_tools_since_last_user_message(&self) -> bool {
914 for message in self.messages.iter().rev() {
915 if self.tool_use.message_has_tool_results(message.id) {
916 return true;
917 } else if message.role == Role::User {
918 return false;
919 }
920 }
921
922 false
923 }
924
925 pub fn to_completion_request(
926 &self,
927 request_kind: RequestKind,
928 cx: &App,
929 ) -> LanguageModelRequest {
930 let mut request = LanguageModelRequest {
931 messages: vec![],
932 tools: Vec::new(),
933 stop: Vec::new(),
934 temperature: None,
935 };
936
937 if let Some(system_prompt_context) = self.system_prompt_context.as_ref() {
938 if let Some(system_prompt) = self
939 .prompt_builder
940 .generate_assistant_system_prompt(system_prompt_context)
941 .context("failed to generate assistant system prompt")
942 .log_err()
943 {
944 request.messages.push(LanguageModelRequestMessage {
945 role: Role::System,
946 content: vec![MessageContent::Text(system_prompt)],
947 cache: true,
948 });
949 }
950 } else {
951 log::error!("system_prompt_context not set.")
952 }
953
954 for message in &self.messages {
955 let mut request_message = LanguageModelRequestMessage {
956 role: message.role,
957 content: Vec::new(),
958 cache: false,
959 };
960
961 match request_kind {
962 RequestKind::Chat => {
963 self.tool_use
964 .attach_tool_results(message.id, &mut request_message);
965 }
966 RequestKind::Summarize => {
967 // We don't care about tool use during summarization.
968 if self.tool_use.message_has_tool_results(message.id) {
969 continue;
970 }
971 }
972 }
973
974 if !message.segments.is_empty() {
975 request_message
976 .content
977 .push(MessageContent::Text(message.to_string()));
978 }
979
980 match request_kind {
981 RequestKind::Chat => {
982 self.tool_use
983 .attach_tool_uses(message.id, &mut request_message);
984 }
985 RequestKind::Summarize => {
986 // We don't care about tool use during summarization.
987 }
988 };
989
990 request.messages.push(request_message);
991 }
992
993 // https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
994 if let Some(last) = request.messages.last_mut() {
995 last.cache = true;
996 }
997
998 self.attached_tracked_files_state(&mut request.messages, cx);
999
1000 // Add reminder to the last user message about code blocks
1001 if let Some(last_user_message) = request
1002 .messages
1003 .iter_mut()
1004 .rev()
1005 .find(|msg| msg.role == Role::User)
1006 {
1007 last_user_message
1008 .content
1009 .push(MessageContent::Text(system_prompt_reminder(
1010 &self.prompt_builder,
1011 )));
1012 }
1013
1014 request
1015 }
1016
1017 fn attached_tracked_files_state(
1018 &self,
1019 messages: &mut Vec<LanguageModelRequestMessage>,
1020 cx: &App,
1021 ) {
1022 const STALE_FILES_HEADER: &str = "These files changed since last read:";
1023
1024 let mut stale_message = String::new();
1025
1026 let action_log = self.action_log.read(cx);
1027
1028 for stale_file in action_log.stale_buffers(cx) {
1029 let Some(file) = stale_file.read(cx).file() else {
1030 continue;
1031 };
1032
1033 if stale_message.is_empty() {
1034 write!(&mut stale_message, "{}\n", STALE_FILES_HEADER).ok();
1035 }
1036
1037 writeln!(&mut stale_message, "- {}", file.path().display()).ok();
1038 }
1039
1040 let mut content = Vec::with_capacity(2);
1041
1042 if !stale_message.is_empty() {
1043 content.push(stale_message.into());
1044 }
1045
1046 if action_log.has_edited_files_since_project_diagnostics_check() {
1047 content.push(
1048 "\n\nWhen you're done making changes, make sure to check project diagnostics \
1049 and fix all errors AND warnings you introduced! \
1050 DO NOT mention you're going to do this until you're done."
1051 .into(),
1052 );
1053 }
1054
1055 if !content.is_empty() {
1056 let context_message = LanguageModelRequestMessage {
1057 role: Role::User,
1058 content,
1059 cache: false,
1060 };
1061
1062 messages.push(context_message);
1063 }
1064 }
1065
1066 pub fn stream_completion(
1067 &mut self,
1068 request: LanguageModelRequest,
1069 model: Arc<dyn LanguageModel>,
1070 cx: &mut Context<Self>,
1071 ) {
1072 let pending_completion_id = post_inc(&mut self.completion_count);
1073
1074 let task = cx.spawn(async move |thread, cx| {
1075 let stream = model.stream_completion(request, &cx);
1076 let initial_token_usage =
1077 thread.read_with(cx, |thread, _cx| thread.cumulative_token_usage.clone());
1078 let stream_completion = async {
1079 let mut events = stream.await?;
1080 let mut stop_reason = StopReason::EndTurn;
1081 let mut current_token_usage = TokenUsage::default();
1082
1083 while let Some(event) = events.next().await {
1084 let event = event?;
1085
1086 thread.update(cx, |thread, cx| {
1087 match event {
1088 LanguageModelCompletionEvent::StartMessage { .. } => {
1089 thread.insert_message(
1090 Role::Assistant,
1091 vec![MessageSegment::Text(String::new())],
1092 cx,
1093 );
1094 }
1095 LanguageModelCompletionEvent::Stop(reason) => {
1096 stop_reason = reason;
1097 }
1098 LanguageModelCompletionEvent::UsageUpdate(token_usage) => {
1099 thread.cumulative_token_usage =
1100 thread.cumulative_token_usage.clone() + token_usage.clone()
1101 - current_token_usage.clone();
1102 current_token_usage = token_usage;
1103 }
1104 LanguageModelCompletionEvent::Text(chunk) => {
1105 if let Some(last_message) = thread.messages.last_mut() {
1106 if last_message.role == Role::Assistant {
1107 last_message.push_text(&chunk);
1108 cx.emit(ThreadEvent::StreamedAssistantText(
1109 last_message.id,
1110 chunk,
1111 ));
1112 } else {
1113 // If we won't have an Assistant message yet, assume this chunk marks the beginning
1114 // of a new Assistant response.
1115 //
1116 // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1117 // will result in duplicating the text of the chunk in the rendered Markdown.
1118 thread.insert_message(
1119 Role::Assistant,
1120 vec![MessageSegment::Text(chunk.to_string())],
1121 cx,
1122 );
1123 };
1124 }
1125 }
1126 LanguageModelCompletionEvent::Thinking(chunk) => {
1127 if let Some(last_message) = thread.messages.last_mut() {
1128 if last_message.role == Role::Assistant {
1129 last_message.push_thinking(&chunk);
1130 cx.emit(ThreadEvent::StreamedAssistantThinking(
1131 last_message.id,
1132 chunk,
1133 ));
1134 } else {
1135 // If we won't have an Assistant message yet, assume this chunk marks the beginning
1136 // of a new Assistant response.
1137 //
1138 // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1139 // will result in duplicating the text of the chunk in the rendered Markdown.
1140 thread.insert_message(
1141 Role::Assistant,
1142 vec![MessageSegment::Thinking(chunk.to_string())],
1143 cx,
1144 );
1145 };
1146 }
1147 }
1148 LanguageModelCompletionEvent::ToolUse(tool_use) => {
1149 let last_assistant_message_id = thread
1150 .messages
1151 .iter_mut()
1152 .rfind(|message| message.role == Role::Assistant)
1153 .map(|message| message.id)
1154 .unwrap_or_else(|| {
1155 thread.insert_message(Role::Assistant, vec![], cx)
1156 });
1157
1158 thread.tool_use.request_tool_use(
1159 last_assistant_message_id,
1160 tool_use,
1161 cx,
1162 );
1163 }
1164 }
1165
1166 thread.touch_updated_at();
1167 cx.emit(ThreadEvent::StreamedCompletion);
1168 cx.notify();
1169 })?;
1170
1171 smol::future::yield_now().await;
1172 }
1173
1174 thread.update(cx, |thread, cx| {
1175 thread
1176 .pending_completions
1177 .retain(|completion| completion.id != pending_completion_id);
1178
1179 if thread.summary.is_none() && thread.messages.len() >= 2 {
1180 thread.summarize(cx);
1181 }
1182 })?;
1183
1184 anyhow::Ok(stop_reason)
1185 };
1186
1187 let result = stream_completion.await;
1188
1189 thread
1190 .update(cx, |thread, cx| {
1191 thread.finalize_pending_checkpoint(cx);
1192 match result.as_ref() {
1193 Ok(stop_reason) => match stop_reason {
1194 StopReason::ToolUse => {
1195 cx.emit(ThreadEvent::UsePendingTools);
1196 }
1197 StopReason::EndTurn => {}
1198 StopReason::MaxTokens => {}
1199 },
1200 Err(error) => {
1201 if error.is::<PaymentRequiredError>() {
1202 cx.emit(ThreadEvent::ShowError(ThreadError::PaymentRequired));
1203 } else if error.is::<MaxMonthlySpendReachedError>() {
1204 cx.emit(ThreadEvent::ShowError(
1205 ThreadError::MaxMonthlySpendReached,
1206 ));
1207 } else {
1208 let error_message = error
1209 .chain()
1210 .map(|err| err.to_string())
1211 .collect::<Vec<_>>()
1212 .join("\n");
1213 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1214 header: "Error interacting with language model".into(),
1215 message: SharedString::from(error_message.clone()),
1216 }));
1217 }
1218
1219 thread.cancel_last_completion(cx);
1220 }
1221 }
1222 cx.emit(ThreadEvent::DoneStreaming);
1223
1224 if let Ok(initial_usage) = initial_token_usage {
1225 let usage = thread.cumulative_token_usage.clone() - initial_usage;
1226
1227 telemetry::event!(
1228 "Assistant Thread Completion",
1229 thread_id = thread.id().to_string(),
1230 model = model.telemetry_id(),
1231 model_provider = model.provider_id().to_string(),
1232 input_tokens = usage.input_tokens,
1233 output_tokens = usage.output_tokens,
1234 cache_creation_input_tokens = usage.cache_creation_input_tokens,
1235 cache_read_input_tokens = usage.cache_read_input_tokens,
1236 );
1237 }
1238 })
1239 .ok();
1240 });
1241
1242 self.pending_completions.push(PendingCompletion {
1243 id: pending_completion_id,
1244 _task: task,
1245 });
1246 }
1247
1248 pub fn summarize(&mut self, cx: &mut Context<Self>) {
1249 let Some(model) = LanguageModelRegistry::read_global(cx).thread_summary_model() else {
1250 return;
1251 };
1252
1253 if !model.provider.is_authenticated(cx) {
1254 return;
1255 }
1256
1257 let mut request = self.to_completion_request(RequestKind::Summarize, cx);
1258 request.messages.push(LanguageModelRequestMessage {
1259 role: Role::User,
1260 content: vec![
1261 "Generate a concise 3-7 word title for this conversation, omitting punctuation. \
1262 Go straight to the title, without any preamble and prefix like `Here's a concise suggestion:...` or `Title:`. \
1263 If the conversation is about a specific subject, include it in the title. \
1264 Be descriptive. DO NOT speak in the first person."
1265 .into(),
1266 ],
1267 cache: false,
1268 });
1269
1270 self.pending_summary = cx.spawn(async move |this, cx| {
1271 async move {
1272 let stream = model.model.stream_completion_text(request, &cx);
1273 let mut messages = stream.await?;
1274
1275 let mut new_summary = String::new();
1276 while let Some(message) = messages.stream.next().await {
1277 let text = message?;
1278 let mut lines = text.lines();
1279 new_summary.extend(lines.next());
1280
1281 // Stop if the LLM generated multiple lines.
1282 if lines.next().is_some() {
1283 break;
1284 }
1285 }
1286
1287 this.update(cx, |this, cx| {
1288 if !new_summary.is_empty() {
1289 this.summary = Some(new_summary.into());
1290 }
1291
1292 cx.emit(ThreadEvent::SummaryGenerated);
1293 })?;
1294
1295 anyhow::Ok(())
1296 }
1297 .log_err()
1298 .await
1299 });
1300 }
1301
1302 pub fn generate_detailed_summary(&mut self, cx: &mut Context<Self>) -> Option<Task<()>> {
1303 let last_message_id = self.messages.last().map(|message| message.id)?;
1304
1305 match &self.detailed_summary_state {
1306 DetailedSummaryState::Generating { message_id, .. }
1307 | DetailedSummaryState::Generated { message_id, .. }
1308 if *message_id == last_message_id =>
1309 {
1310 // Already up-to-date
1311 return None;
1312 }
1313 _ => {}
1314 }
1315
1316 let ConfiguredModel { model, provider } =
1317 LanguageModelRegistry::read_global(cx).thread_summary_model()?;
1318
1319 if !provider.is_authenticated(cx) {
1320 return None;
1321 }
1322
1323 let mut request = self.to_completion_request(RequestKind::Summarize, cx);
1324
1325 request.messages.push(LanguageModelRequestMessage {
1326 role: Role::User,
1327 content: vec![
1328 "Generate a detailed summary of this conversation. Include:\n\
1329 1. A brief overview of what was discussed\n\
1330 2. Key facts or information discovered\n\
1331 3. Outcomes or conclusions reached\n\
1332 4. Any action items or next steps if any\n\
1333 Format it in Markdown with headings and bullet points."
1334 .into(),
1335 ],
1336 cache: false,
1337 });
1338
1339 let task = cx.spawn(async move |thread, cx| {
1340 let stream = model.stream_completion_text(request, &cx);
1341 let Some(mut messages) = stream.await.log_err() else {
1342 thread
1343 .update(cx, |this, _cx| {
1344 this.detailed_summary_state = DetailedSummaryState::NotGenerated;
1345 })
1346 .log_err();
1347
1348 return;
1349 };
1350
1351 let mut new_detailed_summary = String::new();
1352
1353 while let Some(chunk) = messages.stream.next().await {
1354 if let Some(chunk) = chunk.log_err() {
1355 new_detailed_summary.push_str(&chunk);
1356 }
1357 }
1358
1359 thread
1360 .update(cx, |this, _cx| {
1361 this.detailed_summary_state = DetailedSummaryState::Generated {
1362 text: new_detailed_summary.into(),
1363 message_id: last_message_id,
1364 };
1365 })
1366 .log_err();
1367 });
1368
1369 self.detailed_summary_state = DetailedSummaryState::Generating {
1370 message_id: last_message_id,
1371 };
1372
1373 Some(task)
1374 }
1375
1376 pub fn is_generating_detailed_summary(&self) -> bool {
1377 matches!(
1378 self.detailed_summary_state,
1379 DetailedSummaryState::Generating { .. }
1380 )
1381 }
1382
1383 pub fn use_pending_tools(
1384 &mut self,
1385 cx: &mut Context<Self>,
1386 ) -> impl IntoIterator<Item = PendingToolUse> + use<> {
1387 let request = self.to_completion_request(RequestKind::Chat, cx);
1388 let messages = Arc::new(request.messages);
1389 let pending_tool_uses = self
1390 .tool_use
1391 .pending_tool_uses()
1392 .into_iter()
1393 .filter(|tool_use| tool_use.status.is_idle())
1394 .cloned()
1395 .collect::<Vec<_>>();
1396
1397 for tool_use in pending_tool_uses.iter() {
1398 if let Some(tool) = self.tools.tool(&tool_use.name, cx) {
1399 if tool.needs_confirmation(&tool_use.input, cx)
1400 && !AssistantSettings::get_global(cx).always_allow_tool_actions
1401 {
1402 self.tool_use.confirm_tool_use(
1403 tool_use.id.clone(),
1404 tool_use.ui_text.clone(),
1405 tool_use.input.clone(),
1406 messages.clone(),
1407 tool,
1408 );
1409 cx.emit(ThreadEvent::ToolConfirmationNeeded);
1410 } else {
1411 self.run_tool(
1412 tool_use.id.clone(),
1413 tool_use.ui_text.clone(),
1414 tool_use.input.clone(),
1415 &messages,
1416 tool,
1417 cx,
1418 );
1419 }
1420 }
1421 }
1422
1423 pending_tool_uses
1424 }
1425
1426 pub fn run_tool(
1427 &mut self,
1428 tool_use_id: LanguageModelToolUseId,
1429 ui_text: impl Into<SharedString>,
1430 input: serde_json::Value,
1431 messages: &[LanguageModelRequestMessage],
1432 tool: Arc<dyn Tool>,
1433 cx: &mut Context<Thread>,
1434 ) {
1435 let task = self.spawn_tool_use(tool_use_id.clone(), messages, input, tool, cx);
1436 self.tool_use
1437 .run_pending_tool(tool_use_id, ui_text.into(), task);
1438 }
1439
1440 fn spawn_tool_use(
1441 &mut self,
1442 tool_use_id: LanguageModelToolUseId,
1443 messages: &[LanguageModelRequestMessage],
1444 input: serde_json::Value,
1445 tool: Arc<dyn Tool>,
1446 cx: &mut Context<Thread>,
1447 ) -> Task<()> {
1448 let tool_name: Arc<str> = tool.name().into();
1449
1450 let run_tool = if self.tools.is_disabled(&tool.source(), &tool_name) {
1451 Task::ready(Err(anyhow!("tool is disabled: {tool_name}")))
1452 } else {
1453 tool.run(
1454 input,
1455 messages,
1456 self.project.clone(),
1457 self.action_log.clone(),
1458 cx,
1459 )
1460 };
1461
1462 cx.spawn({
1463 async move |thread: WeakEntity<Thread>, cx| {
1464 let output = run_tool.await;
1465
1466 thread
1467 .update(cx, |thread, cx| {
1468 let pending_tool_use = thread.tool_use.insert_tool_output(
1469 tool_use_id.clone(),
1470 tool_name,
1471 output,
1472 cx,
1473 );
1474
1475 cx.emit(ThreadEvent::ToolFinished {
1476 tool_use_id,
1477 pending_tool_use,
1478 canceled: false,
1479 });
1480 })
1481 .ok();
1482 }
1483 })
1484 }
1485
1486 pub fn attach_tool_results(&mut self, cx: &mut Context<Self>) {
1487 // Insert a user message to contain the tool results.
1488 self.insert_user_message(
1489 // TODO: Sending up a user message without any content results in the model sending back
1490 // responses that also don't have any content. We currently don't handle this case well,
1491 // so for now we provide some text to keep the model on track.
1492 "Here are the tool results.",
1493 Vec::new(),
1494 None,
1495 cx,
1496 );
1497 }
1498
1499 /// Cancels the last pending completion, if there are any pending.
1500 ///
1501 /// Returns whether a completion was canceled.
1502 pub fn cancel_last_completion(&mut self, cx: &mut Context<Self>) -> bool {
1503 let canceled = if self.pending_completions.pop().is_some() {
1504 true
1505 } else {
1506 let mut canceled = false;
1507 for pending_tool_use in self.tool_use.cancel_pending() {
1508 canceled = true;
1509 cx.emit(ThreadEvent::ToolFinished {
1510 tool_use_id: pending_tool_use.id.clone(),
1511 pending_tool_use: Some(pending_tool_use),
1512 canceled: true,
1513 });
1514 }
1515 canceled
1516 };
1517 self.finalize_pending_checkpoint(cx);
1518 canceled
1519 }
1520
1521 /// Returns the feedback given to the thread, if any.
1522 pub fn feedback(&self) -> Option<ThreadFeedback> {
1523 self.feedback
1524 }
1525
1526 /// Reports feedback about the thread and stores it in our telemetry backend.
1527 pub fn report_feedback(
1528 &mut self,
1529 feedback: ThreadFeedback,
1530 cx: &mut Context<Self>,
1531 ) -> Task<Result<()>> {
1532 let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
1533 let serialized_thread = self.serialize(cx);
1534 let thread_id = self.id().clone();
1535 let client = self.project.read(cx).client();
1536 self.feedback = Some(feedback);
1537 cx.notify();
1538
1539 cx.background_spawn(async move {
1540 let final_project_snapshot = final_project_snapshot.await;
1541 let serialized_thread = serialized_thread.await?;
1542 let thread_data =
1543 serde_json::to_value(serialized_thread).unwrap_or_else(|_| serde_json::Value::Null);
1544
1545 let rating = match feedback {
1546 ThreadFeedback::Positive => "positive",
1547 ThreadFeedback::Negative => "negative",
1548 };
1549 telemetry::event!(
1550 "Assistant Thread Rated",
1551 rating,
1552 thread_id,
1553 thread_data,
1554 final_project_snapshot
1555 );
1556 client.telemetry().flush_events();
1557
1558 Ok(())
1559 })
1560 }
1561
1562 /// Create a snapshot of the current project state including git information and unsaved buffers.
1563 fn project_snapshot(
1564 project: Entity<Project>,
1565 cx: &mut Context<Self>,
1566 ) -> Task<Arc<ProjectSnapshot>> {
1567 let git_store = project.read(cx).git_store().clone();
1568 let worktree_snapshots: Vec<_> = project
1569 .read(cx)
1570 .visible_worktrees(cx)
1571 .map(|worktree| Self::worktree_snapshot(worktree, git_store.clone(), cx))
1572 .collect();
1573
1574 cx.spawn(async move |_, cx| {
1575 let worktree_snapshots = futures::future::join_all(worktree_snapshots).await;
1576
1577 let mut unsaved_buffers = Vec::new();
1578 cx.update(|app_cx| {
1579 let buffer_store = project.read(app_cx).buffer_store();
1580 for buffer_handle in buffer_store.read(app_cx).buffers() {
1581 let buffer = buffer_handle.read(app_cx);
1582 if buffer.is_dirty() {
1583 if let Some(file) = buffer.file() {
1584 let path = file.path().to_string_lossy().to_string();
1585 unsaved_buffers.push(path);
1586 }
1587 }
1588 }
1589 })
1590 .ok();
1591
1592 Arc::new(ProjectSnapshot {
1593 worktree_snapshots,
1594 unsaved_buffer_paths: unsaved_buffers,
1595 timestamp: Utc::now(),
1596 })
1597 })
1598 }
1599
1600 fn worktree_snapshot(
1601 worktree: Entity<project::Worktree>,
1602 git_store: Entity<GitStore>,
1603 cx: &App,
1604 ) -> Task<WorktreeSnapshot> {
1605 cx.spawn(async move |cx| {
1606 // Get worktree path and snapshot
1607 let worktree_info = cx.update(|app_cx| {
1608 let worktree = worktree.read(app_cx);
1609 let path = worktree.abs_path().to_string_lossy().to_string();
1610 let snapshot = worktree.snapshot();
1611 (path, snapshot)
1612 });
1613
1614 let Ok((worktree_path, _snapshot)) = worktree_info else {
1615 return WorktreeSnapshot {
1616 worktree_path: String::new(),
1617 git_state: None,
1618 };
1619 };
1620
1621 let git_state = git_store
1622 .update(cx, |git_store, cx| {
1623 git_store
1624 .repositories()
1625 .values()
1626 .find(|repo| {
1627 repo.read(cx)
1628 .abs_path_to_repo_path(&worktree.read(cx).abs_path())
1629 .is_some()
1630 })
1631 .cloned()
1632 })
1633 .ok()
1634 .flatten()
1635 .map(|repo| {
1636 repo.update(cx, |repo, _| {
1637 let current_branch =
1638 repo.branch.as_ref().map(|branch| branch.name.to_string());
1639 repo.send_job(None, |state, _| async move {
1640 let RepositoryState::Local { backend, .. } = state else {
1641 return GitState {
1642 remote_url: None,
1643 head_sha: None,
1644 current_branch,
1645 diff: None,
1646 };
1647 };
1648
1649 let remote_url = backend.remote_url("origin");
1650 let head_sha = backend.head_sha();
1651 let diff = backend.diff(DiffType::HeadToWorktree).await.ok();
1652
1653 GitState {
1654 remote_url,
1655 head_sha,
1656 current_branch,
1657 diff,
1658 }
1659 })
1660 })
1661 });
1662
1663 let git_state = match git_state {
1664 Some(git_state) => match git_state.ok() {
1665 Some(git_state) => git_state.await.ok(),
1666 None => None,
1667 },
1668 None => None,
1669 };
1670
1671 WorktreeSnapshot {
1672 worktree_path,
1673 git_state,
1674 }
1675 })
1676 }
1677
1678 pub fn to_markdown(&self, cx: &App) -> Result<String> {
1679 let mut markdown = Vec::new();
1680
1681 if let Some(summary) = self.summary() {
1682 writeln!(markdown, "# {summary}\n")?;
1683 };
1684
1685 for message in self.messages() {
1686 writeln!(
1687 markdown,
1688 "## {role}\n",
1689 role = match message.role {
1690 Role::User => "User",
1691 Role::Assistant => "Assistant",
1692 Role::System => "System",
1693 }
1694 )?;
1695
1696 if !message.context.is_empty() {
1697 writeln!(markdown, "{}", message.context)?;
1698 }
1699
1700 for segment in &message.segments {
1701 match segment {
1702 MessageSegment::Text(text) => writeln!(markdown, "{}\n", text)?,
1703 MessageSegment::Thinking(text) => {
1704 writeln!(markdown, "<think>{}</think>\n", text)?
1705 }
1706 }
1707 }
1708
1709 for tool_use in self.tool_uses_for_message(message.id, cx) {
1710 writeln!(
1711 markdown,
1712 "**Use Tool: {} ({})**",
1713 tool_use.name, tool_use.id
1714 )?;
1715 writeln!(markdown, "```json")?;
1716 writeln!(
1717 markdown,
1718 "{}",
1719 serde_json::to_string_pretty(&tool_use.input)?
1720 )?;
1721 writeln!(markdown, "```")?;
1722 }
1723
1724 for tool_result in self.tool_results_for_message(message.id) {
1725 write!(markdown, "**Tool Results: {}", tool_result.tool_use_id)?;
1726 if tool_result.is_error {
1727 write!(markdown, " (Error)")?;
1728 }
1729
1730 writeln!(markdown, "**\n")?;
1731 writeln!(markdown, "{}", tool_result.content)?;
1732 }
1733 }
1734
1735 Ok(String::from_utf8_lossy(&markdown).to_string())
1736 }
1737
1738 pub fn keep_edits_in_range(
1739 &mut self,
1740 buffer: Entity<language::Buffer>,
1741 buffer_range: Range<language::Anchor>,
1742 cx: &mut Context<Self>,
1743 ) {
1744 self.action_log.update(cx, |action_log, cx| {
1745 action_log.keep_edits_in_range(buffer, buffer_range, cx)
1746 });
1747 }
1748
1749 pub fn keep_all_edits(&mut self, cx: &mut Context<Self>) {
1750 self.action_log
1751 .update(cx, |action_log, cx| action_log.keep_all_edits(cx));
1752 }
1753
1754 pub fn reject_edits_in_range(
1755 &mut self,
1756 buffer: Entity<language::Buffer>,
1757 buffer_range: Range<language::Anchor>,
1758 cx: &mut Context<Self>,
1759 ) -> Task<Result<()>> {
1760 self.action_log.update(cx, |action_log, cx| {
1761 action_log.reject_edits_in_range(buffer, buffer_range, cx)
1762 })
1763 }
1764
1765 pub fn action_log(&self) -> &Entity<ActionLog> {
1766 &self.action_log
1767 }
1768
1769 pub fn project(&self) -> &Entity<Project> {
1770 &self.project
1771 }
1772
1773 pub fn cumulative_token_usage(&self) -> TokenUsage {
1774 self.cumulative_token_usage.clone()
1775 }
1776
1777 pub fn total_token_usage(&self, cx: &App) -> TotalTokenUsage {
1778 let model_registry = LanguageModelRegistry::read_global(cx);
1779 let Some(model) = model_registry.default_model() else {
1780 return TotalTokenUsage::default();
1781 };
1782
1783 let max = model.model.max_token_count();
1784
1785 #[cfg(debug_assertions)]
1786 let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
1787 .unwrap_or("0.8".to_string())
1788 .parse()
1789 .unwrap();
1790 #[cfg(not(debug_assertions))]
1791 let warning_threshold: f32 = 0.8;
1792
1793 let total = self.cumulative_token_usage.total_tokens() as usize;
1794
1795 let ratio = if total >= max {
1796 TokenUsageRatio::Exceeded
1797 } else if total as f32 / max as f32 >= warning_threshold {
1798 TokenUsageRatio::Warning
1799 } else {
1800 TokenUsageRatio::Normal
1801 };
1802
1803 TotalTokenUsage { total, max, ratio }
1804 }
1805
1806 pub fn deny_tool_use(
1807 &mut self,
1808 tool_use_id: LanguageModelToolUseId,
1809 tool_name: Arc<str>,
1810 cx: &mut Context<Self>,
1811 ) {
1812 let err = Err(anyhow::anyhow!(
1813 "Permission to run tool action denied by user"
1814 ));
1815
1816 self.tool_use
1817 .insert_tool_output(tool_use_id.clone(), tool_name, err, cx);
1818
1819 cx.emit(ThreadEvent::ToolFinished {
1820 tool_use_id,
1821 pending_tool_use: None,
1822 canceled: true,
1823 });
1824 }
1825}
1826
1827pub fn system_prompt_reminder(prompt_builder: &prompt_store::PromptBuilder) -> String {
1828 prompt_builder
1829 .generate_assistant_system_prompt_reminder()
1830 .unwrap_or_default()
1831}
1832
1833#[derive(Debug, Clone)]
1834pub enum ThreadError {
1835 PaymentRequired,
1836 MaxMonthlySpendReached,
1837 Message {
1838 header: SharedString,
1839 message: SharedString,
1840 },
1841}
1842
1843#[derive(Debug, Clone)]
1844pub enum ThreadEvent {
1845 ShowError(ThreadError),
1846 StreamedCompletion,
1847 StreamedAssistantText(MessageId, String),
1848 StreamedAssistantThinking(MessageId, String),
1849 DoneStreaming,
1850 MessageAdded(MessageId),
1851 MessageEdited(MessageId),
1852 MessageDeleted(MessageId),
1853 SummaryGenerated,
1854 SummaryChanged,
1855 UsePendingTools,
1856 ToolFinished {
1857 #[allow(unused)]
1858 tool_use_id: LanguageModelToolUseId,
1859 /// The pending tool use that corresponds to this tool.
1860 pending_tool_use: Option<PendingToolUse>,
1861 /// Whether the tool was canceled by the user.
1862 canceled: bool,
1863 },
1864 CheckpointChanged,
1865 ToolConfirmationNeeded,
1866}
1867
1868impl EventEmitter<ThreadEvent> for Thread {}
1869
1870struct PendingCompletion {
1871 id: usize,
1872 _task: Task<()>,
1873}
1874
1875#[cfg(test)]
1876mod tests {
1877 use super::*;
1878 use crate::{ThreadStore, context_store::ContextStore, thread_store};
1879 use assistant_settings::AssistantSettings;
1880 use context_server::ContextServerSettings;
1881 use editor::EditorSettings;
1882 use gpui::TestAppContext;
1883 use project::{FakeFs, Project};
1884 use prompt_store::PromptBuilder;
1885 use serde_json::json;
1886 use settings::{Settings, SettingsStore};
1887 use std::sync::Arc;
1888 use theme::ThemeSettings;
1889 use util::path;
1890 use workspace::Workspace;
1891
1892 #[gpui::test]
1893 async fn test_message_with_context(cx: &mut TestAppContext) {
1894 init_test_settings(cx);
1895
1896 let project = create_test_project(
1897 cx,
1898 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
1899 )
1900 .await;
1901
1902 let (_workspace, _thread_store, thread, context_store, prompt_builder) =
1903 setup_test_environment(cx, project.clone()).await;
1904
1905 add_file_to_context(&project, &context_store, "test/code.rs", cx)
1906 .await
1907 .unwrap();
1908
1909 let context =
1910 context_store.update(cx, |store, _| store.context().first().cloned().unwrap());
1911
1912 // Insert user message with context
1913 let message_id = thread.update(cx, |thread, cx| {
1914 thread.insert_user_message("Please explain this code", vec![context], None, cx)
1915 });
1916
1917 // Check content and context in message object
1918 let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
1919
1920 // Use different path format strings based on platform for the test
1921 #[cfg(windows)]
1922 let path_part = r"test\code.rs";
1923 #[cfg(not(windows))]
1924 let path_part = "test/code.rs";
1925
1926 let expected_context = format!(
1927 r#"
1928<context>
1929The following items were attached by the user. You don't need to use other tools to read them.
1930
1931<files>
1932```rs {path_part}
1933fn main() {{
1934 println!("Hello, world!");
1935}}
1936```
1937</files>
1938</context>
1939"#
1940 );
1941
1942 assert_eq!(message.role, Role::User);
1943 assert_eq!(message.segments.len(), 1);
1944 assert_eq!(
1945 message.segments[0],
1946 MessageSegment::Text("Please explain this code".to_string())
1947 );
1948 assert_eq!(message.context, expected_context);
1949
1950 // Check message in request
1951 let request = thread.read_with(cx, |thread, cx| {
1952 thread.to_completion_request(RequestKind::Chat, cx)
1953 });
1954
1955 assert_eq!(request.messages.len(), 1);
1956 let actual_message = request.messages[0].string_contents();
1957 let expected_content = format!(
1958 "{}Please explain this code{}",
1959 expected_context,
1960 system_prompt_reminder(&prompt_builder)
1961 );
1962
1963 assert_eq!(actual_message, expected_content);
1964 }
1965
1966 #[gpui::test]
1967 async fn test_only_include_new_contexts(cx: &mut TestAppContext) {
1968 init_test_settings(cx);
1969
1970 let project = create_test_project(
1971 cx,
1972 json!({
1973 "file1.rs": "fn function1() {}\n",
1974 "file2.rs": "fn function2() {}\n",
1975 "file3.rs": "fn function3() {}\n",
1976 }),
1977 )
1978 .await;
1979
1980 let (_, _thread_store, thread, context_store, _prompt_builder) =
1981 setup_test_environment(cx, project.clone()).await;
1982
1983 // Open files individually
1984 add_file_to_context(&project, &context_store, "test/file1.rs", cx)
1985 .await
1986 .unwrap();
1987 add_file_to_context(&project, &context_store, "test/file2.rs", cx)
1988 .await
1989 .unwrap();
1990 add_file_to_context(&project, &context_store, "test/file3.rs", cx)
1991 .await
1992 .unwrap();
1993
1994 // Get the context objects
1995 let contexts = context_store.update(cx, |store, _| store.context().clone());
1996 assert_eq!(contexts.len(), 3);
1997
1998 // First message with context 1
1999 let message1_id = thread.update(cx, |thread, cx| {
2000 thread.insert_user_message("Message 1", vec![contexts[0].clone()], None, cx)
2001 });
2002
2003 // Second message with contexts 1 and 2 (context 1 should be skipped as it's already included)
2004 let message2_id = thread.update(cx, |thread, cx| {
2005 thread.insert_user_message(
2006 "Message 2",
2007 vec![contexts[0].clone(), contexts[1].clone()],
2008 None,
2009 cx,
2010 )
2011 });
2012
2013 // Third message with all three contexts (contexts 1 and 2 should be skipped)
2014 let message3_id = thread.update(cx, |thread, cx| {
2015 thread.insert_user_message(
2016 "Message 3",
2017 vec![
2018 contexts[0].clone(),
2019 contexts[1].clone(),
2020 contexts[2].clone(),
2021 ],
2022 None,
2023 cx,
2024 )
2025 });
2026
2027 // Check what contexts are included in each message
2028 let (message1, message2, message3) = thread.read_with(cx, |thread, _| {
2029 (
2030 thread.message(message1_id).unwrap().clone(),
2031 thread.message(message2_id).unwrap().clone(),
2032 thread.message(message3_id).unwrap().clone(),
2033 )
2034 });
2035
2036 // First message should include context 1
2037 assert!(message1.context.contains("file1.rs"));
2038
2039 // Second message should include only context 2 (not 1)
2040 assert!(!message2.context.contains("file1.rs"));
2041 assert!(message2.context.contains("file2.rs"));
2042
2043 // Third message should include only context 3 (not 1 or 2)
2044 assert!(!message3.context.contains("file1.rs"));
2045 assert!(!message3.context.contains("file2.rs"));
2046 assert!(message3.context.contains("file3.rs"));
2047
2048 // Check entire request to make sure all contexts are properly included
2049 let request = thread.read_with(cx, |thread, cx| {
2050 thread.to_completion_request(RequestKind::Chat, cx)
2051 });
2052
2053 // The request should contain all 3 messages
2054 assert_eq!(request.messages.len(), 3);
2055
2056 // Check that the contexts are properly formatted in each message
2057 assert!(request.messages[0].string_contents().contains("file1.rs"));
2058 assert!(!request.messages[0].string_contents().contains("file2.rs"));
2059 assert!(!request.messages[0].string_contents().contains("file3.rs"));
2060
2061 assert!(!request.messages[1].string_contents().contains("file1.rs"));
2062 assert!(request.messages[1].string_contents().contains("file2.rs"));
2063 assert!(!request.messages[1].string_contents().contains("file3.rs"));
2064
2065 assert!(!request.messages[2].string_contents().contains("file1.rs"));
2066 assert!(!request.messages[2].string_contents().contains("file2.rs"));
2067 assert!(request.messages[2].string_contents().contains("file3.rs"));
2068 }
2069
2070 #[gpui::test]
2071 async fn test_message_without_files(cx: &mut TestAppContext) {
2072 init_test_settings(cx);
2073
2074 let project = create_test_project(
2075 cx,
2076 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
2077 )
2078 .await;
2079
2080 let (_, _thread_store, thread, _context_store, prompt_builder) =
2081 setup_test_environment(cx, project.clone()).await;
2082
2083 // Insert user message without any context (empty context vector)
2084 let message_id = thread.update(cx, |thread, cx| {
2085 thread.insert_user_message("What is the best way to learn Rust?", vec![], None, cx)
2086 });
2087
2088 // Check content and context in message object
2089 let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
2090
2091 // Context should be empty when no files are included
2092 assert_eq!(message.role, Role::User);
2093 assert_eq!(message.segments.len(), 1);
2094 assert_eq!(
2095 message.segments[0],
2096 MessageSegment::Text("What is the best way to learn Rust?".to_string())
2097 );
2098 assert_eq!(message.context, "");
2099
2100 // Check message in request
2101 let request = thread.read_with(cx, |thread, cx| {
2102 thread.to_completion_request(RequestKind::Chat, cx)
2103 });
2104
2105 assert_eq!(request.messages.len(), 1);
2106 let actual_message = request.messages[0].string_contents();
2107 let expected_content = format!(
2108 "What is the best way to learn Rust?{}",
2109 system_prompt_reminder(&prompt_builder)
2110 );
2111
2112 assert_eq!(actual_message, expected_content);
2113
2114 // Add second message, also without context
2115 let message2_id = thread.update(cx, |thread, cx| {
2116 thread.insert_user_message("Are there any good books?", vec![], None, cx)
2117 });
2118
2119 let message2 =
2120 thread.read_with(cx, |thread, _| thread.message(message2_id).unwrap().clone());
2121 assert_eq!(message2.context, "");
2122
2123 // Check that both messages appear in the request
2124 let request = thread.read_with(cx, |thread, cx| {
2125 thread.to_completion_request(RequestKind::Chat, cx)
2126 });
2127
2128 assert_eq!(request.messages.len(), 2);
2129 // First message should be the system prompt
2130 assert_eq!(request.messages[0].role, Role::User);
2131
2132 // Second message should be the user message with prompt reminder
2133 let actual_message = request.messages[1].string_contents();
2134 let expected_content = format!(
2135 "Are there any good books?{}",
2136 system_prompt_reminder(&prompt_builder)
2137 );
2138
2139 assert_eq!(actual_message, expected_content);
2140 }
2141
2142 #[gpui::test]
2143 async fn test_stale_buffer_notification(cx: &mut TestAppContext) {
2144 init_test_settings(cx);
2145
2146 let project = create_test_project(
2147 cx,
2148 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
2149 )
2150 .await;
2151
2152 let (_workspace, _thread_store, thread, context_store, prompt_builder) =
2153 setup_test_environment(cx, project.clone()).await;
2154
2155 // Open buffer and add it to context
2156 let buffer = add_file_to_context(&project, &context_store, "test/code.rs", cx)
2157 .await
2158 .unwrap();
2159
2160 let context =
2161 context_store.update(cx, |store, _| store.context().first().cloned().unwrap());
2162
2163 // Insert user message with the buffer as context
2164 thread.update(cx, |thread, cx| {
2165 thread.insert_user_message("Explain this code", vec![context], None, cx)
2166 });
2167
2168 // Create a request and check that it doesn't have a stale buffer warning yet
2169 let initial_request = thread.read_with(cx, |thread, cx| {
2170 thread.to_completion_request(RequestKind::Chat, cx)
2171 });
2172
2173 // Make sure we don't have a stale file warning yet
2174 let has_stale_warning = initial_request.messages.iter().any(|msg| {
2175 msg.string_contents()
2176 .contains("These files changed since last read:")
2177 });
2178 assert!(
2179 !has_stale_warning,
2180 "Should not have stale buffer warning before buffer is modified"
2181 );
2182
2183 // Modify the buffer
2184 buffer.update(cx, |buffer, cx| {
2185 // Find a position at the end of line 1
2186 buffer.edit(
2187 [(1..1, "\n println!(\"Added a new line\");\n")],
2188 None,
2189 cx,
2190 );
2191 });
2192
2193 // Insert another user message without context
2194 thread.update(cx, |thread, cx| {
2195 thread.insert_user_message("What does the code do now?", vec![], None, cx)
2196 });
2197
2198 // Create a new request and check for the stale buffer warning
2199 let new_request = thread.read_with(cx, |thread, cx| {
2200 thread.to_completion_request(RequestKind::Chat, cx)
2201 });
2202
2203 // We should have a stale file warning as the last message
2204 let last_message = new_request
2205 .messages
2206 .last()
2207 .expect("Request should have messages");
2208
2209 // The last message should be the stale buffer notification
2210 assert_eq!(last_message.role, Role::User);
2211
2212 let actual_message = last_message.string_contents();
2213 let expected_content = format!(
2214 "These files changed since last read:\n- code.rs\n{}",
2215 system_prompt_reminder(&prompt_builder)
2216 );
2217
2218 assert_eq!(
2219 actual_message, expected_content,
2220 "Last message should be exactly the stale buffer notification"
2221 );
2222 }
2223
2224 fn init_test_settings(cx: &mut TestAppContext) {
2225 cx.update(|cx| {
2226 let settings_store = SettingsStore::test(cx);
2227 cx.set_global(settings_store);
2228 language::init(cx);
2229 Project::init_settings(cx);
2230 AssistantSettings::register(cx);
2231 thread_store::init(cx);
2232 workspace::init_settings(cx);
2233 ThemeSettings::register(cx);
2234 ContextServerSettings::register(cx);
2235 EditorSettings::register(cx);
2236 });
2237 }
2238
2239 // Helper to create a test project with test files
2240 async fn create_test_project(
2241 cx: &mut TestAppContext,
2242 files: serde_json::Value,
2243 ) -> Entity<Project> {
2244 let fs = FakeFs::new(cx.executor());
2245 fs.insert_tree(path!("/test"), files).await;
2246 Project::test(fs, [path!("/test").as_ref()], cx).await
2247 }
2248
2249 async fn setup_test_environment(
2250 cx: &mut TestAppContext,
2251 project: Entity<Project>,
2252 ) -> (
2253 Entity<Workspace>,
2254 Entity<ThreadStore>,
2255 Entity<Thread>,
2256 Entity<ContextStore>,
2257 Arc<PromptBuilder>,
2258 ) {
2259 let (workspace, cx) =
2260 cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
2261
2262 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2263
2264 let thread_store = cx.update(|_, cx| {
2265 ThreadStore::new(project.clone(), Arc::default(), prompt_builder.clone(), cx).unwrap()
2266 });
2267
2268 let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
2269 let context_store = cx.new(|_cx| ContextStore::new(workspace.downgrade(), None));
2270
2271 (
2272 workspace,
2273 thread_store,
2274 thread,
2275 context_store,
2276 prompt_builder,
2277 )
2278 }
2279
2280 async fn add_file_to_context(
2281 project: &Entity<Project>,
2282 context_store: &Entity<ContextStore>,
2283 path: &str,
2284 cx: &mut TestAppContext,
2285 ) -> Result<Entity<language::Buffer>> {
2286 let buffer_path = project
2287 .read_with(cx, |project, cx| project.find_project_path(path, cx))
2288 .unwrap();
2289
2290 let buffer = project
2291 .update(cx, |project, cx| project.open_buffer(buffer_path, cx))
2292 .await
2293 .unwrap();
2294
2295 context_store
2296 .update(cx, |store, cx| {
2297 store.add_file_from_buffer(buffer.clone(), cx)
2298 })
2299 .await?;
2300
2301 Ok(buffer)
2302 }
2303}