1use std::fmt::Write as _;
2use std::io::Write;
3use std::ops::Range;
4use std::sync::Arc;
5use std::time::Instant;
6
7use agent_settings::{AgentProfileId, AgentSettings, CompletionMode};
8use anyhow::{Result, anyhow};
9use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolWorkingSet};
10use chrono::{DateTime, Utc};
11use collections::HashMap;
12use editor::display_map::CreaseMetadata;
13use feature_flags::{self, FeatureFlagAppExt};
14use futures::future::Shared;
15use futures::{FutureExt, StreamExt as _};
16use git::repository::DiffType;
17use gpui::{
18 AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task,
19 WeakEntity,
20};
21use language_model::{
22 ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
23 LanguageModelId, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest,
24 LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
25 LanguageModelToolResultContent, LanguageModelToolUseId, MessageContent,
26 ModelRequestLimitReachedError, PaymentRequiredError, RequestUsage, Role, SelectedModel,
27 StopReason, TokenUsage,
28};
29use postage::stream::Stream as _;
30use project::Project;
31use project::git_store::{GitStore, GitStoreCheckpoint, RepositoryState};
32use prompt_store::{ModelContext, PromptBuilder};
33use proto::Plan;
34use schemars::JsonSchema;
35use serde::{Deserialize, Serialize};
36use settings::Settings;
37use thiserror::Error;
38use ui::Window;
39use util::{ResultExt as _, post_inc};
40use uuid::Uuid;
41use zed_llm_client::{CompletionIntent, CompletionRequestStatus};
42
43use crate::ThreadStore;
44use crate::agent_profile::AgentProfile;
45use crate::context::{AgentContext, AgentContextHandle, ContextLoadResult, LoadedContext};
46use crate::thread_store::{
47 SerializedCrease, SerializedLanguageModel, SerializedMessage, SerializedMessageSegment,
48 SerializedThread, SerializedToolResult, SerializedToolUse, SharedProjectContext,
49};
50use crate::tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState};
51
52#[derive(
53 Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,
54)]
55pub struct ThreadId(Arc<str>);
56
57impl ThreadId {
58 pub fn new() -> Self {
59 Self(Uuid::new_v4().to_string().into())
60 }
61}
62
63impl std::fmt::Display for ThreadId {
64 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65 write!(f, "{}", self.0)
66 }
67}
68
69impl From<&str> for ThreadId {
70 fn from(value: &str) -> Self {
71 Self(value.into())
72 }
73}
74
75/// The ID of the user prompt that initiated a request.
76///
77/// This equates to the user physically submitting a message to the model (e.g., by pressing the Enter key).
78#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize)]
79pub struct PromptId(Arc<str>);
80
81impl PromptId {
82 pub fn new() -> Self {
83 Self(Uuid::new_v4().to_string().into())
84 }
85}
86
87impl std::fmt::Display for PromptId {
88 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89 write!(f, "{}", self.0)
90 }
91}
92
93#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
94pub struct MessageId(pub(crate) usize);
95
96impl MessageId {
97 fn post_inc(&mut self) -> Self {
98 Self(post_inc(&mut self.0))
99 }
100}
101
102/// Stored information that can be used to resurrect a context crease when creating an editor for a past message.
103#[derive(Clone, Debug)]
104pub struct MessageCrease {
105 pub range: Range<usize>,
106 pub metadata: CreaseMetadata,
107 /// None for a deserialized message, Some otherwise.
108 pub context: Option<AgentContextHandle>,
109}
110
111/// A message in a [`Thread`].
112#[derive(Debug, Clone)]
113pub struct Message {
114 pub id: MessageId,
115 pub role: Role,
116 pub segments: Vec<MessageSegment>,
117 pub loaded_context: LoadedContext,
118 pub creases: Vec<MessageCrease>,
119 pub is_hidden: bool,
120}
121
122impl Message {
123 /// Returns whether the message contains any meaningful text that should be displayed
124 /// The model sometimes runs tool without producing any text or just a marker ([`USING_TOOL_MARKER`])
125 pub fn should_display_content(&self) -> bool {
126 self.segments.iter().all(|segment| segment.should_display())
127 }
128
129 pub fn push_thinking(&mut self, text: &str, signature: Option<String>) {
130 if let Some(MessageSegment::Thinking {
131 text: segment,
132 signature: current_signature,
133 }) = self.segments.last_mut()
134 {
135 if let Some(signature) = signature {
136 *current_signature = Some(signature);
137 }
138 segment.push_str(text);
139 } else {
140 self.segments.push(MessageSegment::Thinking {
141 text: text.to_string(),
142 signature,
143 });
144 }
145 }
146
147 pub fn push_text(&mut self, text: &str) {
148 if let Some(MessageSegment::Text(segment)) = self.segments.last_mut() {
149 segment.push_str(text);
150 } else {
151 self.segments.push(MessageSegment::Text(text.to_string()));
152 }
153 }
154
155 pub fn to_string(&self) -> String {
156 let mut result = String::new();
157
158 if !self.loaded_context.text.is_empty() {
159 result.push_str(&self.loaded_context.text);
160 }
161
162 for segment in &self.segments {
163 match segment {
164 MessageSegment::Text(text) => result.push_str(text),
165 MessageSegment::Thinking { text, .. } => {
166 result.push_str("<think>\n");
167 result.push_str(text);
168 result.push_str("\n</think>");
169 }
170 MessageSegment::RedactedThinking(_) => {}
171 }
172 }
173
174 result
175 }
176}
177
178#[derive(Debug, Clone, PartialEq, Eq)]
179pub enum MessageSegment {
180 Text(String),
181 Thinking {
182 text: String,
183 signature: Option<String>,
184 },
185 RedactedThinking(Vec<u8>),
186}
187
188impl MessageSegment {
189 pub fn should_display(&self) -> bool {
190 match self {
191 Self::Text(text) => text.is_empty(),
192 Self::Thinking { text, .. } => text.is_empty(),
193 Self::RedactedThinking(_) => false,
194 }
195 }
196}
197
198#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
199pub struct ProjectSnapshot {
200 pub worktree_snapshots: Vec<WorktreeSnapshot>,
201 pub unsaved_buffer_paths: Vec<String>,
202 pub timestamp: DateTime<Utc>,
203}
204
205#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
206pub struct WorktreeSnapshot {
207 pub worktree_path: String,
208 pub git_state: Option<GitState>,
209}
210
211#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
212pub struct GitState {
213 pub remote_url: Option<String>,
214 pub head_sha: Option<String>,
215 pub current_branch: Option<String>,
216 pub diff: Option<String>,
217}
218
219#[derive(Clone, Debug)]
220pub struct ThreadCheckpoint {
221 message_id: MessageId,
222 git_checkpoint: GitStoreCheckpoint,
223}
224
225#[derive(Copy, Clone, Debug, PartialEq, Eq)]
226pub enum ThreadFeedback {
227 Positive,
228 Negative,
229}
230
231pub enum LastRestoreCheckpoint {
232 Pending {
233 message_id: MessageId,
234 },
235 Error {
236 message_id: MessageId,
237 error: String,
238 },
239}
240
241impl LastRestoreCheckpoint {
242 pub fn message_id(&self) -> MessageId {
243 match self {
244 LastRestoreCheckpoint::Pending { message_id } => *message_id,
245 LastRestoreCheckpoint::Error { message_id, .. } => *message_id,
246 }
247 }
248}
249
250#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
251pub enum DetailedSummaryState {
252 #[default]
253 NotGenerated,
254 Generating {
255 message_id: MessageId,
256 },
257 Generated {
258 text: SharedString,
259 message_id: MessageId,
260 },
261}
262
263impl DetailedSummaryState {
264 fn text(&self) -> Option<SharedString> {
265 if let Self::Generated { text, .. } = self {
266 Some(text.clone())
267 } else {
268 None
269 }
270 }
271}
272
273#[derive(Default, Debug)]
274pub struct TotalTokenUsage {
275 pub total: u64,
276 pub max: u64,
277}
278
279impl TotalTokenUsage {
280 pub fn ratio(&self) -> TokenUsageRatio {
281 #[cfg(debug_assertions)]
282 let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
283 .unwrap_or("0.8".to_string())
284 .parse()
285 .unwrap();
286 #[cfg(not(debug_assertions))]
287 let warning_threshold: f32 = 0.8;
288
289 // When the maximum is unknown because there is no selected model,
290 // avoid showing the token limit warning.
291 if self.max == 0 {
292 TokenUsageRatio::Normal
293 } else if self.total >= self.max {
294 TokenUsageRatio::Exceeded
295 } else if self.total as f32 / self.max as f32 >= warning_threshold {
296 TokenUsageRatio::Warning
297 } else {
298 TokenUsageRatio::Normal
299 }
300 }
301
302 pub fn add(&self, tokens: u64) -> TotalTokenUsage {
303 TotalTokenUsage {
304 total: self.total + tokens,
305 max: self.max,
306 }
307 }
308}
309
310#[derive(Debug, Default, PartialEq, Eq)]
311pub enum TokenUsageRatio {
312 #[default]
313 Normal,
314 Warning,
315 Exceeded,
316}
317
318#[derive(Debug, Clone, Copy)]
319pub enum QueueState {
320 Sending,
321 Queued { position: usize },
322 Started,
323}
324
325/// A thread of conversation with the LLM.
326pub struct Thread {
327 id: ThreadId,
328 updated_at: DateTime<Utc>,
329 summary: ThreadSummary,
330 pending_summary: Task<Option<()>>,
331 detailed_summary_task: Task<Option<()>>,
332 detailed_summary_tx: postage::watch::Sender<DetailedSummaryState>,
333 detailed_summary_rx: postage::watch::Receiver<DetailedSummaryState>,
334 completion_mode: agent_settings::CompletionMode,
335 messages: Vec<Message>,
336 next_message_id: MessageId,
337 last_prompt_id: PromptId,
338 project_context: SharedProjectContext,
339 checkpoints_by_message: HashMap<MessageId, ThreadCheckpoint>,
340 completion_count: usize,
341 pending_completions: Vec<PendingCompletion>,
342 project: Entity<Project>,
343 prompt_builder: Arc<PromptBuilder>,
344 tools: Entity<ToolWorkingSet>,
345 tool_use: ToolUseState,
346 action_log: Entity<ActionLog>,
347 last_restore_checkpoint: Option<LastRestoreCheckpoint>,
348 pending_checkpoint: Option<ThreadCheckpoint>,
349 initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
350 request_token_usage: Vec<TokenUsage>,
351 cumulative_token_usage: TokenUsage,
352 exceeded_window_error: Option<ExceededWindowError>,
353 last_usage: Option<RequestUsage>,
354 tool_use_limit_reached: bool,
355 feedback: Option<ThreadFeedback>,
356 message_feedback: HashMap<MessageId, ThreadFeedback>,
357 last_auto_capture_at: Option<Instant>,
358 last_received_chunk_at: Option<Instant>,
359 request_callback: Option<
360 Box<dyn FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>])>,
361 >,
362 remaining_turns: u32,
363 configured_model: Option<ConfiguredModel>,
364 profile: AgentProfile,
365}
366
367#[derive(Clone, Debug, PartialEq, Eq)]
368pub enum ThreadSummary {
369 Pending,
370 Generating,
371 Ready(SharedString),
372 Error,
373}
374
375impl ThreadSummary {
376 pub const DEFAULT: SharedString = SharedString::new_static("New Thread");
377
378 pub fn or_default(&self) -> SharedString {
379 self.unwrap_or(Self::DEFAULT)
380 }
381
382 pub fn unwrap_or(&self, message: impl Into<SharedString>) -> SharedString {
383 self.ready().unwrap_or_else(|| message.into())
384 }
385
386 pub fn ready(&self) -> Option<SharedString> {
387 match self {
388 ThreadSummary::Ready(summary) => Some(summary.clone()),
389 ThreadSummary::Pending | ThreadSummary::Generating | ThreadSummary::Error => None,
390 }
391 }
392}
393
394#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
395pub struct ExceededWindowError {
396 /// Model used when last message exceeded context window
397 model_id: LanguageModelId,
398 /// Token count including last message
399 token_count: u64,
400}
401
402impl Thread {
403 pub fn new(
404 project: Entity<Project>,
405 tools: Entity<ToolWorkingSet>,
406 prompt_builder: Arc<PromptBuilder>,
407 system_prompt: SharedProjectContext,
408 cx: &mut Context<Self>,
409 ) -> Self {
410 let (detailed_summary_tx, detailed_summary_rx) = postage::watch::channel();
411 let configured_model = LanguageModelRegistry::read_global(cx).default_model();
412 let profile_id = AgentSettings::get_global(cx).default_profile.clone();
413
414 Self {
415 id: ThreadId::new(),
416 updated_at: Utc::now(),
417 summary: ThreadSummary::Pending,
418 pending_summary: Task::ready(None),
419 detailed_summary_task: Task::ready(None),
420 detailed_summary_tx,
421 detailed_summary_rx,
422 completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,
423 messages: Vec::new(),
424 next_message_id: MessageId(0),
425 last_prompt_id: PromptId::new(),
426 project_context: system_prompt,
427 checkpoints_by_message: HashMap::default(),
428 completion_count: 0,
429 pending_completions: Vec::new(),
430 project: project.clone(),
431 prompt_builder,
432 tools: tools.clone(),
433 last_restore_checkpoint: None,
434 pending_checkpoint: None,
435 tool_use: ToolUseState::new(tools.clone()),
436 action_log: cx.new(|_| ActionLog::new(project.clone())),
437 initial_project_snapshot: {
438 let project_snapshot = Self::project_snapshot(project, cx);
439 cx.foreground_executor()
440 .spawn(async move { Some(project_snapshot.await) })
441 .shared()
442 },
443 request_token_usage: Vec::new(),
444 cumulative_token_usage: TokenUsage::default(),
445 exceeded_window_error: None,
446 last_usage: None,
447 tool_use_limit_reached: false,
448 feedback: None,
449 message_feedback: HashMap::default(),
450 last_auto_capture_at: None,
451 last_received_chunk_at: None,
452 request_callback: None,
453 remaining_turns: u32::MAX,
454 configured_model,
455 profile: AgentProfile::new(profile_id, tools),
456 }
457 }
458
459 pub fn deserialize(
460 id: ThreadId,
461 serialized: SerializedThread,
462 project: Entity<Project>,
463 tools: Entity<ToolWorkingSet>,
464 prompt_builder: Arc<PromptBuilder>,
465 project_context: SharedProjectContext,
466 window: Option<&mut Window>, // None in headless mode
467 cx: &mut Context<Self>,
468 ) -> Self {
469 let next_message_id = MessageId(
470 serialized
471 .messages
472 .last()
473 .map(|message| message.id.0 + 1)
474 .unwrap_or(0),
475 );
476 let tool_use = ToolUseState::from_serialized_messages(
477 tools.clone(),
478 &serialized.messages,
479 project.clone(),
480 window,
481 cx,
482 );
483 let (detailed_summary_tx, detailed_summary_rx) =
484 postage::watch::channel_with(serialized.detailed_summary_state);
485
486 let configured_model = LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
487 serialized
488 .model
489 .and_then(|model| {
490 let model = SelectedModel {
491 provider: model.provider.clone().into(),
492 model: model.model.clone().into(),
493 };
494 registry.select_model(&model, cx)
495 })
496 .or_else(|| registry.default_model())
497 });
498
499 let completion_mode = serialized
500 .completion_mode
501 .unwrap_or_else(|| AgentSettings::get_global(cx).preferred_completion_mode);
502 let profile_id = serialized
503 .profile
504 .unwrap_or_else(|| AgentSettings::get_global(cx).default_profile.clone());
505
506 Self {
507 id,
508 updated_at: serialized.updated_at,
509 summary: ThreadSummary::Ready(serialized.summary),
510 pending_summary: Task::ready(None),
511 detailed_summary_task: Task::ready(None),
512 detailed_summary_tx,
513 detailed_summary_rx,
514 completion_mode,
515 messages: serialized
516 .messages
517 .into_iter()
518 .map(|message| Message {
519 id: message.id,
520 role: message.role,
521 segments: message
522 .segments
523 .into_iter()
524 .map(|segment| match segment {
525 SerializedMessageSegment::Text { text } => MessageSegment::Text(text),
526 SerializedMessageSegment::Thinking { text, signature } => {
527 MessageSegment::Thinking { text, signature }
528 }
529 SerializedMessageSegment::RedactedThinking { data } => {
530 MessageSegment::RedactedThinking(data)
531 }
532 })
533 .collect(),
534 loaded_context: LoadedContext {
535 contexts: Vec::new(),
536 text: message.context,
537 images: Vec::new(),
538 },
539 creases: message
540 .creases
541 .into_iter()
542 .map(|crease| MessageCrease {
543 range: crease.start..crease.end,
544 metadata: CreaseMetadata {
545 icon_path: crease.icon_path,
546 label: crease.label,
547 },
548 context: None,
549 })
550 .collect(),
551 is_hidden: message.is_hidden,
552 })
553 .collect(),
554 next_message_id,
555 last_prompt_id: PromptId::new(),
556 project_context,
557 checkpoints_by_message: HashMap::default(),
558 completion_count: 0,
559 pending_completions: Vec::new(),
560 last_restore_checkpoint: None,
561 pending_checkpoint: None,
562 project: project.clone(),
563 prompt_builder,
564 tools: tools.clone(),
565 tool_use,
566 action_log: cx.new(|_| ActionLog::new(project)),
567 initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
568 request_token_usage: serialized.request_token_usage,
569 cumulative_token_usage: serialized.cumulative_token_usage,
570 exceeded_window_error: None,
571 last_usage: None,
572 tool_use_limit_reached: serialized.tool_use_limit_reached,
573 feedback: None,
574 message_feedback: HashMap::default(),
575 last_auto_capture_at: None,
576 last_received_chunk_at: None,
577 request_callback: None,
578 remaining_turns: u32::MAX,
579 configured_model,
580 profile: AgentProfile::new(profile_id, tools),
581 }
582 }
583
584 pub fn set_request_callback(
585 &mut self,
586 callback: impl 'static
587 + FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>]),
588 ) {
589 self.request_callback = Some(Box::new(callback));
590 }
591
592 pub fn id(&self) -> &ThreadId {
593 &self.id
594 }
595
596 pub fn profile(&self) -> &AgentProfile {
597 &self.profile
598 }
599
600 pub fn set_profile(&mut self, id: AgentProfileId, cx: &mut Context<Self>) {
601 if &id != self.profile.id() {
602 self.profile = AgentProfile::new(id, self.tools.clone());
603 cx.emit(ThreadEvent::ProfileChanged);
604 }
605 }
606
607 pub fn is_empty(&self) -> bool {
608 self.messages.is_empty()
609 }
610
611 pub fn updated_at(&self) -> DateTime<Utc> {
612 self.updated_at
613 }
614
615 pub fn touch_updated_at(&mut self) {
616 self.updated_at = Utc::now();
617 }
618
619 pub fn advance_prompt_id(&mut self) {
620 self.last_prompt_id = PromptId::new();
621 }
622
623 pub fn project_context(&self) -> SharedProjectContext {
624 self.project_context.clone()
625 }
626
627 pub fn get_or_init_configured_model(&mut self, cx: &App) -> Option<ConfiguredModel> {
628 if self.configured_model.is_none() {
629 self.configured_model = LanguageModelRegistry::read_global(cx).default_model();
630 }
631 self.configured_model.clone()
632 }
633
634 pub fn configured_model(&self) -> Option<ConfiguredModel> {
635 self.configured_model.clone()
636 }
637
638 pub fn set_configured_model(&mut self, model: Option<ConfiguredModel>, cx: &mut Context<Self>) {
639 self.configured_model = model;
640 cx.notify();
641 }
642
643 pub fn summary(&self) -> &ThreadSummary {
644 &self.summary
645 }
646
647 pub fn set_summary(&mut self, new_summary: impl Into<SharedString>, cx: &mut Context<Self>) {
648 let current_summary = match &self.summary {
649 ThreadSummary::Pending | ThreadSummary::Generating => return,
650 ThreadSummary::Ready(summary) => summary,
651 ThreadSummary::Error => &ThreadSummary::DEFAULT,
652 };
653
654 let mut new_summary = new_summary.into();
655
656 if new_summary.is_empty() {
657 new_summary = ThreadSummary::DEFAULT;
658 }
659
660 if current_summary != &new_summary {
661 self.summary = ThreadSummary::Ready(new_summary);
662 cx.emit(ThreadEvent::SummaryChanged);
663 }
664 }
665
666 pub fn completion_mode(&self) -> CompletionMode {
667 self.completion_mode
668 }
669
670 pub fn set_completion_mode(&mut self, mode: CompletionMode) {
671 self.completion_mode = mode;
672 }
673
674 pub fn message(&self, id: MessageId) -> Option<&Message> {
675 let index = self
676 .messages
677 .binary_search_by(|message| message.id.cmp(&id))
678 .ok()?;
679
680 self.messages.get(index)
681 }
682
683 pub fn messages(&self) -> impl ExactSizeIterator<Item = &Message> {
684 self.messages.iter()
685 }
686
687 pub fn is_generating(&self) -> bool {
688 !self.pending_completions.is_empty() || !self.all_tools_finished()
689 }
690
691 /// Indicates whether streaming of language model events is stale.
692 /// When `is_generating()` is false, this method returns `None`.
693 pub fn is_generation_stale(&self) -> Option<bool> {
694 const STALE_THRESHOLD: u128 = 250;
695
696 self.last_received_chunk_at
697 .map(|instant| instant.elapsed().as_millis() > STALE_THRESHOLD)
698 }
699
700 fn received_chunk(&mut self) {
701 self.last_received_chunk_at = Some(Instant::now());
702 }
703
704 pub fn queue_state(&self) -> Option<QueueState> {
705 self.pending_completions
706 .first()
707 .map(|pending_completion| pending_completion.queue_state)
708 }
709
710 pub fn tools(&self) -> &Entity<ToolWorkingSet> {
711 &self.tools
712 }
713
714 pub fn pending_tool(&self, id: &LanguageModelToolUseId) -> Option<&PendingToolUse> {
715 self.tool_use
716 .pending_tool_uses()
717 .into_iter()
718 .find(|tool_use| &tool_use.id == id)
719 }
720
721 pub fn tools_needing_confirmation(&self) -> impl Iterator<Item = &PendingToolUse> {
722 self.tool_use
723 .pending_tool_uses()
724 .into_iter()
725 .filter(|tool_use| tool_use.status.needs_confirmation())
726 }
727
728 pub fn has_pending_tool_uses(&self) -> bool {
729 !self.tool_use.pending_tool_uses().is_empty()
730 }
731
732 pub fn checkpoint_for_message(&self, id: MessageId) -> Option<ThreadCheckpoint> {
733 self.checkpoints_by_message.get(&id).cloned()
734 }
735
736 pub fn restore_checkpoint(
737 &mut self,
738 checkpoint: ThreadCheckpoint,
739 cx: &mut Context<Self>,
740 ) -> Task<Result<()>> {
741 self.last_restore_checkpoint = Some(LastRestoreCheckpoint::Pending {
742 message_id: checkpoint.message_id,
743 });
744 cx.emit(ThreadEvent::CheckpointChanged);
745 cx.notify();
746
747 let git_store = self.project().read(cx).git_store().clone();
748 let restore = git_store.update(cx, |git_store, cx| {
749 git_store.restore_checkpoint(checkpoint.git_checkpoint.clone(), cx)
750 });
751
752 cx.spawn(async move |this, cx| {
753 let result = restore.await;
754 this.update(cx, |this, cx| {
755 if let Err(err) = result.as_ref() {
756 this.last_restore_checkpoint = Some(LastRestoreCheckpoint::Error {
757 message_id: checkpoint.message_id,
758 error: err.to_string(),
759 });
760 } else {
761 this.truncate(checkpoint.message_id, cx);
762 this.last_restore_checkpoint = None;
763 }
764 this.pending_checkpoint = None;
765 cx.emit(ThreadEvent::CheckpointChanged);
766 cx.notify();
767 })?;
768 result
769 })
770 }
771
772 fn finalize_pending_checkpoint(&mut self, cx: &mut Context<Self>) {
773 let pending_checkpoint = if self.is_generating() {
774 return;
775 } else if let Some(checkpoint) = self.pending_checkpoint.take() {
776 checkpoint
777 } else {
778 return;
779 };
780
781 self.finalize_checkpoint(pending_checkpoint, cx);
782 }
783
784 fn finalize_checkpoint(
785 &mut self,
786 pending_checkpoint: ThreadCheckpoint,
787 cx: &mut Context<Self>,
788 ) {
789 let git_store = self.project.read(cx).git_store().clone();
790 let final_checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
791 cx.spawn(async move |this, cx| match final_checkpoint.await {
792 Ok(final_checkpoint) => {
793 let equal = git_store
794 .update(cx, |store, cx| {
795 store.compare_checkpoints(
796 pending_checkpoint.git_checkpoint.clone(),
797 final_checkpoint.clone(),
798 cx,
799 )
800 })?
801 .await
802 .unwrap_or(false);
803
804 if !equal {
805 this.update(cx, |this, cx| {
806 this.insert_checkpoint(pending_checkpoint, cx)
807 })?;
808 }
809
810 Ok(())
811 }
812 Err(_) => this.update(cx, |this, cx| {
813 this.insert_checkpoint(pending_checkpoint, cx)
814 }),
815 })
816 .detach();
817 }
818
819 fn insert_checkpoint(&mut self, checkpoint: ThreadCheckpoint, cx: &mut Context<Self>) {
820 self.checkpoints_by_message
821 .insert(checkpoint.message_id, checkpoint);
822 cx.emit(ThreadEvent::CheckpointChanged);
823 cx.notify();
824 }
825
826 pub fn last_restore_checkpoint(&self) -> Option<&LastRestoreCheckpoint> {
827 self.last_restore_checkpoint.as_ref()
828 }
829
830 pub fn truncate(&mut self, message_id: MessageId, cx: &mut Context<Self>) {
831 let Some(message_ix) = self
832 .messages
833 .iter()
834 .rposition(|message| message.id == message_id)
835 else {
836 return;
837 };
838 for deleted_message in self.messages.drain(message_ix..) {
839 self.checkpoints_by_message.remove(&deleted_message.id);
840 }
841 cx.notify();
842 }
843
844 pub fn context_for_message(&self, id: MessageId) -> impl Iterator<Item = &AgentContext> {
845 self.messages
846 .iter()
847 .find(|message| message.id == id)
848 .into_iter()
849 .flat_map(|message| message.loaded_context.contexts.iter())
850 }
851
852 pub fn is_turn_end(&self, ix: usize) -> bool {
853 if self.messages.is_empty() {
854 return false;
855 }
856
857 if !self.is_generating() && ix == self.messages.len() - 1 {
858 return true;
859 }
860
861 let Some(message) = self.messages.get(ix) else {
862 return false;
863 };
864
865 if message.role != Role::Assistant {
866 return false;
867 }
868
869 self.messages
870 .get(ix + 1)
871 .and_then(|message| {
872 self.message(message.id)
873 .map(|next_message| next_message.role == Role::User && !next_message.is_hidden)
874 })
875 .unwrap_or(false)
876 }
877
878 pub fn last_usage(&self) -> Option<RequestUsage> {
879 self.last_usage
880 }
881
882 pub fn tool_use_limit_reached(&self) -> bool {
883 self.tool_use_limit_reached
884 }
885
886 /// Returns whether all of the tool uses have finished running.
887 pub fn all_tools_finished(&self) -> bool {
888 // If the only pending tool uses left are the ones with errors, then
889 // that means that we've finished running all of the pending tools.
890 self.tool_use
891 .pending_tool_uses()
892 .iter()
893 .all(|pending_tool_use| pending_tool_use.status.is_error())
894 }
895
896 /// Returns whether any pending tool uses may perform edits
897 pub fn has_pending_edit_tool_uses(&self) -> bool {
898 self.tool_use
899 .pending_tool_uses()
900 .iter()
901 .filter(|pending_tool_use| !pending_tool_use.status.is_error())
902 .any(|pending_tool_use| pending_tool_use.may_perform_edits)
903 }
904
905 pub fn tool_uses_for_message(&self, id: MessageId, cx: &App) -> Vec<ToolUse> {
906 self.tool_use.tool_uses_for_message(id, cx)
907 }
908
909 pub fn tool_results_for_message(
910 &self,
911 assistant_message_id: MessageId,
912 ) -> Vec<&LanguageModelToolResult> {
913 self.tool_use.tool_results_for_message(assistant_message_id)
914 }
915
916 pub fn tool_result(&self, id: &LanguageModelToolUseId) -> Option<&LanguageModelToolResult> {
917 self.tool_use.tool_result(id)
918 }
919
920 pub fn output_for_tool(&self, id: &LanguageModelToolUseId) -> Option<&Arc<str>> {
921 match &self.tool_use.tool_result(id)?.content {
922 LanguageModelToolResultContent::Text(text) => Some(text),
923 LanguageModelToolResultContent::Image(_) => {
924 // TODO: We should display image
925 None
926 }
927 }
928 }
929
930 pub fn card_for_tool(&self, id: &LanguageModelToolUseId) -> Option<AnyToolCard> {
931 self.tool_use.tool_result_card(id).cloned()
932 }
933
934 /// Return tools that are both enabled and supported by the model
935 pub fn available_tools(
936 &self,
937 cx: &App,
938 model: Arc<dyn LanguageModel>,
939 ) -> Vec<LanguageModelRequestTool> {
940 if model.supports_tools() {
941 self.profile
942 .enabled_tools(cx)
943 .into_iter()
944 .filter_map(|tool| {
945 // Skip tools that cannot be supported
946 let input_schema = tool.input_schema(model.tool_input_format()).ok()?;
947 Some(LanguageModelRequestTool {
948 name: tool.name(),
949 description: tool.description(),
950 input_schema,
951 })
952 })
953 .collect()
954 } else {
955 Vec::default()
956 }
957 }
958
959 pub fn insert_user_message(
960 &mut self,
961 text: impl Into<String>,
962 loaded_context: ContextLoadResult,
963 git_checkpoint: Option<GitStoreCheckpoint>,
964 creases: Vec<MessageCrease>,
965 cx: &mut Context<Self>,
966 ) -> MessageId {
967 if !loaded_context.referenced_buffers.is_empty() {
968 self.action_log.update(cx, |log, cx| {
969 for buffer in loaded_context.referenced_buffers {
970 log.buffer_read(buffer, cx);
971 }
972 });
973 }
974
975 let message_id = self.insert_message(
976 Role::User,
977 vec![MessageSegment::Text(text.into())],
978 loaded_context.loaded_context,
979 creases,
980 false,
981 cx,
982 );
983
984 if let Some(git_checkpoint) = git_checkpoint {
985 self.pending_checkpoint = Some(ThreadCheckpoint {
986 message_id,
987 git_checkpoint,
988 });
989 }
990
991 self.auto_capture_telemetry(cx);
992
993 message_id
994 }
995
996 pub fn insert_invisible_continue_message(&mut self, cx: &mut Context<Self>) -> MessageId {
997 let id = self.insert_message(
998 Role::User,
999 vec![MessageSegment::Text("Continue where you left off".into())],
1000 LoadedContext::default(),
1001 vec![],
1002 true,
1003 cx,
1004 );
1005 self.pending_checkpoint = None;
1006
1007 id
1008 }
1009
1010 pub fn insert_assistant_message(
1011 &mut self,
1012 segments: Vec<MessageSegment>,
1013 cx: &mut Context<Self>,
1014 ) -> MessageId {
1015 self.insert_message(
1016 Role::Assistant,
1017 segments,
1018 LoadedContext::default(),
1019 Vec::new(),
1020 false,
1021 cx,
1022 )
1023 }
1024
1025 pub fn insert_message(
1026 &mut self,
1027 role: Role,
1028 segments: Vec<MessageSegment>,
1029 loaded_context: LoadedContext,
1030 creases: Vec<MessageCrease>,
1031 is_hidden: bool,
1032 cx: &mut Context<Self>,
1033 ) -> MessageId {
1034 let id = self.next_message_id.post_inc();
1035 self.messages.push(Message {
1036 id,
1037 role,
1038 segments,
1039 loaded_context,
1040 creases,
1041 is_hidden,
1042 });
1043 self.touch_updated_at();
1044 cx.emit(ThreadEvent::MessageAdded(id));
1045 id
1046 }
1047
1048 pub fn edit_message(
1049 &mut self,
1050 id: MessageId,
1051 new_role: Role,
1052 new_segments: Vec<MessageSegment>,
1053 creases: Vec<MessageCrease>,
1054 loaded_context: Option<LoadedContext>,
1055 checkpoint: Option<GitStoreCheckpoint>,
1056 cx: &mut Context<Self>,
1057 ) -> bool {
1058 let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
1059 return false;
1060 };
1061 message.role = new_role;
1062 message.segments = new_segments;
1063 message.creases = creases;
1064 if let Some(context) = loaded_context {
1065 message.loaded_context = context;
1066 }
1067 if let Some(git_checkpoint) = checkpoint {
1068 self.checkpoints_by_message.insert(
1069 id,
1070 ThreadCheckpoint {
1071 message_id: id,
1072 git_checkpoint,
1073 },
1074 );
1075 }
1076 self.touch_updated_at();
1077 cx.emit(ThreadEvent::MessageEdited(id));
1078 true
1079 }
1080
1081 pub fn delete_message(&mut self, id: MessageId, cx: &mut Context<Self>) -> bool {
1082 let Some(index) = self.messages.iter().position(|message| message.id == id) else {
1083 return false;
1084 };
1085 self.messages.remove(index);
1086 self.touch_updated_at();
1087 cx.emit(ThreadEvent::MessageDeleted(id));
1088 true
1089 }
1090
1091 /// Returns the representation of this [`Thread`] in a textual form.
1092 ///
1093 /// This is the representation we use when attaching a thread as context to another thread.
1094 pub fn text(&self) -> String {
1095 let mut text = String::new();
1096
1097 for message in &self.messages {
1098 text.push_str(match message.role {
1099 language_model::Role::User => "User:",
1100 language_model::Role::Assistant => "Agent:",
1101 language_model::Role::System => "System:",
1102 });
1103 text.push('\n');
1104
1105 for segment in &message.segments {
1106 match segment {
1107 MessageSegment::Text(content) => text.push_str(content),
1108 MessageSegment::Thinking { text: content, .. } => {
1109 text.push_str(&format!("<think>{}</think>", content))
1110 }
1111 MessageSegment::RedactedThinking(_) => {}
1112 }
1113 }
1114 text.push('\n');
1115 }
1116
1117 text
1118 }
1119
1120 /// Serializes this thread into a format for storage or telemetry.
1121 pub fn serialize(&self, cx: &mut Context<Self>) -> Task<Result<SerializedThread>> {
1122 let initial_project_snapshot = self.initial_project_snapshot.clone();
1123 cx.spawn(async move |this, cx| {
1124 let initial_project_snapshot = initial_project_snapshot.await;
1125 this.read_with(cx, |this, cx| SerializedThread {
1126 version: SerializedThread::VERSION.to_string(),
1127 summary: this.summary().or_default(),
1128 updated_at: this.updated_at(),
1129 messages: this
1130 .messages()
1131 .map(|message| SerializedMessage {
1132 id: message.id,
1133 role: message.role,
1134 segments: message
1135 .segments
1136 .iter()
1137 .map(|segment| match segment {
1138 MessageSegment::Text(text) => {
1139 SerializedMessageSegment::Text { text: text.clone() }
1140 }
1141 MessageSegment::Thinking { text, signature } => {
1142 SerializedMessageSegment::Thinking {
1143 text: text.clone(),
1144 signature: signature.clone(),
1145 }
1146 }
1147 MessageSegment::RedactedThinking(data) => {
1148 SerializedMessageSegment::RedactedThinking {
1149 data: data.clone(),
1150 }
1151 }
1152 })
1153 .collect(),
1154 tool_uses: this
1155 .tool_uses_for_message(message.id, cx)
1156 .into_iter()
1157 .map(|tool_use| SerializedToolUse {
1158 id: tool_use.id,
1159 name: tool_use.name,
1160 input: tool_use.input,
1161 })
1162 .collect(),
1163 tool_results: this
1164 .tool_results_for_message(message.id)
1165 .into_iter()
1166 .map(|tool_result| SerializedToolResult {
1167 tool_use_id: tool_result.tool_use_id.clone(),
1168 is_error: tool_result.is_error,
1169 content: tool_result.content.clone(),
1170 output: tool_result.output.clone(),
1171 })
1172 .collect(),
1173 context: message.loaded_context.text.clone(),
1174 creases: message
1175 .creases
1176 .iter()
1177 .map(|crease| SerializedCrease {
1178 start: crease.range.start,
1179 end: crease.range.end,
1180 icon_path: crease.metadata.icon_path.clone(),
1181 label: crease.metadata.label.clone(),
1182 })
1183 .collect(),
1184 is_hidden: message.is_hidden,
1185 })
1186 .collect(),
1187 initial_project_snapshot,
1188 cumulative_token_usage: this.cumulative_token_usage,
1189 request_token_usage: this.request_token_usage.clone(),
1190 detailed_summary_state: this.detailed_summary_rx.borrow().clone(),
1191 exceeded_window_error: this.exceeded_window_error.clone(),
1192 model: this
1193 .configured_model
1194 .as_ref()
1195 .map(|model| SerializedLanguageModel {
1196 provider: model.provider.id().0.to_string(),
1197 model: model.model.id().0.to_string(),
1198 }),
1199 completion_mode: Some(this.completion_mode),
1200 tool_use_limit_reached: this.tool_use_limit_reached,
1201 profile: Some(this.profile.id().clone()),
1202 })
1203 })
1204 }
1205
1206 pub fn remaining_turns(&self) -> u32 {
1207 self.remaining_turns
1208 }
1209
1210 pub fn set_remaining_turns(&mut self, remaining_turns: u32) {
1211 self.remaining_turns = remaining_turns;
1212 }
1213
1214 pub fn send_to_model(
1215 &mut self,
1216 model: Arc<dyn LanguageModel>,
1217 intent: CompletionIntent,
1218 window: Option<AnyWindowHandle>,
1219 cx: &mut Context<Self>,
1220 ) {
1221 if self.remaining_turns == 0 {
1222 return;
1223 }
1224
1225 self.remaining_turns -= 1;
1226
1227 let request = self.to_completion_request(model.clone(), intent, cx);
1228
1229 self.stream_completion(request, model, window, cx);
1230 }
1231
1232 pub fn used_tools_since_last_user_message(&self) -> bool {
1233 for message in self.messages.iter().rev() {
1234 if self.tool_use.message_has_tool_results(message.id) {
1235 return true;
1236 } else if message.role == Role::User {
1237 return false;
1238 }
1239 }
1240
1241 false
1242 }
1243
1244 pub fn to_completion_request(
1245 &self,
1246 model: Arc<dyn LanguageModel>,
1247 intent: CompletionIntent,
1248 cx: &mut Context<Self>,
1249 ) -> LanguageModelRequest {
1250 let mut request = LanguageModelRequest {
1251 thread_id: Some(self.id.to_string()),
1252 prompt_id: Some(self.last_prompt_id.to_string()),
1253 intent: Some(intent),
1254 mode: None,
1255 messages: vec![],
1256 tools: Vec::new(),
1257 tool_choice: None,
1258 stop: Vec::new(),
1259 temperature: AgentSettings::temperature_for_model(&model, cx),
1260 };
1261
1262 let available_tools = self.available_tools(cx, model.clone());
1263 let available_tool_names = available_tools
1264 .iter()
1265 .map(|tool| tool.name.clone())
1266 .collect();
1267
1268 let model_context = &ModelContext {
1269 available_tools: available_tool_names,
1270 };
1271
1272 if let Some(project_context) = self.project_context.borrow().as_ref() {
1273 match self
1274 .prompt_builder
1275 .generate_assistant_system_prompt(project_context, model_context)
1276 {
1277 Err(err) => {
1278 let message = format!("{err:?}").into();
1279 log::error!("{message}");
1280 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1281 header: "Error generating system prompt".into(),
1282 message,
1283 }));
1284 }
1285 Ok(system_prompt) => {
1286 request.messages.push(LanguageModelRequestMessage {
1287 role: Role::System,
1288 content: vec![MessageContent::Text(system_prompt)],
1289 cache: true,
1290 });
1291 }
1292 }
1293 } else {
1294 let message = "Context for system prompt unexpectedly not ready.".into();
1295 log::error!("{message}");
1296 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1297 header: "Error generating system prompt".into(),
1298 message,
1299 }));
1300 }
1301
1302 let mut message_ix_to_cache = None;
1303 for message in &self.messages {
1304 let mut request_message = LanguageModelRequestMessage {
1305 role: message.role,
1306 content: Vec::new(),
1307 cache: false,
1308 };
1309
1310 message
1311 .loaded_context
1312 .add_to_request_message(&mut request_message);
1313
1314 for segment in &message.segments {
1315 match segment {
1316 MessageSegment::Text(text) => {
1317 if !text.is_empty() {
1318 request_message
1319 .content
1320 .push(MessageContent::Text(text.into()));
1321 }
1322 }
1323 MessageSegment::Thinking { text, signature } => {
1324 if !text.is_empty() {
1325 request_message.content.push(MessageContent::Thinking {
1326 text: text.into(),
1327 signature: signature.clone(),
1328 });
1329 }
1330 }
1331 MessageSegment::RedactedThinking(data) => {
1332 request_message
1333 .content
1334 .push(MessageContent::RedactedThinking(data.clone()));
1335 }
1336 };
1337 }
1338
1339 let mut cache_message = true;
1340 let mut tool_results_message = LanguageModelRequestMessage {
1341 role: Role::User,
1342 content: Vec::new(),
1343 cache: false,
1344 };
1345 for (tool_use, tool_result) in self.tool_use.tool_results(message.id) {
1346 if let Some(tool_result) = tool_result {
1347 request_message
1348 .content
1349 .push(MessageContent::ToolUse(tool_use.clone()));
1350 tool_results_message
1351 .content
1352 .push(MessageContent::ToolResult(LanguageModelToolResult {
1353 tool_use_id: tool_use.id.clone(),
1354 tool_name: tool_result.tool_name.clone(),
1355 is_error: tool_result.is_error,
1356 content: if tool_result.content.is_empty() {
1357 // Surprisingly, the API fails if we return an empty string here.
1358 // It thinks we are sending a tool use without a tool result.
1359 "<Tool returned an empty string>".into()
1360 } else {
1361 tool_result.content.clone()
1362 },
1363 output: None,
1364 }));
1365 } else {
1366 cache_message = false;
1367 log::debug!(
1368 "skipped tool use {:?} because it is still pending",
1369 tool_use
1370 );
1371 }
1372 }
1373
1374 if cache_message {
1375 message_ix_to_cache = Some(request.messages.len());
1376 }
1377 request.messages.push(request_message);
1378
1379 if !tool_results_message.content.is_empty() {
1380 if cache_message {
1381 message_ix_to_cache = Some(request.messages.len());
1382 }
1383 request.messages.push(tool_results_message);
1384 }
1385 }
1386
1387 // https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
1388 if let Some(message_ix_to_cache) = message_ix_to_cache {
1389 request.messages[message_ix_to_cache].cache = true;
1390 }
1391
1392 self.attach_tracked_files_state(&mut request.messages, cx);
1393
1394 request.tools = available_tools;
1395 request.mode = if model.supports_max_mode() {
1396 Some(self.completion_mode.into())
1397 } else {
1398 Some(CompletionMode::Normal.into())
1399 };
1400
1401 request
1402 }
1403
1404 fn to_summarize_request(
1405 &self,
1406 model: &Arc<dyn LanguageModel>,
1407 intent: CompletionIntent,
1408 added_user_message: String,
1409 cx: &App,
1410 ) -> LanguageModelRequest {
1411 let mut request = LanguageModelRequest {
1412 thread_id: None,
1413 prompt_id: None,
1414 intent: Some(intent),
1415 mode: None,
1416 messages: vec![],
1417 tools: Vec::new(),
1418 tool_choice: None,
1419 stop: Vec::new(),
1420 temperature: AgentSettings::temperature_for_model(model, cx),
1421 };
1422
1423 for message in &self.messages {
1424 let mut request_message = LanguageModelRequestMessage {
1425 role: message.role,
1426 content: Vec::new(),
1427 cache: false,
1428 };
1429
1430 for segment in &message.segments {
1431 match segment {
1432 MessageSegment::Text(text) => request_message
1433 .content
1434 .push(MessageContent::Text(text.clone())),
1435 MessageSegment::Thinking { .. } => {}
1436 MessageSegment::RedactedThinking(_) => {}
1437 }
1438 }
1439
1440 if request_message.content.is_empty() {
1441 continue;
1442 }
1443
1444 request.messages.push(request_message);
1445 }
1446
1447 request.messages.push(LanguageModelRequestMessage {
1448 role: Role::User,
1449 content: vec![MessageContent::Text(added_user_message)],
1450 cache: false,
1451 });
1452
1453 request
1454 }
1455
1456 fn attach_tracked_files_state(
1457 &self,
1458 messages: &mut Vec<LanguageModelRequestMessage>,
1459 cx: &App,
1460 ) {
1461 let mut stale_files = String::new();
1462
1463 let action_log = self.action_log.read(cx);
1464
1465 for stale_file in action_log.stale_buffers(cx) {
1466 if let Some(file) = stale_file.read(cx).file() {
1467 writeln!(&mut stale_files, "- {}", file.path().display()).ok();
1468 }
1469 }
1470
1471 if stale_files.is_empty() {
1472 return;
1473 }
1474
1475 // NOTE: Changes to this prompt require a symmetric update in the LLM Worker
1476 const STALE_FILES_HEADER: &str = include_str!("./prompts/stale_files_prompt_header.txt");
1477 let content = MessageContent::Text(
1478 format!("{STALE_FILES_HEADER}{stale_files}").replace("\r\n", "\n"),
1479 );
1480
1481 // Insert our message before the last Assistant message.
1482 // Inserting it to the tail distracts the agent too much
1483 let insert_position = messages
1484 .iter()
1485 .enumerate()
1486 .rfind(|(_, message)| message.role == Role::Assistant)
1487 .map_or(messages.len(), |(i, _)| i);
1488
1489 let request_message = LanguageModelRequestMessage {
1490 role: Role::User,
1491 content: vec![content],
1492 cache: false,
1493 };
1494
1495 messages.insert(insert_position, request_message);
1496
1497 // It makes no sense to cache messages after this one because
1498 // the cache is invalidated when this message is gone.
1499 // Move the cache marker before this message.
1500 let has_cached_messages_after = messages
1501 .iter()
1502 .skip(insert_position + 1)
1503 .any(|message| message.cache);
1504
1505 if has_cached_messages_after {
1506 messages[insert_position - 1].cache = true;
1507 }
1508 }
1509
1510 pub fn stream_completion(
1511 &mut self,
1512 request: LanguageModelRequest,
1513 model: Arc<dyn LanguageModel>,
1514 window: Option<AnyWindowHandle>,
1515 cx: &mut Context<Self>,
1516 ) {
1517 self.tool_use_limit_reached = false;
1518
1519 let pending_completion_id = post_inc(&mut self.completion_count);
1520 let mut request_callback_parameters = if self.request_callback.is_some() {
1521 Some((request.clone(), Vec::new()))
1522 } else {
1523 None
1524 };
1525 let prompt_id = self.last_prompt_id.clone();
1526 let tool_use_metadata = ToolUseMetadata {
1527 model: model.clone(),
1528 thread_id: self.id.clone(),
1529 prompt_id: prompt_id.clone(),
1530 };
1531
1532 self.last_received_chunk_at = Some(Instant::now());
1533
1534 let task = cx.spawn(async move |thread, cx| {
1535 let stream_completion_future = model.stream_completion(request, &cx);
1536 let initial_token_usage =
1537 thread.read_with(cx, |thread, _cx| thread.cumulative_token_usage);
1538 let stream_completion = async {
1539 let mut events = stream_completion_future.await?;
1540
1541 let mut stop_reason = StopReason::EndTurn;
1542 let mut current_token_usage = TokenUsage::default();
1543
1544 thread
1545 .update(cx, |_thread, cx| {
1546 cx.emit(ThreadEvent::NewRequest);
1547 })
1548 .ok();
1549
1550 let mut request_assistant_message_id = None;
1551
1552 while let Some(event) = events.next().await {
1553 if let Some((_, response_events)) = request_callback_parameters.as_mut() {
1554 response_events
1555 .push(event.as_ref().map_err(|error| error.to_string()).cloned());
1556 }
1557
1558 thread.update(cx, |thread, cx| {
1559 let event = match event {
1560 Ok(event) => event,
1561 Err(LanguageModelCompletionError::BadInputJson {
1562 id,
1563 tool_name,
1564 raw_input: invalid_input_json,
1565 json_parse_error,
1566 }) => {
1567 thread.receive_invalid_tool_json(
1568 id,
1569 tool_name,
1570 invalid_input_json,
1571 json_parse_error,
1572 window,
1573 cx,
1574 );
1575 return Ok(());
1576 }
1577 Err(LanguageModelCompletionError::Other(error)) => {
1578 return Err(error);
1579 }
1580 Err(err @ LanguageModelCompletionError::RateLimit(..)) => {
1581 return Err(err.into());
1582 }
1583 };
1584
1585 match event {
1586 LanguageModelCompletionEvent::StartMessage { .. } => {
1587 request_assistant_message_id =
1588 Some(thread.insert_assistant_message(
1589 vec![MessageSegment::Text(String::new())],
1590 cx,
1591 ));
1592 }
1593 LanguageModelCompletionEvent::Stop(reason) => {
1594 stop_reason = reason;
1595 }
1596 LanguageModelCompletionEvent::UsageUpdate(token_usage) => {
1597 thread.update_token_usage_at_last_message(token_usage);
1598 thread.cumulative_token_usage = thread.cumulative_token_usage
1599 + token_usage
1600 - current_token_usage;
1601 current_token_usage = token_usage;
1602 }
1603 LanguageModelCompletionEvent::Text(chunk) => {
1604 thread.received_chunk();
1605
1606 cx.emit(ThreadEvent::ReceivedTextChunk);
1607 if let Some(last_message) = thread.messages.last_mut() {
1608 if last_message.role == Role::Assistant
1609 && !thread.tool_use.has_tool_results(last_message.id)
1610 {
1611 last_message.push_text(&chunk);
1612 cx.emit(ThreadEvent::StreamedAssistantText(
1613 last_message.id,
1614 chunk,
1615 ));
1616 } else {
1617 // If we won't have an Assistant message yet, assume this chunk marks the beginning
1618 // of a new Assistant response.
1619 //
1620 // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1621 // will result in duplicating the text of the chunk in the rendered Markdown.
1622 request_assistant_message_id =
1623 Some(thread.insert_assistant_message(
1624 vec![MessageSegment::Text(chunk.to_string())],
1625 cx,
1626 ));
1627 };
1628 }
1629 }
1630 LanguageModelCompletionEvent::Thinking {
1631 text: chunk,
1632 signature,
1633 } => {
1634 thread.received_chunk();
1635
1636 if let Some(last_message) = thread.messages.last_mut() {
1637 if last_message.role == Role::Assistant
1638 && !thread.tool_use.has_tool_results(last_message.id)
1639 {
1640 last_message.push_thinking(&chunk, signature);
1641 cx.emit(ThreadEvent::StreamedAssistantThinking(
1642 last_message.id,
1643 chunk,
1644 ));
1645 } else {
1646 // If we won't have an Assistant message yet, assume this chunk marks the beginning
1647 // of a new Assistant response.
1648 //
1649 // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1650 // will result in duplicating the text of the chunk in the rendered Markdown.
1651 request_assistant_message_id =
1652 Some(thread.insert_assistant_message(
1653 vec![MessageSegment::Thinking {
1654 text: chunk.to_string(),
1655 signature,
1656 }],
1657 cx,
1658 ));
1659 };
1660 }
1661 }
1662 LanguageModelCompletionEvent::ToolUse(tool_use) => {
1663 let last_assistant_message_id = request_assistant_message_id
1664 .unwrap_or_else(|| {
1665 let new_assistant_message_id =
1666 thread.insert_assistant_message(vec![], cx);
1667 request_assistant_message_id =
1668 Some(new_assistant_message_id);
1669 new_assistant_message_id
1670 });
1671
1672 let tool_use_id = tool_use.id.clone();
1673 let streamed_input = if tool_use.is_input_complete {
1674 None
1675 } else {
1676 Some((&tool_use.input).clone())
1677 };
1678
1679 let ui_text = thread.tool_use.request_tool_use(
1680 last_assistant_message_id,
1681 tool_use,
1682 tool_use_metadata.clone(),
1683 cx,
1684 );
1685
1686 if let Some(input) = streamed_input {
1687 cx.emit(ThreadEvent::StreamedToolUse {
1688 tool_use_id,
1689 ui_text,
1690 input,
1691 });
1692 }
1693 }
1694 LanguageModelCompletionEvent::StatusUpdate(status_update) => {
1695 if let Some(completion) = thread
1696 .pending_completions
1697 .iter_mut()
1698 .find(|completion| completion.id == pending_completion_id)
1699 {
1700 match status_update {
1701 CompletionRequestStatus::Queued {
1702 position,
1703 } => {
1704 completion.queue_state = QueueState::Queued { position };
1705 }
1706 CompletionRequestStatus::Started => {
1707 completion.queue_state = QueueState::Started;
1708 }
1709 CompletionRequestStatus::Failed {
1710 code, message, request_id
1711 } => {
1712 anyhow::bail!("completion request failed. request_id: {request_id}, code: {code}, message: {message}");
1713 }
1714 CompletionRequestStatus::UsageUpdated {
1715 amount, limit
1716 } => {
1717 let usage = RequestUsage { limit, amount: amount as i32 };
1718
1719 thread.last_usage = Some(usage);
1720 }
1721 CompletionRequestStatus::ToolUseLimitReached => {
1722 thread.tool_use_limit_reached = true;
1723 cx.emit(ThreadEvent::ToolUseLimitReached);
1724 }
1725 }
1726 }
1727 }
1728 }
1729
1730 thread.touch_updated_at();
1731 cx.emit(ThreadEvent::StreamedCompletion);
1732 cx.notify();
1733
1734 thread.auto_capture_telemetry(cx);
1735 Ok(())
1736 })??;
1737
1738 smol::future::yield_now().await;
1739 }
1740
1741 thread.update(cx, |thread, cx| {
1742 thread.last_received_chunk_at = None;
1743 thread
1744 .pending_completions
1745 .retain(|completion| completion.id != pending_completion_id);
1746
1747 // If there is a response without tool use, summarize the message. Otherwise,
1748 // allow two tool uses before summarizing.
1749 if matches!(thread.summary, ThreadSummary::Pending)
1750 && thread.messages.len() >= 2
1751 && (!thread.has_pending_tool_uses() || thread.messages.len() >= 6)
1752 {
1753 thread.summarize(cx);
1754 }
1755 })?;
1756
1757 anyhow::Ok(stop_reason)
1758 };
1759
1760 let result = stream_completion.await;
1761
1762 thread
1763 .update(cx, |thread, cx| {
1764 thread.finalize_pending_checkpoint(cx);
1765 match result.as_ref() {
1766 Ok(stop_reason) => match stop_reason {
1767 StopReason::ToolUse => {
1768 let tool_uses = thread.use_pending_tools(window, cx, model.clone());
1769 cx.emit(ThreadEvent::UsePendingTools { tool_uses });
1770 }
1771 StopReason::EndTurn | StopReason::MaxTokens => {
1772 thread.project.update(cx, |project, cx| {
1773 project.set_agent_location(None, cx);
1774 });
1775 }
1776 StopReason::Refusal => {
1777 thread.project.update(cx, |project, cx| {
1778 project.set_agent_location(None, cx);
1779 });
1780
1781 // Remove the turn that was refused.
1782 //
1783 // https://docs.anthropic.com/en/docs/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals#reset-context-after-refusal
1784 {
1785 let mut messages_to_remove = Vec::new();
1786
1787 for (ix, message) in thread.messages.iter().enumerate().rev() {
1788 messages_to_remove.push(message.id);
1789
1790 if message.role == Role::User {
1791 if ix == 0 {
1792 break;
1793 }
1794
1795 if let Some(prev_message) = thread.messages.get(ix - 1) {
1796 if prev_message.role == Role::Assistant {
1797 break;
1798 }
1799 }
1800 }
1801 }
1802
1803 for message_id in messages_to_remove {
1804 thread.delete_message(message_id, cx);
1805 }
1806 }
1807
1808 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1809 header: "Language model refusal".into(),
1810 message: "Model refused to generate content for safety reasons.".into(),
1811 }));
1812 }
1813 },
1814 Err(error) => {
1815 thread.project.update(cx, |project, cx| {
1816 project.set_agent_location(None, cx);
1817 });
1818
1819 if error.is::<PaymentRequiredError>() {
1820 cx.emit(ThreadEvent::ShowError(ThreadError::PaymentRequired));
1821 } else if let Some(error) =
1822 error.downcast_ref::<ModelRequestLimitReachedError>()
1823 {
1824 cx.emit(ThreadEvent::ShowError(
1825 ThreadError::ModelRequestLimitReached { plan: error.plan },
1826 ));
1827 } else if let Some(known_error) =
1828 error.downcast_ref::<LanguageModelKnownError>()
1829 {
1830 match known_error {
1831 LanguageModelKnownError::ContextWindowLimitExceeded {
1832 tokens,
1833 } => {
1834 thread.exceeded_window_error = Some(ExceededWindowError {
1835 model_id: model.id(),
1836 token_count: *tokens,
1837 });
1838 cx.notify();
1839 }
1840 }
1841 } else {
1842 let error_message = error
1843 .chain()
1844 .map(|err| err.to_string())
1845 .collect::<Vec<_>>()
1846 .join("\n");
1847 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1848 header: "Error interacting with language model".into(),
1849 message: SharedString::from(error_message.clone()),
1850 }));
1851 }
1852
1853 thread.cancel_last_completion(window, cx);
1854 }
1855 }
1856
1857 cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new)));
1858
1859 if let Some((request_callback, (request, response_events))) = thread
1860 .request_callback
1861 .as_mut()
1862 .zip(request_callback_parameters.as_ref())
1863 {
1864 request_callback(request, response_events);
1865 }
1866
1867 thread.auto_capture_telemetry(cx);
1868
1869 if let Ok(initial_usage) = initial_token_usage {
1870 let usage = thread.cumulative_token_usage - initial_usage;
1871
1872 telemetry::event!(
1873 "Assistant Thread Completion",
1874 thread_id = thread.id().to_string(),
1875 prompt_id = prompt_id,
1876 model = model.telemetry_id(),
1877 model_provider = model.provider_id().to_string(),
1878 input_tokens = usage.input_tokens,
1879 output_tokens = usage.output_tokens,
1880 cache_creation_input_tokens = usage.cache_creation_input_tokens,
1881 cache_read_input_tokens = usage.cache_read_input_tokens,
1882 );
1883 }
1884 })
1885 .ok();
1886 });
1887
1888 self.pending_completions.push(PendingCompletion {
1889 id: pending_completion_id,
1890 queue_state: QueueState::Sending,
1891 _task: task,
1892 });
1893 }
1894
1895 pub fn summarize(&mut self, cx: &mut Context<Self>) {
1896 let Some(model) = LanguageModelRegistry::read_global(cx).thread_summary_model() else {
1897 println!("No thread summary model");
1898 return;
1899 };
1900
1901 if !model.provider.is_authenticated(cx) {
1902 return;
1903 }
1904
1905 let added_user_message = include_str!("./prompts/summarize_thread_prompt.txt");
1906
1907 let request = self.to_summarize_request(
1908 &model.model,
1909 CompletionIntent::ThreadSummarization,
1910 added_user_message.into(),
1911 cx,
1912 );
1913
1914 self.summary = ThreadSummary::Generating;
1915
1916 self.pending_summary = cx.spawn(async move |this, cx| {
1917 let result = async {
1918 let mut messages = model.model.stream_completion(request, &cx).await?;
1919
1920 let mut new_summary = String::new();
1921 while let Some(event) = messages.next().await {
1922 let Ok(event) = event else {
1923 continue;
1924 };
1925 let text = match event {
1926 LanguageModelCompletionEvent::Text(text) => text,
1927 LanguageModelCompletionEvent::StatusUpdate(
1928 CompletionRequestStatus::UsageUpdated { amount, limit },
1929 ) => {
1930 this.update(cx, |thread, _cx| {
1931 thread.last_usage = Some(RequestUsage {
1932 limit,
1933 amount: amount as i32,
1934 });
1935 })?;
1936 continue;
1937 }
1938 _ => continue,
1939 };
1940
1941 let mut lines = text.lines();
1942 new_summary.extend(lines.next());
1943
1944 // Stop if the LLM generated multiple lines.
1945 if lines.next().is_some() {
1946 break;
1947 }
1948 }
1949
1950 anyhow::Ok(new_summary)
1951 }
1952 .await;
1953
1954 this.update(cx, |this, cx| {
1955 match result {
1956 Ok(new_summary) => {
1957 if new_summary.is_empty() {
1958 this.summary = ThreadSummary::Error;
1959 } else {
1960 this.summary = ThreadSummary::Ready(new_summary.into());
1961 }
1962 }
1963 Err(err) => {
1964 this.summary = ThreadSummary::Error;
1965 log::error!("Failed to generate thread summary: {}", err);
1966 }
1967 }
1968 cx.emit(ThreadEvent::SummaryGenerated);
1969 })
1970 .log_err()?;
1971
1972 Some(())
1973 });
1974 }
1975
1976 pub fn start_generating_detailed_summary_if_needed(
1977 &mut self,
1978 thread_store: WeakEntity<ThreadStore>,
1979 cx: &mut Context<Self>,
1980 ) {
1981 let Some(last_message_id) = self.messages.last().map(|message| message.id) else {
1982 return;
1983 };
1984
1985 match &*self.detailed_summary_rx.borrow() {
1986 DetailedSummaryState::Generating { message_id, .. }
1987 | DetailedSummaryState::Generated { message_id, .. }
1988 if *message_id == last_message_id =>
1989 {
1990 // Already up-to-date
1991 return;
1992 }
1993 _ => {}
1994 }
1995
1996 let Some(ConfiguredModel { model, provider }) =
1997 LanguageModelRegistry::read_global(cx).thread_summary_model()
1998 else {
1999 return;
2000 };
2001
2002 if !provider.is_authenticated(cx) {
2003 return;
2004 }
2005
2006 let added_user_message = include_str!("./prompts/summarize_thread_detailed_prompt.txt");
2007
2008 let request = self.to_summarize_request(
2009 &model,
2010 CompletionIntent::ThreadContextSummarization,
2011 added_user_message.into(),
2012 cx,
2013 );
2014
2015 *self.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generating {
2016 message_id: last_message_id,
2017 };
2018
2019 // Replace the detailed summarization task if there is one, cancelling it. It would probably
2020 // be better to allow the old task to complete, but this would require logic for choosing
2021 // which result to prefer (the old task could complete after the new one, resulting in a
2022 // stale summary).
2023 self.detailed_summary_task = cx.spawn(async move |thread, cx| {
2024 let stream = model.stream_completion_text(request, &cx);
2025 let Some(mut messages) = stream.await.log_err() else {
2026 thread
2027 .update(cx, |thread, _cx| {
2028 *thread.detailed_summary_tx.borrow_mut() =
2029 DetailedSummaryState::NotGenerated;
2030 })
2031 .ok()?;
2032 return None;
2033 };
2034
2035 let mut new_detailed_summary = String::new();
2036
2037 while let Some(chunk) = messages.stream.next().await {
2038 if let Some(chunk) = chunk.log_err() {
2039 new_detailed_summary.push_str(&chunk);
2040 }
2041 }
2042
2043 thread
2044 .update(cx, |thread, _cx| {
2045 *thread.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generated {
2046 text: new_detailed_summary.into(),
2047 message_id: last_message_id,
2048 };
2049 })
2050 .ok()?;
2051
2052 // Save thread so its summary can be reused later
2053 if let Some(thread) = thread.upgrade() {
2054 if let Ok(Ok(save_task)) = cx.update(|cx| {
2055 thread_store
2056 .update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx))
2057 }) {
2058 save_task.await.log_err();
2059 }
2060 }
2061
2062 Some(())
2063 });
2064 }
2065
2066 pub async fn wait_for_detailed_summary_or_text(
2067 this: &Entity<Self>,
2068 cx: &mut AsyncApp,
2069 ) -> Option<SharedString> {
2070 let mut detailed_summary_rx = this
2071 .read_with(cx, |this, _cx| this.detailed_summary_rx.clone())
2072 .ok()?;
2073 loop {
2074 match detailed_summary_rx.recv().await? {
2075 DetailedSummaryState::Generating { .. } => {}
2076 DetailedSummaryState::NotGenerated => {
2077 return this.read_with(cx, |this, _cx| this.text().into()).ok();
2078 }
2079 DetailedSummaryState::Generated { text, .. } => return Some(text),
2080 }
2081 }
2082 }
2083
2084 pub fn latest_detailed_summary_or_text(&self) -> SharedString {
2085 self.detailed_summary_rx
2086 .borrow()
2087 .text()
2088 .unwrap_or_else(|| self.text().into())
2089 }
2090
2091 pub fn is_generating_detailed_summary(&self) -> bool {
2092 matches!(
2093 &*self.detailed_summary_rx.borrow(),
2094 DetailedSummaryState::Generating { .. }
2095 )
2096 }
2097
2098 pub fn use_pending_tools(
2099 &mut self,
2100 window: Option<AnyWindowHandle>,
2101 cx: &mut Context<Self>,
2102 model: Arc<dyn LanguageModel>,
2103 ) -> Vec<PendingToolUse> {
2104 self.auto_capture_telemetry(cx);
2105 let request =
2106 Arc::new(self.to_completion_request(model.clone(), CompletionIntent::ToolResults, cx));
2107 let pending_tool_uses = self
2108 .tool_use
2109 .pending_tool_uses()
2110 .into_iter()
2111 .filter(|tool_use| tool_use.status.is_idle())
2112 .cloned()
2113 .collect::<Vec<_>>();
2114
2115 for tool_use in pending_tool_uses.iter() {
2116 if let Some(tool) = self.tools.read(cx).tool(&tool_use.name, cx) {
2117 if tool.needs_confirmation(&tool_use.input, cx)
2118 && !AgentSettings::get_global(cx).always_allow_tool_actions
2119 {
2120 self.tool_use.confirm_tool_use(
2121 tool_use.id.clone(),
2122 tool_use.ui_text.clone(),
2123 tool_use.input.clone(),
2124 request.clone(),
2125 tool,
2126 );
2127 cx.emit(ThreadEvent::ToolConfirmationNeeded);
2128 } else {
2129 self.run_tool(
2130 tool_use.id.clone(),
2131 tool_use.ui_text.clone(),
2132 tool_use.input.clone(),
2133 request.clone(),
2134 tool,
2135 model.clone(),
2136 window,
2137 cx,
2138 );
2139 }
2140 } else {
2141 self.handle_hallucinated_tool_use(
2142 tool_use.id.clone(),
2143 tool_use.name.clone(),
2144 window,
2145 cx,
2146 );
2147 }
2148 }
2149
2150 pending_tool_uses
2151 }
2152
2153 pub fn handle_hallucinated_tool_use(
2154 &mut self,
2155 tool_use_id: LanguageModelToolUseId,
2156 hallucinated_tool_name: Arc<str>,
2157 window: Option<AnyWindowHandle>,
2158 cx: &mut Context<Thread>,
2159 ) {
2160 let available_tools = self.profile.enabled_tools(cx);
2161
2162 let tool_list = available_tools
2163 .iter()
2164 .map(|tool| format!("- {}: {}", tool.name(), tool.description()))
2165 .collect::<Vec<_>>()
2166 .join("\n");
2167
2168 let error_message = format!(
2169 "The tool '{}' doesn't exist or is not enabled. Available tools:\n{}",
2170 hallucinated_tool_name, tool_list
2171 );
2172
2173 let pending_tool_use = self.tool_use.insert_tool_output(
2174 tool_use_id.clone(),
2175 hallucinated_tool_name,
2176 Err(anyhow!("Missing tool call: {error_message}")),
2177 self.configured_model.as_ref(),
2178 );
2179
2180 cx.emit(ThreadEvent::MissingToolUse {
2181 tool_use_id: tool_use_id.clone(),
2182 ui_text: error_message.into(),
2183 });
2184
2185 self.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2186 }
2187
2188 pub fn receive_invalid_tool_json(
2189 &mut self,
2190 tool_use_id: LanguageModelToolUseId,
2191 tool_name: Arc<str>,
2192 invalid_json: Arc<str>,
2193 error: String,
2194 window: Option<AnyWindowHandle>,
2195 cx: &mut Context<Thread>,
2196 ) {
2197 log::error!("The model returned invalid input JSON: {invalid_json}");
2198
2199 let pending_tool_use = self.tool_use.insert_tool_output(
2200 tool_use_id.clone(),
2201 tool_name,
2202 Err(anyhow!("Error parsing input JSON: {error}")),
2203 self.configured_model.as_ref(),
2204 );
2205 let ui_text = if let Some(pending_tool_use) = &pending_tool_use {
2206 pending_tool_use.ui_text.clone()
2207 } else {
2208 log::error!(
2209 "There was no pending tool use for tool use {tool_use_id}, even though it finished (with invalid input JSON)."
2210 );
2211 format!("Unknown tool {}", tool_use_id).into()
2212 };
2213
2214 cx.emit(ThreadEvent::InvalidToolInput {
2215 tool_use_id: tool_use_id.clone(),
2216 ui_text,
2217 invalid_input_json: invalid_json,
2218 });
2219
2220 self.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2221 }
2222
2223 pub fn run_tool(
2224 &mut self,
2225 tool_use_id: LanguageModelToolUseId,
2226 ui_text: impl Into<SharedString>,
2227 input: serde_json::Value,
2228 request: Arc<LanguageModelRequest>,
2229 tool: Arc<dyn Tool>,
2230 model: Arc<dyn LanguageModel>,
2231 window: Option<AnyWindowHandle>,
2232 cx: &mut Context<Thread>,
2233 ) {
2234 let task =
2235 self.spawn_tool_use(tool_use_id.clone(), request, input, tool, model, window, cx);
2236 self.tool_use
2237 .run_pending_tool(tool_use_id, ui_text.into(), task);
2238 }
2239
2240 fn spawn_tool_use(
2241 &mut self,
2242 tool_use_id: LanguageModelToolUseId,
2243 request: Arc<LanguageModelRequest>,
2244 input: serde_json::Value,
2245 tool: Arc<dyn Tool>,
2246 model: Arc<dyn LanguageModel>,
2247 window: Option<AnyWindowHandle>,
2248 cx: &mut Context<Thread>,
2249 ) -> Task<()> {
2250 let tool_name: Arc<str> = tool.name().into();
2251
2252 let tool_result = tool.run(
2253 input,
2254 request,
2255 self.project.clone(),
2256 self.action_log.clone(),
2257 model,
2258 window,
2259 cx,
2260 );
2261
2262 // Store the card separately if it exists
2263 if let Some(card) = tool_result.card.clone() {
2264 self.tool_use
2265 .insert_tool_result_card(tool_use_id.clone(), card);
2266 }
2267
2268 cx.spawn({
2269 async move |thread: WeakEntity<Thread>, cx| {
2270 let output = tool_result.output.await;
2271
2272 thread
2273 .update(cx, |thread, cx| {
2274 let pending_tool_use = thread.tool_use.insert_tool_output(
2275 tool_use_id.clone(),
2276 tool_name,
2277 output,
2278 thread.configured_model.as_ref(),
2279 );
2280 thread.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2281 })
2282 .ok();
2283 }
2284 })
2285 }
2286
2287 fn tool_finished(
2288 &mut self,
2289 tool_use_id: LanguageModelToolUseId,
2290 pending_tool_use: Option<PendingToolUse>,
2291 canceled: bool,
2292 window: Option<AnyWindowHandle>,
2293 cx: &mut Context<Self>,
2294 ) {
2295 if self.all_tools_finished() {
2296 if let Some(ConfiguredModel { model, .. }) = self.configured_model.as_ref() {
2297 if !canceled {
2298 self.send_to_model(model.clone(), CompletionIntent::ToolResults, window, cx);
2299 }
2300 self.auto_capture_telemetry(cx);
2301 }
2302 }
2303
2304 cx.emit(ThreadEvent::ToolFinished {
2305 tool_use_id,
2306 pending_tool_use,
2307 });
2308 }
2309
2310 /// Cancels the last pending completion, if there are any pending.
2311 ///
2312 /// Returns whether a completion was canceled.
2313 pub fn cancel_last_completion(
2314 &mut self,
2315 window: Option<AnyWindowHandle>,
2316 cx: &mut Context<Self>,
2317 ) -> bool {
2318 let mut canceled = self.pending_completions.pop().is_some();
2319
2320 for pending_tool_use in self.tool_use.cancel_pending() {
2321 canceled = true;
2322 self.tool_finished(
2323 pending_tool_use.id.clone(),
2324 Some(pending_tool_use),
2325 true,
2326 window,
2327 cx,
2328 );
2329 }
2330
2331 if canceled {
2332 cx.emit(ThreadEvent::CompletionCanceled);
2333
2334 // When canceled, we always want to insert the checkpoint.
2335 // (We skip over finalize_pending_checkpoint, because it
2336 // would conclude we didn't have anything to insert here.)
2337 if let Some(checkpoint) = self.pending_checkpoint.take() {
2338 self.insert_checkpoint(checkpoint, cx);
2339 }
2340 } else {
2341 self.finalize_pending_checkpoint(cx);
2342 }
2343
2344 canceled
2345 }
2346
2347 /// Signals that any in-progress editing should be canceled.
2348 ///
2349 /// This method is used to notify listeners (like ActiveThread) that
2350 /// they should cancel any editing operations.
2351 pub fn cancel_editing(&mut self, cx: &mut Context<Self>) {
2352 cx.emit(ThreadEvent::CancelEditing);
2353 }
2354
2355 pub fn feedback(&self) -> Option<ThreadFeedback> {
2356 self.feedback
2357 }
2358
2359 pub fn message_feedback(&self, message_id: MessageId) -> Option<ThreadFeedback> {
2360 self.message_feedback.get(&message_id).copied()
2361 }
2362
2363 pub fn report_message_feedback(
2364 &mut self,
2365 message_id: MessageId,
2366 feedback: ThreadFeedback,
2367 cx: &mut Context<Self>,
2368 ) -> Task<Result<()>> {
2369 if self.message_feedback.get(&message_id) == Some(&feedback) {
2370 return Task::ready(Ok(()));
2371 }
2372
2373 let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
2374 let serialized_thread = self.serialize(cx);
2375 let thread_id = self.id().clone();
2376 let client = self.project.read(cx).client();
2377
2378 let enabled_tool_names: Vec<String> = self
2379 .profile
2380 .enabled_tools(cx)
2381 .iter()
2382 .map(|tool| tool.name())
2383 .collect();
2384
2385 self.message_feedback.insert(message_id, feedback);
2386
2387 cx.notify();
2388
2389 let message_content = self
2390 .message(message_id)
2391 .map(|msg| msg.to_string())
2392 .unwrap_or_default();
2393
2394 cx.background_spawn(async move {
2395 let final_project_snapshot = final_project_snapshot.await;
2396 let serialized_thread = serialized_thread.await?;
2397 let thread_data =
2398 serde_json::to_value(serialized_thread).unwrap_or_else(|_| serde_json::Value::Null);
2399
2400 let rating = match feedback {
2401 ThreadFeedback::Positive => "positive",
2402 ThreadFeedback::Negative => "negative",
2403 };
2404 telemetry::event!(
2405 "Assistant Thread Rated",
2406 rating,
2407 thread_id,
2408 enabled_tool_names,
2409 message_id = message_id.0,
2410 message_content,
2411 thread_data,
2412 final_project_snapshot
2413 );
2414 client.telemetry().flush_events().await;
2415
2416 Ok(())
2417 })
2418 }
2419
2420 pub fn report_feedback(
2421 &mut self,
2422 feedback: ThreadFeedback,
2423 cx: &mut Context<Self>,
2424 ) -> Task<Result<()>> {
2425 let last_assistant_message_id = self
2426 .messages
2427 .iter()
2428 .rev()
2429 .find(|msg| msg.role == Role::Assistant)
2430 .map(|msg| msg.id);
2431
2432 if let Some(message_id) = last_assistant_message_id {
2433 self.report_message_feedback(message_id, feedback, cx)
2434 } else {
2435 let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
2436 let serialized_thread = self.serialize(cx);
2437 let thread_id = self.id().clone();
2438 let client = self.project.read(cx).client();
2439 self.feedback = Some(feedback);
2440 cx.notify();
2441
2442 cx.background_spawn(async move {
2443 let final_project_snapshot = final_project_snapshot.await;
2444 let serialized_thread = serialized_thread.await?;
2445 let thread_data = serde_json::to_value(serialized_thread)
2446 .unwrap_or_else(|_| serde_json::Value::Null);
2447
2448 let rating = match feedback {
2449 ThreadFeedback::Positive => "positive",
2450 ThreadFeedback::Negative => "negative",
2451 };
2452 telemetry::event!(
2453 "Assistant Thread Rated",
2454 rating,
2455 thread_id,
2456 thread_data,
2457 final_project_snapshot
2458 );
2459 client.telemetry().flush_events().await;
2460
2461 Ok(())
2462 })
2463 }
2464 }
2465
2466 /// Create a snapshot of the current project state including git information and unsaved buffers.
2467 fn project_snapshot(
2468 project: Entity<Project>,
2469 cx: &mut Context<Self>,
2470 ) -> Task<Arc<ProjectSnapshot>> {
2471 let git_store = project.read(cx).git_store().clone();
2472 let worktree_snapshots: Vec<_> = project
2473 .read(cx)
2474 .visible_worktrees(cx)
2475 .map(|worktree| Self::worktree_snapshot(worktree, git_store.clone(), cx))
2476 .collect();
2477
2478 cx.spawn(async move |_, cx| {
2479 let worktree_snapshots = futures::future::join_all(worktree_snapshots).await;
2480
2481 let mut unsaved_buffers = Vec::new();
2482 cx.update(|app_cx| {
2483 let buffer_store = project.read(app_cx).buffer_store();
2484 for buffer_handle in buffer_store.read(app_cx).buffers() {
2485 let buffer = buffer_handle.read(app_cx);
2486 if buffer.is_dirty() {
2487 if let Some(file) = buffer.file() {
2488 let path = file.path().to_string_lossy().to_string();
2489 unsaved_buffers.push(path);
2490 }
2491 }
2492 }
2493 })
2494 .ok();
2495
2496 Arc::new(ProjectSnapshot {
2497 worktree_snapshots,
2498 unsaved_buffer_paths: unsaved_buffers,
2499 timestamp: Utc::now(),
2500 })
2501 })
2502 }
2503
2504 fn worktree_snapshot(
2505 worktree: Entity<project::Worktree>,
2506 git_store: Entity<GitStore>,
2507 cx: &App,
2508 ) -> Task<WorktreeSnapshot> {
2509 cx.spawn(async move |cx| {
2510 // Get worktree path and snapshot
2511 let worktree_info = cx.update(|app_cx| {
2512 let worktree = worktree.read(app_cx);
2513 let path = worktree.abs_path().to_string_lossy().to_string();
2514 let snapshot = worktree.snapshot();
2515 (path, snapshot)
2516 });
2517
2518 let Ok((worktree_path, _snapshot)) = worktree_info else {
2519 return WorktreeSnapshot {
2520 worktree_path: String::new(),
2521 git_state: None,
2522 };
2523 };
2524
2525 let git_state = git_store
2526 .update(cx, |git_store, cx| {
2527 git_store
2528 .repositories()
2529 .values()
2530 .find(|repo| {
2531 repo.read(cx)
2532 .abs_path_to_repo_path(&worktree.read(cx).abs_path())
2533 .is_some()
2534 })
2535 .cloned()
2536 })
2537 .ok()
2538 .flatten()
2539 .map(|repo| {
2540 repo.update(cx, |repo, _| {
2541 let current_branch =
2542 repo.branch.as_ref().map(|branch| branch.name().to_owned());
2543 repo.send_job(None, |state, _| async move {
2544 let RepositoryState::Local { backend, .. } = state else {
2545 return GitState {
2546 remote_url: None,
2547 head_sha: None,
2548 current_branch,
2549 diff: None,
2550 };
2551 };
2552
2553 let remote_url = backend.remote_url("origin");
2554 let head_sha = backend.head_sha().await;
2555 let diff = backend.diff(DiffType::HeadToWorktree).await.ok();
2556
2557 GitState {
2558 remote_url,
2559 head_sha,
2560 current_branch,
2561 diff,
2562 }
2563 })
2564 })
2565 });
2566
2567 let git_state = match git_state {
2568 Some(git_state) => match git_state.ok() {
2569 Some(git_state) => git_state.await.ok(),
2570 None => None,
2571 },
2572 None => None,
2573 };
2574
2575 WorktreeSnapshot {
2576 worktree_path,
2577 git_state,
2578 }
2579 })
2580 }
2581
2582 pub fn to_markdown(&self, cx: &App) -> Result<String> {
2583 let mut markdown = Vec::new();
2584
2585 let summary = self.summary().or_default();
2586 writeln!(markdown, "# {summary}\n")?;
2587
2588 for message in self.messages() {
2589 writeln!(
2590 markdown,
2591 "## {role}\n",
2592 role = match message.role {
2593 Role::User => "User",
2594 Role::Assistant => "Agent",
2595 Role::System => "System",
2596 }
2597 )?;
2598
2599 if !message.loaded_context.text.is_empty() {
2600 writeln!(markdown, "{}", message.loaded_context.text)?;
2601 }
2602
2603 if !message.loaded_context.images.is_empty() {
2604 writeln!(
2605 markdown,
2606 "\n{} images attached as context.\n",
2607 message.loaded_context.images.len()
2608 )?;
2609 }
2610
2611 for segment in &message.segments {
2612 match segment {
2613 MessageSegment::Text(text) => writeln!(markdown, "{}\n", text)?,
2614 MessageSegment::Thinking { text, .. } => {
2615 writeln!(markdown, "<think>\n{}\n</think>\n", text)?
2616 }
2617 MessageSegment::RedactedThinking(_) => {}
2618 }
2619 }
2620
2621 for tool_use in self.tool_uses_for_message(message.id, cx) {
2622 writeln!(
2623 markdown,
2624 "**Use Tool: {} ({})**",
2625 tool_use.name, tool_use.id
2626 )?;
2627 writeln!(markdown, "```json")?;
2628 writeln!(
2629 markdown,
2630 "{}",
2631 serde_json::to_string_pretty(&tool_use.input)?
2632 )?;
2633 writeln!(markdown, "```")?;
2634 }
2635
2636 for tool_result in self.tool_results_for_message(message.id) {
2637 write!(markdown, "\n**Tool Results: {}", tool_result.tool_use_id)?;
2638 if tool_result.is_error {
2639 write!(markdown, " (Error)")?;
2640 }
2641
2642 writeln!(markdown, "**\n")?;
2643 match &tool_result.content {
2644 LanguageModelToolResultContent::Text(text) => {
2645 writeln!(markdown, "{text}")?;
2646 }
2647 LanguageModelToolResultContent::Image(image) => {
2648 writeln!(markdown, "", image.source)?;
2649 }
2650 }
2651
2652 if let Some(output) = tool_result.output.as_ref() {
2653 writeln!(
2654 markdown,
2655 "\n\nDebug Output:\n\n```json\n{}\n```\n",
2656 serde_json::to_string_pretty(output)?
2657 )?;
2658 }
2659 }
2660 }
2661
2662 Ok(String::from_utf8_lossy(&markdown).to_string())
2663 }
2664
2665 pub fn keep_edits_in_range(
2666 &mut self,
2667 buffer: Entity<language::Buffer>,
2668 buffer_range: Range<language::Anchor>,
2669 cx: &mut Context<Self>,
2670 ) {
2671 self.action_log.update(cx, |action_log, cx| {
2672 action_log.keep_edits_in_range(buffer, buffer_range, cx)
2673 });
2674 }
2675
2676 pub fn keep_all_edits(&mut self, cx: &mut Context<Self>) {
2677 self.action_log
2678 .update(cx, |action_log, cx| action_log.keep_all_edits(cx));
2679 }
2680
2681 pub fn reject_edits_in_ranges(
2682 &mut self,
2683 buffer: Entity<language::Buffer>,
2684 buffer_ranges: Vec<Range<language::Anchor>>,
2685 cx: &mut Context<Self>,
2686 ) -> Task<Result<()>> {
2687 self.action_log.update(cx, |action_log, cx| {
2688 action_log.reject_edits_in_ranges(buffer, buffer_ranges, cx)
2689 })
2690 }
2691
2692 pub fn action_log(&self) -> &Entity<ActionLog> {
2693 &self.action_log
2694 }
2695
2696 pub fn project(&self) -> &Entity<Project> {
2697 &self.project
2698 }
2699
2700 pub fn auto_capture_telemetry(&mut self, cx: &mut Context<Self>) {
2701 if !cx.has_flag::<feature_flags::ThreadAutoCaptureFeatureFlag>() {
2702 return;
2703 }
2704
2705 let now = Instant::now();
2706 if let Some(last) = self.last_auto_capture_at {
2707 if now.duration_since(last).as_secs() < 10 {
2708 return;
2709 }
2710 }
2711
2712 self.last_auto_capture_at = Some(now);
2713
2714 let thread_id = self.id().clone();
2715 let github_login = self
2716 .project
2717 .read(cx)
2718 .user_store()
2719 .read(cx)
2720 .current_user()
2721 .map(|user| user.github_login.clone());
2722 let client = self.project.read(cx).client();
2723 let serialize_task = self.serialize(cx);
2724
2725 cx.background_executor()
2726 .spawn(async move {
2727 if let Ok(serialized_thread) = serialize_task.await {
2728 if let Ok(thread_data) = serde_json::to_value(serialized_thread) {
2729 telemetry::event!(
2730 "Agent Thread Auto-Captured",
2731 thread_id = thread_id.to_string(),
2732 thread_data = thread_data,
2733 auto_capture_reason = "tracked_user",
2734 github_login = github_login
2735 );
2736
2737 client.telemetry().flush_events().await;
2738 }
2739 }
2740 })
2741 .detach();
2742 }
2743
2744 pub fn cumulative_token_usage(&self) -> TokenUsage {
2745 self.cumulative_token_usage
2746 }
2747
2748 pub fn token_usage_up_to_message(&self, message_id: MessageId) -> TotalTokenUsage {
2749 let Some(model) = self.configured_model.as_ref() else {
2750 return TotalTokenUsage::default();
2751 };
2752
2753 let max = model.model.max_token_count();
2754
2755 let index = self
2756 .messages
2757 .iter()
2758 .position(|msg| msg.id == message_id)
2759 .unwrap_or(0);
2760
2761 if index == 0 {
2762 return TotalTokenUsage { total: 0, max };
2763 }
2764
2765 let token_usage = &self
2766 .request_token_usage
2767 .get(index - 1)
2768 .cloned()
2769 .unwrap_or_default();
2770
2771 TotalTokenUsage {
2772 total: token_usage.total_tokens(),
2773 max,
2774 }
2775 }
2776
2777 pub fn total_token_usage(&self) -> Option<TotalTokenUsage> {
2778 let model = self.configured_model.as_ref()?;
2779
2780 let max = model.model.max_token_count();
2781
2782 if let Some(exceeded_error) = &self.exceeded_window_error {
2783 if model.model.id() == exceeded_error.model_id {
2784 return Some(TotalTokenUsage {
2785 total: exceeded_error.token_count,
2786 max,
2787 });
2788 }
2789 }
2790
2791 let total = self
2792 .token_usage_at_last_message()
2793 .unwrap_or_default()
2794 .total_tokens();
2795
2796 Some(TotalTokenUsage { total, max })
2797 }
2798
2799 fn token_usage_at_last_message(&self) -> Option<TokenUsage> {
2800 self.request_token_usage
2801 .get(self.messages.len().saturating_sub(1))
2802 .or_else(|| self.request_token_usage.last())
2803 .cloned()
2804 }
2805
2806 fn update_token_usage_at_last_message(&mut self, token_usage: TokenUsage) {
2807 let placeholder = self.token_usage_at_last_message().unwrap_or_default();
2808 self.request_token_usage
2809 .resize(self.messages.len(), placeholder);
2810
2811 if let Some(last) = self.request_token_usage.last_mut() {
2812 *last = token_usage;
2813 }
2814 }
2815
2816 pub fn deny_tool_use(
2817 &mut self,
2818 tool_use_id: LanguageModelToolUseId,
2819 tool_name: Arc<str>,
2820 window: Option<AnyWindowHandle>,
2821 cx: &mut Context<Self>,
2822 ) {
2823 let err = Err(anyhow::anyhow!(
2824 "Permission to run tool action denied by user"
2825 ));
2826
2827 self.tool_use.insert_tool_output(
2828 tool_use_id.clone(),
2829 tool_name,
2830 err,
2831 self.configured_model.as_ref(),
2832 );
2833 self.tool_finished(tool_use_id.clone(), None, true, window, cx);
2834 }
2835}
2836
2837#[derive(Debug, Clone, Error)]
2838pub enum ThreadError {
2839 #[error("Payment required")]
2840 PaymentRequired,
2841 #[error("Model request limit reached")]
2842 ModelRequestLimitReached { plan: Plan },
2843 #[error("Message {header}: {message}")]
2844 Message {
2845 header: SharedString,
2846 message: SharedString,
2847 },
2848}
2849
2850#[derive(Debug, Clone)]
2851pub enum ThreadEvent {
2852 ShowError(ThreadError),
2853 StreamedCompletion,
2854 ReceivedTextChunk,
2855 NewRequest,
2856 StreamedAssistantText(MessageId, String),
2857 StreamedAssistantThinking(MessageId, String),
2858 StreamedToolUse {
2859 tool_use_id: LanguageModelToolUseId,
2860 ui_text: Arc<str>,
2861 input: serde_json::Value,
2862 },
2863 MissingToolUse {
2864 tool_use_id: LanguageModelToolUseId,
2865 ui_text: Arc<str>,
2866 },
2867 InvalidToolInput {
2868 tool_use_id: LanguageModelToolUseId,
2869 ui_text: Arc<str>,
2870 invalid_input_json: Arc<str>,
2871 },
2872 Stopped(Result<StopReason, Arc<anyhow::Error>>),
2873 MessageAdded(MessageId),
2874 MessageEdited(MessageId),
2875 MessageDeleted(MessageId),
2876 SummaryGenerated,
2877 SummaryChanged,
2878 UsePendingTools {
2879 tool_uses: Vec<PendingToolUse>,
2880 },
2881 ToolFinished {
2882 #[allow(unused)]
2883 tool_use_id: LanguageModelToolUseId,
2884 /// The pending tool use that corresponds to this tool.
2885 pending_tool_use: Option<PendingToolUse>,
2886 },
2887 CheckpointChanged,
2888 ToolConfirmationNeeded,
2889 ToolUseLimitReached,
2890 CancelEditing,
2891 CompletionCanceled,
2892 ProfileChanged,
2893}
2894
2895impl EventEmitter<ThreadEvent> for Thread {}
2896
2897struct PendingCompletion {
2898 id: usize,
2899 queue_state: QueueState,
2900 _task: Task<()>,
2901}
2902
2903#[cfg(test)]
2904mod tests {
2905 use super::*;
2906 use crate::{ThreadStore, context::load_context, context_store::ContextStore, thread_store};
2907 use agent_settings::{AgentProfileId, AgentSettings, LanguageModelParameters};
2908 use assistant_tool::ToolRegistry;
2909 use editor::EditorSettings;
2910 use gpui::TestAppContext;
2911 use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider};
2912 use project::{FakeFs, Project};
2913 use prompt_store::PromptBuilder;
2914 use serde_json::json;
2915 use settings::{Settings, SettingsStore};
2916 use std::sync::Arc;
2917 use theme::ThemeSettings;
2918 use util::path;
2919 use workspace::Workspace;
2920
2921 #[gpui::test]
2922 async fn test_message_with_context(cx: &mut TestAppContext) {
2923 init_test_settings(cx);
2924
2925 let project = create_test_project(
2926 cx,
2927 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
2928 )
2929 .await;
2930
2931 let (_workspace, _thread_store, thread, context_store, model) =
2932 setup_test_environment(cx, project.clone()).await;
2933
2934 add_file_to_context(&project, &context_store, "test/code.rs", cx)
2935 .await
2936 .unwrap();
2937
2938 let context =
2939 context_store.read_with(cx, |store, _| store.context().next().cloned().unwrap());
2940 let loaded_context = cx
2941 .update(|cx| load_context(vec![context], &project, &None, cx))
2942 .await;
2943
2944 // Insert user message with context
2945 let message_id = thread.update(cx, |thread, cx| {
2946 thread.insert_user_message(
2947 "Please explain this code",
2948 loaded_context,
2949 None,
2950 Vec::new(),
2951 cx,
2952 )
2953 });
2954
2955 // Check content and context in message object
2956 let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
2957
2958 // Use different path format strings based on platform for the test
2959 #[cfg(windows)]
2960 let path_part = r"test\code.rs";
2961 #[cfg(not(windows))]
2962 let path_part = "test/code.rs";
2963
2964 let expected_context = format!(
2965 r#"
2966<context>
2967The following items were attached by the user. They are up-to-date and don't need to be re-read.
2968
2969<files>
2970```rs {path_part}
2971fn main() {{
2972 println!("Hello, world!");
2973}}
2974```
2975</files>
2976</context>
2977"#
2978 );
2979
2980 assert_eq!(message.role, Role::User);
2981 assert_eq!(message.segments.len(), 1);
2982 assert_eq!(
2983 message.segments[0],
2984 MessageSegment::Text("Please explain this code".to_string())
2985 );
2986 assert_eq!(message.loaded_context.text, expected_context);
2987
2988 // Check message in request
2989 let request = thread.update(cx, |thread, cx| {
2990 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
2991 });
2992
2993 assert_eq!(request.messages.len(), 2);
2994 let expected_full_message = format!("{}Please explain this code", expected_context);
2995 assert_eq!(request.messages[1].string_contents(), expected_full_message);
2996 }
2997
2998 #[gpui::test]
2999 async fn test_only_include_new_contexts(cx: &mut TestAppContext) {
3000 init_test_settings(cx);
3001
3002 let project = create_test_project(
3003 cx,
3004 json!({
3005 "file1.rs": "fn function1() {}\n",
3006 "file2.rs": "fn function2() {}\n",
3007 "file3.rs": "fn function3() {}\n",
3008 "file4.rs": "fn function4() {}\n",
3009 }),
3010 )
3011 .await;
3012
3013 let (_, _thread_store, thread, context_store, model) =
3014 setup_test_environment(cx, project.clone()).await;
3015
3016 // First message with context 1
3017 add_file_to_context(&project, &context_store, "test/file1.rs", cx)
3018 .await
3019 .unwrap();
3020 let new_contexts = context_store.update(cx, |store, cx| {
3021 store.new_context_for_thread(thread.read(cx), None)
3022 });
3023 assert_eq!(new_contexts.len(), 1);
3024 let loaded_context = cx
3025 .update(|cx| load_context(new_contexts, &project, &None, cx))
3026 .await;
3027 let message1_id = thread.update(cx, |thread, cx| {
3028 thread.insert_user_message("Message 1", loaded_context, None, Vec::new(), cx)
3029 });
3030
3031 // Second message with contexts 1 and 2 (context 1 should be skipped as it's already included)
3032 add_file_to_context(&project, &context_store, "test/file2.rs", cx)
3033 .await
3034 .unwrap();
3035 let new_contexts = context_store.update(cx, |store, cx| {
3036 store.new_context_for_thread(thread.read(cx), None)
3037 });
3038 assert_eq!(new_contexts.len(), 1);
3039 let loaded_context = cx
3040 .update(|cx| load_context(new_contexts, &project, &None, cx))
3041 .await;
3042 let message2_id = thread.update(cx, |thread, cx| {
3043 thread.insert_user_message("Message 2", loaded_context, None, Vec::new(), cx)
3044 });
3045
3046 // Third message with all three contexts (contexts 1 and 2 should be skipped)
3047 //
3048 add_file_to_context(&project, &context_store, "test/file3.rs", cx)
3049 .await
3050 .unwrap();
3051 let new_contexts = context_store.update(cx, |store, cx| {
3052 store.new_context_for_thread(thread.read(cx), None)
3053 });
3054 assert_eq!(new_contexts.len(), 1);
3055 let loaded_context = cx
3056 .update(|cx| load_context(new_contexts, &project, &None, cx))
3057 .await;
3058 let message3_id = thread.update(cx, |thread, cx| {
3059 thread.insert_user_message("Message 3", loaded_context, None, Vec::new(), cx)
3060 });
3061
3062 // Check what contexts are included in each message
3063 let (message1, message2, message3) = thread.read_with(cx, |thread, _| {
3064 (
3065 thread.message(message1_id).unwrap().clone(),
3066 thread.message(message2_id).unwrap().clone(),
3067 thread.message(message3_id).unwrap().clone(),
3068 )
3069 });
3070
3071 // First message should include context 1
3072 assert!(message1.loaded_context.text.contains("file1.rs"));
3073
3074 // Second message should include only context 2 (not 1)
3075 assert!(!message2.loaded_context.text.contains("file1.rs"));
3076 assert!(message2.loaded_context.text.contains("file2.rs"));
3077
3078 // Third message should include only context 3 (not 1 or 2)
3079 assert!(!message3.loaded_context.text.contains("file1.rs"));
3080 assert!(!message3.loaded_context.text.contains("file2.rs"));
3081 assert!(message3.loaded_context.text.contains("file3.rs"));
3082
3083 // Check entire request to make sure all contexts are properly included
3084 let request = thread.update(cx, |thread, cx| {
3085 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3086 });
3087
3088 // The request should contain all 3 messages
3089 assert_eq!(request.messages.len(), 4);
3090
3091 // Check that the contexts are properly formatted in each message
3092 assert!(request.messages[1].string_contents().contains("file1.rs"));
3093 assert!(!request.messages[1].string_contents().contains("file2.rs"));
3094 assert!(!request.messages[1].string_contents().contains("file3.rs"));
3095
3096 assert!(!request.messages[2].string_contents().contains("file1.rs"));
3097 assert!(request.messages[2].string_contents().contains("file2.rs"));
3098 assert!(!request.messages[2].string_contents().contains("file3.rs"));
3099
3100 assert!(!request.messages[3].string_contents().contains("file1.rs"));
3101 assert!(!request.messages[3].string_contents().contains("file2.rs"));
3102 assert!(request.messages[3].string_contents().contains("file3.rs"));
3103
3104 add_file_to_context(&project, &context_store, "test/file4.rs", cx)
3105 .await
3106 .unwrap();
3107 let new_contexts = context_store.update(cx, |store, cx| {
3108 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3109 });
3110 assert_eq!(new_contexts.len(), 3);
3111 let loaded_context = cx
3112 .update(|cx| load_context(new_contexts, &project, &None, cx))
3113 .await
3114 .loaded_context;
3115
3116 assert!(!loaded_context.text.contains("file1.rs"));
3117 assert!(loaded_context.text.contains("file2.rs"));
3118 assert!(loaded_context.text.contains("file3.rs"));
3119 assert!(loaded_context.text.contains("file4.rs"));
3120
3121 let new_contexts = context_store.update(cx, |store, cx| {
3122 // Remove file4.rs
3123 store.remove_context(&loaded_context.contexts[2].handle(), cx);
3124 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3125 });
3126 assert_eq!(new_contexts.len(), 2);
3127 let loaded_context = cx
3128 .update(|cx| load_context(new_contexts, &project, &None, cx))
3129 .await
3130 .loaded_context;
3131
3132 assert!(!loaded_context.text.contains("file1.rs"));
3133 assert!(loaded_context.text.contains("file2.rs"));
3134 assert!(loaded_context.text.contains("file3.rs"));
3135 assert!(!loaded_context.text.contains("file4.rs"));
3136
3137 let new_contexts = context_store.update(cx, |store, cx| {
3138 // Remove file3.rs
3139 store.remove_context(&loaded_context.contexts[1].handle(), cx);
3140 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3141 });
3142 assert_eq!(new_contexts.len(), 1);
3143 let loaded_context = cx
3144 .update(|cx| load_context(new_contexts, &project, &None, cx))
3145 .await
3146 .loaded_context;
3147
3148 assert!(!loaded_context.text.contains("file1.rs"));
3149 assert!(loaded_context.text.contains("file2.rs"));
3150 assert!(!loaded_context.text.contains("file3.rs"));
3151 assert!(!loaded_context.text.contains("file4.rs"));
3152 }
3153
3154 #[gpui::test]
3155 async fn test_message_without_files(cx: &mut TestAppContext) {
3156 init_test_settings(cx);
3157
3158 let project = create_test_project(
3159 cx,
3160 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3161 )
3162 .await;
3163
3164 let (_, _thread_store, thread, _context_store, model) =
3165 setup_test_environment(cx, project.clone()).await;
3166
3167 // Insert user message without any context (empty context vector)
3168 let message_id = thread.update(cx, |thread, cx| {
3169 thread.insert_user_message(
3170 "What is the best way to learn Rust?",
3171 ContextLoadResult::default(),
3172 None,
3173 Vec::new(),
3174 cx,
3175 )
3176 });
3177
3178 // Check content and context in message object
3179 let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
3180
3181 // Context should be empty when no files are included
3182 assert_eq!(message.role, Role::User);
3183 assert_eq!(message.segments.len(), 1);
3184 assert_eq!(
3185 message.segments[0],
3186 MessageSegment::Text("What is the best way to learn Rust?".to_string())
3187 );
3188 assert_eq!(message.loaded_context.text, "");
3189
3190 // Check message in request
3191 let request = thread.update(cx, |thread, cx| {
3192 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3193 });
3194
3195 assert_eq!(request.messages.len(), 2);
3196 assert_eq!(
3197 request.messages[1].string_contents(),
3198 "What is the best way to learn Rust?"
3199 );
3200
3201 // Add second message, also without context
3202 let message2_id = thread.update(cx, |thread, cx| {
3203 thread.insert_user_message(
3204 "Are there any good books?",
3205 ContextLoadResult::default(),
3206 None,
3207 Vec::new(),
3208 cx,
3209 )
3210 });
3211
3212 let message2 =
3213 thread.read_with(cx, |thread, _| thread.message(message2_id).unwrap().clone());
3214 assert_eq!(message2.loaded_context.text, "");
3215
3216 // Check that both messages appear in the request
3217 let request = thread.update(cx, |thread, cx| {
3218 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3219 });
3220
3221 assert_eq!(request.messages.len(), 3);
3222 assert_eq!(
3223 request.messages[1].string_contents(),
3224 "What is the best way to learn Rust?"
3225 );
3226 assert_eq!(
3227 request.messages[2].string_contents(),
3228 "Are there any good books?"
3229 );
3230 }
3231
3232 #[gpui::test]
3233 async fn test_stale_buffer_notification(cx: &mut TestAppContext) {
3234 init_test_settings(cx);
3235
3236 let project = create_test_project(
3237 cx,
3238 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3239 )
3240 .await;
3241
3242 let (_workspace, _thread_store, thread, context_store, model) =
3243 setup_test_environment(cx, project.clone()).await;
3244
3245 // Open buffer and add it to context
3246 let buffer = add_file_to_context(&project, &context_store, "test/code.rs", cx)
3247 .await
3248 .unwrap();
3249
3250 let context =
3251 context_store.read_with(cx, |store, _| store.context().next().cloned().unwrap());
3252 let loaded_context = cx
3253 .update(|cx| load_context(vec![context], &project, &None, cx))
3254 .await;
3255
3256 // Insert user message with the buffer as context
3257 thread.update(cx, |thread, cx| {
3258 thread.insert_user_message("Explain this code", loaded_context, None, Vec::new(), cx)
3259 });
3260
3261 // Create a request and check that it doesn't have a stale buffer warning yet
3262 let initial_request = thread.update(cx, |thread, cx| {
3263 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3264 });
3265
3266 // Make sure we don't have a stale file warning yet
3267 let has_stale_warning = initial_request.messages.iter().any(|msg| {
3268 msg.string_contents()
3269 .contains("These files changed since last read:")
3270 });
3271 assert!(
3272 !has_stale_warning,
3273 "Should not have stale buffer warning before buffer is modified"
3274 );
3275
3276 // Modify the buffer
3277 buffer.update(cx, |buffer, cx| {
3278 // Find a position at the end of line 1
3279 buffer.edit(
3280 [(1..1, "\n println!(\"Added a new line\");\n")],
3281 None,
3282 cx,
3283 );
3284 });
3285
3286 // Insert another user message without context
3287 thread.update(cx, |thread, cx| {
3288 thread.insert_user_message(
3289 "What does the code do now?",
3290 ContextLoadResult::default(),
3291 None,
3292 Vec::new(),
3293 cx,
3294 )
3295 });
3296
3297 // Create a new request and check for the stale buffer warning
3298 let new_request = thread.update(cx, |thread, cx| {
3299 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3300 });
3301
3302 // We should have a stale file warning as the last message
3303 let last_message = new_request
3304 .messages
3305 .last()
3306 .expect("Request should have messages");
3307
3308 // The last message should be the stale buffer notification
3309 assert_eq!(last_message.role, Role::User);
3310
3311 // Check the exact content of the message
3312 let expected_content = "[The following is an auto-generated notification; do not reply]
3313
3314These files have changed since the last read:
3315- code.rs
3316";
3317 assert_eq!(
3318 last_message.string_contents(),
3319 expected_content,
3320 "Last message should be exactly the stale buffer notification"
3321 );
3322
3323 // The message before the notification should be cached
3324 let index = new_request.messages.len() - 2;
3325 let previous_message = new_request.messages.get(index).unwrap();
3326 assert!(
3327 previous_message.cache,
3328 "Message before the stale buffer notification should be cached"
3329 );
3330 }
3331
3332 #[gpui::test]
3333 async fn test_storing_profile_setting_per_thread(cx: &mut TestAppContext) {
3334 init_test_settings(cx);
3335
3336 let project = create_test_project(
3337 cx,
3338 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3339 )
3340 .await;
3341
3342 let (_workspace, thread_store, thread, _context_store, _model) =
3343 setup_test_environment(cx, project.clone()).await;
3344
3345 // Check that we are starting with the default profile
3346 let profile = cx.read(|cx| thread.read(cx).profile.clone());
3347 let tool_set = cx.read(|cx| thread_store.read(cx).tools());
3348 assert_eq!(
3349 profile,
3350 AgentProfile::new(AgentProfileId::default(), tool_set)
3351 );
3352 }
3353
3354 #[gpui::test]
3355 async fn test_serializing_thread_profile(cx: &mut TestAppContext) {
3356 init_test_settings(cx);
3357
3358 let project = create_test_project(
3359 cx,
3360 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3361 )
3362 .await;
3363
3364 let (_workspace, thread_store, thread, _context_store, _model) =
3365 setup_test_environment(cx, project.clone()).await;
3366
3367 // Profile gets serialized with default values
3368 let serialized = thread
3369 .update(cx, |thread, cx| thread.serialize(cx))
3370 .await
3371 .unwrap();
3372
3373 assert_eq!(serialized.profile, Some(AgentProfileId::default()));
3374
3375 let deserialized = cx.update(|cx| {
3376 thread.update(cx, |thread, cx| {
3377 Thread::deserialize(
3378 thread.id.clone(),
3379 serialized,
3380 thread.project.clone(),
3381 thread.tools.clone(),
3382 thread.prompt_builder.clone(),
3383 thread.project_context.clone(),
3384 None,
3385 cx,
3386 )
3387 })
3388 });
3389 let tool_set = cx.read(|cx| thread_store.read(cx).tools());
3390
3391 assert_eq!(
3392 deserialized.profile,
3393 AgentProfile::new(AgentProfileId::default(), tool_set)
3394 );
3395 }
3396
3397 #[gpui::test]
3398 async fn test_temperature_setting(cx: &mut TestAppContext) {
3399 init_test_settings(cx);
3400
3401 let project = create_test_project(
3402 cx,
3403 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3404 )
3405 .await;
3406
3407 let (_workspace, _thread_store, thread, _context_store, model) =
3408 setup_test_environment(cx, project.clone()).await;
3409
3410 // Both model and provider
3411 cx.update(|cx| {
3412 AgentSettings::override_global(
3413 AgentSettings {
3414 model_parameters: vec![LanguageModelParameters {
3415 provider: Some(model.provider_id().0.to_string().into()),
3416 model: Some(model.id().0.clone()),
3417 temperature: Some(0.66),
3418 }],
3419 ..AgentSettings::get_global(cx).clone()
3420 },
3421 cx,
3422 );
3423 });
3424
3425 let request = thread.update(cx, |thread, cx| {
3426 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3427 });
3428 assert_eq!(request.temperature, Some(0.66));
3429
3430 // Only model
3431 cx.update(|cx| {
3432 AgentSettings::override_global(
3433 AgentSettings {
3434 model_parameters: vec![LanguageModelParameters {
3435 provider: None,
3436 model: Some(model.id().0.clone()),
3437 temperature: Some(0.66),
3438 }],
3439 ..AgentSettings::get_global(cx).clone()
3440 },
3441 cx,
3442 );
3443 });
3444
3445 let request = thread.update(cx, |thread, cx| {
3446 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3447 });
3448 assert_eq!(request.temperature, Some(0.66));
3449
3450 // Only provider
3451 cx.update(|cx| {
3452 AgentSettings::override_global(
3453 AgentSettings {
3454 model_parameters: vec![LanguageModelParameters {
3455 provider: Some(model.provider_id().0.to_string().into()),
3456 model: None,
3457 temperature: Some(0.66),
3458 }],
3459 ..AgentSettings::get_global(cx).clone()
3460 },
3461 cx,
3462 );
3463 });
3464
3465 let request = thread.update(cx, |thread, cx| {
3466 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3467 });
3468 assert_eq!(request.temperature, Some(0.66));
3469
3470 // Same model name, different provider
3471 cx.update(|cx| {
3472 AgentSettings::override_global(
3473 AgentSettings {
3474 model_parameters: vec![LanguageModelParameters {
3475 provider: Some("anthropic".into()),
3476 model: Some(model.id().0.clone()),
3477 temperature: Some(0.66),
3478 }],
3479 ..AgentSettings::get_global(cx).clone()
3480 },
3481 cx,
3482 );
3483 });
3484
3485 let request = thread.update(cx, |thread, cx| {
3486 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3487 });
3488 assert_eq!(request.temperature, None);
3489 }
3490
3491 #[gpui::test]
3492 async fn test_thread_summary(cx: &mut TestAppContext) {
3493 init_test_settings(cx);
3494
3495 let project = create_test_project(cx, json!({})).await;
3496
3497 let (_, _thread_store, thread, _context_store, model) =
3498 setup_test_environment(cx, project.clone()).await;
3499
3500 // Initial state should be pending
3501 thread.read_with(cx, |thread, _| {
3502 assert!(matches!(thread.summary(), ThreadSummary::Pending));
3503 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3504 });
3505
3506 // Manually setting the summary should not be allowed in this state
3507 thread.update(cx, |thread, cx| {
3508 thread.set_summary("This should not work", cx);
3509 });
3510
3511 thread.read_with(cx, |thread, _| {
3512 assert!(matches!(thread.summary(), ThreadSummary::Pending));
3513 });
3514
3515 // Send a message
3516 thread.update(cx, |thread, cx| {
3517 thread.insert_user_message("Hi!", ContextLoadResult::default(), None, vec![], cx);
3518 thread.send_to_model(
3519 model.clone(),
3520 CompletionIntent::ThreadSummarization,
3521 None,
3522 cx,
3523 );
3524 });
3525
3526 let fake_model = model.as_fake();
3527 simulate_successful_response(&fake_model, cx);
3528
3529 // Should start generating summary when there are >= 2 messages
3530 thread.read_with(cx, |thread, _| {
3531 assert_eq!(*thread.summary(), ThreadSummary::Generating);
3532 });
3533
3534 // Should not be able to set the summary while generating
3535 thread.update(cx, |thread, cx| {
3536 thread.set_summary("This should not work either", cx);
3537 });
3538
3539 thread.read_with(cx, |thread, _| {
3540 assert!(matches!(thread.summary(), ThreadSummary::Generating));
3541 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3542 });
3543
3544 cx.run_until_parked();
3545 fake_model.stream_last_completion_response("Brief");
3546 fake_model.stream_last_completion_response(" Introduction");
3547 fake_model.end_last_completion_stream();
3548 cx.run_until_parked();
3549
3550 // Summary should be set
3551 thread.read_with(cx, |thread, _| {
3552 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
3553 assert_eq!(thread.summary().or_default(), "Brief Introduction");
3554 });
3555
3556 // Now we should be able to set a summary
3557 thread.update(cx, |thread, cx| {
3558 thread.set_summary("Brief Intro", cx);
3559 });
3560
3561 thread.read_with(cx, |thread, _| {
3562 assert_eq!(thread.summary().or_default(), "Brief Intro");
3563 });
3564
3565 // Test setting an empty summary (should default to DEFAULT)
3566 thread.update(cx, |thread, cx| {
3567 thread.set_summary("", cx);
3568 });
3569
3570 thread.read_with(cx, |thread, _| {
3571 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
3572 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3573 });
3574 }
3575
3576 #[gpui::test]
3577 async fn test_thread_summary_error_set_manually(cx: &mut TestAppContext) {
3578 init_test_settings(cx);
3579
3580 let project = create_test_project(cx, json!({})).await;
3581
3582 let (_, _thread_store, thread, _context_store, model) =
3583 setup_test_environment(cx, project.clone()).await;
3584
3585 test_summarize_error(&model, &thread, cx);
3586
3587 // Now we should be able to set a summary
3588 thread.update(cx, |thread, cx| {
3589 thread.set_summary("Brief Intro", cx);
3590 });
3591
3592 thread.read_with(cx, |thread, _| {
3593 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
3594 assert_eq!(thread.summary().or_default(), "Brief Intro");
3595 });
3596 }
3597
3598 #[gpui::test]
3599 async fn test_thread_summary_error_retry(cx: &mut TestAppContext) {
3600 init_test_settings(cx);
3601
3602 let project = create_test_project(cx, json!({})).await;
3603
3604 let (_, _thread_store, thread, _context_store, model) =
3605 setup_test_environment(cx, project.clone()).await;
3606
3607 test_summarize_error(&model, &thread, cx);
3608
3609 // Sending another message should not trigger another summarize request
3610 thread.update(cx, |thread, cx| {
3611 thread.insert_user_message(
3612 "How are you?",
3613 ContextLoadResult::default(),
3614 None,
3615 vec![],
3616 cx,
3617 );
3618 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
3619 });
3620
3621 let fake_model = model.as_fake();
3622 simulate_successful_response(&fake_model, cx);
3623
3624 thread.read_with(cx, |thread, _| {
3625 // State is still Error, not Generating
3626 assert!(matches!(thread.summary(), ThreadSummary::Error));
3627 });
3628
3629 // But the summarize request can be invoked manually
3630 thread.update(cx, |thread, cx| {
3631 thread.summarize(cx);
3632 });
3633
3634 thread.read_with(cx, |thread, _| {
3635 assert!(matches!(thread.summary(), ThreadSummary::Generating));
3636 });
3637
3638 cx.run_until_parked();
3639 fake_model.stream_last_completion_response("A successful summary");
3640 fake_model.end_last_completion_stream();
3641 cx.run_until_parked();
3642
3643 thread.read_with(cx, |thread, _| {
3644 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
3645 assert_eq!(thread.summary().or_default(), "A successful summary");
3646 });
3647 }
3648
3649 fn test_summarize_error(
3650 model: &Arc<dyn LanguageModel>,
3651 thread: &Entity<Thread>,
3652 cx: &mut TestAppContext,
3653 ) {
3654 thread.update(cx, |thread, cx| {
3655 thread.insert_user_message("Hi!", ContextLoadResult::default(), None, vec![], cx);
3656 thread.send_to_model(
3657 model.clone(),
3658 CompletionIntent::ThreadSummarization,
3659 None,
3660 cx,
3661 );
3662 });
3663
3664 let fake_model = model.as_fake();
3665 simulate_successful_response(&fake_model, cx);
3666
3667 thread.read_with(cx, |thread, _| {
3668 assert!(matches!(thread.summary(), ThreadSummary::Generating));
3669 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3670 });
3671
3672 // Simulate summary request ending
3673 cx.run_until_parked();
3674 fake_model.end_last_completion_stream();
3675 cx.run_until_parked();
3676
3677 // State is set to Error and default message
3678 thread.read_with(cx, |thread, _| {
3679 assert!(matches!(thread.summary(), ThreadSummary::Error));
3680 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3681 });
3682 }
3683
3684 fn simulate_successful_response(fake_model: &FakeLanguageModel, cx: &mut TestAppContext) {
3685 cx.run_until_parked();
3686 fake_model.stream_last_completion_response("Assistant response");
3687 fake_model.end_last_completion_stream();
3688 cx.run_until_parked();
3689 }
3690
3691 fn init_test_settings(cx: &mut TestAppContext) {
3692 cx.update(|cx| {
3693 let settings_store = SettingsStore::test(cx);
3694 cx.set_global(settings_store);
3695 language::init(cx);
3696 Project::init_settings(cx);
3697 AgentSettings::register(cx);
3698 prompt_store::init(cx);
3699 thread_store::init(cx);
3700 workspace::init_settings(cx);
3701 language_model::init_settings(cx);
3702 ThemeSettings::register(cx);
3703 EditorSettings::register(cx);
3704 ToolRegistry::default_global(cx);
3705 });
3706 }
3707
3708 // Helper to create a test project with test files
3709 async fn create_test_project(
3710 cx: &mut TestAppContext,
3711 files: serde_json::Value,
3712 ) -> Entity<Project> {
3713 let fs = FakeFs::new(cx.executor());
3714 fs.insert_tree(path!("/test"), files).await;
3715 Project::test(fs, [path!("/test").as_ref()], cx).await
3716 }
3717
3718 async fn setup_test_environment(
3719 cx: &mut TestAppContext,
3720 project: Entity<Project>,
3721 ) -> (
3722 Entity<Workspace>,
3723 Entity<ThreadStore>,
3724 Entity<Thread>,
3725 Entity<ContextStore>,
3726 Arc<dyn LanguageModel>,
3727 ) {
3728 let (workspace, cx) =
3729 cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
3730
3731 let thread_store = cx
3732 .update(|_, cx| {
3733 ThreadStore::load(
3734 project.clone(),
3735 cx.new(|_| ToolWorkingSet::default()),
3736 None,
3737 Arc::new(PromptBuilder::new(None).unwrap()),
3738 cx,
3739 )
3740 })
3741 .await
3742 .unwrap();
3743
3744 let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
3745 let context_store = cx.new(|_cx| ContextStore::new(project.downgrade(), None));
3746
3747 let provider = Arc::new(FakeLanguageModelProvider);
3748 let model = provider.test_model();
3749 let model: Arc<dyn LanguageModel> = Arc::new(model);
3750
3751 cx.update(|_, cx| {
3752 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
3753 registry.set_default_model(
3754 Some(ConfiguredModel {
3755 provider: provider.clone(),
3756 model: model.clone(),
3757 }),
3758 cx,
3759 );
3760 registry.set_thread_summary_model(
3761 Some(ConfiguredModel {
3762 provider,
3763 model: model.clone(),
3764 }),
3765 cx,
3766 );
3767 })
3768 });
3769
3770 (workspace, thread_store, thread, context_store, model)
3771 }
3772
3773 async fn add_file_to_context(
3774 project: &Entity<Project>,
3775 context_store: &Entity<ContextStore>,
3776 path: &str,
3777 cx: &mut TestAppContext,
3778 ) -> Result<Entity<language::Buffer>> {
3779 let buffer_path = project
3780 .read_with(cx, |project, cx| project.find_project_path(path, cx))
3781 .unwrap();
3782
3783 let buffer = project
3784 .update(cx, |project, cx| {
3785 project.open_buffer(buffer_path.clone(), cx)
3786 })
3787 .await
3788 .unwrap();
3789
3790 context_store.update(cx, |context_store, cx| {
3791 context_store.add_file_from_buffer(&buffer_path, buffer.clone(), false, cx);
3792 });
3793
3794 Ok(buffer)
3795 }
3796}