1use std::fmt::Write as _;
2use std::io::Write;
3use std::ops::Range;
4use std::sync::Arc;
5use std::time::Instant;
6
7use agent_settings::{AgentSettings, CompletionMode};
8use anyhow::{Result, anyhow};
9use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolWorkingSet};
10use chrono::{DateTime, Utc};
11use collections::HashMap;
12use editor::display_map::CreaseMetadata;
13use feature_flags::{self, FeatureFlagAppExt};
14use futures::future::Shared;
15use futures::{FutureExt, StreamExt as _};
16use git::repository::DiffType;
17use gpui::{
18 AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task,
19 WeakEntity,
20};
21use language_model::{
22 ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
23 LanguageModelId, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest,
24 LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
25 LanguageModelToolResultContent, LanguageModelToolUseId, MessageContent,
26 ModelRequestLimitReachedError, PaymentRequiredError, RequestUsage, Role, SelectedModel,
27 StopReason, TokenUsage,
28};
29use postage::stream::Stream as _;
30use project::Project;
31use project::git_store::{GitStore, GitStoreCheckpoint, RepositoryState};
32use prompt_store::{ModelContext, PromptBuilder};
33use proto::Plan;
34use schemars::JsonSchema;
35use serde::{Deserialize, Serialize};
36use settings::Settings;
37use thiserror::Error;
38use ui::Window;
39use util::{ResultExt as _, post_inc};
40use uuid::Uuid;
41use zed_llm_client::{CompletionIntent, CompletionRequestStatus};
42
43use crate::ThreadStore;
44use crate::context::{AgentContext, AgentContextHandle, ContextLoadResult, LoadedContext};
45use crate::thread_store::{
46 SerializedCrease, SerializedLanguageModel, SerializedMessage, SerializedMessageSegment,
47 SerializedThread, SerializedToolResult, SerializedToolUse, SharedProjectContext,
48};
49use crate::tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState};
50
51#[derive(
52 Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,
53)]
54pub struct ThreadId(Arc<str>);
55
56impl ThreadId {
57 pub fn new() -> Self {
58 Self(Uuid::new_v4().to_string().into())
59 }
60}
61
62impl std::fmt::Display for ThreadId {
63 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64 write!(f, "{}", self.0)
65 }
66}
67
68impl From<&str> for ThreadId {
69 fn from(value: &str) -> Self {
70 Self(value.into())
71 }
72}
73
74/// The ID of the user prompt that initiated a request.
75///
76/// This equates to the user physically submitting a message to the model (e.g., by pressing the Enter key).
77#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize)]
78pub struct PromptId(Arc<str>);
79
80impl PromptId {
81 pub fn new() -> Self {
82 Self(Uuid::new_v4().to_string().into())
83 }
84}
85
86impl std::fmt::Display for PromptId {
87 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88 write!(f, "{}", self.0)
89 }
90}
91
92#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
93pub struct MessageId(pub(crate) usize);
94
95impl MessageId {
96 fn post_inc(&mut self) -> Self {
97 Self(post_inc(&mut self.0))
98 }
99}
100
101/// Stored information that can be used to resurrect a context crease when creating an editor for a past message.
102#[derive(Clone, Debug)]
103pub struct MessageCrease {
104 pub range: Range<usize>,
105 pub metadata: CreaseMetadata,
106 /// None for a deserialized message, Some otherwise.
107 pub context: Option<AgentContextHandle>,
108}
109
110/// A message in a [`Thread`].
111#[derive(Debug, Clone)]
112pub struct Message {
113 pub id: MessageId,
114 pub role: Role,
115 pub segments: Vec<MessageSegment>,
116 pub loaded_context: LoadedContext,
117 pub creases: Vec<MessageCrease>,
118 pub is_hidden: bool,
119}
120
121impl Message {
122 /// Returns whether the message contains any meaningful text that should be displayed
123 /// The model sometimes runs tool without producing any text or just a marker ([`USING_TOOL_MARKER`])
124 pub fn should_display_content(&self) -> bool {
125 self.segments.iter().all(|segment| segment.should_display())
126 }
127
128 pub fn push_thinking(&mut self, text: &str, signature: Option<String>) {
129 if let Some(MessageSegment::Thinking {
130 text: segment,
131 signature: current_signature,
132 }) = self.segments.last_mut()
133 {
134 if let Some(signature) = signature {
135 *current_signature = Some(signature);
136 }
137 segment.push_str(text);
138 } else {
139 self.segments.push(MessageSegment::Thinking {
140 text: text.to_string(),
141 signature,
142 });
143 }
144 }
145
146 pub fn push_text(&mut self, text: &str) {
147 if let Some(MessageSegment::Text(segment)) = self.segments.last_mut() {
148 segment.push_str(text);
149 } else {
150 self.segments.push(MessageSegment::Text(text.to_string()));
151 }
152 }
153
154 pub fn to_string(&self) -> String {
155 let mut result = String::new();
156
157 if !self.loaded_context.text.is_empty() {
158 result.push_str(&self.loaded_context.text);
159 }
160
161 for segment in &self.segments {
162 match segment {
163 MessageSegment::Text(text) => result.push_str(text),
164 MessageSegment::Thinking { text, .. } => {
165 result.push_str("<think>\n");
166 result.push_str(text);
167 result.push_str("\n</think>");
168 }
169 MessageSegment::RedactedThinking(_) => {}
170 }
171 }
172
173 result
174 }
175}
176
177#[derive(Debug, Clone, PartialEq, Eq)]
178pub enum MessageSegment {
179 Text(String),
180 Thinking {
181 text: String,
182 signature: Option<String>,
183 },
184 RedactedThinking(Vec<u8>),
185}
186
187impl MessageSegment {
188 pub fn should_display(&self) -> bool {
189 match self {
190 Self::Text(text) => text.is_empty(),
191 Self::Thinking { text, .. } => text.is_empty(),
192 Self::RedactedThinking(_) => false,
193 }
194 }
195}
196
197#[derive(Debug, Clone, Serialize, Deserialize)]
198pub struct ProjectSnapshot {
199 pub worktree_snapshots: Vec<WorktreeSnapshot>,
200 pub unsaved_buffer_paths: Vec<String>,
201 pub timestamp: DateTime<Utc>,
202}
203
204#[derive(Debug, Clone, Serialize, Deserialize)]
205pub struct WorktreeSnapshot {
206 pub worktree_path: String,
207 pub git_state: Option<GitState>,
208}
209
210#[derive(Debug, Clone, Serialize, Deserialize)]
211pub struct GitState {
212 pub remote_url: Option<String>,
213 pub head_sha: Option<String>,
214 pub current_branch: Option<String>,
215 pub diff: Option<String>,
216}
217
218#[derive(Clone, Debug)]
219pub struct ThreadCheckpoint {
220 message_id: MessageId,
221 git_checkpoint: GitStoreCheckpoint,
222}
223
224#[derive(Copy, Clone, Debug, PartialEq, Eq)]
225pub enum ThreadFeedback {
226 Positive,
227 Negative,
228}
229
230pub enum LastRestoreCheckpoint {
231 Pending {
232 message_id: MessageId,
233 },
234 Error {
235 message_id: MessageId,
236 error: String,
237 },
238}
239
240impl LastRestoreCheckpoint {
241 pub fn message_id(&self) -> MessageId {
242 match self {
243 LastRestoreCheckpoint::Pending { message_id } => *message_id,
244 LastRestoreCheckpoint::Error { message_id, .. } => *message_id,
245 }
246 }
247}
248
249#[derive(Clone, Debug, Default, Serialize, Deserialize)]
250pub enum DetailedSummaryState {
251 #[default]
252 NotGenerated,
253 Generating {
254 message_id: MessageId,
255 },
256 Generated {
257 text: SharedString,
258 message_id: MessageId,
259 },
260}
261
262impl DetailedSummaryState {
263 fn text(&self) -> Option<SharedString> {
264 if let Self::Generated { text, .. } = self {
265 Some(text.clone())
266 } else {
267 None
268 }
269 }
270}
271
272#[derive(Default, Debug)]
273pub struct TotalTokenUsage {
274 pub total: usize,
275 pub max: usize,
276}
277
278impl TotalTokenUsage {
279 pub fn ratio(&self) -> TokenUsageRatio {
280 #[cfg(debug_assertions)]
281 let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
282 .unwrap_or("0.8".to_string())
283 .parse()
284 .unwrap();
285 #[cfg(not(debug_assertions))]
286 let warning_threshold: f32 = 0.8;
287
288 // When the maximum is unknown because there is no selected model,
289 // avoid showing the token limit warning.
290 if self.max == 0 {
291 TokenUsageRatio::Normal
292 } else if self.total >= self.max {
293 TokenUsageRatio::Exceeded
294 } else if self.total as f32 / self.max as f32 >= warning_threshold {
295 TokenUsageRatio::Warning
296 } else {
297 TokenUsageRatio::Normal
298 }
299 }
300
301 pub fn add(&self, tokens: usize) -> TotalTokenUsage {
302 TotalTokenUsage {
303 total: self.total + tokens,
304 max: self.max,
305 }
306 }
307}
308
309#[derive(Debug, Default, PartialEq, Eq)]
310pub enum TokenUsageRatio {
311 #[default]
312 Normal,
313 Warning,
314 Exceeded,
315}
316
317#[derive(Debug, Clone, Copy)]
318pub enum QueueState {
319 Sending,
320 Queued { position: usize },
321 Started,
322}
323
324/// A thread of conversation with the LLM.
325pub struct Thread {
326 id: ThreadId,
327 updated_at: DateTime<Utc>,
328 summary: ThreadSummary,
329 pending_summary: Task<Option<()>>,
330 detailed_summary_task: Task<Option<()>>,
331 detailed_summary_tx: postage::watch::Sender<DetailedSummaryState>,
332 detailed_summary_rx: postage::watch::Receiver<DetailedSummaryState>,
333 completion_mode: agent_settings::CompletionMode,
334 messages: Vec<Message>,
335 next_message_id: MessageId,
336 last_prompt_id: PromptId,
337 project_context: SharedProjectContext,
338 checkpoints_by_message: HashMap<MessageId, ThreadCheckpoint>,
339 completion_count: usize,
340 pending_completions: Vec<PendingCompletion>,
341 project: Entity<Project>,
342 prompt_builder: Arc<PromptBuilder>,
343 tools: Entity<ToolWorkingSet>,
344 tool_use: ToolUseState,
345 action_log: Entity<ActionLog>,
346 last_restore_checkpoint: Option<LastRestoreCheckpoint>,
347 pending_checkpoint: Option<ThreadCheckpoint>,
348 initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
349 request_token_usage: Vec<TokenUsage>,
350 cumulative_token_usage: TokenUsage,
351 exceeded_window_error: Option<ExceededWindowError>,
352 last_usage: Option<RequestUsage>,
353 tool_use_limit_reached: bool,
354 feedback: Option<ThreadFeedback>,
355 message_feedback: HashMap<MessageId, ThreadFeedback>,
356 last_auto_capture_at: Option<Instant>,
357 last_received_chunk_at: Option<Instant>,
358 request_callback: Option<
359 Box<dyn FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>])>,
360 >,
361 remaining_turns: u32,
362 configured_model: Option<ConfiguredModel>,
363}
364
365#[derive(Clone, Debug, PartialEq, Eq)]
366pub enum ThreadSummary {
367 Pending,
368 Generating,
369 Ready(SharedString),
370 Error,
371}
372
373impl ThreadSummary {
374 pub const DEFAULT: SharedString = SharedString::new_static("New Thread");
375
376 pub fn or_default(&self) -> SharedString {
377 self.unwrap_or(Self::DEFAULT)
378 }
379
380 pub fn unwrap_or(&self, message: impl Into<SharedString>) -> SharedString {
381 self.ready().unwrap_or_else(|| message.into())
382 }
383
384 pub fn ready(&self) -> Option<SharedString> {
385 match self {
386 ThreadSummary::Ready(summary) => Some(summary.clone()),
387 ThreadSummary::Pending | ThreadSummary::Generating | ThreadSummary::Error => None,
388 }
389 }
390}
391
392#[derive(Debug, Clone, Serialize, Deserialize)]
393pub struct ExceededWindowError {
394 /// Model used when last message exceeded context window
395 model_id: LanguageModelId,
396 /// Token count including last message
397 token_count: usize,
398}
399
400impl Thread {
401 pub fn new(
402 project: Entity<Project>,
403 tools: Entity<ToolWorkingSet>,
404 prompt_builder: Arc<PromptBuilder>,
405 system_prompt: SharedProjectContext,
406 cx: &mut Context<Self>,
407 ) -> Self {
408 let (detailed_summary_tx, detailed_summary_rx) = postage::watch::channel();
409 let configured_model = LanguageModelRegistry::read_global(cx).default_model();
410
411 Self {
412 id: ThreadId::new(),
413 updated_at: Utc::now(),
414 summary: ThreadSummary::Pending,
415 pending_summary: Task::ready(None),
416 detailed_summary_task: Task::ready(None),
417 detailed_summary_tx,
418 detailed_summary_rx,
419 completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,
420 messages: Vec::new(),
421 next_message_id: MessageId(0),
422 last_prompt_id: PromptId::new(),
423 project_context: system_prompt,
424 checkpoints_by_message: HashMap::default(),
425 completion_count: 0,
426 pending_completions: Vec::new(),
427 project: project.clone(),
428 prompt_builder,
429 tools: tools.clone(),
430 last_restore_checkpoint: None,
431 pending_checkpoint: None,
432 tool_use: ToolUseState::new(tools.clone()),
433 action_log: cx.new(|_| ActionLog::new(project.clone())),
434 initial_project_snapshot: {
435 let project_snapshot = Self::project_snapshot(project, cx);
436 cx.foreground_executor()
437 .spawn(async move { Some(project_snapshot.await) })
438 .shared()
439 },
440 request_token_usage: Vec::new(),
441 cumulative_token_usage: TokenUsage::default(),
442 exceeded_window_error: None,
443 last_usage: None,
444 tool_use_limit_reached: false,
445 feedback: None,
446 message_feedback: HashMap::default(),
447 last_auto_capture_at: None,
448 last_received_chunk_at: None,
449 request_callback: None,
450 remaining_turns: u32::MAX,
451 configured_model,
452 }
453 }
454
455 pub fn deserialize(
456 id: ThreadId,
457 serialized: SerializedThread,
458 project: Entity<Project>,
459 tools: Entity<ToolWorkingSet>,
460 prompt_builder: Arc<PromptBuilder>,
461 project_context: SharedProjectContext,
462 window: Option<&mut Window>, // None in headless mode
463 cx: &mut Context<Self>,
464 ) -> Self {
465 let next_message_id = MessageId(
466 serialized
467 .messages
468 .last()
469 .map(|message| message.id.0 + 1)
470 .unwrap_or(0),
471 );
472 let tool_use = ToolUseState::from_serialized_messages(
473 tools.clone(),
474 &serialized.messages,
475 project.clone(),
476 window,
477 cx,
478 );
479 let (detailed_summary_tx, detailed_summary_rx) =
480 postage::watch::channel_with(serialized.detailed_summary_state);
481
482 let configured_model = LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
483 serialized
484 .model
485 .and_then(|model| {
486 let model = SelectedModel {
487 provider: model.provider.clone().into(),
488 model: model.model.clone().into(),
489 };
490 registry.select_model(&model, cx)
491 })
492 .or_else(|| registry.default_model())
493 });
494
495 let completion_mode = serialized
496 .completion_mode
497 .unwrap_or_else(|| AgentSettings::get_global(cx).preferred_completion_mode);
498
499 Self {
500 id,
501 updated_at: serialized.updated_at,
502 summary: ThreadSummary::Ready(serialized.summary),
503 pending_summary: Task::ready(None),
504 detailed_summary_task: Task::ready(None),
505 detailed_summary_tx,
506 detailed_summary_rx,
507 completion_mode,
508 messages: serialized
509 .messages
510 .into_iter()
511 .map(|message| Message {
512 id: message.id,
513 role: message.role,
514 segments: message
515 .segments
516 .into_iter()
517 .map(|segment| match segment {
518 SerializedMessageSegment::Text { text } => MessageSegment::Text(text),
519 SerializedMessageSegment::Thinking { text, signature } => {
520 MessageSegment::Thinking { text, signature }
521 }
522 SerializedMessageSegment::RedactedThinking { data } => {
523 MessageSegment::RedactedThinking(data)
524 }
525 })
526 .collect(),
527 loaded_context: LoadedContext {
528 contexts: Vec::new(),
529 text: message.context,
530 images: Vec::new(),
531 },
532 creases: message
533 .creases
534 .into_iter()
535 .map(|crease| MessageCrease {
536 range: crease.start..crease.end,
537 metadata: CreaseMetadata {
538 icon_path: crease.icon_path,
539 label: crease.label,
540 },
541 context: None,
542 })
543 .collect(),
544 is_hidden: message.is_hidden,
545 })
546 .collect(),
547 next_message_id,
548 last_prompt_id: PromptId::new(),
549 project_context,
550 checkpoints_by_message: HashMap::default(),
551 completion_count: 0,
552 pending_completions: Vec::new(),
553 last_restore_checkpoint: None,
554 pending_checkpoint: None,
555 project: project.clone(),
556 prompt_builder,
557 tools,
558 tool_use,
559 action_log: cx.new(|_| ActionLog::new(project)),
560 initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
561 request_token_usage: serialized.request_token_usage,
562 cumulative_token_usage: serialized.cumulative_token_usage,
563 exceeded_window_error: None,
564 last_usage: None,
565 tool_use_limit_reached: serialized.tool_use_limit_reached,
566 feedback: None,
567 message_feedback: HashMap::default(),
568 last_auto_capture_at: None,
569 last_received_chunk_at: None,
570 request_callback: None,
571 remaining_turns: u32::MAX,
572 configured_model,
573 }
574 }
575
576 pub fn set_request_callback(
577 &mut self,
578 callback: impl 'static
579 + FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>]),
580 ) {
581 self.request_callback = Some(Box::new(callback));
582 }
583
584 pub fn id(&self) -> &ThreadId {
585 &self.id
586 }
587
588 pub fn is_empty(&self) -> bool {
589 self.messages.is_empty()
590 }
591
592 pub fn updated_at(&self) -> DateTime<Utc> {
593 self.updated_at
594 }
595
596 pub fn touch_updated_at(&mut self) {
597 self.updated_at = Utc::now();
598 }
599
600 pub fn advance_prompt_id(&mut self) {
601 self.last_prompt_id = PromptId::new();
602 }
603
604 pub fn project_context(&self) -> SharedProjectContext {
605 self.project_context.clone()
606 }
607
608 pub fn get_or_init_configured_model(&mut self, cx: &App) -> Option<ConfiguredModel> {
609 if self.configured_model.is_none() {
610 self.configured_model = LanguageModelRegistry::read_global(cx).default_model();
611 }
612 self.configured_model.clone()
613 }
614
615 pub fn configured_model(&self) -> Option<ConfiguredModel> {
616 self.configured_model.clone()
617 }
618
619 pub fn set_configured_model(&mut self, model: Option<ConfiguredModel>, cx: &mut Context<Self>) {
620 self.configured_model = model;
621 cx.notify();
622 }
623
624 pub fn summary(&self) -> &ThreadSummary {
625 &self.summary
626 }
627
628 pub fn set_summary(&mut self, new_summary: impl Into<SharedString>, cx: &mut Context<Self>) {
629 let current_summary = match &self.summary {
630 ThreadSummary::Pending | ThreadSummary::Generating => return,
631 ThreadSummary::Ready(summary) => summary,
632 ThreadSummary::Error => &ThreadSummary::DEFAULT,
633 };
634
635 let mut new_summary = new_summary.into();
636
637 if new_summary.is_empty() {
638 new_summary = ThreadSummary::DEFAULT;
639 }
640
641 if current_summary != &new_summary {
642 self.summary = ThreadSummary::Ready(new_summary);
643 cx.emit(ThreadEvent::SummaryChanged);
644 }
645 }
646
647 pub fn completion_mode(&self) -> CompletionMode {
648 self.completion_mode
649 }
650
651 pub fn set_completion_mode(&mut self, mode: CompletionMode) {
652 self.completion_mode = mode;
653 }
654
655 pub fn message(&self, id: MessageId) -> Option<&Message> {
656 let index = self
657 .messages
658 .binary_search_by(|message| message.id.cmp(&id))
659 .ok()?;
660
661 self.messages.get(index)
662 }
663
664 pub fn messages(&self) -> impl ExactSizeIterator<Item = &Message> {
665 self.messages.iter()
666 }
667
668 pub fn is_generating(&self) -> bool {
669 !self.pending_completions.is_empty() || !self.all_tools_finished()
670 }
671
672 /// Indicates whether streaming of language model events is stale.
673 /// When `is_generating()` is false, this method returns `None`.
674 pub fn is_generation_stale(&self) -> Option<bool> {
675 const STALE_THRESHOLD: u128 = 250;
676
677 self.last_received_chunk_at
678 .map(|instant| instant.elapsed().as_millis() > STALE_THRESHOLD)
679 }
680
681 fn received_chunk(&mut self) {
682 self.last_received_chunk_at = Some(Instant::now());
683 }
684
685 pub fn queue_state(&self) -> Option<QueueState> {
686 self.pending_completions
687 .first()
688 .map(|pending_completion| pending_completion.queue_state)
689 }
690
691 pub fn tools(&self) -> &Entity<ToolWorkingSet> {
692 &self.tools
693 }
694
695 pub fn pending_tool(&self, id: &LanguageModelToolUseId) -> Option<&PendingToolUse> {
696 self.tool_use
697 .pending_tool_uses()
698 .into_iter()
699 .find(|tool_use| &tool_use.id == id)
700 }
701
702 pub fn tools_needing_confirmation(&self) -> impl Iterator<Item = &PendingToolUse> {
703 self.tool_use
704 .pending_tool_uses()
705 .into_iter()
706 .filter(|tool_use| tool_use.status.needs_confirmation())
707 }
708
709 pub fn has_pending_tool_uses(&self) -> bool {
710 !self.tool_use.pending_tool_uses().is_empty()
711 }
712
713 pub fn checkpoint_for_message(&self, id: MessageId) -> Option<ThreadCheckpoint> {
714 self.checkpoints_by_message.get(&id).cloned()
715 }
716
717 pub fn restore_checkpoint(
718 &mut self,
719 checkpoint: ThreadCheckpoint,
720 cx: &mut Context<Self>,
721 ) -> Task<Result<()>> {
722 self.last_restore_checkpoint = Some(LastRestoreCheckpoint::Pending {
723 message_id: checkpoint.message_id,
724 });
725 cx.emit(ThreadEvent::CheckpointChanged);
726 cx.notify();
727
728 let git_store = self.project().read(cx).git_store().clone();
729 let restore = git_store.update(cx, |git_store, cx| {
730 git_store.restore_checkpoint(checkpoint.git_checkpoint.clone(), cx)
731 });
732
733 cx.spawn(async move |this, cx| {
734 let result = restore.await;
735 this.update(cx, |this, cx| {
736 if let Err(err) = result.as_ref() {
737 this.last_restore_checkpoint = Some(LastRestoreCheckpoint::Error {
738 message_id: checkpoint.message_id,
739 error: err.to_string(),
740 });
741 } else {
742 this.truncate(checkpoint.message_id, cx);
743 this.last_restore_checkpoint = None;
744 }
745 this.pending_checkpoint = None;
746 cx.emit(ThreadEvent::CheckpointChanged);
747 cx.notify();
748 })?;
749 result
750 })
751 }
752
753 fn finalize_pending_checkpoint(&mut self, cx: &mut Context<Self>) {
754 let pending_checkpoint = if self.is_generating() {
755 return;
756 } else if let Some(checkpoint) = self.pending_checkpoint.take() {
757 checkpoint
758 } else {
759 return;
760 };
761
762 self.finalize_checkpoint(pending_checkpoint, cx);
763 }
764
765 fn finalize_checkpoint(
766 &mut self,
767 pending_checkpoint: ThreadCheckpoint,
768 cx: &mut Context<Self>,
769 ) {
770 let git_store = self.project.read(cx).git_store().clone();
771 let final_checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
772 cx.spawn(async move |this, cx| match final_checkpoint.await {
773 Ok(final_checkpoint) => {
774 let equal = git_store
775 .update(cx, |store, cx| {
776 store.compare_checkpoints(
777 pending_checkpoint.git_checkpoint.clone(),
778 final_checkpoint.clone(),
779 cx,
780 )
781 })?
782 .await
783 .unwrap_or(false);
784
785 if !equal {
786 this.update(cx, |this, cx| {
787 this.insert_checkpoint(pending_checkpoint, cx)
788 })?;
789 }
790
791 Ok(())
792 }
793 Err(_) => this.update(cx, |this, cx| {
794 this.insert_checkpoint(pending_checkpoint, cx)
795 }),
796 })
797 .detach();
798 }
799
800 fn insert_checkpoint(&mut self, checkpoint: ThreadCheckpoint, cx: &mut Context<Self>) {
801 self.checkpoints_by_message
802 .insert(checkpoint.message_id, checkpoint);
803 cx.emit(ThreadEvent::CheckpointChanged);
804 cx.notify();
805 }
806
807 pub fn last_restore_checkpoint(&self) -> Option<&LastRestoreCheckpoint> {
808 self.last_restore_checkpoint.as_ref()
809 }
810
811 pub fn truncate(&mut self, message_id: MessageId, cx: &mut Context<Self>) {
812 let Some(message_ix) = self
813 .messages
814 .iter()
815 .rposition(|message| message.id == message_id)
816 else {
817 return;
818 };
819 for deleted_message in self.messages.drain(message_ix..) {
820 self.checkpoints_by_message.remove(&deleted_message.id);
821 }
822 cx.notify();
823 }
824
825 pub fn context_for_message(&self, id: MessageId) -> impl Iterator<Item = &AgentContext> {
826 self.messages
827 .iter()
828 .find(|message| message.id == id)
829 .into_iter()
830 .flat_map(|message| message.loaded_context.contexts.iter())
831 }
832
833 pub fn is_turn_end(&self, ix: usize) -> bool {
834 if self.messages.is_empty() {
835 return false;
836 }
837
838 if !self.is_generating() && ix == self.messages.len() - 1 {
839 return true;
840 }
841
842 let Some(message) = self.messages.get(ix) else {
843 return false;
844 };
845
846 if message.role != Role::Assistant {
847 return false;
848 }
849
850 self.messages
851 .get(ix + 1)
852 .and_then(|message| {
853 self.message(message.id)
854 .map(|next_message| next_message.role == Role::User && !next_message.is_hidden)
855 })
856 .unwrap_or(false)
857 }
858
859 pub fn last_usage(&self) -> Option<RequestUsage> {
860 self.last_usage
861 }
862
863 pub fn tool_use_limit_reached(&self) -> bool {
864 self.tool_use_limit_reached
865 }
866
867 /// Returns whether all of the tool uses have finished running.
868 pub fn all_tools_finished(&self) -> bool {
869 // If the only pending tool uses left are the ones with errors, then
870 // that means that we've finished running all of the pending tools.
871 self.tool_use
872 .pending_tool_uses()
873 .iter()
874 .all(|pending_tool_use| pending_tool_use.status.is_error())
875 }
876
877 /// Returns whether any pending tool uses may perform edits
878 pub fn has_pending_edit_tool_uses(&self) -> bool {
879 self.tool_use
880 .pending_tool_uses()
881 .iter()
882 .filter(|pending_tool_use| !pending_tool_use.status.is_error())
883 .any(|pending_tool_use| pending_tool_use.may_perform_edits)
884 }
885
886 pub fn tool_uses_for_message(&self, id: MessageId, cx: &App) -> Vec<ToolUse> {
887 self.tool_use.tool_uses_for_message(id, cx)
888 }
889
890 pub fn tool_results_for_message(
891 &self,
892 assistant_message_id: MessageId,
893 ) -> Vec<&LanguageModelToolResult> {
894 self.tool_use.tool_results_for_message(assistant_message_id)
895 }
896
897 pub fn tool_result(&self, id: &LanguageModelToolUseId) -> Option<&LanguageModelToolResult> {
898 self.tool_use.tool_result(id)
899 }
900
901 pub fn output_for_tool(&self, id: &LanguageModelToolUseId) -> Option<&Arc<str>> {
902 match &self.tool_use.tool_result(id)?.content {
903 LanguageModelToolResultContent::Text(text) => Some(text),
904 LanguageModelToolResultContent::Image(_) => {
905 // TODO: We should display image
906 None
907 }
908 }
909 }
910
911 pub fn card_for_tool(&self, id: &LanguageModelToolUseId) -> Option<AnyToolCard> {
912 self.tool_use.tool_result_card(id).cloned()
913 }
914
915 /// Return tools that are both enabled and supported by the model
916 pub fn available_tools(
917 &self,
918 cx: &App,
919 model: Arc<dyn LanguageModel>,
920 ) -> Vec<LanguageModelRequestTool> {
921 if model.supports_tools() {
922 self.tools()
923 .read(cx)
924 .enabled_tools(cx)
925 .into_iter()
926 .filter_map(|tool| {
927 // Skip tools that cannot be supported
928 let input_schema = tool.input_schema(model.tool_input_format()).ok()?;
929 Some(LanguageModelRequestTool {
930 name: tool.name(),
931 description: tool.description(),
932 input_schema,
933 })
934 })
935 .collect()
936 } else {
937 Vec::default()
938 }
939 }
940
941 pub fn insert_user_message(
942 &mut self,
943 text: impl Into<String>,
944 loaded_context: ContextLoadResult,
945 git_checkpoint: Option<GitStoreCheckpoint>,
946 creases: Vec<MessageCrease>,
947 cx: &mut Context<Self>,
948 ) -> MessageId {
949 if !loaded_context.referenced_buffers.is_empty() {
950 self.action_log.update(cx, |log, cx| {
951 for buffer in loaded_context.referenced_buffers {
952 log.buffer_read(buffer, cx);
953 }
954 });
955 }
956
957 let message_id = self.insert_message(
958 Role::User,
959 vec![MessageSegment::Text(text.into())],
960 loaded_context.loaded_context,
961 creases,
962 false,
963 cx,
964 );
965
966 if let Some(git_checkpoint) = git_checkpoint {
967 self.pending_checkpoint = Some(ThreadCheckpoint {
968 message_id,
969 git_checkpoint,
970 });
971 }
972
973 self.auto_capture_telemetry(cx);
974
975 message_id
976 }
977
978 pub fn insert_invisible_continue_message(&mut self, cx: &mut Context<Self>) -> MessageId {
979 let id = self.insert_message(
980 Role::User,
981 vec![MessageSegment::Text("Continue where you left off".into())],
982 LoadedContext::default(),
983 vec![],
984 true,
985 cx,
986 );
987 self.pending_checkpoint = None;
988
989 id
990 }
991
992 pub fn insert_assistant_message(
993 &mut self,
994 segments: Vec<MessageSegment>,
995 cx: &mut Context<Self>,
996 ) -> MessageId {
997 self.insert_message(
998 Role::Assistant,
999 segments,
1000 LoadedContext::default(),
1001 Vec::new(),
1002 false,
1003 cx,
1004 )
1005 }
1006
1007 pub fn insert_message(
1008 &mut self,
1009 role: Role,
1010 segments: Vec<MessageSegment>,
1011 loaded_context: LoadedContext,
1012 creases: Vec<MessageCrease>,
1013 is_hidden: bool,
1014 cx: &mut Context<Self>,
1015 ) -> MessageId {
1016 let id = self.next_message_id.post_inc();
1017 self.messages.push(Message {
1018 id,
1019 role,
1020 segments,
1021 loaded_context,
1022 creases,
1023 is_hidden,
1024 });
1025 self.touch_updated_at();
1026 cx.emit(ThreadEvent::MessageAdded(id));
1027 id
1028 }
1029
1030 pub fn edit_message(
1031 &mut self,
1032 id: MessageId,
1033 new_role: Role,
1034 new_segments: Vec<MessageSegment>,
1035 loaded_context: Option<LoadedContext>,
1036 checkpoint: Option<GitStoreCheckpoint>,
1037 cx: &mut Context<Self>,
1038 ) -> bool {
1039 let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
1040 return false;
1041 };
1042 message.role = new_role;
1043 message.segments = new_segments;
1044 if let Some(context) = loaded_context {
1045 message.loaded_context = context;
1046 }
1047 if let Some(git_checkpoint) = checkpoint {
1048 self.checkpoints_by_message.insert(
1049 id,
1050 ThreadCheckpoint {
1051 message_id: id,
1052 git_checkpoint,
1053 },
1054 );
1055 }
1056 self.touch_updated_at();
1057 cx.emit(ThreadEvent::MessageEdited(id));
1058 true
1059 }
1060
1061 pub fn delete_message(&mut self, id: MessageId, cx: &mut Context<Self>) -> bool {
1062 let Some(index) = self.messages.iter().position(|message| message.id == id) else {
1063 return false;
1064 };
1065 self.messages.remove(index);
1066 self.touch_updated_at();
1067 cx.emit(ThreadEvent::MessageDeleted(id));
1068 true
1069 }
1070
1071 /// Returns the representation of this [`Thread`] in a textual form.
1072 ///
1073 /// This is the representation we use when attaching a thread as context to another thread.
1074 pub fn text(&self) -> String {
1075 let mut text = String::new();
1076
1077 for message in &self.messages {
1078 text.push_str(match message.role {
1079 language_model::Role::User => "User:",
1080 language_model::Role::Assistant => "Agent:",
1081 language_model::Role::System => "System:",
1082 });
1083 text.push('\n');
1084
1085 for segment in &message.segments {
1086 match segment {
1087 MessageSegment::Text(content) => text.push_str(content),
1088 MessageSegment::Thinking { text: content, .. } => {
1089 text.push_str(&format!("<think>{}</think>", content))
1090 }
1091 MessageSegment::RedactedThinking(_) => {}
1092 }
1093 }
1094 text.push('\n');
1095 }
1096
1097 text
1098 }
1099
1100 /// Serializes this thread into a format for storage or telemetry.
1101 pub fn serialize(&self, cx: &mut Context<Self>) -> Task<Result<SerializedThread>> {
1102 let initial_project_snapshot = self.initial_project_snapshot.clone();
1103 cx.spawn(async move |this, cx| {
1104 let initial_project_snapshot = initial_project_snapshot.await;
1105 this.read_with(cx, |this, cx| SerializedThread {
1106 version: SerializedThread::VERSION.to_string(),
1107 summary: this.summary().or_default(),
1108 updated_at: this.updated_at(),
1109 messages: this
1110 .messages()
1111 .map(|message| SerializedMessage {
1112 id: message.id,
1113 role: message.role,
1114 segments: message
1115 .segments
1116 .iter()
1117 .map(|segment| match segment {
1118 MessageSegment::Text(text) => {
1119 SerializedMessageSegment::Text { text: text.clone() }
1120 }
1121 MessageSegment::Thinking { text, signature } => {
1122 SerializedMessageSegment::Thinking {
1123 text: text.clone(),
1124 signature: signature.clone(),
1125 }
1126 }
1127 MessageSegment::RedactedThinking(data) => {
1128 SerializedMessageSegment::RedactedThinking {
1129 data: data.clone(),
1130 }
1131 }
1132 })
1133 .collect(),
1134 tool_uses: this
1135 .tool_uses_for_message(message.id, cx)
1136 .into_iter()
1137 .map(|tool_use| SerializedToolUse {
1138 id: tool_use.id,
1139 name: tool_use.name,
1140 input: tool_use.input,
1141 })
1142 .collect(),
1143 tool_results: this
1144 .tool_results_for_message(message.id)
1145 .into_iter()
1146 .map(|tool_result| SerializedToolResult {
1147 tool_use_id: tool_result.tool_use_id.clone(),
1148 is_error: tool_result.is_error,
1149 content: tool_result.content.clone(),
1150 output: tool_result.output.clone(),
1151 })
1152 .collect(),
1153 context: message.loaded_context.text.clone(),
1154 creases: message
1155 .creases
1156 .iter()
1157 .map(|crease| SerializedCrease {
1158 start: crease.range.start,
1159 end: crease.range.end,
1160 icon_path: crease.metadata.icon_path.clone(),
1161 label: crease.metadata.label.clone(),
1162 })
1163 .collect(),
1164 is_hidden: message.is_hidden,
1165 })
1166 .collect(),
1167 initial_project_snapshot,
1168 cumulative_token_usage: this.cumulative_token_usage,
1169 request_token_usage: this.request_token_usage.clone(),
1170 detailed_summary_state: this.detailed_summary_rx.borrow().clone(),
1171 exceeded_window_error: this.exceeded_window_error.clone(),
1172 model: this
1173 .configured_model
1174 .as_ref()
1175 .map(|model| SerializedLanguageModel {
1176 provider: model.provider.id().0.to_string(),
1177 model: model.model.id().0.to_string(),
1178 }),
1179 completion_mode: Some(this.completion_mode),
1180 tool_use_limit_reached: this.tool_use_limit_reached,
1181 })
1182 })
1183 }
1184
1185 pub fn remaining_turns(&self) -> u32 {
1186 self.remaining_turns
1187 }
1188
1189 pub fn set_remaining_turns(&mut self, remaining_turns: u32) {
1190 self.remaining_turns = remaining_turns;
1191 }
1192
1193 pub fn send_to_model(
1194 &mut self,
1195 model: Arc<dyn LanguageModel>,
1196 intent: CompletionIntent,
1197 window: Option<AnyWindowHandle>,
1198 cx: &mut Context<Self>,
1199 ) {
1200 if self.remaining_turns == 0 {
1201 return;
1202 }
1203
1204 self.remaining_turns -= 1;
1205
1206 let request = self.to_completion_request(model.clone(), intent, cx);
1207
1208 self.stream_completion(request, model, window, cx);
1209 }
1210
1211 pub fn used_tools_since_last_user_message(&self) -> bool {
1212 for message in self.messages.iter().rev() {
1213 if self.tool_use.message_has_tool_results(message.id) {
1214 return true;
1215 } else if message.role == Role::User {
1216 return false;
1217 }
1218 }
1219
1220 false
1221 }
1222
1223 pub fn to_completion_request(
1224 &self,
1225 model: Arc<dyn LanguageModel>,
1226 intent: CompletionIntent,
1227 cx: &mut Context<Self>,
1228 ) -> LanguageModelRequest {
1229 let mut request = LanguageModelRequest {
1230 thread_id: Some(self.id.to_string()),
1231 prompt_id: Some(self.last_prompt_id.to_string()),
1232 intent: Some(intent),
1233 mode: None,
1234 messages: vec![],
1235 tools: Vec::new(),
1236 tool_choice: None,
1237 stop: Vec::new(),
1238 temperature: AgentSettings::temperature_for_model(&model, cx),
1239 };
1240
1241 let available_tools = self.available_tools(cx, model.clone());
1242 let available_tool_names = available_tools
1243 .iter()
1244 .map(|tool| tool.name.clone())
1245 .collect();
1246
1247 let model_context = &ModelContext {
1248 available_tools: available_tool_names,
1249 };
1250
1251 if let Some(project_context) = self.project_context.borrow().as_ref() {
1252 match self
1253 .prompt_builder
1254 .generate_assistant_system_prompt(project_context, model_context)
1255 {
1256 Err(err) => {
1257 let message = format!("{err:?}").into();
1258 log::error!("{message}");
1259 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1260 header: "Error generating system prompt".into(),
1261 message,
1262 }));
1263 }
1264 Ok(system_prompt) => {
1265 request.messages.push(LanguageModelRequestMessage {
1266 role: Role::System,
1267 content: vec![MessageContent::Text(system_prompt)],
1268 cache: true,
1269 });
1270 }
1271 }
1272 } else {
1273 let message = "Context for system prompt unexpectedly not ready.".into();
1274 log::error!("{message}");
1275 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1276 header: "Error generating system prompt".into(),
1277 message,
1278 }));
1279 }
1280
1281 let mut message_ix_to_cache = None;
1282 for message in &self.messages {
1283 let mut request_message = LanguageModelRequestMessage {
1284 role: message.role,
1285 content: Vec::new(),
1286 cache: false,
1287 };
1288
1289 message
1290 .loaded_context
1291 .add_to_request_message(&mut request_message);
1292
1293 for segment in &message.segments {
1294 match segment {
1295 MessageSegment::Text(text) => {
1296 if !text.is_empty() {
1297 request_message
1298 .content
1299 .push(MessageContent::Text(text.into()));
1300 }
1301 }
1302 MessageSegment::Thinking { text, signature } => {
1303 if !text.is_empty() {
1304 request_message.content.push(MessageContent::Thinking {
1305 text: text.into(),
1306 signature: signature.clone(),
1307 });
1308 }
1309 }
1310 MessageSegment::RedactedThinking(data) => {
1311 request_message
1312 .content
1313 .push(MessageContent::RedactedThinking(data.clone()));
1314 }
1315 };
1316 }
1317
1318 let mut cache_message = true;
1319 let mut tool_results_message = LanguageModelRequestMessage {
1320 role: Role::User,
1321 content: Vec::new(),
1322 cache: false,
1323 };
1324 for (tool_use, tool_result) in self.tool_use.tool_results(message.id) {
1325 if let Some(tool_result) = tool_result {
1326 request_message
1327 .content
1328 .push(MessageContent::ToolUse(tool_use.clone()));
1329 tool_results_message
1330 .content
1331 .push(MessageContent::ToolResult(LanguageModelToolResult {
1332 tool_use_id: tool_use.id.clone(),
1333 tool_name: tool_result.tool_name.clone(),
1334 is_error: tool_result.is_error,
1335 content: if tool_result.content.is_empty() {
1336 // Surprisingly, the API fails if we return an empty string here.
1337 // It thinks we are sending a tool use without a tool result.
1338 "<Tool returned an empty string>".into()
1339 } else {
1340 tool_result.content.clone()
1341 },
1342 output: None,
1343 }));
1344 } else {
1345 cache_message = false;
1346 log::debug!(
1347 "skipped tool use {:?} because it is still pending",
1348 tool_use
1349 );
1350 }
1351 }
1352
1353 if cache_message {
1354 message_ix_to_cache = Some(request.messages.len());
1355 }
1356 request.messages.push(request_message);
1357
1358 if !tool_results_message.content.is_empty() {
1359 if cache_message {
1360 message_ix_to_cache = Some(request.messages.len());
1361 }
1362 request.messages.push(tool_results_message);
1363 }
1364 }
1365
1366 // https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
1367 if let Some(message_ix_to_cache) = message_ix_to_cache {
1368 request.messages[message_ix_to_cache].cache = true;
1369 }
1370
1371 self.attached_tracked_files_state(&mut request.messages, cx);
1372
1373 request.tools = available_tools;
1374 request.mode = if model.supports_max_mode() {
1375 Some(self.completion_mode.into())
1376 } else {
1377 Some(CompletionMode::Normal.into())
1378 };
1379
1380 request
1381 }
1382
1383 fn to_summarize_request(
1384 &self,
1385 model: &Arc<dyn LanguageModel>,
1386 intent: CompletionIntent,
1387 added_user_message: String,
1388 cx: &App,
1389 ) -> LanguageModelRequest {
1390 let mut request = LanguageModelRequest {
1391 thread_id: None,
1392 prompt_id: None,
1393 intent: Some(intent),
1394 mode: None,
1395 messages: vec![],
1396 tools: Vec::new(),
1397 tool_choice: None,
1398 stop: Vec::new(),
1399 temperature: AgentSettings::temperature_for_model(model, cx),
1400 };
1401
1402 for message in &self.messages {
1403 let mut request_message = LanguageModelRequestMessage {
1404 role: message.role,
1405 content: Vec::new(),
1406 cache: false,
1407 };
1408
1409 for segment in &message.segments {
1410 match segment {
1411 MessageSegment::Text(text) => request_message
1412 .content
1413 .push(MessageContent::Text(text.clone())),
1414 MessageSegment::Thinking { .. } => {}
1415 MessageSegment::RedactedThinking(_) => {}
1416 }
1417 }
1418
1419 if request_message.content.is_empty() {
1420 continue;
1421 }
1422
1423 request.messages.push(request_message);
1424 }
1425
1426 request.messages.push(LanguageModelRequestMessage {
1427 role: Role::User,
1428 content: vec![MessageContent::Text(added_user_message)],
1429 cache: false,
1430 });
1431
1432 request
1433 }
1434
1435 fn attached_tracked_files_state(
1436 &self,
1437 messages: &mut Vec<LanguageModelRequestMessage>,
1438 cx: &App,
1439 ) {
1440 const STALE_FILES_HEADER: &str = include_str!("./prompts/stale_files_prompt_header.txt");
1441
1442 let mut stale_message = String::new();
1443
1444 let action_log = self.action_log.read(cx);
1445
1446 for stale_file in action_log.stale_buffers(cx) {
1447 let Some(file) = stale_file.read(cx).file() else {
1448 continue;
1449 };
1450
1451 if stale_message.is_empty() {
1452 write!(&mut stale_message, "{}\n", STALE_FILES_HEADER.trim()).ok();
1453 }
1454
1455 writeln!(&mut stale_message, "- {}", file.path().display()).ok();
1456 }
1457
1458 let mut content = Vec::with_capacity(2);
1459
1460 if !stale_message.is_empty() {
1461 content.push(stale_message.into());
1462 }
1463
1464 if !content.is_empty() {
1465 let context_message = LanguageModelRequestMessage {
1466 role: Role::User,
1467 content,
1468 cache: false,
1469 };
1470
1471 messages.push(context_message);
1472 }
1473 }
1474
1475 pub fn stream_completion(
1476 &mut self,
1477 request: LanguageModelRequest,
1478 model: Arc<dyn LanguageModel>,
1479 window: Option<AnyWindowHandle>,
1480 cx: &mut Context<Self>,
1481 ) {
1482 self.tool_use_limit_reached = false;
1483
1484 let pending_completion_id = post_inc(&mut self.completion_count);
1485 let mut request_callback_parameters = if self.request_callback.is_some() {
1486 Some((request.clone(), Vec::new()))
1487 } else {
1488 None
1489 };
1490 let prompt_id = self.last_prompt_id.clone();
1491 let tool_use_metadata = ToolUseMetadata {
1492 model: model.clone(),
1493 thread_id: self.id.clone(),
1494 prompt_id: prompt_id.clone(),
1495 };
1496
1497 self.last_received_chunk_at = Some(Instant::now());
1498
1499 let task = cx.spawn(async move |thread, cx| {
1500 let stream_completion_future = model.stream_completion(request, &cx);
1501 let initial_token_usage =
1502 thread.read_with(cx, |thread, _cx| thread.cumulative_token_usage);
1503 let stream_completion = async {
1504 let mut events = stream_completion_future.await?;
1505
1506 let mut stop_reason = StopReason::EndTurn;
1507 let mut current_token_usage = TokenUsage::default();
1508
1509 thread
1510 .update(cx, |_thread, cx| {
1511 cx.emit(ThreadEvent::NewRequest);
1512 })
1513 .ok();
1514
1515 let mut request_assistant_message_id = None;
1516
1517 while let Some(event) = events.next().await {
1518 if let Some((_, response_events)) = request_callback_parameters.as_mut() {
1519 response_events
1520 .push(event.as_ref().map_err(|error| error.to_string()).cloned());
1521 }
1522
1523 thread.update(cx, |thread, cx| {
1524 let event = match event {
1525 Ok(event) => event,
1526 Err(LanguageModelCompletionError::BadInputJson {
1527 id,
1528 tool_name,
1529 raw_input: invalid_input_json,
1530 json_parse_error,
1531 }) => {
1532 thread.receive_invalid_tool_json(
1533 id,
1534 tool_name,
1535 invalid_input_json,
1536 json_parse_error,
1537 window,
1538 cx,
1539 );
1540 return Ok(());
1541 }
1542 Err(LanguageModelCompletionError::Other(error)) => {
1543 return Err(error);
1544 }
1545 };
1546
1547 match event {
1548 LanguageModelCompletionEvent::StartMessage { .. } => {
1549 request_assistant_message_id =
1550 Some(thread.insert_assistant_message(
1551 vec![MessageSegment::Text(String::new())],
1552 cx,
1553 ));
1554 }
1555 LanguageModelCompletionEvent::Stop(reason) => {
1556 stop_reason = reason;
1557 }
1558 LanguageModelCompletionEvent::UsageUpdate(token_usage) => {
1559 thread.update_token_usage_at_last_message(token_usage);
1560 thread.cumulative_token_usage = thread.cumulative_token_usage
1561 + token_usage
1562 - current_token_usage;
1563 current_token_usage = token_usage;
1564 }
1565 LanguageModelCompletionEvent::Text(chunk) => {
1566 thread.received_chunk();
1567
1568 cx.emit(ThreadEvent::ReceivedTextChunk);
1569 if let Some(last_message) = thread.messages.last_mut() {
1570 if last_message.role == Role::Assistant
1571 && !thread.tool_use.has_tool_results(last_message.id)
1572 {
1573 last_message.push_text(&chunk);
1574 cx.emit(ThreadEvent::StreamedAssistantText(
1575 last_message.id,
1576 chunk,
1577 ));
1578 } else {
1579 // If we won't have an Assistant message yet, assume this chunk marks the beginning
1580 // of a new Assistant response.
1581 //
1582 // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1583 // will result in duplicating the text of the chunk in the rendered Markdown.
1584 request_assistant_message_id =
1585 Some(thread.insert_assistant_message(
1586 vec![MessageSegment::Text(chunk.to_string())],
1587 cx,
1588 ));
1589 };
1590 }
1591 }
1592 LanguageModelCompletionEvent::Thinking {
1593 text: chunk,
1594 signature,
1595 } => {
1596 thread.received_chunk();
1597
1598 if let Some(last_message) = thread.messages.last_mut() {
1599 if last_message.role == Role::Assistant
1600 && !thread.tool_use.has_tool_results(last_message.id)
1601 {
1602 last_message.push_thinking(&chunk, signature);
1603 cx.emit(ThreadEvent::StreamedAssistantThinking(
1604 last_message.id,
1605 chunk,
1606 ));
1607 } else {
1608 // If we won't have an Assistant message yet, assume this chunk marks the beginning
1609 // of a new Assistant response.
1610 //
1611 // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1612 // will result in duplicating the text of the chunk in the rendered Markdown.
1613 request_assistant_message_id =
1614 Some(thread.insert_assistant_message(
1615 vec![MessageSegment::Thinking {
1616 text: chunk.to_string(),
1617 signature,
1618 }],
1619 cx,
1620 ));
1621 };
1622 }
1623 }
1624 LanguageModelCompletionEvent::ToolUse(tool_use) => {
1625 let last_assistant_message_id = request_assistant_message_id
1626 .unwrap_or_else(|| {
1627 let new_assistant_message_id =
1628 thread.insert_assistant_message(vec![], cx);
1629 request_assistant_message_id =
1630 Some(new_assistant_message_id);
1631 new_assistant_message_id
1632 });
1633
1634 let tool_use_id = tool_use.id.clone();
1635 let streamed_input = if tool_use.is_input_complete {
1636 None
1637 } else {
1638 Some((&tool_use.input).clone())
1639 };
1640
1641 let ui_text = thread.tool_use.request_tool_use(
1642 last_assistant_message_id,
1643 tool_use,
1644 tool_use_metadata.clone(),
1645 cx,
1646 );
1647
1648 if let Some(input) = streamed_input {
1649 cx.emit(ThreadEvent::StreamedToolUse {
1650 tool_use_id,
1651 ui_text,
1652 input,
1653 });
1654 }
1655 }
1656 LanguageModelCompletionEvent::StatusUpdate(status_update) => {
1657 if let Some(completion) = thread
1658 .pending_completions
1659 .iter_mut()
1660 .find(|completion| completion.id == pending_completion_id)
1661 {
1662 match status_update {
1663 CompletionRequestStatus::Queued {
1664 position,
1665 } => {
1666 completion.queue_state = QueueState::Queued { position };
1667 }
1668 CompletionRequestStatus::Started => {
1669 completion.queue_state = QueueState::Started;
1670 }
1671 CompletionRequestStatus::Failed {
1672 code, message, request_id
1673 } => {
1674 anyhow::bail!("completion request failed. request_id: {request_id}, code: {code}, message: {message}");
1675 }
1676 CompletionRequestStatus::UsageUpdated {
1677 amount, limit
1678 } => {
1679 let usage = RequestUsage { limit, amount: amount as i32 };
1680
1681 thread.last_usage = Some(usage);
1682 }
1683 CompletionRequestStatus::ToolUseLimitReached => {
1684 thread.tool_use_limit_reached = true;
1685 cx.emit(ThreadEvent::ToolUseLimitReached);
1686 }
1687 }
1688 }
1689 }
1690 }
1691
1692 thread.touch_updated_at();
1693 cx.emit(ThreadEvent::StreamedCompletion);
1694 cx.notify();
1695
1696 thread.auto_capture_telemetry(cx);
1697 Ok(())
1698 })??;
1699
1700 smol::future::yield_now().await;
1701 }
1702
1703 thread.update(cx, |thread, cx| {
1704 thread.last_received_chunk_at = None;
1705 thread
1706 .pending_completions
1707 .retain(|completion| completion.id != pending_completion_id);
1708
1709 // If there is a response without tool use, summarize the message. Otherwise,
1710 // allow two tool uses before summarizing.
1711 if matches!(thread.summary, ThreadSummary::Pending)
1712 && thread.messages.len() >= 2
1713 && (!thread.has_pending_tool_uses() || thread.messages.len() >= 6)
1714 {
1715 thread.summarize(cx);
1716 }
1717 })?;
1718
1719 anyhow::Ok(stop_reason)
1720 };
1721
1722 let result = stream_completion.await;
1723
1724 thread
1725 .update(cx, |thread, cx| {
1726 thread.finalize_pending_checkpoint(cx);
1727 match result.as_ref() {
1728 Ok(stop_reason) => match stop_reason {
1729 StopReason::ToolUse => {
1730 let tool_uses = thread.use_pending_tools(window, cx, model.clone());
1731 cx.emit(ThreadEvent::UsePendingTools { tool_uses });
1732 }
1733 StopReason::EndTurn | StopReason::MaxTokens => {
1734 thread.project.update(cx, |project, cx| {
1735 project.set_agent_location(None, cx);
1736 });
1737 }
1738 StopReason::Refusal => {
1739 thread.project.update(cx, |project, cx| {
1740 project.set_agent_location(None, cx);
1741 });
1742
1743 // Remove the turn that was refused.
1744 //
1745 // https://docs.anthropic.com/en/docs/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals#reset-context-after-refusal
1746 {
1747 let mut messages_to_remove = Vec::new();
1748
1749 for (ix, message) in thread.messages.iter().enumerate().rev() {
1750 messages_to_remove.push(message.id);
1751
1752 if message.role == Role::User {
1753 if ix == 0 {
1754 break;
1755 }
1756
1757 if let Some(prev_message) = thread.messages.get(ix - 1) {
1758 if prev_message.role == Role::Assistant {
1759 break;
1760 }
1761 }
1762 }
1763 }
1764
1765 for message_id in messages_to_remove {
1766 thread.delete_message(message_id, cx);
1767 }
1768 }
1769
1770 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1771 header: "Language model refusal".into(),
1772 message: "Model refused to generate content for safety reasons.".into(),
1773 }));
1774 }
1775 },
1776 Err(error) => {
1777 thread.project.update(cx, |project, cx| {
1778 project.set_agent_location(None, cx);
1779 });
1780
1781 if error.is::<PaymentRequiredError>() {
1782 cx.emit(ThreadEvent::ShowError(ThreadError::PaymentRequired));
1783 } else if let Some(error) =
1784 error.downcast_ref::<ModelRequestLimitReachedError>()
1785 {
1786 cx.emit(ThreadEvent::ShowError(
1787 ThreadError::ModelRequestLimitReached { plan: error.plan },
1788 ));
1789 } else if let Some(known_error) =
1790 error.downcast_ref::<LanguageModelKnownError>()
1791 {
1792 match known_error {
1793 LanguageModelKnownError::ContextWindowLimitExceeded {
1794 tokens,
1795 } => {
1796 thread.exceeded_window_error = Some(ExceededWindowError {
1797 model_id: model.id(),
1798 token_count: *tokens,
1799 });
1800 cx.notify();
1801 }
1802 }
1803 } else {
1804 let error_message = error
1805 .chain()
1806 .map(|err| err.to_string())
1807 .collect::<Vec<_>>()
1808 .join("\n");
1809 cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1810 header: "Error interacting with language model".into(),
1811 message: SharedString::from(error_message.clone()),
1812 }));
1813 }
1814
1815 thread.cancel_last_completion(window, cx);
1816 }
1817 }
1818
1819 cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new)));
1820
1821 if let Some((request_callback, (request, response_events))) = thread
1822 .request_callback
1823 .as_mut()
1824 .zip(request_callback_parameters.as_ref())
1825 {
1826 request_callback(request, response_events);
1827 }
1828
1829 thread.auto_capture_telemetry(cx);
1830
1831 if let Ok(initial_usage) = initial_token_usage {
1832 let usage = thread.cumulative_token_usage - initial_usage;
1833
1834 telemetry::event!(
1835 "Assistant Thread Completion",
1836 thread_id = thread.id().to_string(),
1837 prompt_id = prompt_id,
1838 model = model.telemetry_id(),
1839 model_provider = model.provider_id().to_string(),
1840 input_tokens = usage.input_tokens,
1841 output_tokens = usage.output_tokens,
1842 cache_creation_input_tokens = usage.cache_creation_input_tokens,
1843 cache_read_input_tokens = usage.cache_read_input_tokens,
1844 );
1845 }
1846 })
1847 .ok();
1848 });
1849
1850 self.pending_completions.push(PendingCompletion {
1851 id: pending_completion_id,
1852 queue_state: QueueState::Sending,
1853 _task: task,
1854 });
1855 }
1856
1857 pub fn summarize(&mut self, cx: &mut Context<Self>) {
1858 let Some(model) = LanguageModelRegistry::read_global(cx).thread_summary_model() else {
1859 println!("No thread summary model");
1860 return;
1861 };
1862
1863 if !model.provider.is_authenticated(cx) {
1864 return;
1865 }
1866
1867 let added_user_message = include_str!("./prompts/summarize_thread_prompt.txt");
1868
1869 let request = self.to_summarize_request(
1870 &model.model,
1871 CompletionIntent::ThreadSummarization,
1872 added_user_message.into(),
1873 cx,
1874 );
1875
1876 self.summary = ThreadSummary::Generating;
1877
1878 self.pending_summary = cx.spawn(async move |this, cx| {
1879 let result = async {
1880 let mut messages = model.model.stream_completion(request, &cx).await?;
1881
1882 let mut new_summary = String::new();
1883 while let Some(event) = messages.next().await {
1884 let Ok(event) = event else {
1885 continue;
1886 };
1887 let text = match event {
1888 LanguageModelCompletionEvent::Text(text) => text,
1889 LanguageModelCompletionEvent::StatusUpdate(
1890 CompletionRequestStatus::UsageUpdated { amount, limit },
1891 ) => {
1892 this.update(cx, |thread, _cx| {
1893 thread.last_usage = Some(RequestUsage {
1894 limit,
1895 amount: amount as i32,
1896 });
1897 })?;
1898 continue;
1899 }
1900 _ => continue,
1901 };
1902
1903 let mut lines = text.lines();
1904 new_summary.extend(lines.next());
1905
1906 // Stop if the LLM generated multiple lines.
1907 if lines.next().is_some() {
1908 break;
1909 }
1910 }
1911
1912 anyhow::Ok(new_summary)
1913 }
1914 .await;
1915
1916 this.update(cx, |this, cx| {
1917 match result {
1918 Ok(new_summary) => {
1919 if new_summary.is_empty() {
1920 this.summary = ThreadSummary::Error;
1921 } else {
1922 this.summary = ThreadSummary::Ready(new_summary.into());
1923 }
1924 }
1925 Err(err) => {
1926 this.summary = ThreadSummary::Error;
1927 log::error!("Failed to generate thread summary: {}", err);
1928 }
1929 }
1930 cx.emit(ThreadEvent::SummaryGenerated);
1931 })
1932 .log_err()?;
1933
1934 Some(())
1935 });
1936 }
1937
1938 pub fn start_generating_detailed_summary_if_needed(
1939 &mut self,
1940 thread_store: WeakEntity<ThreadStore>,
1941 cx: &mut Context<Self>,
1942 ) {
1943 let Some(last_message_id) = self.messages.last().map(|message| message.id) else {
1944 return;
1945 };
1946
1947 match &*self.detailed_summary_rx.borrow() {
1948 DetailedSummaryState::Generating { message_id, .. }
1949 | DetailedSummaryState::Generated { message_id, .. }
1950 if *message_id == last_message_id =>
1951 {
1952 // Already up-to-date
1953 return;
1954 }
1955 _ => {}
1956 }
1957
1958 let Some(ConfiguredModel { model, provider }) =
1959 LanguageModelRegistry::read_global(cx).thread_summary_model()
1960 else {
1961 return;
1962 };
1963
1964 if !provider.is_authenticated(cx) {
1965 return;
1966 }
1967
1968 let added_user_message = include_str!("./prompts/summarize_thread_detailed_prompt.txt");
1969
1970 let request = self.to_summarize_request(
1971 &model,
1972 CompletionIntent::ThreadContextSummarization,
1973 added_user_message.into(),
1974 cx,
1975 );
1976
1977 *self.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generating {
1978 message_id: last_message_id,
1979 };
1980
1981 // Replace the detailed summarization task if there is one, cancelling it. It would probably
1982 // be better to allow the old task to complete, but this would require logic for choosing
1983 // which result to prefer (the old task could complete after the new one, resulting in a
1984 // stale summary).
1985 self.detailed_summary_task = cx.spawn(async move |thread, cx| {
1986 let stream = model.stream_completion_text(request, &cx);
1987 let Some(mut messages) = stream.await.log_err() else {
1988 thread
1989 .update(cx, |thread, _cx| {
1990 *thread.detailed_summary_tx.borrow_mut() =
1991 DetailedSummaryState::NotGenerated;
1992 })
1993 .ok()?;
1994 return None;
1995 };
1996
1997 let mut new_detailed_summary = String::new();
1998
1999 while let Some(chunk) = messages.stream.next().await {
2000 if let Some(chunk) = chunk.log_err() {
2001 new_detailed_summary.push_str(&chunk);
2002 }
2003 }
2004
2005 thread
2006 .update(cx, |thread, _cx| {
2007 *thread.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generated {
2008 text: new_detailed_summary.into(),
2009 message_id: last_message_id,
2010 };
2011 })
2012 .ok()?;
2013
2014 // Save thread so its summary can be reused later
2015 if let Some(thread) = thread.upgrade() {
2016 if let Ok(Ok(save_task)) = cx.update(|cx| {
2017 thread_store
2018 .update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx))
2019 }) {
2020 save_task.await.log_err();
2021 }
2022 }
2023
2024 Some(())
2025 });
2026 }
2027
2028 pub async fn wait_for_detailed_summary_or_text(
2029 this: &Entity<Self>,
2030 cx: &mut AsyncApp,
2031 ) -> Option<SharedString> {
2032 let mut detailed_summary_rx = this
2033 .read_with(cx, |this, _cx| this.detailed_summary_rx.clone())
2034 .ok()?;
2035 loop {
2036 match detailed_summary_rx.recv().await? {
2037 DetailedSummaryState::Generating { .. } => {}
2038 DetailedSummaryState::NotGenerated => {
2039 return this.read_with(cx, |this, _cx| this.text().into()).ok();
2040 }
2041 DetailedSummaryState::Generated { text, .. } => return Some(text),
2042 }
2043 }
2044 }
2045
2046 pub fn latest_detailed_summary_or_text(&self) -> SharedString {
2047 self.detailed_summary_rx
2048 .borrow()
2049 .text()
2050 .unwrap_or_else(|| self.text().into())
2051 }
2052
2053 pub fn is_generating_detailed_summary(&self) -> bool {
2054 matches!(
2055 &*self.detailed_summary_rx.borrow(),
2056 DetailedSummaryState::Generating { .. }
2057 )
2058 }
2059
2060 pub fn use_pending_tools(
2061 &mut self,
2062 window: Option<AnyWindowHandle>,
2063 cx: &mut Context<Self>,
2064 model: Arc<dyn LanguageModel>,
2065 ) -> Vec<PendingToolUse> {
2066 self.auto_capture_telemetry(cx);
2067 let request =
2068 Arc::new(self.to_completion_request(model.clone(), CompletionIntent::ToolResults, cx));
2069 let pending_tool_uses = self
2070 .tool_use
2071 .pending_tool_uses()
2072 .into_iter()
2073 .filter(|tool_use| tool_use.status.is_idle())
2074 .cloned()
2075 .collect::<Vec<_>>();
2076
2077 for tool_use in pending_tool_uses.iter() {
2078 if let Some(tool) = self.tools.read(cx).tool(&tool_use.name, cx) {
2079 if tool.needs_confirmation(&tool_use.input, cx)
2080 && !AgentSettings::get_global(cx).always_allow_tool_actions
2081 {
2082 self.tool_use.confirm_tool_use(
2083 tool_use.id.clone(),
2084 tool_use.ui_text.clone(),
2085 tool_use.input.clone(),
2086 request.clone(),
2087 tool,
2088 );
2089 cx.emit(ThreadEvent::ToolConfirmationNeeded);
2090 } else {
2091 self.run_tool(
2092 tool_use.id.clone(),
2093 tool_use.ui_text.clone(),
2094 tool_use.input.clone(),
2095 request.clone(),
2096 tool,
2097 model.clone(),
2098 window,
2099 cx,
2100 );
2101 }
2102 } else {
2103 self.handle_hallucinated_tool_use(
2104 tool_use.id.clone(),
2105 tool_use.name.clone(),
2106 window,
2107 cx,
2108 );
2109 }
2110 }
2111
2112 pending_tool_uses
2113 }
2114
2115 pub fn handle_hallucinated_tool_use(
2116 &mut self,
2117 tool_use_id: LanguageModelToolUseId,
2118 hallucinated_tool_name: Arc<str>,
2119 window: Option<AnyWindowHandle>,
2120 cx: &mut Context<Thread>,
2121 ) {
2122 let available_tools = self.tools.read(cx).enabled_tools(cx);
2123
2124 let tool_list = available_tools
2125 .iter()
2126 .map(|tool| format!("- {}: {}", tool.name(), tool.description()))
2127 .collect::<Vec<_>>()
2128 .join("\n");
2129
2130 let error_message = format!(
2131 "The tool '{}' doesn't exist or is not enabled. Available tools:\n{}",
2132 hallucinated_tool_name, tool_list
2133 );
2134
2135 let pending_tool_use = self.tool_use.insert_tool_output(
2136 tool_use_id.clone(),
2137 hallucinated_tool_name,
2138 Err(anyhow!("Missing tool call: {error_message}")),
2139 self.configured_model.as_ref(),
2140 );
2141
2142 cx.emit(ThreadEvent::MissingToolUse {
2143 tool_use_id: tool_use_id.clone(),
2144 ui_text: error_message.into(),
2145 });
2146
2147 self.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2148 }
2149
2150 pub fn receive_invalid_tool_json(
2151 &mut self,
2152 tool_use_id: LanguageModelToolUseId,
2153 tool_name: Arc<str>,
2154 invalid_json: Arc<str>,
2155 error: String,
2156 window: Option<AnyWindowHandle>,
2157 cx: &mut Context<Thread>,
2158 ) {
2159 log::error!("The model returned invalid input JSON: {invalid_json}");
2160
2161 let pending_tool_use = self.tool_use.insert_tool_output(
2162 tool_use_id.clone(),
2163 tool_name,
2164 Err(anyhow!("Error parsing input JSON: {error}")),
2165 self.configured_model.as_ref(),
2166 );
2167 let ui_text = if let Some(pending_tool_use) = &pending_tool_use {
2168 pending_tool_use.ui_text.clone()
2169 } else {
2170 log::error!(
2171 "There was no pending tool use for tool use {tool_use_id}, even though it finished (with invalid input JSON)."
2172 );
2173 format!("Unknown tool {}", tool_use_id).into()
2174 };
2175
2176 cx.emit(ThreadEvent::InvalidToolInput {
2177 tool_use_id: tool_use_id.clone(),
2178 ui_text,
2179 invalid_input_json: invalid_json,
2180 });
2181
2182 self.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2183 }
2184
2185 pub fn run_tool(
2186 &mut self,
2187 tool_use_id: LanguageModelToolUseId,
2188 ui_text: impl Into<SharedString>,
2189 input: serde_json::Value,
2190 request: Arc<LanguageModelRequest>,
2191 tool: Arc<dyn Tool>,
2192 model: Arc<dyn LanguageModel>,
2193 window: Option<AnyWindowHandle>,
2194 cx: &mut Context<Thread>,
2195 ) {
2196 let task =
2197 self.spawn_tool_use(tool_use_id.clone(), request, input, tool, model, window, cx);
2198 self.tool_use
2199 .run_pending_tool(tool_use_id, ui_text.into(), task);
2200 }
2201
2202 fn spawn_tool_use(
2203 &mut self,
2204 tool_use_id: LanguageModelToolUseId,
2205 request: Arc<LanguageModelRequest>,
2206 input: serde_json::Value,
2207 tool: Arc<dyn Tool>,
2208 model: Arc<dyn LanguageModel>,
2209 window: Option<AnyWindowHandle>,
2210 cx: &mut Context<Thread>,
2211 ) -> Task<()> {
2212 let tool_name: Arc<str> = tool.name().into();
2213
2214 let tool_result = if self.tools.read(cx).is_disabled(&tool.source(), &tool_name) {
2215 Task::ready(Err(anyhow!("tool is disabled: {tool_name}"))).into()
2216 } else {
2217 tool.run(
2218 input,
2219 request,
2220 self.project.clone(),
2221 self.action_log.clone(),
2222 model,
2223 window,
2224 cx,
2225 )
2226 };
2227
2228 // Store the card separately if it exists
2229 if let Some(card) = tool_result.card.clone() {
2230 self.tool_use
2231 .insert_tool_result_card(tool_use_id.clone(), card);
2232 }
2233
2234 cx.spawn({
2235 async move |thread: WeakEntity<Thread>, cx| {
2236 let output = tool_result.output.await;
2237
2238 thread
2239 .update(cx, |thread, cx| {
2240 let pending_tool_use = thread.tool_use.insert_tool_output(
2241 tool_use_id.clone(),
2242 tool_name,
2243 output,
2244 thread.configured_model.as_ref(),
2245 );
2246 thread.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2247 })
2248 .ok();
2249 }
2250 })
2251 }
2252
2253 fn tool_finished(
2254 &mut self,
2255 tool_use_id: LanguageModelToolUseId,
2256 pending_tool_use: Option<PendingToolUse>,
2257 canceled: bool,
2258 window: Option<AnyWindowHandle>,
2259 cx: &mut Context<Self>,
2260 ) {
2261 if self.all_tools_finished() {
2262 if let Some(ConfiguredModel { model, .. }) = self.configured_model.as_ref() {
2263 if !canceled {
2264 self.send_to_model(model.clone(), CompletionIntent::ToolResults, window, cx);
2265 }
2266 self.auto_capture_telemetry(cx);
2267 }
2268 }
2269
2270 cx.emit(ThreadEvent::ToolFinished {
2271 tool_use_id,
2272 pending_tool_use,
2273 });
2274 }
2275
2276 /// Cancels the last pending completion, if there are any pending.
2277 ///
2278 /// Returns whether a completion was canceled.
2279 pub fn cancel_last_completion(
2280 &mut self,
2281 window: Option<AnyWindowHandle>,
2282 cx: &mut Context<Self>,
2283 ) -> bool {
2284 let mut canceled = self.pending_completions.pop().is_some();
2285
2286 for pending_tool_use in self.tool_use.cancel_pending() {
2287 canceled = true;
2288 self.tool_finished(
2289 pending_tool_use.id.clone(),
2290 Some(pending_tool_use),
2291 true,
2292 window,
2293 cx,
2294 );
2295 }
2296
2297 if canceled {
2298 cx.emit(ThreadEvent::CompletionCanceled);
2299
2300 // When canceled, we always want to insert the checkpoint.
2301 // (We skip over finalize_pending_checkpoint, because it
2302 // would conclude we didn't have anything to insert here.)
2303 if let Some(checkpoint) = self.pending_checkpoint.take() {
2304 self.insert_checkpoint(checkpoint, cx);
2305 }
2306 } else {
2307 self.finalize_pending_checkpoint(cx);
2308 }
2309
2310 canceled
2311 }
2312
2313 /// Signals that any in-progress editing should be canceled.
2314 ///
2315 /// This method is used to notify listeners (like ActiveThread) that
2316 /// they should cancel any editing operations.
2317 pub fn cancel_editing(&mut self, cx: &mut Context<Self>) {
2318 cx.emit(ThreadEvent::CancelEditing);
2319 }
2320
2321 pub fn feedback(&self) -> Option<ThreadFeedback> {
2322 self.feedback
2323 }
2324
2325 pub fn message_feedback(&self, message_id: MessageId) -> Option<ThreadFeedback> {
2326 self.message_feedback.get(&message_id).copied()
2327 }
2328
2329 pub fn report_message_feedback(
2330 &mut self,
2331 message_id: MessageId,
2332 feedback: ThreadFeedback,
2333 cx: &mut Context<Self>,
2334 ) -> Task<Result<()>> {
2335 if self.message_feedback.get(&message_id) == Some(&feedback) {
2336 return Task::ready(Ok(()));
2337 }
2338
2339 let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
2340 let serialized_thread = self.serialize(cx);
2341 let thread_id = self.id().clone();
2342 let client = self.project.read(cx).client();
2343
2344 let enabled_tool_names: Vec<String> = self
2345 .tools()
2346 .read(cx)
2347 .enabled_tools(cx)
2348 .iter()
2349 .map(|tool| tool.name())
2350 .collect();
2351
2352 self.message_feedback.insert(message_id, feedback);
2353
2354 cx.notify();
2355
2356 let message_content = self
2357 .message(message_id)
2358 .map(|msg| msg.to_string())
2359 .unwrap_or_default();
2360
2361 cx.background_spawn(async move {
2362 let final_project_snapshot = final_project_snapshot.await;
2363 let serialized_thread = serialized_thread.await?;
2364 let thread_data =
2365 serde_json::to_value(serialized_thread).unwrap_or_else(|_| serde_json::Value::Null);
2366
2367 let rating = match feedback {
2368 ThreadFeedback::Positive => "positive",
2369 ThreadFeedback::Negative => "negative",
2370 };
2371 telemetry::event!(
2372 "Assistant Thread Rated",
2373 rating,
2374 thread_id,
2375 enabled_tool_names,
2376 message_id = message_id.0,
2377 message_content,
2378 thread_data,
2379 final_project_snapshot
2380 );
2381 client.telemetry().flush_events().await;
2382
2383 Ok(())
2384 })
2385 }
2386
2387 pub fn report_feedback(
2388 &mut self,
2389 feedback: ThreadFeedback,
2390 cx: &mut Context<Self>,
2391 ) -> Task<Result<()>> {
2392 let last_assistant_message_id = self
2393 .messages
2394 .iter()
2395 .rev()
2396 .find(|msg| msg.role == Role::Assistant)
2397 .map(|msg| msg.id);
2398
2399 if let Some(message_id) = last_assistant_message_id {
2400 self.report_message_feedback(message_id, feedback, cx)
2401 } else {
2402 let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
2403 let serialized_thread = self.serialize(cx);
2404 let thread_id = self.id().clone();
2405 let client = self.project.read(cx).client();
2406 self.feedback = Some(feedback);
2407 cx.notify();
2408
2409 cx.background_spawn(async move {
2410 let final_project_snapshot = final_project_snapshot.await;
2411 let serialized_thread = serialized_thread.await?;
2412 let thread_data = serde_json::to_value(serialized_thread)
2413 .unwrap_or_else(|_| serde_json::Value::Null);
2414
2415 let rating = match feedback {
2416 ThreadFeedback::Positive => "positive",
2417 ThreadFeedback::Negative => "negative",
2418 };
2419 telemetry::event!(
2420 "Assistant Thread Rated",
2421 rating,
2422 thread_id,
2423 thread_data,
2424 final_project_snapshot
2425 );
2426 client.telemetry().flush_events().await;
2427
2428 Ok(())
2429 })
2430 }
2431 }
2432
2433 /// Create a snapshot of the current project state including git information and unsaved buffers.
2434 fn project_snapshot(
2435 project: Entity<Project>,
2436 cx: &mut Context<Self>,
2437 ) -> Task<Arc<ProjectSnapshot>> {
2438 let git_store = project.read(cx).git_store().clone();
2439 let worktree_snapshots: Vec<_> = project
2440 .read(cx)
2441 .visible_worktrees(cx)
2442 .map(|worktree| Self::worktree_snapshot(worktree, git_store.clone(), cx))
2443 .collect();
2444
2445 cx.spawn(async move |_, cx| {
2446 let worktree_snapshots = futures::future::join_all(worktree_snapshots).await;
2447
2448 let mut unsaved_buffers = Vec::new();
2449 cx.update(|app_cx| {
2450 let buffer_store = project.read(app_cx).buffer_store();
2451 for buffer_handle in buffer_store.read(app_cx).buffers() {
2452 let buffer = buffer_handle.read(app_cx);
2453 if buffer.is_dirty() {
2454 if let Some(file) = buffer.file() {
2455 let path = file.path().to_string_lossy().to_string();
2456 unsaved_buffers.push(path);
2457 }
2458 }
2459 }
2460 })
2461 .ok();
2462
2463 Arc::new(ProjectSnapshot {
2464 worktree_snapshots,
2465 unsaved_buffer_paths: unsaved_buffers,
2466 timestamp: Utc::now(),
2467 })
2468 })
2469 }
2470
2471 fn worktree_snapshot(
2472 worktree: Entity<project::Worktree>,
2473 git_store: Entity<GitStore>,
2474 cx: &App,
2475 ) -> Task<WorktreeSnapshot> {
2476 cx.spawn(async move |cx| {
2477 // Get worktree path and snapshot
2478 let worktree_info = cx.update(|app_cx| {
2479 let worktree = worktree.read(app_cx);
2480 let path = worktree.abs_path().to_string_lossy().to_string();
2481 let snapshot = worktree.snapshot();
2482 (path, snapshot)
2483 });
2484
2485 let Ok((worktree_path, _snapshot)) = worktree_info else {
2486 return WorktreeSnapshot {
2487 worktree_path: String::new(),
2488 git_state: None,
2489 };
2490 };
2491
2492 let git_state = git_store
2493 .update(cx, |git_store, cx| {
2494 git_store
2495 .repositories()
2496 .values()
2497 .find(|repo| {
2498 repo.read(cx)
2499 .abs_path_to_repo_path(&worktree.read(cx).abs_path())
2500 .is_some()
2501 })
2502 .cloned()
2503 })
2504 .ok()
2505 .flatten()
2506 .map(|repo| {
2507 repo.update(cx, |repo, _| {
2508 let current_branch =
2509 repo.branch.as_ref().map(|branch| branch.name().to_owned());
2510 repo.send_job(None, |state, _| async move {
2511 let RepositoryState::Local { backend, .. } = state else {
2512 return GitState {
2513 remote_url: None,
2514 head_sha: None,
2515 current_branch,
2516 diff: None,
2517 };
2518 };
2519
2520 let remote_url = backend.remote_url("origin");
2521 let head_sha = backend.head_sha().await;
2522 let diff = backend.diff(DiffType::HeadToWorktree).await.ok();
2523
2524 GitState {
2525 remote_url,
2526 head_sha,
2527 current_branch,
2528 diff,
2529 }
2530 })
2531 })
2532 });
2533
2534 let git_state = match git_state {
2535 Some(git_state) => match git_state.ok() {
2536 Some(git_state) => git_state.await.ok(),
2537 None => None,
2538 },
2539 None => None,
2540 };
2541
2542 WorktreeSnapshot {
2543 worktree_path,
2544 git_state,
2545 }
2546 })
2547 }
2548
2549 pub fn to_markdown(&self, cx: &App) -> Result<String> {
2550 let mut markdown = Vec::new();
2551
2552 let summary = self.summary().or_default();
2553 writeln!(markdown, "# {summary}\n")?;
2554
2555 for message in self.messages() {
2556 writeln!(
2557 markdown,
2558 "## {role}\n",
2559 role = match message.role {
2560 Role::User => "User",
2561 Role::Assistant => "Agent",
2562 Role::System => "System",
2563 }
2564 )?;
2565
2566 if !message.loaded_context.text.is_empty() {
2567 writeln!(markdown, "{}", message.loaded_context.text)?;
2568 }
2569
2570 if !message.loaded_context.images.is_empty() {
2571 writeln!(
2572 markdown,
2573 "\n{} images attached as context.\n",
2574 message.loaded_context.images.len()
2575 )?;
2576 }
2577
2578 for segment in &message.segments {
2579 match segment {
2580 MessageSegment::Text(text) => writeln!(markdown, "{}\n", text)?,
2581 MessageSegment::Thinking { text, .. } => {
2582 writeln!(markdown, "<think>\n{}\n</think>\n", text)?
2583 }
2584 MessageSegment::RedactedThinking(_) => {}
2585 }
2586 }
2587
2588 for tool_use in self.tool_uses_for_message(message.id, cx) {
2589 writeln!(
2590 markdown,
2591 "**Use Tool: {} ({})**",
2592 tool_use.name, tool_use.id
2593 )?;
2594 writeln!(markdown, "```json")?;
2595 writeln!(
2596 markdown,
2597 "{}",
2598 serde_json::to_string_pretty(&tool_use.input)?
2599 )?;
2600 writeln!(markdown, "```")?;
2601 }
2602
2603 for tool_result in self.tool_results_for_message(message.id) {
2604 write!(markdown, "\n**Tool Results: {}", tool_result.tool_use_id)?;
2605 if tool_result.is_error {
2606 write!(markdown, " (Error)")?;
2607 }
2608
2609 writeln!(markdown, "**\n")?;
2610 match &tool_result.content {
2611 LanguageModelToolResultContent::Text(text) => {
2612 writeln!(markdown, "{text}")?;
2613 }
2614 LanguageModelToolResultContent::Image(image) => {
2615 writeln!(markdown, "", image.source)?;
2616 }
2617 }
2618
2619 if let Some(output) = tool_result.output.as_ref() {
2620 writeln!(
2621 markdown,
2622 "\n\nDebug Output:\n\n```json\n{}\n```\n",
2623 serde_json::to_string_pretty(output)?
2624 )?;
2625 }
2626 }
2627 }
2628
2629 Ok(String::from_utf8_lossy(&markdown).to_string())
2630 }
2631
2632 pub fn keep_edits_in_range(
2633 &mut self,
2634 buffer: Entity<language::Buffer>,
2635 buffer_range: Range<language::Anchor>,
2636 cx: &mut Context<Self>,
2637 ) {
2638 self.action_log.update(cx, |action_log, cx| {
2639 action_log.keep_edits_in_range(buffer, buffer_range, cx)
2640 });
2641 }
2642
2643 pub fn keep_all_edits(&mut self, cx: &mut Context<Self>) {
2644 self.action_log
2645 .update(cx, |action_log, cx| action_log.keep_all_edits(cx));
2646 }
2647
2648 pub fn reject_edits_in_ranges(
2649 &mut self,
2650 buffer: Entity<language::Buffer>,
2651 buffer_ranges: Vec<Range<language::Anchor>>,
2652 cx: &mut Context<Self>,
2653 ) -> Task<Result<()>> {
2654 self.action_log.update(cx, |action_log, cx| {
2655 action_log.reject_edits_in_ranges(buffer, buffer_ranges, cx)
2656 })
2657 }
2658
2659 pub fn action_log(&self) -> &Entity<ActionLog> {
2660 &self.action_log
2661 }
2662
2663 pub fn project(&self) -> &Entity<Project> {
2664 &self.project
2665 }
2666
2667 pub fn auto_capture_telemetry(&mut self, cx: &mut Context<Self>) {
2668 if !cx.has_flag::<feature_flags::ThreadAutoCaptureFeatureFlag>() {
2669 return;
2670 }
2671
2672 let now = Instant::now();
2673 if let Some(last) = self.last_auto_capture_at {
2674 if now.duration_since(last).as_secs() < 10 {
2675 return;
2676 }
2677 }
2678
2679 self.last_auto_capture_at = Some(now);
2680
2681 let thread_id = self.id().clone();
2682 let github_login = self
2683 .project
2684 .read(cx)
2685 .user_store()
2686 .read(cx)
2687 .current_user()
2688 .map(|user| user.github_login.clone());
2689 let client = self.project.read(cx).client();
2690 let serialize_task = self.serialize(cx);
2691
2692 cx.background_executor()
2693 .spawn(async move {
2694 if let Ok(serialized_thread) = serialize_task.await {
2695 if let Ok(thread_data) = serde_json::to_value(serialized_thread) {
2696 telemetry::event!(
2697 "Agent Thread Auto-Captured",
2698 thread_id = thread_id.to_string(),
2699 thread_data = thread_data,
2700 auto_capture_reason = "tracked_user",
2701 github_login = github_login
2702 );
2703
2704 client.telemetry().flush_events().await;
2705 }
2706 }
2707 })
2708 .detach();
2709 }
2710
2711 pub fn cumulative_token_usage(&self) -> TokenUsage {
2712 self.cumulative_token_usage
2713 }
2714
2715 pub fn token_usage_up_to_message(&self, message_id: MessageId) -> TotalTokenUsage {
2716 let Some(model) = self.configured_model.as_ref() else {
2717 return TotalTokenUsage::default();
2718 };
2719
2720 let max = model.model.max_token_count();
2721
2722 let index = self
2723 .messages
2724 .iter()
2725 .position(|msg| msg.id == message_id)
2726 .unwrap_or(0);
2727
2728 if index == 0 {
2729 return TotalTokenUsage { total: 0, max };
2730 }
2731
2732 let token_usage = &self
2733 .request_token_usage
2734 .get(index - 1)
2735 .cloned()
2736 .unwrap_or_default();
2737
2738 TotalTokenUsage {
2739 total: token_usage.total_tokens() as usize,
2740 max,
2741 }
2742 }
2743
2744 pub fn total_token_usage(&self) -> Option<TotalTokenUsage> {
2745 let model = self.configured_model.as_ref()?;
2746
2747 let max = model.model.max_token_count();
2748
2749 if let Some(exceeded_error) = &self.exceeded_window_error {
2750 if model.model.id() == exceeded_error.model_id {
2751 return Some(TotalTokenUsage {
2752 total: exceeded_error.token_count,
2753 max,
2754 });
2755 }
2756 }
2757
2758 let total = self
2759 .token_usage_at_last_message()
2760 .unwrap_or_default()
2761 .total_tokens() as usize;
2762
2763 Some(TotalTokenUsage { total, max })
2764 }
2765
2766 fn token_usage_at_last_message(&self) -> Option<TokenUsage> {
2767 self.request_token_usage
2768 .get(self.messages.len().saturating_sub(1))
2769 .or_else(|| self.request_token_usage.last())
2770 .cloned()
2771 }
2772
2773 fn update_token_usage_at_last_message(&mut self, token_usage: TokenUsage) {
2774 let placeholder = self.token_usage_at_last_message().unwrap_or_default();
2775 self.request_token_usage
2776 .resize(self.messages.len(), placeholder);
2777
2778 if let Some(last) = self.request_token_usage.last_mut() {
2779 *last = token_usage;
2780 }
2781 }
2782
2783 pub fn deny_tool_use(
2784 &mut self,
2785 tool_use_id: LanguageModelToolUseId,
2786 tool_name: Arc<str>,
2787 window: Option<AnyWindowHandle>,
2788 cx: &mut Context<Self>,
2789 ) {
2790 let err = Err(anyhow::anyhow!(
2791 "Permission to run tool action denied by user"
2792 ));
2793
2794 self.tool_use.insert_tool_output(
2795 tool_use_id.clone(),
2796 tool_name,
2797 err,
2798 self.configured_model.as_ref(),
2799 );
2800 self.tool_finished(tool_use_id.clone(), None, true, window, cx);
2801 }
2802}
2803
2804#[derive(Debug, Clone, Error)]
2805pub enum ThreadError {
2806 #[error("Payment required")]
2807 PaymentRequired,
2808 #[error("Model request limit reached")]
2809 ModelRequestLimitReached { plan: Plan },
2810 #[error("Message {header}: {message}")]
2811 Message {
2812 header: SharedString,
2813 message: SharedString,
2814 },
2815}
2816
2817#[derive(Debug, Clone)]
2818pub enum ThreadEvent {
2819 ShowError(ThreadError),
2820 StreamedCompletion,
2821 ReceivedTextChunk,
2822 NewRequest,
2823 StreamedAssistantText(MessageId, String),
2824 StreamedAssistantThinking(MessageId, String),
2825 StreamedToolUse {
2826 tool_use_id: LanguageModelToolUseId,
2827 ui_text: Arc<str>,
2828 input: serde_json::Value,
2829 },
2830 MissingToolUse {
2831 tool_use_id: LanguageModelToolUseId,
2832 ui_text: Arc<str>,
2833 },
2834 InvalidToolInput {
2835 tool_use_id: LanguageModelToolUseId,
2836 ui_text: Arc<str>,
2837 invalid_input_json: Arc<str>,
2838 },
2839 Stopped(Result<StopReason, Arc<anyhow::Error>>),
2840 MessageAdded(MessageId),
2841 MessageEdited(MessageId),
2842 MessageDeleted(MessageId),
2843 SummaryGenerated,
2844 SummaryChanged,
2845 UsePendingTools {
2846 tool_uses: Vec<PendingToolUse>,
2847 },
2848 ToolFinished {
2849 #[allow(unused)]
2850 tool_use_id: LanguageModelToolUseId,
2851 /// The pending tool use that corresponds to this tool.
2852 pending_tool_use: Option<PendingToolUse>,
2853 },
2854 CheckpointChanged,
2855 ToolConfirmationNeeded,
2856 ToolUseLimitReached,
2857 CancelEditing,
2858 CompletionCanceled,
2859}
2860
2861impl EventEmitter<ThreadEvent> for Thread {}
2862
2863struct PendingCompletion {
2864 id: usize,
2865 queue_state: QueueState,
2866 _task: Task<()>,
2867}
2868
2869#[cfg(test)]
2870mod tests {
2871 use super::*;
2872 use crate::{ThreadStore, context::load_context, context_store::ContextStore, thread_store};
2873 use agent_settings::{AgentSettings, LanguageModelParameters};
2874 use assistant_tool::ToolRegistry;
2875 use editor::EditorSettings;
2876 use gpui::TestAppContext;
2877 use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider};
2878 use project::{FakeFs, Project};
2879 use prompt_store::PromptBuilder;
2880 use serde_json::json;
2881 use settings::{Settings, SettingsStore};
2882 use std::sync::Arc;
2883 use theme::ThemeSettings;
2884 use util::path;
2885 use workspace::Workspace;
2886
2887 #[gpui::test]
2888 async fn test_message_with_context(cx: &mut TestAppContext) {
2889 init_test_settings(cx);
2890
2891 let project = create_test_project(
2892 cx,
2893 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
2894 )
2895 .await;
2896
2897 let (_workspace, _thread_store, thread, context_store, model) =
2898 setup_test_environment(cx, project.clone()).await;
2899
2900 add_file_to_context(&project, &context_store, "test/code.rs", cx)
2901 .await
2902 .unwrap();
2903
2904 let context =
2905 context_store.read_with(cx, |store, _| store.context().next().cloned().unwrap());
2906 let loaded_context = cx
2907 .update(|cx| load_context(vec![context], &project, &None, cx))
2908 .await;
2909
2910 // Insert user message with context
2911 let message_id = thread.update(cx, |thread, cx| {
2912 thread.insert_user_message(
2913 "Please explain this code",
2914 loaded_context,
2915 None,
2916 Vec::new(),
2917 cx,
2918 )
2919 });
2920
2921 // Check content and context in message object
2922 let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
2923
2924 // Use different path format strings based on platform for the test
2925 #[cfg(windows)]
2926 let path_part = r"test\code.rs";
2927 #[cfg(not(windows))]
2928 let path_part = "test/code.rs";
2929
2930 let expected_context = format!(
2931 r#"
2932<context>
2933The following items were attached by the user. They are up-to-date and don't need to be re-read.
2934
2935<files>
2936```rs {path_part}
2937fn main() {{
2938 println!("Hello, world!");
2939}}
2940```
2941</files>
2942</context>
2943"#
2944 );
2945
2946 assert_eq!(message.role, Role::User);
2947 assert_eq!(message.segments.len(), 1);
2948 assert_eq!(
2949 message.segments[0],
2950 MessageSegment::Text("Please explain this code".to_string())
2951 );
2952 assert_eq!(message.loaded_context.text, expected_context);
2953
2954 // Check message in request
2955 let request = thread.update(cx, |thread, cx| {
2956 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
2957 });
2958
2959 assert_eq!(request.messages.len(), 2);
2960 let expected_full_message = format!("{}Please explain this code", expected_context);
2961 assert_eq!(request.messages[1].string_contents(), expected_full_message);
2962 }
2963
2964 #[gpui::test]
2965 async fn test_only_include_new_contexts(cx: &mut TestAppContext) {
2966 init_test_settings(cx);
2967
2968 let project = create_test_project(
2969 cx,
2970 json!({
2971 "file1.rs": "fn function1() {}\n",
2972 "file2.rs": "fn function2() {}\n",
2973 "file3.rs": "fn function3() {}\n",
2974 "file4.rs": "fn function4() {}\n",
2975 }),
2976 )
2977 .await;
2978
2979 let (_, _thread_store, thread, context_store, model) =
2980 setup_test_environment(cx, project.clone()).await;
2981
2982 // First message with context 1
2983 add_file_to_context(&project, &context_store, "test/file1.rs", cx)
2984 .await
2985 .unwrap();
2986 let new_contexts = context_store.update(cx, |store, cx| {
2987 store.new_context_for_thread(thread.read(cx), None)
2988 });
2989 assert_eq!(new_contexts.len(), 1);
2990 let loaded_context = cx
2991 .update(|cx| load_context(new_contexts, &project, &None, cx))
2992 .await;
2993 let message1_id = thread.update(cx, |thread, cx| {
2994 thread.insert_user_message("Message 1", loaded_context, None, Vec::new(), cx)
2995 });
2996
2997 // Second message with contexts 1 and 2 (context 1 should be skipped as it's already included)
2998 add_file_to_context(&project, &context_store, "test/file2.rs", cx)
2999 .await
3000 .unwrap();
3001 let new_contexts = context_store.update(cx, |store, cx| {
3002 store.new_context_for_thread(thread.read(cx), None)
3003 });
3004 assert_eq!(new_contexts.len(), 1);
3005 let loaded_context = cx
3006 .update(|cx| load_context(new_contexts, &project, &None, cx))
3007 .await;
3008 let message2_id = thread.update(cx, |thread, cx| {
3009 thread.insert_user_message("Message 2", loaded_context, None, Vec::new(), cx)
3010 });
3011
3012 // Third message with all three contexts (contexts 1 and 2 should be skipped)
3013 //
3014 add_file_to_context(&project, &context_store, "test/file3.rs", cx)
3015 .await
3016 .unwrap();
3017 let new_contexts = context_store.update(cx, |store, cx| {
3018 store.new_context_for_thread(thread.read(cx), None)
3019 });
3020 assert_eq!(new_contexts.len(), 1);
3021 let loaded_context = cx
3022 .update(|cx| load_context(new_contexts, &project, &None, cx))
3023 .await;
3024 let message3_id = thread.update(cx, |thread, cx| {
3025 thread.insert_user_message("Message 3", loaded_context, None, Vec::new(), cx)
3026 });
3027
3028 // Check what contexts are included in each message
3029 let (message1, message2, message3) = thread.read_with(cx, |thread, _| {
3030 (
3031 thread.message(message1_id).unwrap().clone(),
3032 thread.message(message2_id).unwrap().clone(),
3033 thread.message(message3_id).unwrap().clone(),
3034 )
3035 });
3036
3037 // First message should include context 1
3038 assert!(message1.loaded_context.text.contains("file1.rs"));
3039
3040 // Second message should include only context 2 (not 1)
3041 assert!(!message2.loaded_context.text.contains("file1.rs"));
3042 assert!(message2.loaded_context.text.contains("file2.rs"));
3043
3044 // Third message should include only context 3 (not 1 or 2)
3045 assert!(!message3.loaded_context.text.contains("file1.rs"));
3046 assert!(!message3.loaded_context.text.contains("file2.rs"));
3047 assert!(message3.loaded_context.text.contains("file3.rs"));
3048
3049 // Check entire request to make sure all contexts are properly included
3050 let request = thread.update(cx, |thread, cx| {
3051 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3052 });
3053
3054 // The request should contain all 3 messages
3055 assert_eq!(request.messages.len(), 4);
3056
3057 // Check that the contexts are properly formatted in each message
3058 assert!(request.messages[1].string_contents().contains("file1.rs"));
3059 assert!(!request.messages[1].string_contents().contains("file2.rs"));
3060 assert!(!request.messages[1].string_contents().contains("file3.rs"));
3061
3062 assert!(!request.messages[2].string_contents().contains("file1.rs"));
3063 assert!(request.messages[2].string_contents().contains("file2.rs"));
3064 assert!(!request.messages[2].string_contents().contains("file3.rs"));
3065
3066 assert!(!request.messages[3].string_contents().contains("file1.rs"));
3067 assert!(!request.messages[3].string_contents().contains("file2.rs"));
3068 assert!(request.messages[3].string_contents().contains("file3.rs"));
3069
3070 add_file_to_context(&project, &context_store, "test/file4.rs", cx)
3071 .await
3072 .unwrap();
3073 let new_contexts = context_store.update(cx, |store, cx| {
3074 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3075 });
3076 assert_eq!(new_contexts.len(), 3);
3077 let loaded_context = cx
3078 .update(|cx| load_context(new_contexts, &project, &None, cx))
3079 .await
3080 .loaded_context;
3081
3082 assert!(!loaded_context.text.contains("file1.rs"));
3083 assert!(loaded_context.text.contains("file2.rs"));
3084 assert!(loaded_context.text.contains("file3.rs"));
3085 assert!(loaded_context.text.contains("file4.rs"));
3086
3087 let new_contexts = context_store.update(cx, |store, cx| {
3088 // Remove file4.rs
3089 store.remove_context(&loaded_context.contexts[2].handle(), cx);
3090 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3091 });
3092 assert_eq!(new_contexts.len(), 2);
3093 let loaded_context = cx
3094 .update(|cx| load_context(new_contexts, &project, &None, cx))
3095 .await
3096 .loaded_context;
3097
3098 assert!(!loaded_context.text.contains("file1.rs"));
3099 assert!(loaded_context.text.contains("file2.rs"));
3100 assert!(loaded_context.text.contains("file3.rs"));
3101 assert!(!loaded_context.text.contains("file4.rs"));
3102
3103 let new_contexts = context_store.update(cx, |store, cx| {
3104 // Remove file3.rs
3105 store.remove_context(&loaded_context.contexts[1].handle(), cx);
3106 store.new_context_for_thread(thread.read(cx), Some(message2_id))
3107 });
3108 assert_eq!(new_contexts.len(), 1);
3109 let loaded_context = cx
3110 .update(|cx| load_context(new_contexts, &project, &None, cx))
3111 .await
3112 .loaded_context;
3113
3114 assert!(!loaded_context.text.contains("file1.rs"));
3115 assert!(loaded_context.text.contains("file2.rs"));
3116 assert!(!loaded_context.text.contains("file3.rs"));
3117 assert!(!loaded_context.text.contains("file4.rs"));
3118 }
3119
3120 #[gpui::test]
3121 async fn test_message_without_files(cx: &mut TestAppContext) {
3122 init_test_settings(cx);
3123
3124 let project = create_test_project(
3125 cx,
3126 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3127 )
3128 .await;
3129
3130 let (_, _thread_store, thread, _context_store, model) =
3131 setup_test_environment(cx, project.clone()).await;
3132
3133 // Insert user message without any context (empty context vector)
3134 let message_id = thread.update(cx, |thread, cx| {
3135 thread.insert_user_message(
3136 "What is the best way to learn Rust?",
3137 ContextLoadResult::default(),
3138 None,
3139 Vec::new(),
3140 cx,
3141 )
3142 });
3143
3144 // Check content and context in message object
3145 let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
3146
3147 // Context should be empty when no files are included
3148 assert_eq!(message.role, Role::User);
3149 assert_eq!(message.segments.len(), 1);
3150 assert_eq!(
3151 message.segments[0],
3152 MessageSegment::Text("What is the best way to learn Rust?".to_string())
3153 );
3154 assert_eq!(message.loaded_context.text, "");
3155
3156 // Check message in request
3157 let request = thread.update(cx, |thread, cx| {
3158 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3159 });
3160
3161 assert_eq!(request.messages.len(), 2);
3162 assert_eq!(
3163 request.messages[1].string_contents(),
3164 "What is the best way to learn Rust?"
3165 );
3166
3167 // Add second message, also without context
3168 let message2_id = thread.update(cx, |thread, cx| {
3169 thread.insert_user_message(
3170 "Are there any good books?",
3171 ContextLoadResult::default(),
3172 None,
3173 Vec::new(),
3174 cx,
3175 )
3176 });
3177
3178 let message2 =
3179 thread.read_with(cx, |thread, _| thread.message(message2_id).unwrap().clone());
3180 assert_eq!(message2.loaded_context.text, "");
3181
3182 // Check that both messages appear in the request
3183 let request = thread.update(cx, |thread, cx| {
3184 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3185 });
3186
3187 assert_eq!(request.messages.len(), 3);
3188 assert_eq!(
3189 request.messages[1].string_contents(),
3190 "What is the best way to learn Rust?"
3191 );
3192 assert_eq!(
3193 request.messages[2].string_contents(),
3194 "Are there any good books?"
3195 );
3196 }
3197
3198 #[gpui::test]
3199 async fn test_stale_buffer_notification(cx: &mut TestAppContext) {
3200 init_test_settings(cx);
3201
3202 let project = create_test_project(
3203 cx,
3204 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3205 )
3206 .await;
3207
3208 let (_workspace, _thread_store, thread, context_store, model) =
3209 setup_test_environment(cx, project.clone()).await;
3210
3211 // Open buffer and add it to context
3212 let buffer = add_file_to_context(&project, &context_store, "test/code.rs", cx)
3213 .await
3214 .unwrap();
3215
3216 let context =
3217 context_store.read_with(cx, |store, _| store.context().next().cloned().unwrap());
3218 let loaded_context = cx
3219 .update(|cx| load_context(vec![context], &project, &None, cx))
3220 .await;
3221
3222 // Insert user message with the buffer as context
3223 thread.update(cx, |thread, cx| {
3224 thread.insert_user_message("Explain this code", loaded_context, None, Vec::new(), cx)
3225 });
3226
3227 // Create a request and check that it doesn't have a stale buffer warning yet
3228 let initial_request = thread.update(cx, |thread, cx| {
3229 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3230 });
3231
3232 // Make sure we don't have a stale file warning yet
3233 let has_stale_warning = initial_request.messages.iter().any(|msg| {
3234 msg.string_contents()
3235 .contains("These files changed since last read:")
3236 });
3237 assert!(
3238 !has_stale_warning,
3239 "Should not have stale buffer warning before buffer is modified"
3240 );
3241
3242 // Modify the buffer
3243 buffer.update(cx, |buffer, cx| {
3244 // Find a position at the end of line 1
3245 buffer.edit(
3246 [(1..1, "\n println!(\"Added a new line\");\n")],
3247 None,
3248 cx,
3249 );
3250 });
3251
3252 // Insert another user message without context
3253 thread.update(cx, |thread, cx| {
3254 thread.insert_user_message(
3255 "What does the code do now?",
3256 ContextLoadResult::default(),
3257 None,
3258 Vec::new(),
3259 cx,
3260 )
3261 });
3262
3263 // Create a new request and check for the stale buffer warning
3264 let new_request = thread.update(cx, |thread, cx| {
3265 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3266 });
3267
3268 // We should have a stale file warning as the last message
3269 let last_message = new_request
3270 .messages
3271 .last()
3272 .expect("Request should have messages");
3273
3274 // The last message should be the stale buffer notification
3275 assert_eq!(last_message.role, Role::User);
3276
3277 // Check the exact content of the message
3278 let expected_content = "These files changed since last read:\n- code.rs\n";
3279 assert_eq!(
3280 last_message.string_contents(),
3281 expected_content,
3282 "Last message should be exactly the stale buffer notification"
3283 );
3284 }
3285
3286 #[gpui::test]
3287 async fn test_temperature_setting(cx: &mut TestAppContext) {
3288 init_test_settings(cx);
3289
3290 let project = create_test_project(
3291 cx,
3292 json!({"code.rs": "fn main() {\n println!(\"Hello, world!\");\n}"}),
3293 )
3294 .await;
3295
3296 let (_workspace, _thread_store, thread, _context_store, model) =
3297 setup_test_environment(cx, project.clone()).await;
3298
3299 // Both model and provider
3300 cx.update(|cx| {
3301 AgentSettings::override_global(
3302 AgentSettings {
3303 model_parameters: vec![LanguageModelParameters {
3304 provider: Some(model.provider_id().0.to_string().into()),
3305 model: Some(model.id().0.clone()),
3306 temperature: Some(0.66),
3307 }],
3308 ..AgentSettings::get_global(cx).clone()
3309 },
3310 cx,
3311 );
3312 });
3313
3314 let request = thread.update(cx, |thread, cx| {
3315 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3316 });
3317 assert_eq!(request.temperature, Some(0.66));
3318
3319 // Only model
3320 cx.update(|cx| {
3321 AgentSettings::override_global(
3322 AgentSettings {
3323 model_parameters: vec![LanguageModelParameters {
3324 provider: None,
3325 model: Some(model.id().0.clone()),
3326 temperature: Some(0.66),
3327 }],
3328 ..AgentSettings::get_global(cx).clone()
3329 },
3330 cx,
3331 );
3332 });
3333
3334 let request = thread.update(cx, |thread, cx| {
3335 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3336 });
3337 assert_eq!(request.temperature, Some(0.66));
3338
3339 // Only provider
3340 cx.update(|cx| {
3341 AgentSettings::override_global(
3342 AgentSettings {
3343 model_parameters: vec![LanguageModelParameters {
3344 provider: Some(model.provider_id().0.to_string().into()),
3345 model: None,
3346 temperature: Some(0.66),
3347 }],
3348 ..AgentSettings::get_global(cx).clone()
3349 },
3350 cx,
3351 );
3352 });
3353
3354 let request = thread.update(cx, |thread, cx| {
3355 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3356 });
3357 assert_eq!(request.temperature, Some(0.66));
3358
3359 // Same model name, different provider
3360 cx.update(|cx| {
3361 AgentSettings::override_global(
3362 AgentSettings {
3363 model_parameters: vec![LanguageModelParameters {
3364 provider: Some("anthropic".into()),
3365 model: Some(model.id().0.clone()),
3366 temperature: Some(0.66),
3367 }],
3368 ..AgentSettings::get_global(cx).clone()
3369 },
3370 cx,
3371 );
3372 });
3373
3374 let request = thread.update(cx, |thread, cx| {
3375 thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3376 });
3377 assert_eq!(request.temperature, None);
3378 }
3379
3380 #[gpui::test]
3381 async fn test_thread_summary(cx: &mut TestAppContext) {
3382 init_test_settings(cx);
3383
3384 let project = create_test_project(cx, json!({})).await;
3385
3386 let (_, _thread_store, thread, _context_store, model) =
3387 setup_test_environment(cx, project.clone()).await;
3388
3389 // Initial state should be pending
3390 thread.read_with(cx, |thread, _| {
3391 assert!(matches!(thread.summary(), ThreadSummary::Pending));
3392 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3393 });
3394
3395 // Manually setting the summary should not be allowed in this state
3396 thread.update(cx, |thread, cx| {
3397 thread.set_summary("This should not work", cx);
3398 });
3399
3400 thread.read_with(cx, |thread, _| {
3401 assert!(matches!(thread.summary(), ThreadSummary::Pending));
3402 });
3403
3404 // Send a message
3405 thread.update(cx, |thread, cx| {
3406 thread.insert_user_message("Hi!", ContextLoadResult::default(), None, vec![], cx);
3407 thread.send_to_model(
3408 model.clone(),
3409 CompletionIntent::ThreadSummarization,
3410 None,
3411 cx,
3412 );
3413 });
3414
3415 let fake_model = model.as_fake();
3416 simulate_successful_response(&fake_model, cx);
3417
3418 // Should start generating summary when there are >= 2 messages
3419 thread.read_with(cx, |thread, _| {
3420 assert_eq!(*thread.summary(), ThreadSummary::Generating);
3421 });
3422
3423 // Should not be able to set the summary while generating
3424 thread.update(cx, |thread, cx| {
3425 thread.set_summary("This should not work either", cx);
3426 });
3427
3428 thread.read_with(cx, |thread, _| {
3429 assert!(matches!(thread.summary(), ThreadSummary::Generating));
3430 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3431 });
3432
3433 cx.run_until_parked();
3434 fake_model.stream_last_completion_response("Brief");
3435 fake_model.stream_last_completion_response(" Introduction");
3436 fake_model.end_last_completion_stream();
3437 cx.run_until_parked();
3438
3439 // Summary should be set
3440 thread.read_with(cx, |thread, _| {
3441 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
3442 assert_eq!(thread.summary().or_default(), "Brief Introduction");
3443 });
3444
3445 // Now we should be able to set a summary
3446 thread.update(cx, |thread, cx| {
3447 thread.set_summary("Brief Intro", cx);
3448 });
3449
3450 thread.read_with(cx, |thread, _| {
3451 assert_eq!(thread.summary().or_default(), "Brief Intro");
3452 });
3453
3454 // Test setting an empty summary (should default to DEFAULT)
3455 thread.update(cx, |thread, cx| {
3456 thread.set_summary("", cx);
3457 });
3458
3459 thread.read_with(cx, |thread, _| {
3460 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
3461 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3462 });
3463 }
3464
3465 #[gpui::test]
3466 async fn test_thread_summary_error_set_manually(cx: &mut TestAppContext) {
3467 init_test_settings(cx);
3468
3469 let project = create_test_project(cx, json!({})).await;
3470
3471 let (_, _thread_store, thread, _context_store, model) =
3472 setup_test_environment(cx, project.clone()).await;
3473
3474 test_summarize_error(&model, &thread, cx);
3475
3476 // Now we should be able to set a summary
3477 thread.update(cx, |thread, cx| {
3478 thread.set_summary("Brief Intro", cx);
3479 });
3480
3481 thread.read_with(cx, |thread, _| {
3482 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
3483 assert_eq!(thread.summary().or_default(), "Brief Intro");
3484 });
3485 }
3486
3487 #[gpui::test]
3488 async fn test_thread_summary_error_retry(cx: &mut TestAppContext) {
3489 init_test_settings(cx);
3490
3491 let project = create_test_project(cx, json!({})).await;
3492
3493 let (_, _thread_store, thread, _context_store, model) =
3494 setup_test_environment(cx, project.clone()).await;
3495
3496 test_summarize_error(&model, &thread, cx);
3497
3498 // Sending another message should not trigger another summarize request
3499 thread.update(cx, |thread, cx| {
3500 thread.insert_user_message(
3501 "How are you?",
3502 ContextLoadResult::default(),
3503 None,
3504 vec![],
3505 cx,
3506 );
3507 thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
3508 });
3509
3510 let fake_model = model.as_fake();
3511 simulate_successful_response(&fake_model, cx);
3512
3513 thread.read_with(cx, |thread, _| {
3514 // State is still Error, not Generating
3515 assert!(matches!(thread.summary(), ThreadSummary::Error));
3516 });
3517
3518 // But the summarize request can be invoked manually
3519 thread.update(cx, |thread, cx| {
3520 thread.summarize(cx);
3521 });
3522
3523 thread.read_with(cx, |thread, _| {
3524 assert!(matches!(thread.summary(), ThreadSummary::Generating));
3525 });
3526
3527 cx.run_until_parked();
3528 fake_model.stream_last_completion_response("A successful summary");
3529 fake_model.end_last_completion_stream();
3530 cx.run_until_parked();
3531
3532 thread.read_with(cx, |thread, _| {
3533 assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
3534 assert_eq!(thread.summary().or_default(), "A successful summary");
3535 });
3536 }
3537
3538 fn test_summarize_error(
3539 model: &Arc<dyn LanguageModel>,
3540 thread: &Entity<Thread>,
3541 cx: &mut TestAppContext,
3542 ) {
3543 thread.update(cx, |thread, cx| {
3544 thread.insert_user_message("Hi!", ContextLoadResult::default(), None, vec![], cx);
3545 thread.send_to_model(
3546 model.clone(),
3547 CompletionIntent::ThreadSummarization,
3548 None,
3549 cx,
3550 );
3551 });
3552
3553 let fake_model = model.as_fake();
3554 simulate_successful_response(&fake_model, cx);
3555
3556 thread.read_with(cx, |thread, _| {
3557 assert!(matches!(thread.summary(), ThreadSummary::Generating));
3558 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3559 });
3560
3561 // Simulate summary request ending
3562 cx.run_until_parked();
3563 fake_model.end_last_completion_stream();
3564 cx.run_until_parked();
3565
3566 // State is set to Error and default message
3567 thread.read_with(cx, |thread, _| {
3568 assert!(matches!(thread.summary(), ThreadSummary::Error));
3569 assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
3570 });
3571 }
3572
3573 fn simulate_successful_response(fake_model: &FakeLanguageModel, cx: &mut TestAppContext) {
3574 cx.run_until_parked();
3575 fake_model.stream_last_completion_response("Assistant response");
3576 fake_model.end_last_completion_stream();
3577 cx.run_until_parked();
3578 }
3579
3580 fn init_test_settings(cx: &mut TestAppContext) {
3581 cx.update(|cx| {
3582 let settings_store = SettingsStore::test(cx);
3583 cx.set_global(settings_store);
3584 language::init(cx);
3585 Project::init_settings(cx);
3586 AgentSettings::register(cx);
3587 prompt_store::init(cx);
3588 thread_store::init(cx);
3589 workspace::init_settings(cx);
3590 language_model::init_settings(cx);
3591 ThemeSettings::register(cx);
3592 EditorSettings::register(cx);
3593 ToolRegistry::default_global(cx);
3594 });
3595 }
3596
3597 // Helper to create a test project with test files
3598 async fn create_test_project(
3599 cx: &mut TestAppContext,
3600 files: serde_json::Value,
3601 ) -> Entity<Project> {
3602 let fs = FakeFs::new(cx.executor());
3603 fs.insert_tree(path!("/test"), files).await;
3604 Project::test(fs, [path!("/test").as_ref()], cx).await
3605 }
3606
3607 async fn setup_test_environment(
3608 cx: &mut TestAppContext,
3609 project: Entity<Project>,
3610 ) -> (
3611 Entity<Workspace>,
3612 Entity<ThreadStore>,
3613 Entity<Thread>,
3614 Entity<ContextStore>,
3615 Arc<dyn LanguageModel>,
3616 ) {
3617 let (workspace, cx) =
3618 cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
3619
3620 let thread_store = cx
3621 .update(|_, cx| {
3622 ThreadStore::load(
3623 project.clone(),
3624 cx.new(|_| ToolWorkingSet::default()),
3625 None,
3626 Arc::new(PromptBuilder::new(None).unwrap()),
3627 cx,
3628 )
3629 })
3630 .await
3631 .unwrap();
3632
3633 let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
3634 let context_store = cx.new(|_cx| ContextStore::new(project.downgrade(), None));
3635
3636 let provider = Arc::new(FakeLanguageModelProvider);
3637 let model = provider.test_model();
3638 let model: Arc<dyn LanguageModel> = Arc::new(model);
3639
3640 cx.update(|_, cx| {
3641 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
3642 registry.set_default_model(
3643 Some(ConfiguredModel {
3644 provider: provider.clone(),
3645 model: model.clone(),
3646 }),
3647 cx,
3648 );
3649 registry.set_thread_summary_model(
3650 Some(ConfiguredModel {
3651 provider,
3652 model: model.clone(),
3653 }),
3654 cx,
3655 );
3656 })
3657 });
3658
3659 (workspace, thread_store, thread, context_store, model)
3660 }
3661
3662 async fn add_file_to_context(
3663 project: &Entity<Project>,
3664 context_store: &Entity<ContextStore>,
3665 path: &str,
3666 cx: &mut TestAppContext,
3667 ) -> Result<Entity<language::Buffer>> {
3668 let buffer_path = project
3669 .read_with(cx, |project, cx| project.find_project_path(path, cx))
3670 .unwrap();
3671
3672 let buffer = project
3673 .update(cx, |project, cx| {
3674 project.open_buffer(buffer_path.clone(), cx)
3675 })
3676 .await
3677 .unwrap();
3678
3679 context_store.update(cx, |context_store, cx| {
3680 context_store.add_file_from_buffer(&buffer_path, buffer.clone(), false, cx);
3681 });
3682
3683 Ok(buffer)
3684 }
3685}