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