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