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