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