1use crate::{
2 agent_profile::AgentProfile,
3 context::{AgentContext, AgentContextHandle, ContextLoadResult, LoadedContext},
4 thread_store::{
5 SerializedCrease, SerializedLanguageModel, SerializedMessage, SerializedMessageSegment,
6 SerializedThread, SerializedToolResult, SerializedToolUse, SharedProjectContext,
7 ThreadStore,
8 },
9 tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState},
10};
11use action_log::ActionLog;
12use agent_settings::{AgentProfileId, AgentSettings, CompletionMode, SUMMARIZE_THREAD_PROMPT};
13use anyhow::{Result, anyhow};
14use assistant_tool::{AnyToolCard, Tool, ToolWorkingSet};
15use chrono::{DateTime, Utc};
16use client::{ModelRequestUsage, RequestUsage};
17use cloud_llm_client::{CompletionIntent, CompletionRequestStatus, Plan, UsageLimit};
18use collections::HashMap;
19use feature_flags::{self, FeatureFlagAppExt};
20use futures::{FutureExt, StreamExt as _, future::Shared};
21use git::repository::DiffType;
22use gpui::{
23 AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task,
24 WeakEntity, Window,
25};
26use http_client::StatusCode;
27use language_model::{
28 ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
29 LanguageModelExt as _, LanguageModelId, LanguageModelRegistry, LanguageModelRequest,
30 LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
31 LanguageModelToolResultContent, LanguageModelToolUse, LanguageModelToolUseId, MessageContent,
32 ModelRequestLimitReachedError, PaymentRequiredError, Role, SelectedModel, StopReason,
33 TokenUsage,
34};
35use postage::stream::Stream as _;
36use project::{
37 Project,
38 git_store::{GitStore, GitStoreCheckpoint, RepositoryState},
39};
40use prompt_store::{ModelContext, PromptBuilder};
41use schemars::JsonSchema;
42use serde::{Deserialize, Serialize};
43use settings::Settings;
44use std::{
45 io::Write,
46 ops::Range,
47 sync::Arc,
48 time::{Duration, Instant},
49};
50use thiserror::Error;
51use util::{ResultExt as _, post_inc};
52use uuid::Uuid;
53
54const MAX_RETRY_ATTEMPTS: u8 = 4;
55const BASE_RETRY_DELAY: Duration = Duration::from_secs(5);
56
57#[derive(Debug, Clone)]
58enum RetryStrategy {
59 ExponentialBackoff {
60 initial_delay: Duration,
61 max_attempts: u8,
62 },
63 Fixed {
64 delay: Duration,
65 max_attempts: u8,
66 },
67}
68
69#[derive(
70 Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,
71)]
72pub struct ThreadId(Arc<str>);
73
74impl ThreadId {
75 pub fn new() -> Self {
76 Self(Uuid::new_v4().to_string().into())
77 }
78}
79
80impl std::fmt::Display for ThreadId {
81 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82 write!(f, "{}", self.0)
83 }
84}
85
86impl From<&str> for ThreadId {
87 fn from(value: &str) -> Self {
88 Self(value.into())
89 }
90}
91
92/// The ID of the user prompt that initiated a request.
93///
94/// This equates to the user physically submitting a message to the model (e.g., by pressing the Enter key).
95#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize)]
96pub struct PromptId(Arc<str>);
97
98impl PromptId {
99 pub fn new() -> Self {
100 Self(Uuid::new_v4().to_string().into())
101 }
102}
103
104impl std::fmt::Display for PromptId {
105 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106 write!(f, "{}", self.0)
107 }
108}
109
110#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
111pub struct MessageId(pub(crate) usize);
112
113impl MessageId {
114 fn post_inc(&mut self) -> Self {
115 Self(post_inc(&mut self.0))
116 }
117
118 pub fn as_usize(&self) -> usize {
119 self.0
120 }
121}
122
123/// Stored information that can be used to resurrect a context crease when creating an editor for a past message.
124#[derive(Clone, Debug)]
125pub struct MessageCrease {
126 pub range: Range<usize>,
127 pub icon_path: SharedString,
128 pub label: SharedString,
129 /// None for a deserialized message, Some otherwise.
130 pub context: Option<AgentContextHandle>,
131}
132
133/// A message in a [`Thread`].
134#[derive(Debug, Clone)]
135pub struct Message {
136 pub id: MessageId,
137 pub role: Role,
138 pub segments: Vec<MessageSegment>,
139 pub loaded_context: LoadedContext,
140 pub creases: Vec<MessageCrease>,
141 pub is_hidden: bool,
142 pub ui_only: bool,
143}
144
145impl Message {
146 /// Returns whether the message contains any meaningful text that should be displayed
147 /// The model sometimes runs tool without producing any text or just a marker ([`USING_TOOL_MARKER`])
148 pub fn should_display_content(&self) -> bool {
149 self.segments.iter().all(|segment| segment.should_display())
150 }
151
152 pub fn push_thinking(&mut self, text: &str, signature: Option<String>) {
153 if let Some(MessageSegment::Thinking {
154 text: segment,
155 signature: current_signature,
156 }) = self.segments.last_mut()
157 {
158 if let Some(signature) = signature {
159 *current_signature = Some(signature);
160 }
161 segment.push_str(text);
162 } else {
163 self.segments.push(MessageSegment::Thinking {
164 text: text.to_string(),
165 signature,
166 });
167 }
168 }
169
170 pub fn push_redacted_thinking(&mut self, data: String) {
171 self.segments.push(MessageSegment::RedactedThinking(data));
172 }
173
174 pub fn push_text(&mut self, text: &str) {
175 if let Some(MessageSegment::Text(segment)) = self.segments.last_mut() {
176 segment.push_str(text);
177 } else {
178 self.segments.push(MessageSegment::Text(text.to_string()));
179 }
180 }
181
182 pub fn to_string(&self) -> String {
183 let mut result = String::new();
184
185 if !self.loaded_context.text.is_empty() {
186 result.push_str(&self.loaded_context.text);
187 }
188
189 for segment in &self.segments {
190 match segment {
191 MessageSegment::Text(text) => result.push_str(text),
192 MessageSegment::Thinking { text, .. } => {
193 result.push_str("<think>\n");
194 result.push_str(text);
195 result.push_str("\n</think>");
196 }
197 MessageSegment::RedactedThinking(_) => {}
198 }
199 }
200
201 result
202 }
203}
204
205#[derive(Debug, Clone, PartialEq, Eq)]
206pub enum MessageSegment {
207 Text(String),
208 Thinking {
209 text: String,
210 signature: Option<String>,
211 },
212 RedactedThinking(String),
213}
214
215impl MessageSegment {
216 pub fn should_display(&self) -> bool {
217 match self {
218 Self::Text(text) => text.is_empty(),
219 Self::Thinking { text, .. } => text.is_empty(),
220 Self::RedactedThinking(_) => false,
221 }
222 }
223
224 pub fn text(&self) -> Option<&str> {
225 match self {
226 MessageSegment::Text(text) => Some(text),
227 _ => None,
228 }
229 }
230}
231
232#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
233pub struct ProjectSnapshot {
234 pub worktree_snapshots: Vec<WorktreeSnapshot>,
235 pub unsaved_buffer_paths: Vec<String>,
236 pub timestamp: DateTime<Utc>,
237}
238
239#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
240pub struct WorktreeSnapshot {
241 pub worktree_path: String,
242 pub git_state: Option<GitState>,
243}
244
245#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
246pub struct GitState {
247 pub remote_url: Option<String>,
248 pub head_sha: Option<String>,
249 pub current_branch: Option<String>,
250 pub diff: Option<String>,
251}
252
253#[derive(Clone, Debug)]
254pub struct ThreadCheckpoint {
255 message_id: MessageId,
256 git_checkpoint: GitStoreCheckpoint,
257}
258
259#[derive(Copy, Clone, Debug, PartialEq, Eq)]
260pub enum ThreadFeedback {
261 Positive,
262 Negative,
263}
264
265pub enum LastRestoreCheckpoint {
266 Pending {
267 message_id: MessageId,
268 },
269 Error {
270 message_id: MessageId,
271 error: String,
272 },
273}
274
275impl LastRestoreCheckpoint {
276 pub fn message_id(&self) -> MessageId {
277 match self {
278 LastRestoreCheckpoint::Pending { message_id } => *message_id,
279 LastRestoreCheckpoint::Error { message_id, .. } => *message_id,
280 }
281 }
282}
283
284#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
285pub enum DetailedSummaryState {
286 #[default]
287 NotGenerated,
288 Generating {
289 message_id: MessageId,
290 },
291 Generated {
292 text: SharedString,
293 message_id: MessageId,
294 },
295}
296
297impl DetailedSummaryState {
298 fn text(&self) -> Option<SharedString> {
299 if let Self::Generated { text, .. } = self {
300 Some(text.clone())
301 } else {
302 None
303 }
304 }
305}
306
307#[derive(Default, Debug)]
308pub struct TotalTokenUsage {
309 pub total: u64,
310 pub max: u64,
311}
312
313impl TotalTokenUsage {
314 pub fn ratio(&self) -> TokenUsageRatio {
315 #[cfg(debug_assertions)]
316 let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
317 .unwrap_or("0.8".to_string())
318 .parse()
319 .unwrap();
320 #[cfg(not(debug_assertions))]
321 let warning_threshold: f32 = 0.8;
322
323 // When the maximum is unknown because there is no selected model,
324 // avoid showing the token limit warning.
325 if self.max == 0 {
326 TokenUsageRatio::Normal
327 } else if self.total >= self.max {
328 TokenUsageRatio::Exceeded
329 } else if self.total as f32 / self.max as f32 >= warning_threshold {
330 TokenUsageRatio::Warning
331 } else {
332 TokenUsageRatio::Normal
333 }
334 }
335
336 pub fn add(&self, tokens: u64) -> TotalTokenUsage {
337 TotalTokenUsage {
338 total: self.total + tokens,
339 max: self.max,
340 }
341 }
342}
343
344#[derive(Debug, Default, PartialEq, Eq)]
345pub enum TokenUsageRatio {
346 #[default]
347 Normal,
348 Warning,
349 Exceeded,
350}
351
352#[derive(Debug, Clone, Copy)]
353pub enum QueueState {
354 Sending,
355 Queued { position: usize },
356 Started,
357}
358
359/// A thread of conversation with the LLM.
360pub struct Thread {
361 id: ThreadId,
362 updated_at: DateTime<Utc>,
363 summary: ThreadSummary,
364 pending_summary: Task<Option<()>>,
365 detailed_summary_task: Task<Option<()>>,
366 detailed_summary_tx: postage::watch::Sender<DetailedSummaryState>,
367 detailed_summary_rx: postage::watch::Receiver<DetailedSummaryState>,
368 completion_mode: agent_settings::CompletionMode,
369 messages: Vec<Message>,
370 next_message_id: MessageId,
371 last_prompt_id: PromptId,
372 project_context: SharedProjectContext,
373 checkpoints_by_message: HashMap<MessageId, ThreadCheckpoint>,
374 completion_count: usize,
375 pending_completions: Vec<PendingCompletion>,
376 project: Entity<Project>,
377 prompt_builder: Arc<PromptBuilder>,
378 tools: Entity<ToolWorkingSet>,
379 tool_use: ToolUseState,
380 action_log: Entity<ActionLog>,
381 last_restore_checkpoint: Option<LastRestoreCheckpoint>,
382 pending_checkpoint: Option<ThreadCheckpoint>,
383 initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
384 request_token_usage: Vec<TokenUsage>,
385 cumulative_token_usage: TokenUsage,
386 exceeded_window_error: Option<ExceededWindowError>,
387 tool_use_limit_reached: bool,
388 feedback: Option<ThreadFeedback>,
389 retry_state: Option<RetryState>,
390 message_feedback: HashMap<MessageId, ThreadFeedback>,
391 last_auto_capture_at: Option<Instant>,
392 last_received_chunk_at: Option<Instant>,
393 request_callback: Option<
394 Box<dyn FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>])>,
395 >,
396 remaining_turns: u32,
397 configured_model: Option<ConfiguredModel>,
398 profile: AgentProfile,
399 last_error_context: Option<(Arc<dyn LanguageModel>, CompletionIntent)>,
400}
401
402#[derive(Clone, Debug)]
403struct RetryState {
404 attempt: u8,
405 max_attempts: u8,
406 intent: CompletionIntent,
407}
408
409#[derive(Clone, Debug, PartialEq, Eq)]
410pub enum ThreadSummary {
411 Pending,
412 Generating,
413 Ready(SharedString),
414 Error,
415}
416
417impl ThreadSummary {
418 pub const DEFAULT: SharedString = SharedString::new_static("New Thread");
419
420 pub fn or_default(&self) -> SharedString {
421 self.unwrap_or(Self::DEFAULT)
422 }
423
424 pub fn unwrap_or(&self, message: impl Into<SharedString>) -> SharedString {
425 self.ready().unwrap_or_else(|| message.into())
426 }
427
428 pub fn ready(&self) -> Option<SharedString> {
429 match self {
430 ThreadSummary::Ready(summary) => Some(summary.clone()),
431 ThreadSummary::Pending | ThreadSummary::Generating | ThreadSummary::Error => None,
432 }
433 }
434}
435
436#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
437pub struct ExceededWindowError {
438 /// Model used when last message exceeded context window
439 model_id: LanguageModelId,
440 /// Token count including last message
441 token_count: u64,
442}
443
444impl Thread {
445 pub fn new(
446 project: Entity<Project>,
447 tools: Entity<ToolWorkingSet>,
448 prompt_builder: Arc<PromptBuilder>,
449 system_prompt: SharedProjectContext,
450 cx: &mut Context<Self>,
451 ) -> Self {
452 let (detailed_summary_tx, detailed_summary_rx) = postage::watch::channel();
453 let configured_model = LanguageModelRegistry::read_global(cx).default_model();
454 let profile_id = AgentSettings::get_global(cx).default_profile.clone();
455
456 Self {
457 id: ThreadId::new(),
458 updated_at: Utc::now(),
459 summary: ThreadSummary::Pending,
460 pending_summary: Task::ready(None),
461 detailed_summary_task: Task::ready(None),
462 detailed_summary_tx,
463 detailed_summary_rx,
464 completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,
465 messages: Vec::new(),
466 next_message_id: MessageId(0),
467 last_prompt_id: PromptId::new(),
468 project_context: system_prompt,
469 checkpoints_by_message: HashMap::default(),
470 completion_count: 0,
471 pending_completions: Vec::new(),
472 project: project.clone(),
473 prompt_builder,
474 tools: tools.clone(),
475 last_restore_checkpoint: None,
476 pending_checkpoint: None,
477 tool_use: ToolUseState::new(tools.clone()),
478 action_log: cx.new(|_| ActionLog::new(project.clone())),
479 initial_project_snapshot: {
480 let project_snapshot = Self::project_snapshot(project, cx);
481 cx.foreground_executor()
482 .spawn(async move { Some(project_snapshot.await) })
483 .shared()
484 },
485 request_token_usage: Vec::new(),
486 cumulative_token_usage: TokenUsage::default(),
487 exceeded_window_error: None,
488 tool_use_limit_reached: false,
489 feedback: None,
490 retry_state: None,
491 message_feedback: HashMap::default(),
492 last_auto_capture_at: None,
493 last_error_context: None,
494 last_received_chunk_at: None,
495 request_callback: None,
496 remaining_turns: u32::MAX,
497 configured_model: configured_model.clone(),
498 profile: AgentProfile::new(profile_id, tools),
499 }
500 }
501
502 pub fn deserialize(
503 id: ThreadId,
504 serialized: SerializedThread,
505 project: Entity<Project>,
506 tools: Entity<ToolWorkingSet>,
507 prompt_builder: Arc<PromptBuilder>,
508 project_context: SharedProjectContext,
509 window: Option<&mut Window>, // None in headless mode
510 cx: &mut Context<Self>,
511 ) -> Self {
512 let next_message_id = MessageId(
513 serialized
514 .messages
515 .last()
516 .map(|message| message.id.0 + 1)
517 .unwrap_or(0),
518 );
519 let tool_use = ToolUseState::from_serialized_messages(
520 tools.clone(),
521 &serialized.messages,
522 project.clone(),
523 window,
524 cx,
525 );
526 let (detailed_summary_tx, detailed_summary_rx) =
527 postage::watch::channel_with(serialized.detailed_summary_state);
528
529 let configured_model = LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
530 serialized
531 .model
532 .and_then(|model| {
533 let model = SelectedModel {
534 provider: model.provider.clone().into(),
535 model: model.model.clone().into(),
536 };
537 registry.select_model(&model, cx)
538 })
539 .or_else(|| registry.default_model())
540 });
541
542 let completion_mode = serialized
543 .completion_mode
544 .unwrap_or_else(|| AgentSettings::get_global(cx).preferred_completion_mode);
545 let profile_id = serialized
546 .profile
547 .unwrap_or_else(|| AgentSettings::get_global(cx).default_profile.clone());
548
549 Self {
550 id,
551 updated_at: serialized.updated_at,
552 summary: ThreadSummary::Ready(serialized.summary),
553 pending_summary: Task::ready(None),
554 detailed_summary_task: Task::ready(None),
555 detailed_summary_tx,
556 detailed_summary_rx,
557 completion_mode,
558 retry_state: None,
559 messages: serialized
560 .messages
561 .into_iter()
562 .map(|message| Message {
563 id: message.id,
564 role: message.role,
565 segments: message
566 .segments
567 .into_iter()
568 .map(|segment| match segment {
569 SerializedMessageSegment::Text { text } => MessageSegment::Text(text),
570 SerializedMessageSegment::Thinking { text, signature } => {
571 MessageSegment::Thinking { text, signature }
572 }
573 SerializedMessageSegment::RedactedThinking { data } => {
574 MessageSegment::RedactedThinking(data)
575 }
576 })
577 .collect(),
578 loaded_context: LoadedContext {
579 contexts: Vec::new(),
580 text: message.context,
581 images: Vec::new(),
582 },
583 creases: message
584 .creases
585 .into_iter()
586 .map(|crease| MessageCrease {
587 range: crease.start..crease.end,
588 icon_path: crease.icon_path,
589 label: crease.label,
590 context: None,
591 })
592 .collect(),
593 is_hidden: message.is_hidden,
594 ui_only: false, // UI-only messages are not persisted
595 })
596 .collect(),
597 next_message_id,
598 last_prompt_id: PromptId::new(),
599 project_context,
600 checkpoints_by_message: HashMap::default(),
601 completion_count: 0,
602 pending_completions: Vec::new(),
603 last_restore_checkpoint: None,
604 pending_checkpoint: None,
605 project: project.clone(),
606 prompt_builder,
607 tools: tools.clone(),
608 tool_use,
609 action_log: cx.new(|_| ActionLog::new(project)),
610 initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
611 request_token_usage: serialized.request_token_usage,
612 cumulative_token_usage: serialized.cumulative_token_usage,
613 exceeded_window_error: None,
614 tool_use_limit_reached: serialized.tool_use_limit_reached,
615 feedback: None,
616 message_feedback: HashMap::default(),
617 last_auto_capture_at: None,
618 last_error_context: None,
619 last_received_chunk_at: None,
620 request_callback: None,
621 remaining_turns: u32::MAX,
622 configured_model,
623 profile: AgentProfile::new(profile_id, tools),
624 }
625 }
626
627 pub fn set_request_callback(
628 &mut self,
629 callback: impl 'static
630 + FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>]),
631 ) {
632 self.request_callback = Some(Box::new(callback));
633 }
634
635 pub fn id(&self) -> &ThreadId {
636 &self.id
637 }
638
639 pub fn profile(&self) -> &AgentProfile {
640 &self.profile
641 }
642
643 pub fn set_profile(&mut self, id: AgentProfileId, cx: &mut Context<Self>) {
644 if &id != self.profile.id() {
645 self.profile = AgentProfile::new(id, self.tools.clone());
646 cx.emit(ThreadEvent::ProfileChanged);
647 }
648 }
649
650 pub fn is_empty(&self) -> bool {
651 self.messages.is_empty()
652 }
653
654 pub fn updated_at(&self) -> DateTime<Utc> {
655 self.updated_at
656 }
657
658 pub fn touch_updated_at(&mut self) {
659 self.updated_at = Utc::now();
660 }
661
662 pub fn advance_prompt_id(&mut self) {
663 self.last_prompt_id = PromptId::new();
664 }
665
666 pub fn project_context(&self) -> SharedProjectContext {
667 self.project_context.clone()
668 }
669
670 pub fn get_or_init_configured_model(&mut self, cx: &App) -> Option<ConfiguredModel> {
671 if self.configured_model.is_none() {
672 self.configured_model = LanguageModelRegistry::read_global(cx).default_model();
673 }
674 self.configured_model.clone()
675 }
676
677 pub fn configured_model(&self) -> Option<ConfiguredModel> {
678 self.configured_model.clone()
679 }
680
681 pub fn set_configured_model(&mut self, model: Option<ConfiguredModel>, cx: &mut Context<Self>) {
682 self.configured_model = model;
683 cx.notify();
684 }
685
686 pub fn summary(&self) -> &ThreadSummary {
687 &self.summary
688 }
689
690 pub fn set_summary(&mut self, new_summary: impl Into<SharedString>, cx: &mut Context<Self>) {
691 let current_summary = match &self.summary {
692 ThreadSummary::Pending | ThreadSummary::Generating => return,
693 ThreadSummary::Ready(summary) => summary,
694 ThreadSummary::Error => &ThreadSummary::DEFAULT,
695 };
696
697 let mut new_summary = new_summary.into();
698
699 if new_summary.is_empty() {
700 new_summary = ThreadSummary::DEFAULT;
701 }
702
703 if current_summary != &new_summary {
704 self.summary = ThreadSummary::Ready(new_summary);
705 cx.emit(ThreadEvent::SummaryChanged);
706 }
707 }
708
709 pub fn completion_mode(&self) -> CompletionMode {
710 self.completion_mode
711 }
712
713 pub fn set_completion_mode(&mut self, mode: CompletionMode) {
714 self.completion_mode = mode;
715 }
716
717 pub fn message(&self, id: MessageId) -> Option<&Message> {
718 let index = self
719 .messages
720 .binary_search_by(|message| message.id.cmp(&id))
721 .ok()?;
722
723 self.messages.get(index)
724 }
725
726 pub fn messages(&self) -> impl ExactSizeIterator<Item = &Message> {
727 self.messages.iter()
728 }
729
730 pub fn is_generating(&self) -> bool {
731 !self.pending_completions.is_empty() || !self.all_tools_finished()
732 }
733
734 /// Indicates whether streaming of language model events is stale.
735 /// When `is_generating()` is false, this method returns `None`.
736 pub fn is_generation_stale(&self) -> Option<bool> {
737 const STALE_THRESHOLD: u128 = 250;
738
739 self.last_received_chunk_at
740 .map(|instant| instant.elapsed().as_millis() > STALE_THRESHOLD)
741 }
742
743 fn received_chunk(&mut self) {
744 self.last_received_chunk_at = Some(Instant::now());
745 }
746
747 pub fn queue_state(&self) -> Option<QueueState> {
748 self.pending_completions
749 .first()
750 .map(|pending_completion| pending_completion.queue_state)
751 }
752
753 pub fn tools(&self) -> &Entity<ToolWorkingSet> {
754 &self.tools
755 }
756
757 pub fn pending_tool(&self, id: &LanguageModelToolUseId) -> Option<&PendingToolUse> {
758 self.tool_use
759 .pending_tool_uses()
760 .into_iter()
761 .find(|tool_use| &tool_use.id == id)
762 }
763
764 pub fn tools_needing_confirmation(&self) -> impl Iterator<Item = &PendingToolUse> {
765 self.tool_use
766 .pending_tool_uses()
767 .into_iter()
768 .filter(|tool_use| tool_use.status.needs_confirmation())
769 }
770
771 pub fn has_pending_tool_uses(&self) -> bool {
772 !self.tool_use.pending_tool_uses().is_empty()
773 }
774
775 pub fn checkpoint_for_message(&self, id: MessageId) -> Option<ThreadCheckpoint> {
776 self.checkpoints_by_message.get(&id).cloned()
777 }
778
779 pub fn restore_checkpoint(
780 &mut self,
781 checkpoint: ThreadCheckpoint,
782 cx: &mut Context<Self>,
783 ) -> Task<Result<()>> {
784 self.last_restore_checkpoint = Some(LastRestoreCheckpoint::Pending {
785 message_id: checkpoint.message_id,
786 });
787 cx.emit(ThreadEvent::CheckpointChanged);
788 cx.notify();
789
790 let git_store = self.project().read(cx).git_store().clone();
791 let restore = git_store.update(cx, |git_store, cx| {
792 git_store.restore_checkpoint(checkpoint.git_checkpoint.clone(), cx)
793 });
794
795 cx.spawn(async move |this, cx| {
796 let result = restore.await;
797 this.update(cx, |this, cx| {
798 if let Err(err) = result.as_ref() {
799 this.last_restore_checkpoint = Some(LastRestoreCheckpoint::Error {
800 message_id: checkpoint.message_id,
801 error: err.to_string(),
802 });
803 } else {
804 this.truncate(checkpoint.message_id, cx);
805 this.last_restore_checkpoint = None;
806 }
807 this.pending_checkpoint = None;
808 cx.emit(ThreadEvent::CheckpointChanged);
809 cx.notify();
810 })?;
811 result
812 })
813 }
814
815 fn finalize_pending_checkpoint(&mut self, cx: &mut Context<Self>) {
816 let pending_checkpoint = if self.is_generating() {
817 return;
818 } else if let Some(checkpoint) = self.pending_checkpoint.take() {
819 checkpoint
820 } else {
821 return;
822 };
823
824 self.finalize_checkpoint(pending_checkpoint, cx);
825 }
826
827 fn finalize_checkpoint(
828 &mut self,
829 pending_checkpoint: ThreadCheckpoint,
830 cx: &mut Context<Self>,
831 ) {
832 let git_store = self.project.read(cx).git_store().clone();
833 let final_checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
834 cx.spawn(async move |this, cx| match final_checkpoint.await {
835 Ok(final_checkpoint) => {
836 let equal = git_store
837 .update(cx, |store, cx| {
838 store.compare_checkpoints(
839 pending_checkpoint.git_checkpoint.clone(),
840 final_checkpoint.clone(),
841 cx,
842 )
843 })?
844 .await
845 .unwrap_or(false);
846
847 this.update(cx, |this, cx| {
848 this.pending_checkpoint = if equal {
849 Some(pending_checkpoint)
850 } else {
851 this.insert_checkpoint(pending_checkpoint, cx);
852 Some(ThreadCheckpoint {
853 message_id: this.next_message_id,
854 git_checkpoint: final_checkpoint,
855 })
856 }
857 })?;
858
859 Ok(())
860 }
861 Err(_) => this.update(cx, |this, cx| {
862 this.insert_checkpoint(pending_checkpoint, cx)
863 }),
864 })
865 .detach();
866 }
867
868 fn insert_checkpoint(&mut self, checkpoint: ThreadCheckpoint, cx: &mut Context<Self>) {
869 self.checkpoints_by_message
870 .insert(checkpoint.message_id, checkpoint);
871 cx.emit(ThreadEvent::CheckpointChanged);
872 cx.notify();
873 }
874
875 pub fn last_restore_checkpoint(&self) -> Option<&LastRestoreCheckpoint> {
876 self.last_restore_checkpoint.as_ref()
877 }
878
879 pub fn truncate(&mut self, message_id: MessageId, cx: &mut Context<Self>) {
880 let Some(message_ix) = self
881 .messages
882 .iter()
883 .rposition(|message| message.id == message_id)
884 else {
885 return;
886 };
887 for deleted_message in self.messages.drain(message_ix..) {
888 self.checkpoints_by_message.remove(&deleted_message.id);
889 }
890 cx.notify();
891 }
892
893 pub fn context_for_message(&self, id: MessageId) -> impl Iterator<Item = &AgentContext> {
894 self.messages
895 .iter()
896 .find(|message| message.id == id)
897 .into_iter()
898 .flat_map(|message| message.loaded_context.contexts.iter())
899 }
900
901 pub fn is_turn_end(&self, ix: usize) -> bool {
902 if self.messages.is_empty() {
903 return false;
904 }
905
906 if !self.is_generating() && ix == self.messages.len() - 1 {
907 return true;
908 }
909
910 let Some(message) = self.messages.get(ix) else {
911 return false;
912 };
913
914 if message.role != Role::Assistant {
915 return false;
916 }
917
918 self.messages
919 .get(ix + 1)
920 .and_then(|message| {
921 self.message(message.id)
922 .map(|next_message| next_message.role == Role::User && !next_message.is_hidden)
923 })
924 .unwrap_or(false)
925 }
926
927 pub fn tool_use_limit_reached(&self) -> bool {
928 self.tool_use_limit_reached
929 }
930
931 /// Returns whether all of the tool uses have finished running.
932 pub fn all_tools_finished(&self) -> bool {
933 // If the only pending tool uses left are the ones with errors, then
934 // that means that we've finished running all of the pending tools.
935 self.tool_use
936 .pending_tool_uses()
937 .iter()
938 .all(|pending_tool_use| pending_tool_use.status.is_error())
939 }
940
941 /// Returns whether any pending tool uses may perform edits
942 pub fn has_pending_edit_tool_uses(&self) -> bool {
943 self.tool_use
944 .pending_tool_uses()
945 .iter()
946 .filter(|pending_tool_use| !pending_tool_use.status.is_error())
947 .any(|pending_tool_use| pending_tool_use.may_perform_edits)
948 }
949
950 pub fn tool_uses_for_message(&self, id: MessageId, cx: &App) -> Vec<ToolUse> {
951 self.tool_use.tool_uses_for_message(id, &self.project, cx)
952 }
953
954 pub fn tool_results_for_message(
955 &self,
956 assistant_message_id: MessageId,
957 ) -> Vec<&LanguageModelToolResult> {
958 self.tool_use.tool_results_for_message(assistant_message_id)
959 }
960
961 pub fn tool_result(&self, id: &LanguageModelToolUseId) -> Option<&LanguageModelToolResult> {
962 self.tool_use.tool_result(id)
963 }
964
965 pub fn output_for_tool(&self, id: &LanguageModelToolUseId) -> Option<&Arc<str>> {
966 match &self.tool_use.tool_result(id)?.content {
967 LanguageModelToolResultContent::Text(text) => Some(text),
968 LanguageModelToolResultContent::Image(_) => {
969 // TODO: We should display image
970 None
971 }
972 }
973 }
974
975 pub fn card_for_tool(&self, id: &LanguageModelToolUseId) -> Option<AnyToolCard> {
976 self.tool_use.tool_result_card(id).cloned()
977 }
978
979 /// Return tools that are both enabled and supported by the model
980 pub fn available_tools(
981 &self,
982 cx: &App,
983 model: Arc<dyn LanguageModel>,
984 ) -> Vec<LanguageModelRequestTool> {
985 if model.supports_tools() {
986 self.profile
987 .enabled_tools(cx)
988 .into_iter()
989 .filter_map(|(name, tool)| {
990 // Skip tools that cannot be supported
991 let input_schema = tool.input_schema(model.tool_input_format()).ok()?;
992 Some(LanguageModelRequestTool {
993 name: name.into(),
994 description: tool.description(),
995 input_schema,
996 })
997 })
998 .collect()
999 } else {
1000 Vec::default()
1001 }
1002 }
1003
1004 pub fn insert_user_message(
1005 &mut self,
1006 text: impl Into<String>,
1007 loaded_context: ContextLoadResult,
1008 git_checkpoint: Option<GitStoreCheckpoint>,
1009 creases: Vec<MessageCrease>,
1010 cx: &mut Context<Self>,
1011 ) -> MessageId {
1012 if !loaded_context.referenced_buffers.is_empty() {
1013 self.action_log.update(cx, |log, cx| {
1014 for buffer in loaded_context.referenced_buffers {
1015 log.buffer_read(buffer, cx);
1016 }
1017 });
1018 }
1019
1020 let message_id = self.insert_message(
1021 Role::User,
1022 vec![MessageSegment::Text(text.into())],
1023 loaded_context.loaded_context,
1024 creases,
1025 false,
1026 cx,
1027 );
1028
1029 if let Some(git_checkpoint) = git_checkpoint {
1030 self.pending_checkpoint = Some(ThreadCheckpoint {
1031 message_id,
1032 git_checkpoint,
1033 });
1034 }
1035
1036 self.auto_capture_telemetry(cx);
1037
1038 message_id
1039 }
1040
1041 pub fn insert_invisible_continue_message(&mut self, cx: &mut Context<Self>) -> MessageId {
1042 let id = self.insert_message(
1043 Role::User,
1044 vec![MessageSegment::Text("Continue where you left off".into())],
1045 LoadedContext::default(),
1046 vec![],
1047 true,
1048 cx,
1049 );
1050 self.pending_checkpoint = None;
1051
1052 id
1053 }
1054
1055 pub fn insert_assistant_message(
1056 &mut self,
1057 segments: Vec<MessageSegment>,
1058 cx: &mut Context<Self>,
1059 ) -> MessageId {
1060 self.insert_message(
1061 Role::Assistant,
1062 segments,
1063 LoadedContext::default(),
1064 Vec::new(),
1065 false,
1066 cx,
1067 )
1068 }
1069
1070 pub fn insert_message(
1071 &mut self,
1072 role: Role,
1073 segments: Vec<MessageSegment>,
1074 loaded_context: LoadedContext,
1075 creases: Vec<MessageCrease>,
1076 is_hidden: bool,
1077 cx: &mut Context<Self>,
1078 ) -> MessageId {
1079 let id = self.next_message_id.post_inc();
1080 self.messages.push(Message {
1081 id,
1082 role,
1083 segments,
1084 loaded_context,
1085 creases,
1086 is_hidden,
1087 ui_only: false,
1088 });
1089 self.touch_updated_at();
1090 cx.emit(ThreadEvent::MessageAdded(id));
1091 id
1092 }
1093
1094 pub fn edit_message(
1095 &mut self,
1096 id: MessageId,
1097 new_role: Role,
1098 new_segments: Vec<MessageSegment>,
1099 creases: Vec<MessageCrease>,
1100 loaded_context: Option<LoadedContext>,
1101 checkpoint: Option<GitStoreCheckpoint>,
1102 cx: &mut Context<Self>,
1103 ) -> bool {
1104 let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
1105 return false;
1106 };
1107 message.role = new_role;
1108 message.segments = new_segments;
1109 message.creases = creases;
1110 if let Some(context) = loaded_context {
1111 message.loaded_context = context;
1112 }
1113 if let Some(git_checkpoint) = checkpoint {
1114 self.checkpoints_by_message.insert(
1115 id,
1116 ThreadCheckpoint {
1117 message_id: id,
1118 git_checkpoint,
1119 },
1120 );
1121 }
1122 self.touch_updated_at();
1123 cx.emit(ThreadEvent::MessageEdited(id));
1124 true
1125 }
1126
1127 pub fn delete_message(&mut self, id: MessageId, cx: &mut Context<Self>) -> bool {
1128 let Some(index) = self.messages.iter().position(|message| message.id == id) else {
1129 return false;
1130 };
1131 self.messages.remove(index);
1132 self.touch_updated_at();
1133 cx.emit(ThreadEvent::MessageDeleted(id));
1134 true
1135 }
1136
1137 /// Returns the representation of this [`Thread`] in a textual form.
1138 ///
1139 /// This is the representation we use when attaching a thread as context to another thread.
1140 pub fn text(&self) -> String {
1141 let mut text = String::new();
1142
1143 for message in &self.messages {
1144 text.push_str(match message.role {
1145 language_model::Role::User => "User:",
1146 language_model::Role::Assistant => "Agent:",
1147 language_model::Role::System => "System:",
1148 });
1149 text.push('\n');
1150
1151 for segment in &message.segments {
1152 match segment {
1153 MessageSegment::Text(content) => text.push_str(content),
1154 MessageSegment::Thinking { text: content, .. } => {
1155 text.push_str(&format!("<think>{}</think>", content))
1156 }
1157 MessageSegment::RedactedThinking(_) => {}
1158 }
1159 }
1160 text.push('\n');
1161 }
1162
1163 text
1164 }
1165
1166 /// Serializes this thread into a format for storage or telemetry.
1167 pub fn serialize(&self, cx: &mut Context<Self>) -> Task<Result<SerializedThread>> {
1168 let initial_project_snapshot = self.initial_project_snapshot.clone();
1169 cx.spawn(async move |this, cx| {
1170 let initial_project_snapshot = initial_project_snapshot.await;
1171 this.read_with(cx, |this, cx| SerializedThread {
1172 version: SerializedThread::VERSION.to_string(),
1173 summary: this.summary().or_default(),
1174 updated_at: this.updated_at(),
1175 messages: this
1176 .messages()
1177 .filter(|message| !message.ui_only)
1178 .map(|message| SerializedMessage {
1179 id: message.id,
1180 role: message.role,
1181 segments: message
1182 .segments
1183 .iter()
1184 .map(|segment| match segment {
1185 MessageSegment::Text(text) => {
1186 SerializedMessageSegment::Text { text: text.clone() }
1187 }
1188 MessageSegment::Thinking { text, signature } => {
1189 SerializedMessageSegment::Thinking {
1190 text: text.clone(),
1191 signature: signature.clone(),
1192 }
1193 }
1194 MessageSegment::RedactedThinking(data) => {
1195 SerializedMessageSegment::RedactedThinking {
1196 data: data.clone(),
1197 }
1198 }
1199 })
1200 .collect(),
1201 tool_uses: this
1202 .tool_uses_for_message(message.id, cx)
1203 .into_iter()
1204 .map(|tool_use| SerializedToolUse {
1205 id: tool_use.id,
1206 name: tool_use.name,
1207 input: tool_use.input,
1208 })
1209 .collect(),
1210 tool_results: this
1211 .tool_results_for_message(message.id)
1212 .into_iter()
1213 .map(|tool_result| SerializedToolResult {
1214 tool_use_id: tool_result.tool_use_id.clone(),
1215 is_error: tool_result.is_error,
1216 content: tool_result.content.clone(),
1217 output: tool_result.output.clone(),
1218 })
1219 .collect(),
1220 context: message.loaded_context.text.clone(),
1221 creases: message
1222 .creases
1223 .iter()
1224 .map(|crease| SerializedCrease {
1225 start: crease.range.start,
1226 end: crease.range.end,
1227 icon_path: crease.icon_path.clone(),
1228 label: crease.label.clone(),
1229 })
1230 .collect(),
1231 is_hidden: message.is_hidden,
1232 })
1233 .collect(),
1234 initial_project_snapshot,
1235 cumulative_token_usage: this.cumulative_token_usage,
1236 request_token_usage: this.request_token_usage.clone(),
1237 detailed_summary_state: this.detailed_summary_rx.borrow().clone(),
1238 exceeded_window_error: this.exceeded_window_error.clone(),
1239 model: this
1240 .configured_model
1241 .as_ref()
1242 .map(|model| SerializedLanguageModel {
1243 provider: model.provider.id().0.to_string(),
1244 model: model.model.id().0.to_string(),
1245 }),
1246 completion_mode: Some(this.completion_mode),
1247 tool_use_limit_reached: this.tool_use_limit_reached,
1248 profile: Some(this.profile.id().clone()),
1249 })
1250 })
1251 }
1252
1253 pub fn remaining_turns(&self) -> u32 {
1254 self.remaining_turns
1255 }
1256
1257 pub fn set_remaining_turns(&mut self, remaining_turns: u32) {
1258 self.remaining_turns = remaining_turns;
1259 }
1260
1261 pub fn send_to_model(
1262 &mut self,
1263 model: Arc<dyn LanguageModel>,
1264 intent: CompletionIntent,
1265 window: Option<AnyWindowHandle>,
1266 cx: &mut Context<Self>,
1267 ) {
1268 if self.remaining_turns == 0 {
1269 return;
1270 }
1271
1272 self.remaining_turns -= 1;
1273
1274 self.flush_notifications(model.clone(), intent, cx);
1275
1276 let _checkpoint = self.finalize_pending_checkpoint(cx);
1277 self.stream_completion(
1278 self.to_completion_request(model.clone(), intent, cx),
1279 model,
1280 intent,
1281 window,
1282 cx,
1283 );
1284 }
1285
1286 pub fn retry_last_completion(
1287 &mut self,
1288 window: Option<AnyWindowHandle>,
1289 cx: &mut Context<Self>,
1290 ) {
1291 // Clear any existing error state
1292 self.retry_state = None;
1293
1294 // Use the last error context if available, otherwise fall back to configured model
1295 let (model, intent) = if let Some((model, intent)) = self.last_error_context.take() {
1296 (model, intent)
1297 } else if let Some(configured_model) = self.configured_model.as_ref() {
1298 let model = configured_model.model.clone();
1299 let intent = if self.has_pending_tool_uses() {
1300 CompletionIntent::ToolResults
1301 } else {
1302 CompletionIntent::UserPrompt
1303 };
1304 (model, intent)
1305 } else if let Some(configured_model) = self.get_or_init_configured_model(cx) {
1306 let model = configured_model.model.clone();
1307 let intent = if self.has_pending_tool_uses() {
1308 CompletionIntent::ToolResults
1309 } else {
1310 CompletionIntent::UserPrompt
1311 };
1312 (model, intent)
1313 } else {
1314 return;
1315 };
1316
1317 self.send_to_model(model, intent, window, cx);
1318 }
1319
1320 pub fn enable_burn_mode_and_retry(
1321 &mut self,
1322 window: Option<AnyWindowHandle>,
1323 cx: &mut Context<Self>,
1324 ) {
1325 self.completion_mode = CompletionMode::Burn;
1326 cx.emit(ThreadEvent::ProfileChanged);
1327 self.retry_last_completion(window, cx);
1328 }
1329
1330 pub fn used_tools_since_last_user_message(&self) -> bool {
1331 for message in self.messages.iter().rev() {
1332 if self.tool_use.message_has_tool_results(message.id) {
1333 return true;
1334 } else if message.role == Role::User {
1335 return false;
1336 }
1337 }
1338
1339 false
1340 }
1341
1342 pub fn to_completion_request(
1343 &self,
1344 model: Arc<dyn LanguageModel>,
1345 intent: CompletionIntent,
1346 cx: &mut Context<Self>,
1347 ) -> LanguageModelRequest {
1348 let mut request = LanguageModelRequest {
1349 thread_id: Some(self.id.to_string()),
1350 prompt_id: Some(self.last_prompt_id.to_string()),
1351 intent: Some(intent),
1352 mode: None,
1353 messages: vec![],
1354 tools: Vec::new(),
1355 tool_choice: None,
1356 stop: Vec::new(),
1357 temperature: AgentSettings::temperature_for_model(&model, cx),
1358 thinking_allowed: true,
1359 };
1360
1361 let available_tools = self.available_tools(cx, model.clone());
1362 let available_tool_names = available_tools
1363 .iter()
1364 .map(|tool| tool.name.clone())
1365 .collect();
1366
1367 let model_context = &ModelContext {
1368 available_tools: available_tool_names,
1369 };
1370
1371 if let Some(project_context) = self.project_context.borrow().as_ref() {
1372 match self
1373 .prompt_builder
1374 .generate_assistant_system_prompt(project_context, model_context)
1375 {
1376 Err(err) => {
1377 let message = format!("{err:?}").into();
1378 log::error!("{message}");
1379 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1380 header: "Error generating system prompt".into(),
1381 message,
1382 }));
1383 }
1384 Ok(system_prompt) => {
1385 request.messages.push(LanguageModelRequestMessage {
1386 role: Role::System,
1387 content: vec![MessageContent::Text(system_prompt)],
1388 cache: true,
1389 });
1390 }
1391 }
1392 } else {
1393 let message = "Context for system prompt unexpectedly not ready.".into();
1394 log::error!("{message}");
1395 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1396 header: "Error generating system prompt".into(),
1397 message,
1398 }));
1399 }
1400
1401 let mut message_ix_to_cache = None;
1402 for message in &self.messages {
1403 // ui_only messages are for the UI only, not for the model
1404 if message.ui_only {
1405 continue;
1406 }
1407
1408 let mut request_message = LanguageModelRequestMessage {
1409 role: message.role,
1410 content: Vec::new(),
1411 cache: false,
1412 };
1413
1414 message
1415 .loaded_context
1416 .add_to_request_message(&mut request_message);
1417
1418 for segment in &message.segments {
1419 match segment {
1420 MessageSegment::Text(text) => {
1421 let text = text.trim_end();
1422 if !text.is_empty() {
1423 request_message
1424 .content
1425 .push(MessageContent::Text(text.into()));
1426 }
1427 }
1428 MessageSegment::Thinking { text, signature } => {
1429 if !text.is_empty() {
1430 request_message.content.push(MessageContent::Thinking {
1431 text: text.into(),
1432 signature: signature.clone(),
1433 });
1434 }
1435 }
1436 MessageSegment::RedactedThinking(data) => {
1437 request_message
1438 .content
1439 .push(MessageContent::RedactedThinking(data.clone()));
1440 }
1441 };
1442 }
1443
1444 let mut cache_message = true;
1445 let mut tool_results_message = LanguageModelRequestMessage {
1446 role: Role::User,
1447 content: Vec::new(),
1448 cache: false,
1449 };
1450 for (tool_use, tool_result) in self.tool_use.tool_results(message.id) {
1451 if let Some(tool_result) = tool_result {
1452 request_message
1453 .content
1454 .push(MessageContent::ToolUse(tool_use.clone()));
1455 tool_results_message
1456 .content
1457 .push(MessageContent::ToolResult(LanguageModelToolResult {
1458 tool_use_id: tool_use.id.clone(),
1459 tool_name: tool_result.tool_name.clone(),
1460 is_error: tool_result.is_error,
1461 content: if tool_result.content.is_empty() {
1462 // Surprisingly, the API fails if we return an empty string here.
1463 // It thinks we are sending a tool use without a tool result.
1464 "<Tool returned an empty string>".into()
1465 } else {
1466 tool_result.content.clone()
1467 },
1468 output: None,
1469 }));
1470 } else {
1471 cache_message = false;
1472 log::debug!(
1473 "skipped tool use {:?} because it is still pending",
1474 tool_use
1475 );
1476 }
1477 }
1478
1479 if cache_message {
1480 message_ix_to_cache = Some(request.messages.len());
1481 }
1482 request.messages.push(request_message);
1483
1484 if !tool_results_message.content.is_empty() {
1485 if cache_message {
1486 message_ix_to_cache = Some(request.messages.len());
1487 }
1488 request.messages.push(tool_results_message);
1489 }
1490 }
1491
1492 // https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
1493 if let Some(message_ix_to_cache) = message_ix_to_cache {
1494 request.messages[message_ix_to_cache].cache = true;
1495 }
1496
1497 request.tools = available_tools;
1498 request.mode = if model.supports_burn_mode() {
1499 Some(self.completion_mode.into())
1500 } else {
1501 Some(CompletionMode::Normal.into())
1502 };
1503
1504 request
1505 }
1506
1507 fn to_summarize_request(
1508 &self,
1509 model: &Arc<dyn LanguageModel>,
1510 intent: CompletionIntent,
1511 added_user_message: String,
1512 cx: &App,
1513 ) -> LanguageModelRequest {
1514 let mut request = LanguageModelRequest {
1515 thread_id: None,
1516 prompt_id: None,
1517 intent: Some(intent),
1518 mode: None,
1519 messages: vec![],
1520 tools: Vec::new(),
1521 tool_choice: None,
1522 stop: Vec::new(),
1523 temperature: AgentSettings::temperature_for_model(model, cx),
1524 thinking_allowed: false,
1525 };
1526
1527 for message in &self.messages {
1528 let mut request_message = LanguageModelRequestMessage {
1529 role: message.role,
1530 content: Vec::new(),
1531 cache: false,
1532 };
1533
1534 for segment in &message.segments {
1535 match segment {
1536 MessageSegment::Text(text) => request_message
1537 .content
1538 .push(MessageContent::Text(text.clone())),
1539 MessageSegment::Thinking { .. } => {}
1540 MessageSegment::RedactedThinking(_) => {}
1541 }
1542 }
1543
1544 if request_message.content.is_empty() {
1545 continue;
1546 }
1547
1548 request.messages.push(request_message);
1549 }
1550
1551 request.messages.push(LanguageModelRequestMessage {
1552 role: Role::User,
1553 content: vec![MessageContent::Text(added_user_message)],
1554 cache: false,
1555 });
1556
1557 request
1558 }
1559
1560 /// Insert auto-generated notifications (if any) to the thread
1561 fn flush_notifications(
1562 &mut self,
1563 model: Arc<dyn LanguageModel>,
1564 intent: CompletionIntent,
1565 cx: &mut Context<Self>,
1566 ) {
1567 match intent {
1568 CompletionIntent::UserPrompt | CompletionIntent::ToolResults => {
1569 if let Some(pending_tool_use) = self.attach_tracked_files_state(model, cx) {
1570 cx.emit(ThreadEvent::ToolFinished {
1571 tool_use_id: pending_tool_use.id.clone(),
1572 pending_tool_use: Some(pending_tool_use),
1573 });
1574 }
1575 }
1576 CompletionIntent::ThreadSummarization
1577 | CompletionIntent::ThreadContextSummarization
1578 | CompletionIntent::CreateFile
1579 | CompletionIntent::EditFile
1580 | CompletionIntent::InlineAssist
1581 | CompletionIntent::TerminalInlineAssist
1582 | CompletionIntent::GenerateGitCommitMessage => {}
1583 };
1584 }
1585
1586 fn attach_tracked_files_state(
1587 &mut self,
1588 model: Arc<dyn LanguageModel>,
1589 cx: &mut App,
1590 ) -> Option<PendingToolUse> {
1591 // Represent notification as a simulated `project_notifications` tool call
1592 let tool_name = Arc::from("project_notifications");
1593 let tool = self.tools.read(cx).tool(&tool_name, cx)?;
1594
1595 if !self.profile.is_tool_enabled(tool.source(), tool.name(), cx) {
1596 return None;
1597 }
1598
1599 if self
1600 .action_log
1601 .update(cx, |log, cx| log.unnotified_user_edits(cx).is_none())
1602 {
1603 return None;
1604 }
1605
1606 let input = serde_json::json!({});
1607 let request = Arc::new(LanguageModelRequest::default()); // unused
1608 let window = None;
1609 let tool_result = tool.run(
1610 input,
1611 request,
1612 self.project.clone(),
1613 self.action_log.clone(),
1614 model.clone(),
1615 window,
1616 cx,
1617 );
1618
1619 let tool_use_id =
1620 LanguageModelToolUseId::from(format!("project_notifications_{}", self.messages.len()));
1621
1622 let tool_use = LanguageModelToolUse {
1623 id: tool_use_id.clone(),
1624 name: tool_name.clone(),
1625 raw_input: "{}".to_string(),
1626 input: serde_json::json!({}),
1627 is_input_complete: true,
1628 };
1629
1630 let tool_output = cx.background_executor().block(tool_result.output);
1631
1632 // Attach a project_notification tool call to the latest existing
1633 // Assistant message. We cannot create a new Assistant message
1634 // because thinking models require a `thinking` block that we
1635 // cannot mock. We cannot send a notification as a normal
1636 // (non-tool-use) User message because this distracts Agent
1637 // too much.
1638 let tool_message_id = self
1639 .messages
1640 .iter()
1641 .enumerate()
1642 .rfind(|(_, message)| message.role == Role::Assistant)
1643 .map(|(_, message)| message.id)?;
1644
1645 let tool_use_metadata = ToolUseMetadata {
1646 model: model.clone(),
1647 thread_id: self.id.clone(),
1648 prompt_id: self.last_prompt_id.clone(),
1649 };
1650
1651 self.tool_use
1652 .request_tool_use(tool_message_id, tool_use, tool_use_metadata.clone(), cx);
1653
1654 let pending_tool_use = self.tool_use.insert_tool_output(
1655 tool_use_id.clone(),
1656 tool_name,
1657 tool_output,
1658 self.configured_model.as_ref(),
1659 self.completion_mode,
1660 );
1661
1662 pending_tool_use
1663 }
1664
1665 pub fn stream_completion(
1666 &mut self,
1667 request: LanguageModelRequest,
1668 model: Arc<dyn LanguageModel>,
1669 intent: CompletionIntent,
1670 window: Option<AnyWindowHandle>,
1671 cx: &mut Context<Self>,
1672 ) {
1673 self.tool_use_limit_reached = false;
1674
1675 let pending_completion_id = post_inc(&mut self.completion_count);
1676 let mut request_callback_parameters = if self.request_callback.is_some() {
1677 Some((request.clone(), Vec::new()))
1678 } else {
1679 None
1680 };
1681 let prompt_id = self.last_prompt_id.clone();
1682 let tool_use_metadata = ToolUseMetadata {
1683 model: model.clone(),
1684 thread_id: self.id.clone(),
1685 prompt_id: prompt_id.clone(),
1686 };
1687
1688 let completion_mode = request
1689 .mode
1690 .unwrap_or(cloud_llm_client::CompletionMode::Normal);
1691
1692 self.last_received_chunk_at = Some(Instant::now());
1693
1694 let task = cx.spawn(async move |thread, cx| {
1695 let stream_completion_future = model.stream_completion(request, &cx);
1696 let initial_token_usage =
1697 thread.read_with(cx, |thread, _cx| thread.cumulative_token_usage);
1698 let stream_completion = async {
1699 let mut events = stream_completion_future.await?;
1700
1701 let mut stop_reason = StopReason::EndTurn;
1702 let mut current_token_usage = TokenUsage::default();
1703
1704 thread
1705 .update(cx, |_thread, cx| {
1706 cx.emit(ThreadEvent::NewRequest);
1707 })
1708 .ok();
1709
1710 let mut request_assistant_message_id = None;
1711
1712 while let Some(event) = events.next().await {
1713 if let Some((_, response_events)) = request_callback_parameters.as_mut() {
1714 response_events
1715 .push(event.as_ref().map_err(|error| error.to_string()).cloned());
1716 }
1717
1718 thread.update(cx, |thread, cx| {
1719 match event? {
1720 LanguageModelCompletionEvent::StartMessage { .. } => {
1721 request_assistant_message_id =
1722 Some(thread.insert_assistant_message(
1723 vec![MessageSegment::Text(String::new())],
1724 cx,
1725 ));
1726 }
1727 LanguageModelCompletionEvent::Stop(reason) => {
1728 stop_reason = reason;
1729 }
1730 LanguageModelCompletionEvent::UsageUpdate(token_usage) => {
1731 thread.update_token_usage_at_last_message(token_usage);
1732 thread.cumulative_token_usage = thread.cumulative_token_usage
1733 + token_usage
1734 - current_token_usage;
1735 current_token_usage = token_usage;
1736 }
1737 LanguageModelCompletionEvent::Text(chunk) => {
1738 thread.received_chunk();
1739
1740 cx.emit(ThreadEvent::ReceivedTextChunk);
1741 if let Some(last_message) = thread.messages.last_mut() {
1742 if last_message.role == Role::Assistant
1743 && !thread.tool_use.has_tool_results(last_message.id)
1744 {
1745 last_message.push_text(&chunk);
1746 cx.emit(ThreadEvent::StreamedAssistantText(
1747 last_message.id,
1748 chunk,
1749 ));
1750 } else {
1751 // If we won't have an Assistant message yet, assume this chunk marks the beginning
1752 // of a new Assistant response.
1753 //
1754 // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1755 // will result in duplicating the text of the chunk in the rendered Markdown.
1756 request_assistant_message_id =
1757 Some(thread.insert_assistant_message(
1758 vec![MessageSegment::Text(chunk.to_string())],
1759 cx,
1760 ));
1761 };
1762 }
1763 }
1764 LanguageModelCompletionEvent::Thinking {
1765 text: chunk,
1766 signature,
1767 } => {
1768 thread.received_chunk();
1769
1770 if let Some(last_message) = thread.messages.last_mut() {
1771 if last_message.role == Role::Assistant
1772 && !thread.tool_use.has_tool_results(last_message.id)
1773 {
1774 last_message.push_thinking(&chunk, signature);
1775 cx.emit(ThreadEvent::StreamedAssistantThinking(
1776 last_message.id,
1777 chunk,
1778 ));
1779 } else {
1780 // If we won't have an Assistant message yet, assume this chunk marks the beginning
1781 // of a new Assistant response.
1782 //
1783 // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1784 // will result in duplicating the text of the chunk in the rendered Markdown.
1785 request_assistant_message_id =
1786 Some(thread.insert_assistant_message(
1787 vec![MessageSegment::Thinking {
1788 text: chunk.to_string(),
1789 signature,
1790 }],
1791 cx,
1792 ));
1793 };
1794 }
1795 }
1796 LanguageModelCompletionEvent::RedactedThinking { data } => {
1797 thread.received_chunk();
1798
1799 if let Some(last_message) = thread.messages.last_mut() {
1800 if last_message.role == Role::Assistant
1801 && !thread.tool_use.has_tool_results(last_message.id)
1802 {
1803 last_message.push_redacted_thinking(data);
1804 } else {
1805 request_assistant_message_id =
1806 Some(thread.insert_assistant_message(
1807 vec![MessageSegment::RedactedThinking(data)],
1808 cx,
1809 ));
1810 };
1811 }
1812 }
1813 LanguageModelCompletionEvent::ToolUse(tool_use) => {
1814 let last_assistant_message_id = request_assistant_message_id
1815 .unwrap_or_else(|| {
1816 let new_assistant_message_id =
1817 thread.insert_assistant_message(vec![], cx);
1818 request_assistant_message_id =
1819 Some(new_assistant_message_id);
1820 new_assistant_message_id
1821 });
1822
1823 let tool_use_id = tool_use.id.clone();
1824 let streamed_input = if tool_use.is_input_complete {
1825 None
1826 } else {
1827 Some((&tool_use.input).clone())
1828 };
1829
1830 let ui_text = thread.tool_use.request_tool_use(
1831 last_assistant_message_id,
1832 tool_use,
1833 tool_use_metadata.clone(),
1834 cx,
1835 );
1836
1837 if let Some(input) = streamed_input {
1838 cx.emit(ThreadEvent::StreamedToolUse {
1839 tool_use_id,
1840 ui_text,
1841 input,
1842 });
1843 }
1844 }
1845 LanguageModelCompletionEvent::ToolUseJsonParseError {
1846 id,
1847 tool_name,
1848 raw_input: invalid_input_json,
1849 json_parse_error,
1850 } => {
1851 thread.receive_invalid_tool_json(
1852 id,
1853 tool_name,
1854 invalid_input_json,
1855 json_parse_error,
1856 window,
1857 cx,
1858 );
1859 }
1860 LanguageModelCompletionEvent::StatusUpdate(status_update) => {
1861 if let Some(completion) = thread
1862 .pending_completions
1863 .iter_mut()
1864 .find(|completion| completion.id == pending_completion_id)
1865 {
1866 match status_update {
1867 CompletionRequestStatus::Queued { position } => {
1868 completion.queue_state =
1869 QueueState::Queued { position };
1870 }
1871 CompletionRequestStatus::Started => {
1872 completion.queue_state = QueueState::Started;
1873 }
1874 CompletionRequestStatus::Failed {
1875 code,
1876 message,
1877 request_id: _,
1878 retry_after,
1879 } => {
1880 return Err(
1881 LanguageModelCompletionError::from_cloud_failure(
1882 model.upstream_provider_name(),
1883 code,
1884 message,
1885 retry_after.map(Duration::from_secs_f64),
1886 ),
1887 );
1888 }
1889 CompletionRequestStatus::UsageUpdated { amount, limit } => {
1890 thread.update_model_request_usage(
1891 amount as u32,
1892 limit,
1893 cx,
1894 );
1895 }
1896 CompletionRequestStatus::ToolUseLimitReached => {
1897 thread.tool_use_limit_reached = true;
1898 cx.emit(ThreadEvent::ToolUseLimitReached);
1899 }
1900 }
1901 }
1902 }
1903 }
1904
1905 thread.touch_updated_at();
1906 cx.emit(ThreadEvent::StreamedCompletion);
1907 cx.notify();
1908
1909 thread.auto_capture_telemetry(cx);
1910 Ok(())
1911 })??;
1912
1913 smol::future::yield_now().await;
1914 }
1915
1916 thread.update(cx, |thread, cx| {
1917 thread.last_received_chunk_at = None;
1918 thread
1919 .pending_completions
1920 .retain(|completion| completion.id != pending_completion_id);
1921
1922 // If there is a response without tool use, summarize the message. Otherwise,
1923 // allow two tool uses before summarizing.
1924 if matches!(thread.summary, ThreadSummary::Pending)
1925 && thread.messages.len() >= 2
1926 && (!thread.has_pending_tool_uses() || thread.messages.len() >= 6)
1927 {
1928 thread.summarize(cx);
1929 }
1930 })?;
1931
1932 anyhow::Ok(stop_reason)
1933 };
1934
1935 let result = stream_completion.await;
1936 let mut retry_scheduled = false;
1937
1938 thread
1939 .update(cx, |thread, cx| {
1940 thread.finalize_pending_checkpoint(cx);
1941 match result.as_ref() {
1942 Ok(stop_reason) => {
1943 match stop_reason {
1944 StopReason::ToolUse => {
1945 let tool_uses =
1946 thread.use_pending_tools(window, model.clone(), cx);
1947 cx.emit(ThreadEvent::UsePendingTools { tool_uses });
1948 }
1949 StopReason::EndTurn | StopReason::MaxTokens => {
1950 thread.project.update(cx, |project, cx| {
1951 project.set_agent_location(None, cx);
1952 });
1953 }
1954 StopReason::Refusal => {
1955 thread.project.update(cx, |project, cx| {
1956 project.set_agent_location(None, cx);
1957 });
1958
1959 // Remove the turn that was refused.
1960 //
1961 // https://docs.anthropic.com/en/docs/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals#reset-context-after-refusal
1962 {
1963 let mut messages_to_remove = Vec::new();
1964
1965 for (ix, message) in
1966 thread.messages.iter().enumerate().rev()
1967 {
1968 messages_to_remove.push(message.id);
1969
1970 if message.role == Role::User {
1971 if ix == 0 {
1972 break;
1973 }
1974
1975 if let Some(prev_message) =
1976 thread.messages.get(ix - 1)
1977 {
1978 if prev_message.role == Role::Assistant {
1979 break;
1980 }
1981 }
1982 }
1983 }
1984
1985 for message_id in messages_to_remove {
1986 thread.delete_message(message_id, cx);
1987 }
1988 }
1989
1990 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1991 header: "Language model refusal".into(),
1992 message:
1993 "Model refused to generate content for safety reasons."
1994 .into(),
1995 }));
1996 }
1997 }
1998
1999 // We successfully completed, so cancel any remaining retries.
2000 thread.retry_state = None;
2001 }
2002 Err(error) => {
2003 thread.project.update(cx, |project, cx| {
2004 project.set_agent_location(None, cx);
2005 });
2006
2007 if error.is::<PaymentRequiredError>() {
2008 cx.emit(ThreadEvent::ShowError(ThreadError::PaymentRequired));
2009 } else if let Some(error) =
2010 error.downcast_ref::<ModelRequestLimitReachedError>()
2011 {
2012 cx.emit(ThreadEvent::ShowError(
2013 ThreadError::ModelRequestLimitReached { plan: error.plan },
2014 ));
2015 } else if let Some(completion_error) =
2016 error.downcast_ref::<LanguageModelCompletionError>()
2017 {
2018 match &completion_error {
2019 LanguageModelCompletionError::PromptTooLarge {
2020 tokens, ..
2021 } => {
2022 let tokens = tokens.unwrap_or_else(|| {
2023 // We didn't get an exact token count from the API, so fall back on our estimate.
2024 thread
2025 .total_token_usage()
2026 .map(|usage| usage.total)
2027 .unwrap_or(0)
2028 // We know the context window was exceeded in practice, so if our estimate was
2029 // lower than max tokens, the estimate was wrong; return that we exceeded by 1.
2030 .max(
2031 model
2032 .max_token_count_for_mode(completion_mode)
2033 .saturating_add(1),
2034 )
2035 });
2036 thread.exceeded_window_error = Some(ExceededWindowError {
2037 model_id: model.id(),
2038 token_count: tokens,
2039 });
2040 cx.notify();
2041 }
2042 _ => {
2043 if let Some(retry_strategy) =
2044 Thread::get_retry_strategy(completion_error)
2045 {
2046 log::info!(
2047 "Retrying with {:?} for language model completion error {:?}",
2048 retry_strategy,
2049 completion_error
2050 );
2051
2052 retry_scheduled = thread
2053 .handle_retryable_error_with_delay(
2054 &completion_error,
2055 Some(retry_strategy),
2056 model.clone(),
2057 intent,
2058 window,
2059 cx,
2060 );
2061 }
2062 }
2063 }
2064 }
2065
2066 if !retry_scheduled {
2067 thread.cancel_last_completion(window, cx);
2068 }
2069 }
2070 }
2071
2072 if !retry_scheduled {
2073 cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new)));
2074 }
2075
2076 if let Some((request_callback, (request, response_events))) = thread
2077 .request_callback
2078 .as_mut()
2079 .zip(request_callback_parameters.as_ref())
2080 {
2081 request_callback(request, response_events);
2082 }
2083
2084 thread.auto_capture_telemetry(cx);
2085
2086 if let Ok(initial_usage) = initial_token_usage {
2087 let usage = thread.cumulative_token_usage - initial_usage;
2088
2089 telemetry::event!(
2090 "Assistant Thread Completion",
2091 thread_id = thread.id().to_string(),
2092 prompt_id = prompt_id,
2093 model = model.telemetry_id(),
2094 model_provider = model.provider_id().to_string(),
2095 input_tokens = usage.input_tokens,
2096 output_tokens = usage.output_tokens,
2097 cache_creation_input_tokens = usage.cache_creation_input_tokens,
2098 cache_read_input_tokens = usage.cache_read_input_tokens,
2099 );
2100 }
2101 })
2102 .ok();
2103 });
2104
2105 self.pending_completions.push(PendingCompletion {
2106 id: pending_completion_id,
2107 queue_state: QueueState::Sending,
2108 _task: task,
2109 });
2110 }
2111
2112 pub fn summarize(&mut self, cx: &mut Context<Self>) {
2113 let Some(model) = LanguageModelRegistry::read_global(cx).thread_summary_model() else {
2114 println!("No thread summary model");
2115 return;
2116 };
2117
2118 if !model.provider.is_authenticated(cx) {
2119 return;
2120 }
2121
2122 let request = self.to_summarize_request(
2123 &model.model,
2124 CompletionIntent::ThreadSummarization,
2125 SUMMARIZE_THREAD_PROMPT.into(),
2126 cx,
2127 );
2128
2129 self.summary = ThreadSummary::Generating;
2130
2131 self.pending_summary = cx.spawn(async move |this, cx| {
2132 let result = async {
2133 let mut messages = model.model.stream_completion(request, &cx).await?;
2134
2135 let mut new_summary = String::new();
2136 while let Some(event) = messages.next().await {
2137 let Ok(event) = event else {
2138 continue;
2139 };
2140 let text = match event {
2141 LanguageModelCompletionEvent::Text(text) => text,
2142 LanguageModelCompletionEvent::StatusUpdate(
2143 CompletionRequestStatus::UsageUpdated { amount, limit },
2144 ) => {
2145 this.update(cx, |thread, cx| {
2146 thread.update_model_request_usage(amount as u32, limit, cx);
2147 })?;
2148 continue;
2149 }
2150 _ => continue,
2151 };
2152
2153 let mut lines = text.lines();
2154 new_summary.extend(lines.next());
2155
2156 // Stop if the LLM generated multiple lines.
2157 if lines.next().is_some() {
2158 break;
2159 }
2160 }
2161
2162 anyhow::Ok(new_summary)
2163 }
2164 .await;
2165
2166 this.update(cx, |this, cx| {
2167 match result {
2168 Ok(new_summary) => {
2169 if new_summary.is_empty() {
2170 this.summary = ThreadSummary::Error;
2171 } else {
2172 this.summary = ThreadSummary::Ready(new_summary.into());
2173 }
2174 }
2175 Err(err) => {
2176 this.summary = ThreadSummary::Error;
2177 log::error!("Failed to generate thread summary: {}", err);
2178 }
2179 }
2180 cx.emit(ThreadEvent::SummaryGenerated);
2181 })
2182 .log_err()?;
2183
2184 Some(())
2185 });
2186 }
2187
2188 fn get_retry_strategy(error: &LanguageModelCompletionError) -> Option<RetryStrategy> {
2189 use LanguageModelCompletionError::*;
2190
2191 // General strategy here:
2192 // - If retrying won't help (e.g. invalid API key or payload too large), return None so we don't retry at all.
2193 // - If it's a time-based issue (e.g. server overloaded, rate limit exceeded), retry up to 4 times with exponential backoff.
2194 // - If it's an issue that *might* be fixed by retrying (e.g. internal server error), retry up to 3 times.
2195 match error {
2196 HttpResponseError {
2197 status_code: StatusCode::TOO_MANY_REQUESTS,
2198 ..
2199 } => Some(RetryStrategy::ExponentialBackoff {
2200 initial_delay: BASE_RETRY_DELAY,
2201 max_attempts: MAX_RETRY_ATTEMPTS,
2202 }),
2203 ServerOverloaded { retry_after, .. } | RateLimitExceeded { retry_after, .. } => {
2204 Some(RetryStrategy::Fixed {
2205 delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2206 max_attempts: MAX_RETRY_ATTEMPTS,
2207 })
2208 }
2209 UpstreamProviderError {
2210 status,
2211 retry_after,
2212 ..
2213 } => match *status {
2214 StatusCode::TOO_MANY_REQUESTS | StatusCode::SERVICE_UNAVAILABLE => {
2215 Some(RetryStrategy::Fixed {
2216 delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2217 max_attempts: MAX_RETRY_ATTEMPTS,
2218 })
2219 }
2220 StatusCode::INTERNAL_SERVER_ERROR => Some(RetryStrategy::Fixed {
2221 delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2222 // Internal Server Error could be anything, retry up to 3 times.
2223 max_attempts: 3,
2224 }),
2225 status => {
2226 // There is no StatusCode variant for the unofficial HTTP 529 ("The service is overloaded"),
2227 // but we frequently get them in practice. See https://http.dev/529
2228 if status.as_u16() == 529 {
2229 Some(RetryStrategy::Fixed {
2230 delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2231 max_attempts: MAX_RETRY_ATTEMPTS,
2232 })
2233 } else {
2234 Some(RetryStrategy::Fixed {
2235 delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2236 max_attempts: 2,
2237 })
2238 }
2239 }
2240 },
2241 ApiInternalServerError { .. } => Some(RetryStrategy::Fixed {
2242 delay: BASE_RETRY_DELAY,
2243 max_attempts: 3,
2244 }),
2245 ApiReadResponseError { .. }
2246 | HttpSend { .. }
2247 | DeserializeResponse { .. }
2248 | BadRequestFormat { .. } => Some(RetryStrategy::Fixed {
2249 delay: BASE_RETRY_DELAY,
2250 max_attempts: 3,
2251 }),
2252 // Retrying these errors definitely shouldn't help.
2253 HttpResponseError {
2254 status_code:
2255 StatusCode::PAYLOAD_TOO_LARGE | StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED,
2256 ..
2257 }
2258 | AuthenticationError { .. }
2259 | PermissionError { .. }
2260 | NoApiKey { .. }
2261 | ApiEndpointNotFound { .. }
2262 | PromptTooLarge { .. } => None,
2263 // These errors might be transient, so retry them
2264 SerializeRequest { .. } | BuildRequestBody { .. } => Some(RetryStrategy::Fixed {
2265 delay: BASE_RETRY_DELAY,
2266 max_attempts: 1,
2267 }),
2268 // Retry all other 4xx and 5xx errors once.
2269 HttpResponseError { status_code, .. }
2270 if status_code.is_client_error() || status_code.is_server_error() =>
2271 {
2272 Some(RetryStrategy::Fixed {
2273 delay: BASE_RETRY_DELAY,
2274 max_attempts: 3,
2275 })
2276 }
2277 Other(err)
2278 if err.is::<PaymentRequiredError>()
2279 || err.is::<ModelRequestLimitReachedError>() =>
2280 {
2281 // Retrying won't help for Payment Required or Model Request Limit errors (where
2282 // the user must upgrade to usage-based billing to get more requests, or else wait
2283 // for a significant amount of time for the request limit to reset).
2284 None
2285 }
2286 // Conservatively assume that any other errors are non-retryable
2287 HttpResponseError { .. } | Other(..) => Some(RetryStrategy::Fixed {
2288 delay: BASE_RETRY_DELAY,
2289 max_attempts: 2,
2290 }),
2291 }
2292 }
2293
2294 fn handle_retryable_error_with_delay(
2295 &mut self,
2296 error: &LanguageModelCompletionError,
2297 strategy: Option<RetryStrategy>,
2298 model: Arc<dyn LanguageModel>,
2299 intent: CompletionIntent,
2300 window: Option<AnyWindowHandle>,
2301 cx: &mut Context<Self>,
2302 ) -> bool {
2303 // Store context for the Retry button
2304 self.last_error_context = Some((model.clone(), intent));
2305
2306 // Only auto-retry if Burn Mode is enabled
2307 if self.completion_mode != CompletionMode::Burn {
2308 // Show error with retry options
2309 cx.emit(ThreadEvent::ShowError(ThreadError::RetryableError {
2310 message: format!(
2311 "{}\n\nTo automatically retry when similar errors happen, enable Burn Mode.",
2312 error
2313 )
2314 .into(),
2315 can_enable_burn_mode: true,
2316 }));
2317 return false;
2318 }
2319
2320 let Some(strategy) = strategy.or_else(|| Self::get_retry_strategy(error)) else {
2321 return false;
2322 };
2323
2324 let max_attempts = match &strategy {
2325 RetryStrategy::ExponentialBackoff { max_attempts, .. } => *max_attempts,
2326 RetryStrategy::Fixed { max_attempts, .. } => *max_attempts,
2327 };
2328
2329 let retry_state = self.retry_state.get_or_insert(RetryState {
2330 attempt: 0,
2331 max_attempts,
2332 intent,
2333 });
2334
2335 retry_state.attempt += 1;
2336 let attempt = retry_state.attempt;
2337 let max_attempts = retry_state.max_attempts;
2338 let intent = retry_state.intent;
2339
2340 if attempt <= max_attempts {
2341 let delay = match &strategy {
2342 RetryStrategy::ExponentialBackoff { initial_delay, .. } => {
2343 let delay_secs = initial_delay.as_secs() * 2u64.pow((attempt - 1) as u32);
2344 Duration::from_secs(delay_secs)
2345 }
2346 RetryStrategy::Fixed { delay, .. } => *delay,
2347 };
2348
2349 // Add a transient message to inform the user
2350 let delay_secs = delay.as_secs();
2351 let retry_message = if max_attempts == 1 {
2352 format!("{error}. Retrying in {delay_secs} seconds...")
2353 } else {
2354 format!(
2355 "{error}. Retrying (attempt {attempt} of {max_attempts}) \
2356 in {delay_secs} seconds..."
2357 )
2358 };
2359 log::warn!(
2360 "Retrying completion request (attempt {attempt} of {max_attempts}) \
2361 in {delay_secs} seconds: {error:?}",
2362 );
2363
2364 // Add a UI-only message instead of a regular message
2365 let id = self.next_message_id.post_inc();
2366 self.messages.push(Message {
2367 id,
2368 role: Role::System,
2369 segments: vec![MessageSegment::Text(retry_message)],
2370 loaded_context: LoadedContext::default(),
2371 creases: Vec::new(),
2372 is_hidden: false,
2373 ui_only: true,
2374 });
2375 cx.emit(ThreadEvent::MessageAdded(id));
2376
2377 // Schedule the retry
2378 let thread_handle = cx.entity().downgrade();
2379
2380 cx.spawn(async move |_thread, cx| {
2381 cx.background_executor().timer(delay).await;
2382
2383 thread_handle
2384 .update(cx, |thread, cx| {
2385 // Retry the completion
2386 thread.send_to_model(model, intent, window, cx);
2387 })
2388 .log_err();
2389 })
2390 .detach();
2391
2392 true
2393 } else {
2394 // Max retries exceeded
2395 self.retry_state = None;
2396
2397 // Stop generating since we're giving up on retrying.
2398 self.pending_completions.clear();
2399
2400 // Show error alongside a Retry button, but no
2401 // Enable Burn Mode button (since it's already enabled)
2402 cx.emit(ThreadEvent::ShowError(ThreadError::RetryableError {
2403 message: format!("Failed after retrying: {}", error).into(),
2404 can_enable_burn_mode: false,
2405 }));
2406
2407 false
2408 }
2409 }
2410
2411 pub fn start_generating_detailed_summary_if_needed(
2412 &mut self,
2413 thread_store: WeakEntity<ThreadStore>,
2414 cx: &mut Context<Self>,
2415 ) {
2416 let Some(last_message_id) = self.messages.last().map(|message| message.id) else {
2417 return;
2418 };
2419
2420 match &*self.detailed_summary_rx.borrow() {
2421 DetailedSummaryState::Generating { message_id, .. }
2422 | DetailedSummaryState::Generated { message_id, .. }
2423 if *message_id == last_message_id =>
2424 {
2425 // Already up-to-date
2426 return;
2427 }
2428 _ => {}
2429 }
2430
2431 let Some(ConfiguredModel { model, provider }) =
2432 LanguageModelRegistry::read_global(cx).thread_summary_model()
2433 else {
2434 return;
2435 };
2436
2437 if !provider.is_authenticated(cx) {
2438 return;
2439 }
2440
2441 let added_user_message = include_str!("./prompts/summarize_thread_detailed_prompt.txt");
2442
2443 let request = self.to_summarize_request(
2444 &model,
2445 CompletionIntent::ThreadContextSummarization,
2446 added_user_message.into(),
2447 cx,
2448 );
2449
2450 *self.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generating {
2451 message_id: last_message_id,
2452 };
2453
2454 // Replace the detailed summarization task if there is one, cancelling it. It would probably
2455 // be better to allow the old task to complete, but this would require logic for choosing
2456 // which result to prefer (the old task could complete after the new one, resulting in a
2457 // stale summary).
2458 self.detailed_summary_task = cx.spawn(async move |thread, cx| {
2459 let stream = model.stream_completion_text(request, &cx);
2460 let Some(mut messages) = stream.await.log_err() else {
2461 thread
2462 .update(cx, |thread, _cx| {
2463 *thread.detailed_summary_tx.borrow_mut() =
2464 DetailedSummaryState::NotGenerated;
2465 })
2466 .ok()?;
2467 return None;
2468 };
2469
2470 let mut new_detailed_summary = String::new();
2471
2472 while let Some(chunk) = messages.stream.next().await {
2473 if let Some(chunk) = chunk.log_err() {
2474 new_detailed_summary.push_str(&chunk);
2475 }
2476 }
2477
2478 thread
2479 .update(cx, |thread, _cx| {
2480 *thread.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generated {
2481 text: new_detailed_summary.into(),
2482 message_id: last_message_id,
2483 };
2484 })
2485 .ok()?;
2486
2487 // Save thread so its summary can be reused later
2488 if let Some(thread) = thread.upgrade() {
2489 if let Ok(Ok(save_task)) = cx.update(|cx| {
2490 thread_store
2491 .update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx))
2492 }) {
2493 save_task.await.log_err();
2494 }
2495 }
2496
2497 Some(())
2498 });
2499 }
2500
2501 pub async fn wait_for_detailed_summary_or_text(
2502 this: &Entity<Self>,
2503 cx: &mut AsyncApp,
2504 ) -> Option<SharedString> {
2505 let mut detailed_summary_rx = this
2506 .read_with(cx, |this, _cx| this.detailed_summary_rx.clone())
2507 .ok()?;
2508 loop {
2509 match detailed_summary_rx.recv().await? {
2510 DetailedSummaryState::Generating { .. } => {}
2511 DetailedSummaryState::NotGenerated => {
2512 return this.read_with(cx, |this, _cx| this.text().into()).ok();
2513 }
2514 DetailedSummaryState::Generated { text, .. } => return Some(text),
2515 }
2516 }
2517 }
2518
2519 pub fn latest_detailed_summary_or_text(&self) -> SharedString {
2520 self.detailed_summary_rx
2521 .borrow()
2522 .text()
2523 .unwrap_or_else(|| self.text().into())
2524 }
2525
2526 pub fn is_generating_detailed_summary(&self) -> bool {
2527 matches!(
2528 &*self.detailed_summary_rx.borrow(),
2529 DetailedSummaryState::Generating { .. }
2530 )
2531 }
2532
2533 pub fn use_pending_tools(
2534 &mut self,
2535 window: Option<AnyWindowHandle>,
2536 model: Arc<dyn LanguageModel>,
2537 cx: &mut Context<Self>,
2538 ) -> Vec<PendingToolUse> {
2539 self.auto_capture_telemetry(cx);
2540 let request =
2541 Arc::new(self.to_completion_request(model.clone(), CompletionIntent::ToolResults, cx));
2542 let pending_tool_uses = self
2543 .tool_use
2544 .pending_tool_uses()
2545 .into_iter()
2546 .filter(|tool_use| tool_use.status.is_idle())
2547 .cloned()
2548 .collect::<Vec<_>>();
2549
2550 for tool_use in pending_tool_uses.iter() {
2551 self.use_pending_tool(tool_use.clone(), request.clone(), model.clone(), window, cx);
2552 }
2553
2554 pending_tool_uses
2555 }
2556
2557 fn use_pending_tool(
2558 &mut self,
2559 tool_use: PendingToolUse,
2560 request: Arc<LanguageModelRequest>,
2561 model: Arc<dyn LanguageModel>,
2562 window: Option<AnyWindowHandle>,
2563 cx: &mut Context<Self>,
2564 ) {
2565 let Some(tool) = self.tools.read(cx).tool(&tool_use.name, cx) else {
2566 return self.handle_hallucinated_tool_use(tool_use.id, tool_use.name, window, cx);
2567 };
2568
2569 if !self.profile.is_tool_enabled(tool.source(), tool.name(), cx) {
2570 return self.handle_hallucinated_tool_use(tool_use.id, tool_use.name, window, cx);
2571 }
2572
2573 if tool.needs_confirmation(&tool_use.input, &self.project, cx)
2574 && !AgentSettings::get_global(cx).always_allow_tool_actions
2575 {
2576 self.tool_use.confirm_tool_use(
2577 tool_use.id,
2578 tool_use.ui_text,
2579 tool_use.input,
2580 request,
2581 tool,
2582 );
2583 cx.emit(ThreadEvent::ToolConfirmationNeeded);
2584 } else {
2585 self.run_tool(
2586 tool_use.id,
2587 tool_use.ui_text,
2588 tool_use.input,
2589 request,
2590 tool,
2591 model,
2592 window,
2593 cx,
2594 );
2595 }
2596 }
2597
2598 pub fn handle_hallucinated_tool_use(
2599 &mut self,
2600 tool_use_id: LanguageModelToolUseId,
2601 hallucinated_tool_name: Arc<str>,
2602 window: Option<AnyWindowHandle>,
2603 cx: &mut Context<Thread>,
2604 ) {
2605 let available_tools = self.profile.enabled_tools(cx);
2606
2607 let tool_list = available_tools
2608 .iter()
2609 .map(|(name, tool)| format!("- {}: {}", name, tool.description()))
2610 .collect::<Vec<_>>()
2611 .join("\n");
2612
2613 let error_message = format!(
2614 "The tool '{}' doesn't exist or is not enabled. Available tools:\n{}",
2615 hallucinated_tool_name, tool_list
2616 );
2617
2618 let pending_tool_use = self.tool_use.insert_tool_output(
2619 tool_use_id.clone(),
2620 hallucinated_tool_name,
2621 Err(anyhow!("Missing tool call: {error_message}")),
2622 self.configured_model.as_ref(),
2623 self.completion_mode,
2624 );
2625
2626 cx.emit(ThreadEvent::MissingToolUse {
2627 tool_use_id: tool_use_id.clone(),
2628 ui_text: error_message.into(),
2629 });
2630
2631 self.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2632 }
2633
2634 pub fn receive_invalid_tool_json(
2635 &mut self,
2636 tool_use_id: LanguageModelToolUseId,
2637 tool_name: Arc<str>,
2638 invalid_json: Arc<str>,
2639 error: String,
2640 window: Option<AnyWindowHandle>,
2641 cx: &mut Context<Thread>,
2642 ) {
2643 log::error!("The model returned invalid input JSON: {invalid_json}");
2644
2645 let pending_tool_use = self.tool_use.insert_tool_output(
2646 tool_use_id.clone(),
2647 tool_name,
2648 Err(anyhow!("Error parsing input JSON: {error}")),
2649 self.configured_model.as_ref(),
2650 self.completion_mode,
2651 );
2652 let ui_text = if let Some(pending_tool_use) = &pending_tool_use {
2653 pending_tool_use.ui_text.clone()
2654 } else {
2655 log::error!(
2656 "There was no pending tool use for tool use {tool_use_id}, even though it finished (with invalid input JSON)."
2657 );
2658 format!("Unknown tool {}", tool_use_id).into()
2659 };
2660
2661 cx.emit(ThreadEvent::InvalidToolInput {
2662 tool_use_id: tool_use_id.clone(),
2663 ui_text,
2664 invalid_input_json: invalid_json,
2665 });
2666
2667 self.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2668 }
2669
2670 pub fn run_tool(
2671 &mut self,
2672 tool_use_id: LanguageModelToolUseId,
2673 ui_text: impl Into<SharedString>,
2674 input: serde_json::Value,
2675 request: Arc<LanguageModelRequest>,
2676 tool: Arc<dyn Tool>,
2677 model: Arc<dyn LanguageModel>,
2678 window: Option<AnyWindowHandle>,
2679 cx: &mut Context<Thread>,
2680 ) {
2681 let task =
2682 self.spawn_tool_use(tool_use_id.clone(), request, input, tool, model, window, cx);
2683 self.tool_use
2684 .run_pending_tool(tool_use_id, ui_text.into(), task);
2685 }
2686
2687 fn spawn_tool_use(
2688 &mut self,
2689 tool_use_id: LanguageModelToolUseId,
2690 request: Arc<LanguageModelRequest>,
2691 input: serde_json::Value,
2692 tool: Arc<dyn Tool>,
2693 model: Arc<dyn LanguageModel>,
2694 window: Option<AnyWindowHandle>,
2695 cx: &mut Context<Thread>,
2696 ) -> Task<()> {
2697 let tool_name: Arc<str> = tool.name().into();
2698
2699 let tool_result = tool.run(
2700 input,
2701 request,
2702 self.project.clone(),
2703 self.action_log.clone(),
2704 model,
2705 window,
2706 cx,
2707 );
2708
2709 // Store the card separately if it exists
2710 if let Some(card) = tool_result.card.clone() {
2711 self.tool_use
2712 .insert_tool_result_card(tool_use_id.clone(), card);
2713 }
2714
2715 cx.spawn({
2716 async move |thread: WeakEntity<Thread>, cx| {
2717 let output = tool_result.output.await;
2718
2719 thread
2720 .update(cx, |thread, cx| {
2721 let pending_tool_use = thread.tool_use.insert_tool_output(
2722 tool_use_id.clone(),
2723 tool_name,
2724 output,
2725 thread.configured_model.as_ref(),
2726 thread.completion_mode,
2727 );
2728 thread.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2729 })
2730 .ok();
2731 }
2732 })
2733 }
2734
2735 fn tool_finished(
2736 &mut self,
2737 tool_use_id: LanguageModelToolUseId,
2738 pending_tool_use: Option<PendingToolUse>,
2739 canceled: bool,
2740 window: Option<AnyWindowHandle>,
2741 cx: &mut Context<Self>,
2742 ) {
2743 if self.all_tools_finished() {
2744 if let Some(ConfiguredModel { model, .. }) = self.configured_model.as_ref() {
2745 if !canceled {
2746 self.send_to_model(model.clone(), CompletionIntent::ToolResults, window, cx);
2747 }
2748 self.auto_capture_telemetry(cx);
2749 }
2750 }
2751
2752 cx.emit(ThreadEvent::ToolFinished {
2753 tool_use_id,
2754 pending_tool_use,
2755 });
2756 }
2757
2758 /// Cancels the last pending completion, if there are any pending.
2759 ///
2760 /// Returns whether a completion was canceled.
2761 pub fn cancel_last_completion(
2762 &mut self,
2763 window: Option<AnyWindowHandle>,
2764 cx: &mut Context<Self>,
2765 ) -> bool {
2766 let mut canceled = self.pending_completions.pop().is_some() || self.retry_state.is_some();
2767
2768 self.retry_state = None;
2769
2770 for pending_tool_use in self.tool_use.cancel_pending() {
2771 canceled = true;
2772 self.tool_finished(
2773 pending_tool_use.id.clone(),
2774 Some(pending_tool_use),
2775 true,
2776 window,
2777 cx,
2778 );
2779 }
2780
2781 if canceled {
2782 cx.emit(ThreadEvent::CompletionCanceled);
2783
2784 // When canceled, we always want to insert the checkpoint.
2785 // (We skip over finalize_pending_checkpoint, because it
2786 // would conclude we didn't have anything to insert here.)
2787 if let Some(checkpoint) = self.pending_checkpoint.take() {
2788 self.insert_checkpoint(checkpoint, cx);
2789 }
2790 } else {
2791 self.finalize_pending_checkpoint(cx);
2792 }
2793
2794 canceled
2795 }
2796
2797 /// Signals that any in-progress editing should be canceled.
2798 ///
2799 /// This method is used to notify listeners (like ActiveThread) that
2800 /// they should cancel any editing operations.
2801 pub fn cancel_editing(&mut self, cx: &mut Context<Self>) {
2802 cx.emit(ThreadEvent::CancelEditing);
2803 }
2804
2805 pub fn feedback(&self) -> Option<ThreadFeedback> {
2806 self.feedback
2807 }
2808
2809 pub fn message_feedback(&self, message_id: MessageId) -> Option<ThreadFeedback> {
2810 self.message_feedback.get(&message_id).copied()
2811 }
2812
2813 pub fn report_message_feedback(
2814 &mut self,
2815 message_id: MessageId,
2816 feedback: ThreadFeedback,
2817 cx: &mut Context<Self>,
2818 ) -> Task<Result<()>> {
2819 if self.message_feedback.get(&message_id) == Some(&feedback) {
2820 return Task::ready(Ok(()));
2821 }
2822
2823 let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
2824 let serialized_thread = self.serialize(cx);
2825 let thread_id = self.id().clone();
2826 let client = self.project.read(cx).client();
2827
2828 let enabled_tool_names: Vec<String> = self
2829 .profile
2830 .enabled_tools(cx)
2831 .iter()
2832 .map(|(name, _)| name.clone().into())
2833 .collect();
2834
2835 self.message_feedback.insert(message_id, feedback);
2836
2837 cx.notify();
2838
2839 let message_content = self
2840 .message(message_id)
2841 .map(|msg| msg.to_string())
2842 .unwrap_or_default();
2843
2844 cx.background_spawn(async move {
2845 let final_project_snapshot = final_project_snapshot.await;
2846 let serialized_thread = serialized_thread.await?;
2847 let thread_data =
2848 serde_json::to_value(serialized_thread).unwrap_or_else(|_| serde_json::Value::Null);
2849
2850 let rating = match feedback {
2851 ThreadFeedback::Positive => "positive",
2852 ThreadFeedback::Negative => "negative",
2853 };
2854 telemetry::event!(
2855 "Assistant Thread Rated",
2856 rating,
2857 thread_id,
2858 enabled_tool_names,
2859 message_id = message_id.0,
2860 message_content,
2861 thread_data,
2862 final_project_snapshot
2863 );
2864 client.telemetry().flush_events().await;
2865
2866 Ok(())
2867 })
2868 }
2869
2870 pub fn report_feedback(
2871 &mut self,
2872 feedback: ThreadFeedback,
2873 cx: &mut Context<Self>,
2874 ) -> Task<Result<()>> {
2875 let last_assistant_message_id = self
2876 .messages
2877 .iter()
2878 .rev()
2879 .find(|msg| msg.role == Role::Assistant)
2880 .map(|msg| msg.id);
2881
2882 if let Some(message_id) = last_assistant_message_id {
2883 self.report_message_feedback(message_id, feedback, cx)
2884 } else {
2885 let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
2886 let serialized_thread = self.serialize(cx);
2887 let thread_id = self.id().clone();
2888 let client = self.project.read(cx).client();
2889 self.feedback = Some(feedback);
2890 cx.notify();
2891
2892 cx.background_spawn(async move {
2893 let final_project_snapshot = final_project_snapshot.await;
2894 let serialized_thread = serialized_thread.await?;
2895 let thread_data = serde_json::to_value(serialized_thread)
2896 .unwrap_or_else(|_| serde_json::Value::Null);
2897
2898 let rating = match feedback {
2899 ThreadFeedback::Positive => "positive",
2900 ThreadFeedback::Negative => "negative",
2901 };
2902 telemetry::event!(
2903 "Assistant Thread Rated",
2904 rating,
2905 thread_id,
2906 thread_data,
2907 final_project_snapshot
2908 );
2909 client.telemetry().flush_events().await;
2910
2911 Ok(())
2912 })
2913 }
2914 }
2915
2916 /// Create a snapshot of the current project state including git information and unsaved buffers.
2917 fn project_snapshot(
2918 project: Entity<Project>,
2919 cx: &mut Context<Self>,
2920 ) -> Task<Arc<ProjectSnapshot>> {
2921 let git_store = project.read(cx).git_store().clone();
2922 let worktree_snapshots: Vec<_> = project
2923 .read(cx)
2924 .visible_worktrees(cx)
2925 .map(|worktree| Self::worktree_snapshot(worktree, git_store.clone(), cx))
2926 .collect();
2927
2928 cx.spawn(async move |_, cx| {
2929 let worktree_snapshots = futures::future::join_all(worktree_snapshots).await;
2930
2931 let mut unsaved_buffers = Vec::new();
2932 cx.update(|app_cx| {
2933 let buffer_store = project.read(app_cx).buffer_store();
2934 for buffer_handle in buffer_store.read(app_cx).buffers() {
2935 let buffer = buffer_handle.read(app_cx);
2936 if buffer.is_dirty() {
2937 if let Some(file) = buffer.file() {
2938 let path = file.path().to_string_lossy().to_string();
2939 unsaved_buffers.push(path);
2940 }
2941 }
2942 }
2943 })
2944 .ok();
2945
2946 Arc::new(ProjectSnapshot {
2947 worktree_snapshots,
2948 unsaved_buffer_paths: unsaved_buffers,
2949 timestamp: Utc::now(),
2950 })
2951 })
2952 }
2953
2954 fn worktree_snapshot(
2955 worktree: Entity<project::Worktree>,
2956 git_store: Entity<GitStore>,
2957 cx: &App,
2958 ) -> Task<WorktreeSnapshot> {
2959 cx.spawn(async move |cx| {
2960 // Get worktree path and snapshot
2961 let worktree_info = cx.update(|app_cx| {
2962 let worktree = worktree.read(app_cx);
2963 let path = worktree.abs_path().to_string_lossy().to_string();
2964 let snapshot = worktree.snapshot();
2965 (path, snapshot)
2966 });
2967
2968 let Ok((worktree_path, _snapshot)) = worktree_info else {
2969 return WorktreeSnapshot {
2970 worktree_path: String::new(),
2971 git_state: None,
2972 };
2973 };
2974
2975 let git_state = git_store
2976 .update(cx, |git_store, cx| {
2977 git_store
2978 .repositories()
2979 .values()
2980 .find(|repo| {
2981 repo.read(cx)
2982 .abs_path_to_repo_path(&worktree.read(cx).abs_path())
2983 .is_some()
2984 })
2985 .cloned()
2986 })
2987 .ok()
2988 .flatten()
2989 .map(|repo| {
2990 repo.update(cx, |repo, _| {
2991 let current_branch =
2992 repo.branch.as_ref().map(|branch| branch.name().to_owned());
2993 repo.send_job(None, |state, _| async move {
2994 let RepositoryState::Local { backend, .. } = state else {
2995 return GitState {
2996 remote_url: None,
2997 head_sha: None,
2998 current_branch,
2999 diff: None,
3000 };
3001 };
3002
3003 let remote_url = backend.remote_url("origin");
3004 let head_sha = backend.head_sha().await;
3005 let diff = backend.diff(DiffType::HeadToWorktree).await.ok();
3006
3007 GitState {
3008 remote_url,
3009 head_sha,
3010 current_branch,
3011 diff,
3012 }
3013 })
3014 })
3015 });
3016
3017 let git_state = match git_state {
3018 Some(git_state) => match git_state.ok() {
3019 Some(git_state) => git_state.await.ok(),
3020 None => None,
3021 },
3022 None => None,
3023 };
3024
3025 WorktreeSnapshot {
3026 worktree_path,
3027 git_state,
3028 }
3029 })
3030 }
3031
3032 pub fn to_markdown(&self, cx: &App) -> Result<String> {
3033 let mut markdown = Vec::new();
3034
3035 let summary = self.summary().or_default();
3036 writeln!(markdown, "# {summary}\n")?;
3037
3038 for message in self.messages() {
3039 writeln!(
3040 markdown,
3041 "## {role}\n",
3042 role = match message.role {
3043 Role::User => "User",
3044 Role::Assistant => "Agent",
3045 Role::System => "System",
3046 }
3047 )?;
3048
3049 if !message.loaded_context.text.is_empty() {
3050 writeln!(markdown, "{}", message.loaded_context.text)?;
3051 }
3052
3053 if !message.loaded_context.images.is_empty() {
3054 writeln!(
3055 markdown,
3056 "\n{} images attached as context.\n",
3057 message.loaded_context.images.len()
3058 )?;
3059 }
3060
3061 for segment in &message.segments {
3062 match segment {
3063 MessageSegment::Text(text) => writeln!(markdown, "{}\n", text)?,
3064 MessageSegment::Thinking { text, .. } => {
3065 writeln!(markdown, "<think>\n{}\n</think>\n", text)?
3066 }
3067 MessageSegment::RedactedThinking(_) => {}
3068 }
3069 }
3070
3071 for tool_use in self.tool_uses_for_message(message.id, cx) {
3072 writeln!(
3073 markdown,
3074 "**Use Tool: {} ({})**",
3075 tool_use.name, tool_use.id
3076 )?;
3077 writeln!(markdown, "```json")?;
3078 writeln!(
3079 markdown,
3080 "{}",
3081 serde_json::to_string_pretty(&tool_use.input)?
3082 )?;
3083 writeln!(markdown, "```")?;
3084 }
3085
3086 for tool_result in self.tool_results_for_message(message.id) {
3087 write!(markdown, "\n**Tool Results: {}", tool_result.tool_use_id)?;
3088 if tool_result.is_error {
3089 write!(markdown, " (Error)")?;
3090 }
3091
3092 writeln!(markdown, "**\n")?;
3093 match &tool_result.content {
3094 LanguageModelToolResultContent::Text(text) => {
3095 writeln!(markdown, "{text}")?;
3096 }
3097 LanguageModelToolResultContent::Image(image) => {
3098 writeln!(markdown, "", image.source)?;
3099 }
3100 }
3101
3102 if let Some(output) = tool_result.output.as_ref() {
3103 writeln!(
3104 markdown,
3105 "\n\nDebug Output:\n\n```json\n{}\n```\n",
3106 serde_json::to_string_pretty(output)?
3107 )?;
3108 }
3109 }
3110 }
3111
3112 Ok(String::from_utf8_lossy(&markdown).to_string())
3113 }
3114
3115 pub fn keep_edits_in_range(
3116 &mut self,
3117 buffer: Entity<language::Buffer>,
3118 buffer_range: Range<language::Anchor>,
3119 cx: &mut Context<Self>,
3120 ) {
3121 self.action_log.update(cx, |action_log, cx| {
3122 action_log.keep_edits_in_range(buffer, buffer_range, cx)
3123 });
3124 }
3125
3126 pub fn keep_all_edits(&mut self, cx: &mut Context<Self>) {
3127 self.action_log
3128 .update(cx, |action_log, cx| action_log.keep_all_edits(cx));
3129 }
3130
3131 pub fn reject_edits_in_ranges(
3132 &mut self,
3133 buffer: Entity<language::Buffer>,
3134 buffer_ranges: Vec<Range<language::Anchor>>,
3135 cx: &mut Context<Self>,
3136 ) -> Task<Result<()>> {
3137 self.action_log.update(cx, |action_log, cx| {
3138 action_log.reject_edits_in_ranges(buffer, buffer_ranges, cx)
3139 })
3140 }
3141
3142 pub fn action_log(&self) -> &Entity<ActionLog> {
3143 &self.action_log
3144 }
3145
3146 pub fn project(&self) -> &Entity<Project> {
3147 &self.project
3148 }
3149
3150 pub fn auto_capture_telemetry(&mut self, cx: &mut Context<Self>) {
3151 if !cx.has_flag::<feature_flags::ThreadAutoCaptureFeatureFlag>() {
3152 return;
3153 }
3154
3155 let now = Instant::now();
3156 if let Some(last) = self.last_auto_capture_at {
3157 if now.duration_since(last).as_secs() < 10 {
3158 return;
3159 }
3160 }
3161
3162 self.last_auto_capture_at = Some(now);
3163
3164 let thread_id = self.id().clone();
3165 let github_login = self
3166 .project
3167 .read(cx)
3168 .user_store()
3169 .read(cx)
3170 .current_user()
3171 .map(|user| user.github_login.clone());
3172 let client = self.project.read(cx).client();
3173 let serialize_task = self.serialize(cx);
3174
3175 cx.background_executor()
3176 .spawn(async move {
3177 if let Ok(serialized_thread) = serialize_task.await {
3178 if let Ok(thread_data) = serde_json::to_value(serialized_thread) {
3179 telemetry::event!(
3180 "Agent Thread Auto-Captured",
3181 thread_id = thread_id.to_string(),
3182 thread_data = thread_data,
3183 auto_capture_reason = "tracked_user",
3184 github_login = github_login
3185 );
3186
3187 client.telemetry().flush_events().await;
3188 }
3189 }
3190 })
3191 .detach();
3192 }
3193
3194 pub fn cumulative_token_usage(&self) -> TokenUsage {
3195 self.cumulative_token_usage
3196 }
3197
3198 pub fn token_usage_up_to_message(&self, message_id: MessageId) -> TotalTokenUsage {
3199 let Some(model) = self.configured_model.as_ref() else {
3200 return TotalTokenUsage::default();
3201 };
3202
3203 let max = model
3204 .model
3205 .max_token_count_for_mode(self.completion_mode().into());
3206
3207 let index = self
3208 .messages
3209 .iter()
3210 .position(|msg| msg.id == message_id)
3211 .unwrap_or(0);
3212
3213 if index == 0 {
3214 return TotalTokenUsage { total: 0, max };
3215 }
3216
3217 let token_usage = &self
3218 .request_token_usage
3219 .get(index - 1)
3220 .cloned()
3221 .unwrap_or_default();
3222
3223 TotalTokenUsage {
3224 total: token_usage.total_tokens(),
3225 max,
3226 }
3227 }
3228
3229 pub fn total_token_usage(&self) -> Option<TotalTokenUsage> {
3230 let model = self.configured_model.as_ref()?;
3231
3232 let max = model
3233 .model
3234 .max_token_count_for_mode(self.completion_mode().into());
3235
3236 if let Some(exceeded_error) = &self.exceeded_window_error {
3237 if model.model.id() == exceeded_error.model_id {
3238 return Some(TotalTokenUsage {
3239 total: exceeded_error.token_count,
3240 max,
3241 });
3242 }
3243 }
3244
3245 let total = self
3246 .token_usage_at_last_message()
3247 .unwrap_or_default()
3248 .total_tokens();
3249
3250 Some(TotalTokenUsage { total, max })
3251 }
3252
3253 fn token_usage_at_last_message(&self) -> Option<TokenUsage> {
3254 self.request_token_usage
3255 .get(self.messages.len().saturating_sub(1))
3256 .or_else(|| self.request_token_usage.last())
3257 .cloned()
3258 }
3259
3260 fn update_token_usage_at_last_message(&mut self, token_usage: TokenUsage) {
3261 let placeholder = self.token_usage_at_last_message().unwrap_or_default();
3262 self.request_token_usage
3263 .resize(self.messages.len(), placeholder);
3264
3265 if let Some(last) = self.request_token_usage.last_mut() {
3266 *last = token_usage;
3267 }
3268 }
3269
3270 fn update_model_request_usage(&self, amount: u32, limit: UsageLimit, cx: &mut Context<Self>) {
3271 self.project
3272 .read(cx)
3273 .user_store()
3274 .update(cx, |user_store, cx| {
3275 user_store.update_model_request_usage(
3276 ModelRequestUsage(RequestUsage {
3277 amount: amount as i32,
3278 limit,
3279 }),
3280 cx,
3281 )
3282 });
3283 }
3284
3285 pub fn deny_tool_use(
3286 &mut self,
3287 tool_use_id: LanguageModelToolUseId,
3288 tool_name: Arc<str>,
3289 window: Option<AnyWindowHandle>,
3290 cx: &mut Context<Self>,
3291 ) {
3292 let err = Err(anyhow::anyhow!(
3293 "Permission to run tool action denied by user"
3294 ));
3295
3296 self.tool_use.insert_tool_output(
3297 tool_use_id.clone(),
3298 tool_name,
3299 err,
3300 self.configured_model.as_ref(),
3301 self.completion_mode,
3302 );
3303 self.tool_finished(tool_use_id.clone(), None, true, window, cx);
3304 }
3305}
3306
3307#[derive(Debug, Clone, Error)]
3308pub enum ThreadError {
3309 #[error("Payment required")]
3310 PaymentRequired,
3311 #[error("Model request limit reached")]
3312 ModelRequestLimitReached { plan: Plan },
3313 #[error("Message {header}: {message}")]
3314 Message {
3315 header: SharedString,
3316 message: SharedString,
3317 },
3318 #[error("Retryable error: {message}")]
3319 RetryableError {
3320 message: SharedString,
3321 can_enable_burn_mode: bool,
3322 },
3323}
3324
3325#[derive(Debug, Clone)]
3326pub enum ThreadEvent {
3327 ShowError(ThreadError),
3328 StreamedCompletion,
3329 ReceivedTextChunk,
3330 NewRequest,
3331 StreamedAssistantText(MessageId, String),
3332 StreamedAssistantThinking(MessageId, String),
3333 StreamedToolUse {
3334 tool_use_id: LanguageModelToolUseId,
3335 ui_text: Arc<str>,
3336 input: serde_json::Value,
3337 },
3338 MissingToolUse {
3339 tool_use_id: LanguageModelToolUseId,
3340 ui_text: Arc<str>,
3341 },
3342 InvalidToolInput {
3343 tool_use_id: LanguageModelToolUseId,
3344 ui_text: Arc<str>,
3345 invalid_input_json: Arc<str>,
3346 },
3347 Stopped(Result<StopReason, Arc<anyhow::Error>>),
3348 MessageAdded(MessageId),
3349 MessageEdited(MessageId),
3350 MessageDeleted(MessageId),
3351 SummaryGenerated,
3352 SummaryChanged,
3353 UsePendingTools {
3354 tool_uses: Vec<PendingToolUse>,
3355 },
3356 ToolFinished {
3357 #[allow(unused)]
3358 tool_use_id: LanguageModelToolUseId,
3359 /// The pending tool use that corresponds to this tool.
3360 pending_tool_use: Option<PendingToolUse>,
3361 },
3362 CheckpointChanged,
3363 ToolConfirmationNeeded,
3364 ToolUseLimitReached,
3365 CancelEditing,
3366 CompletionCanceled,
3367 ProfileChanged,
3368}
3369
3370impl EventEmitter<ThreadEvent> for Thread {}
3371
3372struct PendingCompletion {
3373 id: usize,
3374 queue_state: QueueState,
3375 _task: Task<()>,
3376}
3377
3378#[cfg(test)]
3379mod tests {
3380 use super::*;
3381 use crate::{
3382 context::load_context, context_store::ContextStore, thread_store, thread_store::ThreadStore,
3383 };
3384
3385 // Test-specific constants
3386 const TEST_RATE_LIMIT_RETRY_SECS: u64 = 30;
3387 use agent_settings::{AgentProfileId, AgentSettings, LanguageModelParameters};
3388 use assistant_tool::ToolRegistry;
3389 use assistant_tools;
3390 use futures::StreamExt;
3391 use futures::future::BoxFuture;
3392 use futures::stream::BoxStream;
3393 use gpui::TestAppContext;
3394 use http_client;
3395 use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider};
3396 use language_model::{
3397 LanguageModelCompletionError, LanguageModelName, LanguageModelProviderId,
3398 LanguageModelProviderName, LanguageModelToolChoice,
3399 };
3400 use parking_lot::Mutex;
3401 use project::{FakeFs, Project};
3402 use prompt_store::PromptBuilder;
3403 use serde_json::json;
3404 use settings::{Settings, SettingsStore};
3405 use std::sync::Arc;
3406 use std::time::Duration;
3407 use theme::ThemeSettings;
3408 use util::path;
3409 use workspace::Workspace;
3410
3411 #[gpui::test]
3412 async fn test_message_with_context(cx: &mut TestAppContext) {
3413 init_test_settings(cx);
3414
3415 let project = create_test_project(
3416 cx,
3417 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3418 )
3419 .await;
3420
3421 let (_workspace, _thread_store, thread, context_store, model) =
3422 setup_test_environment(cx, project.clone()).await;
3423
3424 add_file_to_context(&project, &context_store, "test/code.rs", cx)
3425 .await
3426 .unwrap();
3427
3428 let context =
3429 context_store.read_with(cx, |store, _| store.context().next().cloned().unwrap());
3430 let loaded_context = cx
3431 .update(|cx| load_context(vec![context], &project, &None, cx))
3432 .await;
3433
3434 // Insert user message with context
3435 let message_id = thread.update(cx, |thread, cx| {
3436 thread.insert_user_message(
3437 "Please explain this code",
3438 loaded_context,
3439 None,
3440 Vec::new(),
3441 cx,
3442 )
3443 });
3444
3445 // Check content and context in message object
3446 let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
3447
3448 // Use different path format strings based on platform for the test
3449 #[cfg(windows)]
3450 let path_part = r"test\code.rs";
3451 #[cfg(not(windows))]
3452 let path_part = "test/code.rs";
3453
3454 let expected_context = format!(
3455 r#"
3456<context>
3457The following items were attached by the user. They are up-to-date and don't need to be re-read.
3458
3459<files>
3460```rs {path_part}
3461fn main() {{
3462 println!("Hello, world!");
3463}}
3464```
3465</files>
3466</context>
3467"#
3468 );
3469
3470 assert_eq!(message.role, Role::User);
3471 assert_eq!(message.segments.len(), 1);
3472 assert_eq!(
3473 message.segments[0],
3474 MessageSegment::Text("Please explain this code".to_string())
3475 );
3476 assert_eq!(message.loaded_context.text, expected_context);
3477
3478 // Check message in request
3479 let request = thread.update(cx, |thread, cx| {
3480 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3481 });
3482
3483 assert_eq!(request.messages.len(), 2);
3484 let expected_full_message = format!("{}Please explain this code", expected_context);
3485 assert_eq!(request.messages[1].string_contents(), expected_full_message);
3486 }
3487
3488 #[gpui::test]
3489 async fn test_only_include_new_contexts(cx: &mut TestAppContext) {
3490 init_test_settings(cx);
3491
3492 let project = create_test_project(
3493 cx,
3494 json!({
3495 "file1.rs": "fn function1() {}\n",
3496 "file2.rs": "fn function2() {}\n",
3497 "file3.rs": "fn function3() {}\n",
3498 "file4.rs": "fn function4() {}\n",
3499 }),
3500 )
3501 .await;
3502
3503 let (_, _thread_store, thread, context_store, model) =
3504 setup_test_environment(cx, project.clone()).await;
3505
3506 // First message with context 1
3507 add_file_to_context(&project, &context_store, "test/file1.rs", cx)
3508 .await
3509 .unwrap();
3510 let new_contexts = context_store.update(cx, |store, cx| {
3511 store.new_context_for_thread(thread.read(cx), None)
3512 });
3513 assert_eq!(new_contexts.len(), 1);
3514 let loaded_context = cx
3515 .update(|cx| load_context(new_contexts, &project, &None, cx))
3516 .await;
3517 let message1_id = thread.update(cx, |thread, cx| {
3518 thread.insert_user_message("Message 1", loaded_context, None, Vec::new(), cx)
3519 });
3520
3521 // Second message with contexts 1 and 2 (context 1 should be skipped as it's already included)
3522 add_file_to_context(&project, &context_store, "test/file2.rs", cx)
3523 .await
3524 .unwrap();
3525 let new_contexts = context_store.update(cx, |store, cx| {
3526 store.new_context_for_thread(thread.read(cx), None)
3527 });
3528 assert_eq!(new_contexts.len(), 1);
3529 let loaded_context = cx
3530 .update(|cx| load_context(new_contexts, &project, &None, cx))
3531 .await;
3532 let message2_id = thread.update(cx, |thread, cx| {
3533 thread.insert_user_message("Message 2", loaded_context, None, Vec::new(), cx)
3534 });
3535
3536 // Third message with all three contexts (contexts 1 and 2 should be skipped)
3537 //
3538 add_file_to_context(&project, &context_store, "test/file3.rs", cx)
3539 .await
3540 .unwrap();
3541 let new_contexts = context_store.update(cx, |store, cx| {
3542 store.new_context_for_thread(thread.read(cx), None)
3543 });
3544 assert_eq!(new_contexts.len(), 1);
3545 let loaded_context = cx
3546 .update(|cx| load_context(new_contexts, &project, &None, cx))
3547 .await;
3548 let message3_id = thread.update(cx, |thread, cx| {
3549 thread.insert_user_message("Message 3", loaded_context, None, Vec::new(), cx)
3550 });
3551
3552 // Check what contexts are included in each message
3553 let (message1, message2, message3) = thread.read_with(cx, |thread, _| {
3554 (
3555 thread.message(message1_id).unwrap().clone(),
3556 thread.message(message2_id).unwrap().clone(),
3557 thread.message(message3_id).unwrap().clone(),
3558 )
3559 });
3560
3561 // First message should include context 1
3562 assert!(message1.loaded_context.text.contains("file1.rs"));
3563
3564 // Second message should include only context 2 (not 1)
3565 assert!(!message2.loaded_context.text.contains("file1.rs"));
3566 assert!(message2.loaded_context.text.contains("file2.rs"));
3567
3568 // Third message should include only context 3 (not 1 or 2)
3569 assert!(!message3.loaded_context.text.contains("file1.rs"));
3570 assert!(!message3.loaded_context.text.contains("file2.rs"));
3571 assert!(message3.loaded_context.text.contains("file3.rs"));
3572
3573 // Check entire request to make sure all contexts are properly included
3574 let request = thread.update(cx, |thread, cx| {
3575 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3576 });
3577
3578 // The request should contain all 3 messages
3579 assert_eq!(request.messages.len(), 4);
3580
3581 // Check that the contexts are properly formatted in each message
3582 assert!(request.messages[1].string_contents().contains("file1.rs"));
3583 assert!(!request.messages[1].string_contents().contains("file2.rs"));
3584 assert!(!request.messages[1].string_contents().contains("file3.rs"));
3585
3586 assert!(!request.messages[2].string_contents().contains("file1.rs"));
3587 assert!(request.messages[2].string_contents().contains("file2.rs"));
3588 assert!(!request.messages[2].string_contents().contains("file3.rs"));
3589
3590 assert!(!request.messages[3].string_contents().contains("file1.rs"));
3591 assert!(!request.messages[3].string_contents().contains("file2.rs"));
3592 assert!(request.messages[3].string_contents().contains("file3.rs"));
3593
3594 add_file_to_context(&project, &context_store, "test/file4.rs", cx)
3595 .await
3596 .unwrap();
3597 let new_contexts = context_store.update(cx, |store, cx| {
3598 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3599 });
3600 assert_eq!(new_contexts.len(), 3);
3601 let loaded_context = cx
3602 .update(|cx| load_context(new_contexts, &project, &None, cx))
3603 .await
3604 .loaded_context;
3605
3606 assert!(!loaded_context.text.contains("file1.rs"));
3607 assert!(loaded_context.text.contains("file2.rs"));
3608 assert!(loaded_context.text.contains("file3.rs"));
3609 assert!(loaded_context.text.contains("file4.rs"));
3610
3611 let new_contexts = context_store.update(cx, |store, cx| {
3612 // Remove file4.rs
3613 store.remove_context(&loaded_context.contexts[2].handle(), cx);
3614 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3615 });
3616 assert_eq!(new_contexts.len(), 2);
3617 let loaded_context = cx
3618 .update(|cx| load_context(new_contexts, &project, &None, cx))
3619 .await
3620 .loaded_context;
3621
3622 assert!(!loaded_context.text.contains("file1.rs"));
3623 assert!(loaded_context.text.contains("file2.rs"));
3624 assert!(loaded_context.text.contains("file3.rs"));
3625 assert!(!loaded_context.text.contains("file4.rs"));
3626
3627 let new_contexts = context_store.update(cx, |store, cx| {
3628 // Remove file3.rs
3629 store.remove_context(&loaded_context.contexts[1].handle(), cx);
3630 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3631 });
3632 assert_eq!(new_contexts.len(), 1);
3633 let loaded_context = cx
3634 .update(|cx| load_context(new_contexts, &project, &None, cx))
3635 .await
3636 .loaded_context;
3637
3638 assert!(!loaded_context.text.contains("file1.rs"));
3639 assert!(loaded_context.text.contains("file2.rs"));
3640 assert!(!loaded_context.text.contains("file3.rs"));
3641 assert!(!loaded_context.text.contains("file4.rs"));
3642 }
3643
3644 #[gpui::test]
3645 async fn test_message_without_files(cx: &mut TestAppContext) {
3646 init_test_settings(cx);
3647
3648 let project = create_test_project(
3649 cx,
3650 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3651 )
3652 .await;
3653
3654 let (_, _thread_store, thread, _context_store, model) =
3655 setup_test_environment(cx, project.clone()).await;
3656
3657 // Insert user message without any context (empty context vector)
3658 let message_id = thread.update(cx, |thread, cx| {
3659 thread.insert_user_message(
3660 "What is the best way to learn Rust?",
3661 ContextLoadResult::default(),
3662 None,
3663 Vec::new(),
3664 cx,
3665 )
3666 });
3667
3668 // Check content and context in message object
3669 let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
3670
3671 // Context should be empty when no files are included
3672 assert_eq!(message.role, Role::User);
3673 assert_eq!(message.segments.len(), 1);
3674 assert_eq!(
3675 message.segments[0],
3676 MessageSegment::Text("What is the best way to learn Rust?".to_string())
3677 );
3678 assert_eq!(message.loaded_context.text, "");
3679
3680 // Check message in request
3681 let request = thread.update(cx, |thread, cx| {
3682 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3683 });
3684
3685 assert_eq!(request.messages.len(), 2);
3686 assert_eq!(
3687 request.messages[1].string_contents(),
3688 "What is the best way to learn Rust?"
3689 );
3690
3691 // Add second message, also without context
3692 let message2_id = thread.update(cx, |thread, cx| {
3693 thread.insert_user_message(
3694 "Are there any good books?",
3695 ContextLoadResult::default(),
3696 None,
3697 Vec::new(),
3698 cx,
3699 )
3700 });
3701
3702 let message2 =
3703 thread.read_with(cx, |thread, _| thread.message(message2_id).unwrap().clone());
3704 assert_eq!(message2.loaded_context.text, "");
3705
3706 // Check that both messages appear in the request
3707 let request = thread.update(cx, |thread, cx| {
3708 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3709 });
3710
3711 assert_eq!(request.messages.len(), 3);
3712 assert_eq!(
3713 request.messages[1].string_contents(),
3714 "What is the best way to learn Rust?"
3715 );
3716 assert_eq!(
3717 request.messages[2].string_contents(),
3718 "Are there any good books?"
3719 );
3720 }
3721
3722 #[gpui::test]
3723 #[ignore] // turn this test on when project_notifications tool is re-enabled
3724 async fn test_stale_buffer_notification(cx: &mut TestAppContext) {
3725 init_test_settings(cx);
3726
3727 let project = create_test_project(
3728 cx,
3729 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3730 )
3731 .await;
3732
3733 let (_workspace, _thread_store, thread, context_store, model) =
3734 setup_test_environment(cx, project.clone()).await;
3735
3736 // Add a buffer to the context. This will be a tracked buffer
3737 let buffer = add_file_to_context(&project, &context_store, "test/code.rs", cx)
3738 .await
3739 .unwrap();
3740
3741 let context = context_store
3742 .read_with(cx, |store, _| store.context().next().cloned())
3743 .unwrap();
3744 let loaded_context = cx
3745 .update(|cx| load_context(vec![context], &project, &None, cx))
3746 .await;
3747
3748 // Insert user message and assistant response
3749 thread.update(cx, |thread, cx| {
3750 thread.insert_user_message("Explain this code", loaded_context, None, Vec::new(), cx);
3751 thread.insert_assistant_message(
3752 vec![MessageSegment::Text("This code prints 42.".into())],
3753 cx,
3754 );
3755 });
3756 cx.run_until_parked();
3757
3758 // We shouldn't have a stale buffer notification yet
3759 let notifications = thread.read_with(cx, |thread, _| {
3760 find_tool_uses(thread, "project_notifications")
3761 });
3762 assert!(
3763 notifications.is_empty(),
3764 "Should not have stale buffer notification before buffer is modified"
3765 );
3766
3767 // Modify the buffer
3768 buffer.update(cx, |buffer, cx| {
3769 buffer.edit(
3770 [(1..1, "\n println!(\"Added a new line\");\n")],
3771 None,
3772 cx,
3773 );
3774 });
3775
3776 // Insert another user message
3777 thread.update(cx, |thread, cx| {
3778 thread.insert_user_message(
3779 "What does the code do now?",
3780 ContextLoadResult::default(),
3781 None,
3782 Vec::new(),
3783 cx,
3784 )
3785 });
3786 cx.run_until_parked();
3787
3788 // Check for the stale buffer warning
3789 thread.update(cx, |thread, cx| {
3790 thread.flush_notifications(model.clone(), CompletionIntent::UserPrompt, cx)
3791 });
3792 cx.run_until_parked();
3793
3794 let notifications = thread.read_with(cx, |thread, _cx| {
3795 find_tool_uses(thread, "project_notifications")
3796 });
3797
3798 let [notification] = notifications.as_slice() else {
3799 panic!("Should have a `project_notifications` tool use");
3800 };
3801
3802 let Some(notification_content) = notification.content.to_str() else {
3803 panic!("`project_notifications` should return text");
3804 };
3805
3806 assert!(notification_content.contains("These files have changed since the last read:"));
3807 assert!(notification_content.contains("code.rs"));
3808
3809 // Insert another user message and flush notifications again
3810 thread.update(cx, |thread, cx| {
3811 thread.insert_user_message(
3812 "Can you tell me more?",
3813 ContextLoadResult::default(),
3814 None,
3815 Vec::new(),
3816 cx,
3817 )
3818 });
3819
3820 thread.update(cx, |thread, cx| {
3821 thread.flush_notifications(model.clone(), CompletionIntent::UserPrompt, cx)
3822 });
3823 cx.run_until_parked();
3824
3825 // There should be no new notifications (we already flushed one)
3826 let notifications = thread.read_with(cx, |thread, _cx| {
3827 find_tool_uses(thread, "project_notifications")
3828 });
3829
3830 assert_eq!(
3831 notifications.len(),
3832 1,
3833 "Should still have only one notification after second flush - no duplicates"
3834 );
3835 }
3836
3837 fn find_tool_uses(thread: &Thread, tool_name: &str) -> Vec<LanguageModelToolResult> {
3838 thread
3839 .messages()
3840 .flat_map(|message| {
3841 thread
3842 .tool_results_for_message(message.id)
3843 .into_iter()
3844 .filter(|result| result.tool_name == tool_name.into())
3845 .cloned()
3846 .collect::<Vec<_>>()
3847 })
3848 .collect()
3849 }
3850
3851 #[gpui::test]
3852 async fn test_storing_profile_setting_per_thread(cx: &mut TestAppContext) {
3853 init_test_settings(cx);
3854
3855 let project = create_test_project(
3856 cx,
3857 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3858 )
3859 .await;
3860
3861 let (_workspace, thread_store, thread, _context_store, _model) =
3862 setup_test_environment(cx, project.clone()).await;
3863
3864 // Check that we are starting with the default profile
3865 let profile = cx.read(|cx| thread.read(cx).profile.clone());
3866 let tool_set = cx.read(|cx| thread_store.read(cx).tools());
3867 assert_eq!(
3868 profile,
3869 AgentProfile::new(AgentProfileId::default(), tool_set)
3870 );
3871 }
3872
3873 #[gpui::test]
3874 async fn test_serializing_thread_profile(cx: &mut TestAppContext) {
3875 init_test_settings(cx);
3876
3877 let project = create_test_project(
3878 cx,
3879 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3880 )
3881 .await;
3882
3883 let (_workspace, thread_store, thread, _context_store, _model) =
3884 setup_test_environment(cx, project.clone()).await;
3885
3886 // Profile gets serialized with default values
3887 let serialized = thread
3888 .update(cx, |thread, cx| thread.serialize(cx))
3889 .await
3890 .unwrap();
3891
3892 assert_eq!(serialized.profile, Some(AgentProfileId::default()));
3893
3894 let deserialized = cx.update(|cx| {
3895 thread.update(cx, |thread, cx| {
3896 Thread::deserialize(
3897 thread.id.clone(),
3898 serialized,
3899 thread.project.clone(),
3900 thread.tools.clone(),
3901 thread.prompt_builder.clone(),
3902 thread.project_context.clone(),
3903 None,
3904 cx,
3905 )
3906 })
3907 });
3908 let tool_set = cx.read(|cx| thread_store.read(cx).tools());
3909
3910 assert_eq!(
3911 deserialized.profile,
3912 AgentProfile::new(AgentProfileId::default(), tool_set)
3913 );
3914 }
3915
3916 #[gpui::test]
3917 async fn test_temperature_setting(cx: &mut TestAppContext) {
3918 init_test_settings(cx);
3919
3920 let project = create_test_project(
3921 cx,
3922 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3923 )
3924 .await;
3925
3926 let (_workspace, _thread_store, thread, _context_store, model) =
3927 setup_test_environment(cx, project.clone()).await;
3928
3929 // Both model and provider
3930 cx.update(|cx| {
3931 AgentSettings::override_global(
3932 AgentSettings {
3933 model_parameters: vec![LanguageModelParameters {
3934 provider: Some(model.provider_id().0.to_string().into()),
3935 model: Some(model.id().0.clone()),
3936 temperature: Some(0.66),
3937 }],
3938 ..AgentSettings::get_global(cx).clone()
3939 },
3940 cx,
3941 );
3942 });
3943
3944 let request = thread.update(cx, |thread, cx| {
3945 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3946 });
3947 assert_eq!(request.temperature, Some(0.66));
3948
3949 // Only model
3950 cx.update(|cx| {
3951 AgentSettings::override_global(
3952 AgentSettings {
3953 model_parameters: vec![LanguageModelParameters {
3954 provider: None,
3955 model: Some(model.id().0.clone()),
3956 temperature: Some(0.66),
3957 }],
3958 ..AgentSettings::get_global(cx).clone()
3959 },
3960 cx,
3961 );
3962 });
3963
3964 let request = thread.update(cx, |thread, cx| {
3965 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3966 });
3967 assert_eq!(request.temperature, Some(0.66));
3968
3969 // Only provider
3970 cx.update(|cx| {
3971 AgentSettings::override_global(
3972 AgentSettings {
3973 model_parameters: vec![LanguageModelParameters {
3974 provider: Some(model.provider_id().0.to_string().into()),
3975 model: None,
3976 temperature: Some(0.66),
3977 }],
3978 ..AgentSettings::get_global(cx).clone()
3979 },
3980 cx,
3981 );
3982 });
3983
3984 let request = thread.update(cx, |thread, cx| {
3985 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3986 });
3987 assert_eq!(request.temperature, Some(0.66));
3988
3989 // Same model name, different provider
3990 cx.update(|cx| {
3991 AgentSettings::override_global(
3992 AgentSettings {
3993 model_parameters: vec![LanguageModelParameters {
3994 provider: Some("anthropic".into()),
3995 model: Some(model.id().0.clone()),
3996 temperature: Some(0.66),
3997 }],
3998 ..AgentSettings::get_global(cx).clone()
3999 },
4000 cx,
4001 );
4002 });
4003
4004 let request = thread.update(cx, |thread, cx| {
4005 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
4006 });
4007 assert_eq!(request.temperature, None);
4008 }
4009
4010 #[gpui::test]
4011 async fn test_thread_summary(cx: &mut TestAppContext) {
4012 init_test_settings(cx);
4013
4014 let project = create_test_project(cx, json!({})).await;
4015
4016 let (_, _thread_store, thread, _context_store, model) =
4017 setup_test_environment(cx, project.clone()).await;
4018
4019 // Initial state should be pending
4020 thread.read_with(cx, |thread, _| {
4021 assert!(matches!(thread.summary(), ThreadSummary::Pending));
4022 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
4023 });
4024
4025 // Manually setting the summary should not be allowed in this state
4026 thread.update(cx, |thread, cx| {
4027 thread.set_summary("This should not work", cx);
4028 });
4029
4030 thread.read_with(cx, |thread, _| {
4031 assert!(matches!(thread.summary(), ThreadSummary::Pending));
4032 });
4033
4034 // Send a message
4035 thread.update(cx, |thread, cx| {
4036 thread.insert_user_message("Hi!", ContextLoadResult::default(), None, vec![], cx);
4037 thread.send_to_model(
4038 model.clone(),
4039 CompletionIntent::ThreadSummarization,
4040 None,
4041 cx,
4042 );
4043 });
4044
4045 let fake_model = model.as_fake();
4046 simulate_successful_response(&fake_model, cx);
4047
4048 // Should start generating summary when there are >= 2 messages
4049 thread.read_with(cx, |thread, _| {
4050 assert_eq!(*thread.summary(), ThreadSummary::Generating);
4051 });
4052
4053 // Should not be able to set the summary while generating
4054 thread.update(cx, |thread, cx| {
4055 thread.set_summary("This should not work either", cx);
4056 });
4057
4058 thread.read_with(cx, |thread, _| {
4059 assert!(matches!(thread.summary(), ThreadSummary::Generating));
4060 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
4061 });
4062
4063 cx.run_until_parked();
4064 fake_model.send_last_completion_stream_text_chunk("Brief");
4065 fake_model.send_last_completion_stream_text_chunk(" Introduction");
4066 fake_model.end_last_completion_stream();
4067 cx.run_until_parked();
4068
4069 // Summary should be set
4070 thread.read_with(cx, |thread, _| {
4071 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
4072 assert_eq!(thread.summary().or_default(), "Brief Introduction");
4073 });
4074
4075 // Now we should be able to set a summary
4076 thread.update(cx, |thread, cx| {
4077 thread.set_summary("Brief Intro", cx);
4078 });
4079
4080 thread.read_with(cx, |thread, _| {
4081 assert_eq!(thread.summary().or_default(), "Brief Intro");
4082 });
4083
4084 // Test setting an empty summary (should default to DEFAULT)
4085 thread.update(cx, |thread, cx| {
4086 thread.set_summary("", cx);
4087 });
4088
4089 thread.read_with(cx, |thread, _| {
4090 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
4091 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
4092 });
4093 }
4094
4095 #[gpui::test]
4096 async fn test_thread_summary_error_set_manually(cx: &mut TestAppContext) {
4097 init_test_settings(cx);
4098
4099 let project = create_test_project(cx, json!({})).await;
4100
4101 let (_, _thread_store, thread, _context_store, model) =
4102 setup_test_environment(cx, project.clone()).await;
4103
4104 test_summarize_error(&model, &thread, cx);
4105
4106 // Now we should be able to set a summary
4107 thread.update(cx, |thread, cx| {
4108 thread.set_summary("Brief Intro", cx);
4109 });
4110
4111 thread.read_with(cx, |thread, _| {
4112 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
4113 assert_eq!(thread.summary().or_default(), "Brief Intro");
4114 });
4115 }
4116
4117 #[gpui::test]
4118 async fn test_thread_summary_error_retry(cx: &mut TestAppContext) {
4119 init_test_settings(cx);
4120
4121 let project = create_test_project(cx, json!({})).await;
4122
4123 let (_, _thread_store, thread, _context_store, model) =
4124 setup_test_environment(cx, project.clone()).await;
4125
4126 test_summarize_error(&model, &thread, cx);
4127
4128 // Sending another message should not trigger another summarize request
4129 thread.update(cx, |thread, cx| {
4130 thread.insert_user_message(
4131 "How are you?",
4132 ContextLoadResult::default(),
4133 None,
4134 vec![],
4135 cx,
4136 );
4137 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4138 });
4139
4140 let fake_model = model.as_fake();
4141 simulate_successful_response(&fake_model, cx);
4142
4143 thread.read_with(cx, |thread, _| {
4144 // State is still Error, not Generating
4145 assert!(matches!(thread.summary(), ThreadSummary::Error));
4146 });
4147
4148 // But the summarize request can be invoked manually
4149 thread.update(cx, |thread, cx| {
4150 thread.summarize(cx);
4151 });
4152
4153 thread.read_with(cx, |thread, _| {
4154 assert!(matches!(thread.summary(), ThreadSummary::Generating));
4155 });
4156
4157 cx.run_until_parked();
4158 fake_model.send_last_completion_stream_text_chunk("A successful summary");
4159 fake_model.end_last_completion_stream();
4160 cx.run_until_parked();
4161
4162 thread.read_with(cx, |thread, _| {
4163 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
4164 assert_eq!(thread.summary().or_default(), "A successful summary");
4165 });
4166 }
4167
4168 // Helper to create a model that returns errors
4169 enum TestError {
4170 Overloaded,
4171 InternalServerError,
4172 }
4173
4174 struct ErrorInjector {
4175 inner: Arc<FakeLanguageModel>,
4176 error_type: TestError,
4177 }
4178
4179 impl ErrorInjector {
4180 fn new(error_type: TestError) -> Self {
4181 Self {
4182 inner: Arc::new(FakeLanguageModel::default()),
4183 error_type,
4184 }
4185 }
4186 }
4187
4188 impl LanguageModel for ErrorInjector {
4189 fn id(&self) -> LanguageModelId {
4190 self.inner.id()
4191 }
4192
4193 fn name(&self) -> LanguageModelName {
4194 self.inner.name()
4195 }
4196
4197 fn provider_id(&self) -> LanguageModelProviderId {
4198 self.inner.provider_id()
4199 }
4200
4201 fn provider_name(&self) -> LanguageModelProviderName {
4202 self.inner.provider_name()
4203 }
4204
4205 fn supports_tools(&self) -> bool {
4206 self.inner.supports_tools()
4207 }
4208
4209 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
4210 self.inner.supports_tool_choice(choice)
4211 }
4212
4213 fn supports_images(&self) -> bool {
4214 self.inner.supports_images()
4215 }
4216
4217 fn telemetry_id(&self) -> String {
4218 self.inner.telemetry_id()
4219 }
4220
4221 fn max_token_count(&self) -> u64 {
4222 self.inner.max_token_count()
4223 }
4224
4225 fn count_tokens(
4226 &self,
4227 request: LanguageModelRequest,
4228 cx: &App,
4229 ) -> BoxFuture<'static, Result<u64>> {
4230 self.inner.count_tokens(request, cx)
4231 }
4232
4233 fn stream_completion(
4234 &self,
4235 _request: LanguageModelRequest,
4236 _cx: &AsyncApp,
4237 ) -> BoxFuture<
4238 'static,
4239 Result<
4240 BoxStream<
4241 'static,
4242 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
4243 >,
4244 LanguageModelCompletionError,
4245 >,
4246 > {
4247 let error = match self.error_type {
4248 TestError::Overloaded => LanguageModelCompletionError::ServerOverloaded {
4249 provider: self.provider_name(),
4250 retry_after: None,
4251 },
4252 TestError::InternalServerError => {
4253 LanguageModelCompletionError::ApiInternalServerError {
4254 provider: self.provider_name(),
4255 message: "I'm a teapot orbiting the sun".to_string(),
4256 }
4257 }
4258 };
4259 async move {
4260 let stream = futures::stream::once(async move { Err(error) });
4261 Ok(stream.boxed())
4262 }
4263 .boxed()
4264 }
4265
4266 fn as_fake(&self) -> &FakeLanguageModel {
4267 &self.inner
4268 }
4269 }
4270
4271 #[gpui::test]
4272 async fn test_retry_on_overloaded_error(cx: &mut TestAppContext) {
4273 init_test_settings(cx);
4274
4275 let project = create_test_project(cx, json!({})).await;
4276 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4277
4278 // Enable Burn Mode to allow retries
4279 thread.update(cx, |thread, _| {
4280 thread.set_completion_mode(CompletionMode::Burn);
4281 });
4282
4283 // Create model that returns overloaded error
4284 let model = Arc::new(ErrorInjector::new(TestError::Overloaded));
4285
4286 // Insert a user message
4287 thread.update(cx, |thread, cx| {
4288 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4289 });
4290
4291 // Start completion
4292 thread.update(cx, |thread, cx| {
4293 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4294 });
4295
4296 cx.run_until_parked();
4297
4298 thread.read_with(cx, |thread, _| {
4299 assert!(thread.retry_state.is_some(), "Should have retry state");
4300 let retry_state = thread.retry_state.as_ref().unwrap();
4301 assert_eq!(retry_state.attempt, 1, "Should be first retry attempt");
4302 assert_eq!(
4303 retry_state.max_attempts, MAX_RETRY_ATTEMPTS,
4304 "Should retry MAX_RETRY_ATTEMPTS times for overloaded errors"
4305 );
4306 });
4307
4308 // Check that a retry message was added
4309 thread.read_with(cx, |thread, _| {
4310 let mut messages = thread.messages();
4311 assert!(
4312 messages.any(|msg| {
4313 msg.role == Role::System
4314 && msg.ui_only
4315 && msg.segments.iter().any(|seg| {
4316 if let MessageSegment::Text(text) = seg {
4317 text.contains("overloaded")
4318 && text
4319 .contains(&format!("attempt 1 of {}", MAX_RETRY_ATTEMPTS))
4320 } else {
4321 false
4322 }
4323 })
4324 }),
4325 "Should have added a system retry message"
4326 );
4327 });
4328
4329 let retry_count = thread.update(cx, |thread, _| {
4330 thread
4331 .messages
4332 .iter()
4333 .filter(|m| {
4334 m.ui_only
4335 && m.segments.iter().any(|s| {
4336 if let MessageSegment::Text(text) = s {
4337 text.contains("Retrying") && text.contains("seconds")
4338 } else {
4339 false
4340 }
4341 })
4342 })
4343 .count()
4344 });
4345
4346 assert_eq!(retry_count, 1, "Should have one retry message");
4347 }
4348
4349 #[gpui::test]
4350 async fn test_retry_on_internal_server_error(cx: &mut TestAppContext) {
4351 init_test_settings(cx);
4352
4353 let project = create_test_project(cx, json!({})).await;
4354 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4355
4356 // Enable Burn Mode to allow retries
4357 thread.update(cx, |thread, _| {
4358 thread.set_completion_mode(CompletionMode::Burn);
4359 });
4360
4361 // Create model that returns internal server error
4362 let model = Arc::new(ErrorInjector::new(TestError::InternalServerError));
4363
4364 // Insert a user message
4365 thread.update(cx, |thread, cx| {
4366 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4367 });
4368
4369 // Start completion
4370 thread.update(cx, |thread, cx| {
4371 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4372 });
4373
4374 cx.run_until_parked();
4375
4376 // Check retry state on thread
4377 thread.read_with(cx, |thread, _| {
4378 assert!(thread.retry_state.is_some(), "Should have retry state");
4379 let retry_state = thread.retry_state.as_ref().unwrap();
4380 assert_eq!(retry_state.attempt, 1, "Should be first retry attempt");
4381 assert_eq!(
4382 retry_state.max_attempts, 3,
4383 "Should have correct max attempts"
4384 );
4385 });
4386
4387 // Check that a retry message was added with provider name
4388 thread.read_with(cx, |thread, _| {
4389 let mut messages = thread.messages();
4390 assert!(
4391 messages.any(|msg| {
4392 msg.role == Role::System
4393 && msg.ui_only
4394 && msg.segments.iter().any(|seg| {
4395 if let MessageSegment::Text(text) = seg {
4396 text.contains("internal")
4397 && text.contains("Fake")
4398 && text.contains("Retrying")
4399 && text.contains("attempt 1 of 3")
4400 && text.contains("seconds")
4401 } else {
4402 false
4403 }
4404 })
4405 }),
4406 "Should have added a system retry message with provider name"
4407 );
4408 });
4409
4410 // Count retry messages
4411 let retry_count = thread.update(cx, |thread, _| {
4412 thread
4413 .messages
4414 .iter()
4415 .filter(|m| {
4416 m.ui_only
4417 && m.segments.iter().any(|s| {
4418 if let MessageSegment::Text(text) = s {
4419 text.contains("Retrying") && text.contains("seconds")
4420 } else {
4421 false
4422 }
4423 })
4424 })
4425 .count()
4426 });
4427
4428 assert_eq!(retry_count, 1, "Should have one retry message");
4429 }
4430
4431 #[gpui::test]
4432 async fn test_exponential_backoff_on_retries(cx: &mut TestAppContext) {
4433 init_test_settings(cx);
4434
4435 let project = create_test_project(cx, json!({})).await;
4436 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4437
4438 // Enable Burn Mode to allow retries
4439 thread.update(cx, |thread, _| {
4440 thread.set_completion_mode(CompletionMode::Burn);
4441 });
4442
4443 // Create model that returns internal server error
4444 let model = Arc::new(ErrorInjector::new(TestError::InternalServerError));
4445
4446 // Insert a user message
4447 thread.update(cx, |thread, cx| {
4448 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4449 });
4450
4451 // Track retry events and completion count
4452 // Track completion events
4453 let completion_count = Arc::new(Mutex::new(0));
4454 let completion_count_clone = completion_count.clone();
4455
4456 let _subscription = thread.update(cx, |_, cx| {
4457 cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| {
4458 if let ThreadEvent::NewRequest = event {
4459 *completion_count_clone.lock() += 1;
4460 }
4461 })
4462 });
4463
4464 // First attempt
4465 thread.update(cx, |thread, cx| {
4466 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4467 });
4468 cx.run_until_parked();
4469
4470 // Should have scheduled first retry - count retry messages
4471 let retry_count = thread.update(cx, |thread, _| {
4472 thread
4473 .messages
4474 .iter()
4475 .filter(|m| {
4476 m.ui_only
4477 && m.segments.iter().any(|s| {
4478 if let MessageSegment::Text(text) = s {
4479 text.contains("Retrying") && text.contains("seconds")
4480 } else {
4481 false
4482 }
4483 })
4484 })
4485 .count()
4486 });
4487 assert_eq!(retry_count, 1, "Should have scheduled first retry");
4488
4489 // Check retry state
4490 thread.read_with(cx, |thread, _| {
4491 assert!(thread.retry_state.is_some(), "Should have retry state");
4492 let retry_state = thread.retry_state.as_ref().unwrap();
4493 assert_eq!(retry_state.attempt, 1, "Should be first retry attempt");
4494 assert_eq!(
4495 retry_state.max_attempts, 3,
4496 "Internal server errors should retry up to 3 times"
4497 );
4498 });
4499
4500 // Advance clock for first retry
4501 cx.executor().advance_clock(BASE_RETRY_DELAY);
4502 cx.run_until_parked();
4503
4504 // Advance clock for second retry
4505 cx.executor().advance_clock(BASE_RETRY_DELAY);
4506 cx.run_until_parked();
4507
4508 // Advance clock for third retry
4509 cx.executor().advance_clock(BASE_RETRY_DELAY);
4510 cx.run_until_parked();
4511
4512 // Should have completed all retries - count retry messages
4513 let retry_count = thread.update(cx, |thread, _| {
4514 thread
4515 .messages
4516 .iter()
4517 .filter(|m| {
4518 m.ui_only
4519 && m.segments.iter().any(|s| {
4520 if let MessageSegment::Text(text) = s {
4521 text.contains("Retrying") && text.contains("seconds")
4522 } else {
4523 false
4524 }
4525 })
4526 })
4527 .count()
4528 });
4529 assert_eq!(
4530 retry_count, 3,
4531 "Should have 3 retries for internal server errors"
4532 );
4533
4534 // For internal server errors, we retry 3 times and then give up
4535 // Check that retry_state is cleared after all retries
4536 thread.read_with(cx, |thread, _| {
4537 assert!(
4538 thread.retry_state.is_none(),
4539 "Retry state should be cleared after all retries"
4540 );
4541 });
4542
4543 // Verify total attempts (1 initial + 3 retries)
4544 assert_eq!(
4545 *completion_count.lock(),
4546 4,
4547 "Should have attempted once plus 3 retries"
4548 );
4549 }
4550
4551 #[gpui::test]
4552 async fn test_max_retries_exceeded(cx: &mut TestAppContext) {
4553 init_test_settings(cx);
4554
4555 let project = create_test_project(cx, json!({})).await;
4556 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4557
4558 // Enable Burn Mode to allow retries
4559 thread.update(cx, |thread, _| {
4560 thread.set_completion_mode(CompletionMode::Burn);
4561 });
4562
4563 // Create model that returns overloaded error
4564 let model = Arc::new(ErrorInjector::new(TestError::Overloaded));
4565
4566 // Insert a user message
4567 thread.update(cx, |thread, cx| {
4568 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4569 });
4570
4571 // Track events
4572 let stopped_with_error = Arc::new(Mutex::new(false));
4573 let stopped_with_error_clone = stopped_with_error.clone();
4574
4575 let _subscription = thread.update(cx, |_, cx| {
4576 cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| {
4577 if let ThreadEvent::Stopped(Err(_)) = event {
4578 *stopped_with_error_clone.lock() = true;
4579 }
4580 })
4581 });
4582
4583 // Start initial completion
4584 thread.update(cx, |thread, cx| {
4585 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4586 });
4587 cx.run_until_parked();
4588
4589 // Advance through all retries
4590 for _ in 0..MAX_RETRY_ATTEMPTS {
4591 cx.executor().advance_clock(BASE_RETRY_DELAY);
4592 cx.run_until_parked();
4593 }
4594
4595 let retry_count = thread.update(cx, |thread, _| {
4596 thread
4597 .messages
4598 .iter()
4599 .filter(|m| {
4600 m.ui_only
4601 && m.segments.iter().any(|s| {
4602 if let MessageSegment::Text(text) = s {
4603 text.contains("Retrying") && text.contains("seconds")
4604 } else {
4605 false
4606 }
4607 })
4608 })
4609 .count()
4610 });
4611
4612 // After max retries, should emit Stopped(Err(...)) event
4613 assert_eq!(
4614 retry_count, MAX_RETRY_ATTEMPTS as usize,
4615 "Should have attempted MAX_RETRY_ATTEMPTS retries for overloaded errors"
4616 );
4617 assert!(
4618 *stopped_with_error.lock(),
4619 "Should emit Stopped(Err(...)) event after max retries exceeded"
4620 );
4621
4622 // Retry state should be cleared
4623 thread.read_with(cx, |thread, _| {
4624 assert!(
4625 thread.retry_state.is_none(),
4626 "Retry state should be cleared after max retries"
4627 );
4628
4629 // Verify we have the expected number of retry messages
4630 let retry_messages = thread
4631 .messages
4632 .iter()
4633 .filter(|msg| msg.ui_only && msg.role == Role::System)
4634 .count();
4635 assert_eq!(
4636 retry_messages, MAX_RETRY_ATTEMPTS as usize,
4637 "Should have MAX_RETRY_ATTEMPTS retry messages for overloaded errors"
4638 );
4639 });
4640 }
4641
4642 #[gpui::test]
4643 async fn test_retry_message_removed_on_retry(cx: &mut TestAppContext) {
4644 init_test_settings(cx);
4645
4646 let project = create_test_project(cx, json!({})).await;
4647 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4648
4649 // Enable Burn Mode to allow retries
4650 thread.update(cx, |thread, _| {
4651 thread.set_completion_mode(CompletionMode::Burn);
4652 });
4653
4654 // We'll use a wrapper to switch behavior after first failure
4655 struct RetryTestModel {
4656 inner: Arc<FakeLanguageModel>,
4657 failed_once: Arc<Mutex<bool>>,
4658 }
4659
4660 impl LanguageModel for RetryTestModel {
4661 fn id(&self) -> LanguageModelId {
4662 self.inner.id()
4663 }
4664
4665 fn name(&self) -> LanguageModelName {
4666 self.inner.name()
4667 }
4668
4669 fn provider_id(&self) -> LanguageModelProviderId {
4670 self.inner.provider_id()
4671 }
4672
4673 fn provider_name(&self) -> LanguageModelProviderName {
4674 self.inner.provider_name()
4675 }
4676
4677 fn supports_tools(&self) -> bool {
4678 self.inner.supports_tools()
4679 }
4680
4681 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
4682 self.inner.supports_tool_choice(choice)
4683 }
4684
4685 fn supports_images(&self) -> bool {
4686 self.inner.supports_images()
4687 }
4688
4689 fn telemetry_id(&self) -> String {
4690 self.inner.telemetry_id()
4691 }
4692
4693 fn max_token_count(&self) -> u64 {
4694 self.inner.max_token_count()
4695 }
4696
4697 fn count_tokens(
4698 &self,
4699 request: LanguageModelRequest,
4700 cx: &App,
4701 ) -> BoxFuture<'static, Result<u64>> {
4702 self.inner.count_tokens(request, cx)
4703 }
4704
4705 fn stream_completion(
4706 &self,
4707 request: LanguageModelRequest,
4708 cx: &AsyncApp,
4709 ) -> BoxFuture<
4710 'static,
4711 Result<
4712 BoxStream<
4713 'static,
4714 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
4715 >,
4716 LanguageModelCompletionError,
4717 >,
4718 > {
4719 if !*self.failed_once.lock() {
4720 *self.failed_once.lock() = true;
4721 let provider = self.provider_name();
4722 // Return error on first attempt
4723 let stream = futures::stream::once(async move {
4724 Err(LanguageModelCompletionError::ServerOverloaded {
4725 provider,
4726 retry_after: None,
4727 })
4728 });
4729 async move { Ok(stream.boxed()) }.boxed()
4730 } else {
4731 // Succeed on retry
4732 self.inner.stream_completion(request, cx)
4733 }
4734 }
4735
4736 fn as_fake(&self) -> &FakeLanguageModel {
4737 &self.inner
4738 }
4739 }
4740
4741 let model = Arc::new(RetryTestModel {
4742 inner: Arc::new(FakeLanguageModel::default()),
4743 failed_once: Arc::new(Mutex::new(false)),
4744 });
4745
4746 // Insert a user message
4747 thread.update(cx, |thread, cx| {
4748 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4749 });
4750
4751 // Track message deletions
4752 // Track when retry completes successfully
4753 let retry_completed = Arc::new(Mutex::new(false));
4754 let retry_completed_clone = retry_completed.clone();
4755
4756 let _subscription = thread.update(cx, |_, cx| {
4757 cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| {
4758 if let ThreadEvent::StreamedCompletion = event {
4759 *retry_completed_clone.lock() = true;
4760 }
4761 })
4762 });
4763
4764 // Start completion
4765 thread.update(cx, |thread, cx| {
4766 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4767 });
4768 cx.run_until_parked();
4769
4770 // Get the retry message ID
4771 let retry_message_id = thread.read_with(cx, |thread, _| {
4772 thread
4773 .messages()
4774 .find(|msg| msg.role == Role::System && msg.ui_only)
4775 .map(|msg| msg.id)
4776 .expect("Should have a retry message")
4777 });
4778
4779 // Wait for retry
4780 cx.executor().advance_clock(BASE_RETRY_DELAY);
4781 cx.run_until_parked();
4782
4783 // Stream some successful content
4784 let fake_model = model.as_fake();
4785 // After the retry, there should be a new pending completion
4786 let pending = fake_model.pending_completions();
4787 assert!(
4788 !pending.is_empty(),
4789 "Should have a pending completion after retry"
4790 );
4791 fake_model.send_completion_stream_text_chunk(&pending[0], "Success!");
4792 fake_model.end_completion_stream(&pending[0]);
4793 cx.run_until_parked();
4794
4795 // Check that the retry completed successfully
4796 assert!(
4797 *retry_completed.lock(),
4798 "Retry should have completed successfully"
4799 );
4800
4801 // Retry message should still exist but be marked as ui_only
4802 thread.read_with(cx, |thread, _| {
4803 let retry_msg = thread
4804 .message(retry_message_id)
4805 .expect("Retry message should still exist");
4806 assert!(retry_msg.ui_only, "Retry message should be ui_only");
4807 assert_eq!(
4808 retry_msg.role,
4809 Role::System,
4810 "Retry message should have System role"
4811 );
4812 });
4813 }
4814
4815 #[gpui::test]
4816 async fn test_successful_completion_clears_retry_state(cx: &mut TestAppContext) {
4817 init_test_settings(cx);
4818
4819 let project = create_test_project(cx, json!({})).await;
4820 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4821
4822 // Enable Burn Mode to allow retries
4823 thread.update(cx, |thread, _| {
4824 thread.set_completion_mode(CompletionMode::Burn);
4825 });
4826
4827 // Create a model that fails once then succeeds
4828 struct FailOnceModel {
4829 inner: Arc<FakeLanguageModel>,
4830 failed_once: Arc<Mutex<bool>>,
4831 }
4832
4833 impl LanguageModel for FailOnceModel {
4834 fn id(&self) -> LanguageModelId {
4835 self.inner.id()
4836 }
4837
4838 fn name(&self) -> LanguageModelName {
4839 self.inner.name()
4840 }
4841
4842 fn provider_id(&self) -> LanguageModelProviderId {
4843 self.inner.provider_id()
4844 }
4845
4846 fn provider_name(&self) -> LanguageModelProviderName {
4847 self.inner.provider_name()
4848 }
4849
4850 fn supports_tools(&self) -> bool {
4851 self.inner.supports_tools()
4852 }
4853
4854 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
4855 self.inner.supports_tool_choice(choice)
4856 }
4857
4858 fn supports_images(&self) -> bool {
4859 self.inner.supports_images()
4860 }
4861
4862 fn telemetry_id(&self) -> String {
4863 self.inner.telemetry_id()
4864 }
4865
4866 fn max_token_count(&self) -> u64 {
4867 self.inner.max_token_count()
4868 }
4869
4870 fn count_tokens(
4871 &self,
4872 request: LanguageModelRequest,
4873 cx: &App,
4874 ) -> BoxFuture<'static, Result<u64>> {
4875 self.inner.count_tokens(request, cx)
4876 }
4877
4878 fn stream_completion(
4879 &self,
4880 request: LanguageModelRequest,
4881 cx: &AsyncApp,
4882 ) -> BoxFuture<
4883 'static,
4884 Result<
4885 BoxStream<
4886 'static,
4887 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
4888 >,
4889 LanguageModelCompletionError,
4890 >,
4891 > {
4892 if !*self.failed_once.lock() {
4893 *self.failed_once.lock() = true;
4894 let provider = self.provider_name();
4895 // Return error on first attempt
4896 let stream = futures::stream::once(async move {
4897 Err(LanguageModelCompletionError::ServerOverloaded {
4898 provider,
4899 retry_after: None,
4900 })
4901 });
4902 async move { Ok(stream.boxed()) }.boxed()
4903 } else {
4904 // Succeed on retry
4905 self.inner.stream_completion(request, cx)
4906 }
4907 }
4908 }
4909
4910 let fail_once_model = Arc::new(FailOnceModel {
4911 inner: Arc::new(FakeLanguageModel::default()),
4912 failed_once: Arc::new(Mutex::new(false)),
4913 });
4914
4915 // Insert a user message
4916 thread.update(cx, |thread, cx| {
4917 thread.insert_user_message(
4918 "Test message",
4919 ContextLoadResult::default(),
4920 None,
4921 vec![],
4922 cx,
4923 );
4924 });
4925
4926 // Start completion with fail-once model
4927 thread.update(cx, |thread, cx| {
4928 thread.send_to_model(
4929 fail_once_model.clone(),
4930 CompletionIntent::UserPrompt,
4931 None,
4932 cx,
4933 );
4934 });
4935
4936 cx.run_until_parked();
4937
4938 // Verify retry state exists after first failure
4939 thread.read_with(cx, |thread, _| {
4940 assert!(
4941 thread.retry_state.is_some(),
4942 "Should have retry state after failure"
4943 );
4944 });
4945
4946 // Wait for retry delay
4947 cx.executor().advance_clock(BASE_RETRY_DELAY);
4948 cx.run_until_parked();
4949
4950 // The retry should now use our FailOnceModel which should succeed
4951 // We need to help the FakeLanguageModel complete the stream
4952 let inner_fake = fail_once_model.inner.clone();
4953
4954 // Wait a bit for the retry to start
4955 cx.run_until_parked();
4956
4957 // Check for pending completions and complete them
4958 if let Some(pending) = inner_fake.pending_completions().first() {
4959 inner_fake.send_completion_stream_text_chunk(pending, "Success!");
4960 inner_fake.end_completion_stream(pending);
4961 }
4962 cx.run_until_parked();
4963
4964 thread.read_with(cx, |thread, _| {
4965 assert!(
4966 thread.retry_state.is_none(),
4967 "Retry state should be cleared after successful completion"
4968 );
4969
4970 let has_assistant_message = thread
4971 .messages
4972 .iter()
4973 .any(|msg| msg.role == Role::Assistant && !msg.ui_only);
4974 assert!(
4975 has_assistant_message,
4976 "Should have an assistant message after successful retry"
4977 );
4978 });
4979 }
4980
4981 #[gpui::test]
4982 async fn test_rate_limit_retry_single_attempt(cx: &mut TestAppContext) {
4983 init_test_settings(cx);
4984
4985 let project = create_test_project(cx, json!({})).await;
4986 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4987
4988 // Enable Burn Mode to allow retries
4989 thread.update(cx, |thread, _| {
4990 thread.set_completion_mode(CompletionMode::Burn);
4991 });
4992
4993 // Create a model that returns rate limit error with retry_after
4994 struct RateLimitModel {
4995 inner: Arc<FakeLanguageModel>,
4996 }
4997
4998 impl LanguageModel for RateLimitModel {
4999 fn id(&self) -> LanguageModelId {
5000 self.inner.id()
5001 }
5002
5003 fn name(&self) -> LanguageModelName {
5004 self.inner.name()
5005 }
5006
5007 fn provider_id(&self) -> LanguageModelProviderId {
5008 self.inner.provider_id()
5009 }
5010
5011 fn provider_name(&self) -> LanguageModelProviderName {
5012 self.inner.provider_name()
5013 }
5014
5015 fn supports_tools(&self) -> bool {
5016 self.inner.supports_tools()
5017 }
5018
5019 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
5020 self.inner.supports_tool_choice(choice)
5021 }
5022
5023 fn supports_images(&self) -> bool {
5024 self.inner.supports_images()
5025 }
5026
5027 fn telemetry_id(&self) -> String {
5028 self.inner.telemetry_id()
5029 }
5030
5031 fn max_token_count(&self) -> u64 {
5032 self.inner.max_token_count()
5033 }
5034
5035 fn count_tokens(
5036 &self,
5037 request: LanguageModelRequest,
5038 cx: &App,
5039 ) -> BoxFuture<'static, Result<u64>> {
5040 self.inner.count_tokens(request, cx)
5041 }
5042
5043 fn stream_completion(
5044 &self,
5045 _request: LanguageModelRequest,
5046 _cx: &AsyncApp,
5047 ) -> BoxFuture<
5048 'static,
5049 Result<
5050 BoxStream<
5051 'static,
5052 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
5053 >,
5054 LanguageModelCompletionError,
5055 >,
5056 > {
5057 let provider = self.provider_name();
5058 async move {
5059 let stream = futures::stream::once(async move {
5060 Err(LanguageModelCompletionError::RateLimitExceeded {
5061 provider,
5062 retry_after: Some(Duration::from_secs(TEST_RATE_LIMIT_RETRY_SECS)),
5063 })
5064 });
5065 Ok(stream.boxed())
5066 }
5067 .boxed()
5068 }
5069
5070 fn as_fake(&self) -> &FakeLanguageModel {
5071 &self.inner
5072 }
5073 }
5074
5075 let model = Arc::new(RateLimitModel {
5076 inner: Arc::new(FakeLanguageModel::default()),
5077 });
5078
5079 // Insert a user message
5080 thread.update(cx, |thread, cx| {
5081 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
5082 });
5083
5084 // Start completion
5085 thread.update(cx, |thread, cx| {
5086 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
5087 });
5088
5089 cx.run_until_parked();
5090
5091 let retry_count = thread.update(cx, |thread, _| {
5092 thread
5093 .messages
5094 .iter()
5095 .filter(|m| {
5096 m.ui_only
5097 && m.segments.iter().any(|s| {
5098 if let MessageSegment::Text(text) = s {
5099 text.contains("rate limit exceeded")
5100 } else {
5101 false
5102 }
5103 })
5104 })
5105 .count()
5106 });
5107 assert_eq!(retry_count, 1, "Should have scheduled one retry");
5108
5109 thread.read_with(cx, |thread, _| {
5110 assert!(
5111 thread.retry_state.is_some(),
5112 "Rate limit errors should set retry_state"
5113 );
5114 if let Some(retry_state) = &thread.retry_state {
5115 assert_eq!(
5116 retry_state.max_attempts, MAX_RETRY_ATTEMPTS,
5117 "Rate limit errors should use MAX_RETRY_ATTEMPTS"
5118 );
5119 }
5120 });
5121
5122 // Verify we have one retry message
5123 thread.read_with(cx, |thread, _| {
5124 let retry_messages = thread
5125 .messages
5126 .iter()
5127 .filter(|msg| {
5128 msg.ui_only
5129 && msg.segments.iter().any(|seg| {
5130 if let MessageSegment::Text(text) = seg {
5131 text.contains("rate limit exceeded")
5132 } else {
5133 false
5134 }
5135 })
5136 })
5137 .count();
5138 assert_eq!(
5139 retry_messages, 1,
5140 "Should have one rate limit retry message"
5141 );
5142 });
5143
5144 // Check that retry message doesn't include attempt count
5145 thread.read_with(cx, |thread, _| {
5146 let retry_message = thread
5147 .messages
5148 .iter()
5149 .find(|msg| msg.role == Role::System && msg.ui_only)
5150 .expect("Should have a retry message");
5151
5152 // Check that the message contains attempt count since we use retry_state
5153 if let Some(MessageSegment::Text(text)) = retry_message.segments.first() {
5154 assert!(
5155 text.contains(&format!("attempt 1 of {}", MAX_RETRY_ATTEMPTS)),
5156 "Rate limit retry message should contain attempt count with MAX_RETRY_ATTEMPTS"
5157 );
5158 assert!(
5159 text.contains("Retrying"),
5160 "Rate limit retry message should contain retry text"
5161 );
5162 }
5163 });
5164 }
5165
5166 #[gpui::test]
5167 async fn test_ui_only_messages_not_sent_to_model(cx: &mut TestAppContext) {
5168 init_test_settings(cx);
5169
5170 let project = create_test_project(cx, json!({})).await;
5171 let (_, _, thread, _, model) = setup_test_environment(cx, project.clone()).await;
5172
5173 // Insert a regular user message
5174 thread.update(cx, |thread, cx| {
5175 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
5176 });
5177
5178 // Insert a UI-only message (like our retry notifications)
5179 thread.update(cx, |thread, cx| {
5180 let id = thread.next_message_id.post_inc();
5181 thread.messages.push(Message {
5182 id,
5183 role: Role::System,
5184 segments: vec![MessageSegment::Text(
5185 "This is a UI-only message that should not be sent to the model".to_string(),
5186 )],
5187 loaded_context: LoadedContext::default(),
5188 creases: Vec::new(),
5189 is_hidden: true,
5190 ui_only: true,
5191 });
5192 cx.emit(ThreadEvent::MessageAdded(id));
5193 });
5194
5195 // Insert another regular message
5196 thread.update(cx, |thread, cx| {
5197 thread.insert_user_message(
5198 "How are you?",
5199 ContextLoadResult::default(),
5200 None,
5201 vec![],
5202 cx,
5203 );
5204 });
5205
5206 // Generate the completion request
5207 let request = thread.update(cx, |thread, cx| {
5208 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
5209 });
5210
5211 // Verify that the request only contains non-UI-only messages
5212 // Should have system prompt + 2 user messages, but not the UI-only message
5213 let user_messages: Vec<_> = request
5214 .messages
5215 .iter()
5216 .filter(|msg| msg.role == Role::User)
5217 .collect();
5218 assert_eq!(
5219 user_messages.len(),
5220 2,
5221 "Should have exactly 2 user messages"
5222 );
5223
5224 // Verify the UI-only content is not present anywhere in the request
5225 let request_text = request
5226 .messages
5227 .iter()
5228 .flat_map(|msg| &msg.content)
5229 .filter_map(|content| match content {
5230 MessageContent::Text(text) => Some(text.as_str()),
5231 _ => None,
5232 })
5233 .collect::<String>();
5234
5235 assert!(
5236 !request_text.contains("UI-only message"),
5237 "UI-only message content should not be in the request"
5238 );
5239
5240 // Verify the thread still has all 3 messages (including UI-only)
5241 thread.read_with(cx, |thread, _| {
5242 assert_eq!(
5243 thread.messages().count(),
5244 3,
5245 "Thread should have 3 messages"
5246 );
5247 assert_eq!(
5248 thread.messages().filter(|m| m.ui_only).count(),
5249 1,
5250 "Thread should have 1 UI-only message"
5251 );
5252 });
5253
5254 // Verify that UI-only messages are not serialized
5255 let serialized = thread
5256 .update(cx, |thread, cx| thread.serialize(cx))
5257 .await
5258 .unwrap();
5259 assert_eq!(
5260 serialized.messages.len(),
5261 2,
5262 "Serialized thread should only have 2 messages (no UI-only)"
5263 );
5264 }
5265
5266 #[gpui::test]
5267 async fn test_no_retry_without_burn_mode(cx: &mut TestAppContext) {
5268 init_test_settings(cx);
5269
5270 let project = create_test_project(cx, json!({})).await;
5271 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
5272
5273 // Ensure we're in Normal mode (not Burn mode)
5274 thread.update(cx, |thread, _| {
5275 thread.set_completion_mode(CompletionMode::Normal);
5276 });
5277
5278 // Track error events
5279 let error_events = Arc::new(Mutex::new(Vec::new()));
5280 let error_events_clone = error_events.clone();
5281
5282 let _subscription = thread.update(cx, |_, cx| {
5283 cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| {
5284 if let ThreadEvent::ShowError(error) = event {
5285 error_events_clone.lock().push(error.clone());
5286 }
5287 })
5288 });
5289
5290 // Create model that returns overloaded error
5291 let model = Arc::new(ErrorInjector::new(TestError::Overloaded));
5292
5293 // Insert a user message
5294 thread.update(cx, |thread, cx| {
5295 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
5296 });
5297
5298 // Start completion
5299 thread.update(cx, |thread, cx| {
5300 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
5301 });
5302
5303 cx.run_until_parked();
5304
5305 // Verify no retry state was created
5306 thread.read_with(cx, |thread, _| {
5307 assert!(
5308 thread.retry_state.is_none(),
5309 "Should not have retry state in Normal mode"
5310 );
5311 });
5312
5313 // Check that a retryable error was reported
5314 let errors = error_events.lock();
5315 assert!(!errors.is_empty(), "Should have received an error event");
5316
5317 if let ThreadError::RetryableError {
5318 message: _,
5319 can_enable_burn_mode,
5320 } = &errors[0]
5321 {
5322 assert!(
5323 *can_enable_burn_mode,
5324 "Error should indicate burn mode can be enabled"
5325 );
5326 } else {
5327 panic!("Expected RetryableError, got {:?}", errors[0]);
5328 }
5329
5330 // Verify the thread is no longer generating
5331 thread.read_with(cx, |thread, _| {
5332 assert!(
5333 !thread.is_generating(),
5334 "Should not be generating after error without retry"
5335 );
5336 });
5337 }
5338
5339 #[gpui::test]
5340 async fn test_retry_cancelled_on_stop(cx: &mut TestAppContext) {
5341 init_test_settings(cx);
5342
5343 let project = create_test_project(cx, json!({})).await;
5344 let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
5345
5346 // Enable Burn Mode to allow retries
5347 thread.update(cx, |thread, _| {
5348 thread.set_completion_mode(CompletionMode::Burn);
5349 });
5350
5351 // Create model that returns overloaded error
5352 let model = Arc::new(ErrorInjector::new(TestError::Overloaded));
5353
5354 // Insert a user message
5355 thread.update(cx, |thread, cx| {
5356 thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
5357 });
5358
5359 // Start completion
5360 thread.update(cx, |thread, cx| {
5361 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
5362 });
5363
5364 cx.run_until_parked();
5365
5366 // Verify retry was scheduled by checking for retry message
5367 let has_retry_message = thread.read_with(cx, |thread, _| {
5368 thread.messages.iter().any(|m| {
5369 m.ui_only
5370 && m.segments.iter().any(|s| {
5371 if let MessageSegment::Text(text) = s {
5372 text.contains("Retrying") && text.contains("seconds")
5373 } else {
5374 false
5375 }
5376 })
5377 })
5378 });
5379 assert!(has_retry_message, "Should have scheduled a retry");
5380
5381 // Cancel the completion before the retry happens
5382 thread.update(cx, |thread, cx| {
5383 thread.cancel_last_completion(None, cx);
5384 });
5385
5386 cx.run_until_parked();
5387
5388 // The retry should not have happened - no pending completions
5389 let fake_model = model.as_fake();
5390 assert_eq!(
5391 fake_model.pending_completions().len(),
5392 0,
5393 "Should have no pending completions after cancellation"
5394 );
5395
5396 // Verify the retry was cancelled by checking retry state
5397 thread.read_with(cx, |thread, _| {
5398 if let Some(retry_state) = &thread.retry_state {
5399 panic!(
5400 "retry_state should be cleared after cancellation, but found: attempt={}, max_attempts={}, intent={:?}",
5401 retry_state.attempt, retry_state.max_attempts, retry_state.intent
5402 );
5403 }
5404 });
5405 }
5406
5407 fn test_summarize_error(
5408 model: &Arc<dyn LanguageModel>,
5409 thread: &Entity<Thread>,
5410 cx: &mut TestAppContext,
5411 ) {
5412 thread.update(cx, |thread, cx| {
5413 thread.insert_user_message("Hi!", ContextLoadResult::default(), None, vec![], cx);
5414 thread.send_to_model(
5415 model.clone(),
5416 CompletionIntent::ThreadSummarization,
5417 None,
5418 cx,
5419 );
5420 });
5421
5422 let fake_model = model.as_fake();
5423 simulate_successful_response(&fake_model, cx);
5424
5425 thread.read_with(cx, |thread, _| {
5426 assert!(matches!(thread.summary(), ThreadSummary::Generating));
5427 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
5428 });
5429
5430 // Simulate summary request ending
5431 cx.run_until_parked();
5432 fake_model.end_last_completion_stream();
5433 cx.run_until_parked();
5434
5435 // State is set to Error and default message
5436 thread.read_with(cx, |thread, _| {
5437 assert!(matches!(thread.summary(), ThreadSummary::Error));
5438 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
5439 });
5440 }
5441
5442 fn simulate_successful_response(fake_model: &FakeLanguageModel, cx: &mut TestAppContext) {
5443 cx.run_until_parked();
5444 fake_model.send_last_completion_stream_text_chunk("Assistant response");
5445 fake_model.end_last_completion_stream();
5446 cx.run_until_parked();
5447 }
5448
5449 fn init_test_settings(cx: &mut TestAppContext) {
5450 cx.update(|cx| {
5451 let settings_store = SettingsStore::test(cx);
5452 cx.set_global(settings_store);
5453 language::init(cx);
5454 Project::init_settings(cx);
5455 AgentSettings::register(cx);
5456 prompt_store::init(cx);
5457 thread_store::init(cx);
5458 workspace::init_settings(cx);
5459 language_model::init_settings(cx);
5460 ThemeSettings::register(cx);
5461 ToolRegistry::default_global(cx);
5462 assistant_tool::init(cx);
5463
5464 let http_client = Arc::new(http_client::HttpClientWithUrl::new(
5465 http_client::FakeHttpClient::with_200_response(),
5466 "http://localhost".to_string(),
5467 None,
5468 ));
5469 assistant_tools::init(http_client, cx);
5470 });
5471 }
5472
5473 // Helper to create a test project with test files
5474 async fn create_test_project(
5475 cx: &mut TestAppContext,
5476 files: serde_json::Value,
5477 ) -> Entity<Project> {
5478 let fs = FakeFs::new(cx.executor());
5479 fs.insert_tree(path!("/test"), files).await;
5480 Project::test(fs, [path!("/test").as_ref()], cx).await
5481 }
5482
5483 async fn setup_test_environment(
5484 cx: &mut TestAppContext,
5485 project: Entity<Project>,
5486 ) -> (
5487 Entity<Workspace>,
5488 Entity<ThreadStore>,
5489 Entity<Thread>,
5490 Entity<ContextStore>,
5491 Arc<dyn LanguageModel>,
5492 ) {
5493 let (workspace, cx) =
5494 cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
5495
5496 let thread_store = cx
5497 .update(|_, cx| {
5498 ThreadStore::load(
5499 project.clone(),
5500 cx.new(|_| ToolWorkingSet::default()),
5501 None,
5502 Arc::new(PromptBuilder::new(None).unwrap()),
5503 cx,
5504 )
5505 })
5506 .await
5507 .unwrap();
5508
5509 let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
5510 let context_store = cx.new(|_cx| ContextStore::new(project.downgrade(), None));
5511
5512 let provider = Arc::new(FakeLanguageModelProvider::default());
5513 let model = provider.test_model();
5514 let model: Arc<dyn LanguageModel> = Arc::new(model);
5515
5516 cx.update(|_, cx| {
5517 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
5518 registry.set_default_model(
5519 Some(ConfiguredModel {
5520 provider: provider.clone(),
5521 model: model.clone(),
5522 }),
5523 cx,
5524 );
5525 registry.set_thread_summary_model(
5526 Some(ConfiguredModel {
5527 provider,
5528 model: model.clone(),
5529 }),
5530 cx,
5531 );
5532 })
5533 });
5534
5535 (workspace, thread_store, thread, context_store, model)
5536 }
5537
5538 async fn add_file_to_context(
5539 project: &Entity<Project>,
5540 context_store: &Entity<ContextStore>,
5541 path: &str,
5542 cx: &mut TestAppContext,
5543 ) -> Result<Entity<language::Buffer>> {
5544 let buffer_path = project
5545 .read_with(cx, |project, cx| project.find_project_path(path, cx))
5546 .unwrap();
5547
5548 let buffer = project
5549 .update(cx, |project, cx| {
5550 project.open_buffer(buffer_path.clone(), cx)
5551 })
5552 .await
5553 .unwrap();
5554
5555 context_store.update(cx, |context_store, cx| {
5556 context_store.add_file_from_buffer(&buffer_path, buffer.clone(), false, cx);
5557 });
5558
5559 Ok(buffer)
5560 }
5561}