1mod db;
2mod edit_agent;
3mod legacy_thread;
4mod native_agent_server;
5pub mod outline;
6mod pattern_extraction;
7mod templates;
8#[cfg(test)]
9mod tests;
10mod thread;
11mod thread_store;
12mod tool_permissions;
13mod tools;
14
15use context_server::ContextServerId;
16pub use db::*;
17use itertools::Itertools;
18pub use native_agent_server::NativeAgentServer;
19pub use pattern_extraction::*;
20pub use shell_command_parser::extract_commands;
21pub use templates::*;
22pub use thread::*;
23pub use thread_store::*;
24pub use tool_permissions::*;
25pub use tools::*;
26
27use acp_thread::{
28 AcpThread, AgentModelSelector, AgentSessionInfo, AgentSessionList, AgentSessionListRequest,
29 AgentSessionListResponse, TokenUsageRatio, UserMessageId,
30};
31use agent_client_protocol as acp;
32use anyhow::{Context as _, Result, anyhow};
33use chrono::{DateTime, Utc};
34use collections::{HashMap, HashSet, IndexMap};
35use fs::Fs;
36use futures::channel::{mpsc, oneshot};
37use futures::future::Shared;
38use futures::{FutureExt as _, StreamExt as _, future};
39use gpui::{
40 App, AppContext, AsyncApp, Context, Entity, EntityId, SharedString, Subscription, Task,
41 WeakEntity,
42};
43use language_model::{IconOrSvg, LanguageModel, LanguageModelProvider, LanguageModelRegistry};
44use project::{AgentId, Project, ProjectItem, ProjectPath, Worktree};
45use prompt_store::{
46 ProjectContext, PromptStore, RULES_FILE_NAMES, RulesFileContext, UserRulesContext,
47 WorktreeContext,
48};
49use serde::{Deserialize, Serialize};
50use settings::{LanguageModelSelection, update_settings_file};
51use std::any::Any;
52use std::path::PathBuf;
53use std::rc::Rc;
54use std::sync::{Arc, LazyLock};
55use util::ResultExt;
56use util::path_list::PathList;
57use util::rel_path::RelPath;
58
59#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
60pub struct ProjectSnapshot {
61 pub worktree_snapshots: Vec<project::telemetry_snapshot::TelemetryWorktreeSnapshot>,
62 pub timestamp: DateTime<Utc>,
63}
64
65pub struct RulesLoadingError {
66 pub message: SharedString,
67}
68
69struct ProjectState {
70 project: Entity<Project>,
71 project_context: Entity<ProjectContext>,
72 project_context_needs_refresh: watch::Sender<()>,
73 _maintain_project_context: Task<Result<()>>,
74 context_server_registry: Entity<ContextServerRegistry>,
75 _subscriptions: Vec<Subscription>,
76}
77
78/// Holds both the internal Thread and the AcpThread for a session
79struct Session {
80 /// The internal thread that processes messages
81 thread: Entity<Thread>,
82 /// The ACP thread that handles protocol communication
83 acp_thread: Entity<acp_thread::AcpThread>,
84 project_id: EntityId,
85 pending_save: Task<Result<()>>,
86 _subscriptions: Vec<Subscription>,
87 ref_count: usize,
88}
89
90struct PendingSession {
91 task: Shared<Task<Result<Entity<AcpThread>, Arc<anyhow::Error>>>>,
92 ref_count: usize,
93}
94
95pub struct LanguageModels {
96 /// Access language model by ID
97 models: HashMap<acp::ModelId, Arc<dyn LanguageModel>>,
98 /// Cached list for returning language model information
99 model_list: acp_thread::AgentModelList,
100 refresh_models_rx: watch::Receiver<()>,
101 refresh_models_tx: watch::Sender<()>,
102 _authenticate_all_providers_task: Task<()>,
103}
104
105impl LanguageModels {
106 fn new(cx: &mut App) -> Self {
107 let (refresh_models_tx, refresh_models_rx) = watch::channel(());
108
109 let mut this = Self {
110 models: HashMap::default(),
111 model_list: acp_thread::AgentModelList::Grouped(IndexMap::default()),
112 refresh_models_rx,
113 refresh_models_tx,
114 _authenticate_all_providers_task: Self::authenticate_all_language_model_providers(cx),
115 };
116 this.refresh_list(cx);
117 this
118 }
119
120 fn refresh_list(&mut self, cx: &App) {
121 let providers = LanguageModelRegistry::global(cx)
122 .read(cx)
123 .visible_providers()
124 .into_iter()
125 .filter(|provider| provider.is_authenticated(cx))
126 .collect::<Vec<_>>();
127
128 let mut language_model_list = IndexMap::default();
129 let mut recommended_models = HashSet::default();
130
131 let mut recommended = Vec::new();
132 for provider in &providers {
133 for model in provider.recommended_models(cx) {
134 recommended_models.insert((model.provider_id(), model.id()));
135 recommended.push(Self::map_language_model_to_info(&model, provider));
136 }
137 }
138 if !recommended.is_empty() {
139 language_model_list.insert(
140 acp_thread::AgentModelGroupName("Recommended".into()),
141 recommended,
142 );
143 }
144
145 let mut models = HashMap::default();
146 for provider in providers {
147 let mut provider_models = Vec::new();
148 for model in provider.provided_models(cx) {
149 let model_info = Self::map_language_model_to_info(&model, &provider);
150 let model_id = model_info.id.clone();
151 provider_models.push(model_info);
152 models.insert(model_id, model);
153 }
154 if !provider_models.is_empty() {
155 language_model_list.insert(
156 acp_thread::AgentModelGroupName(provider.name().0.clone()),
157 provider_models,
158 );
159 }
160 }
161
162 self.models = models;
163 self.model_list = acp_thread::AgentModelList::Grouped(language_model_list);
164 self.refresh_models_tx.send(()).ok();
165 }
166
167 fn watch(&self) -> watch::Receiver<()> {
168 self.refresh_models_rx.clone()
169 }
170
171 pub fn model_from_id(&self, model_id: &acp::ModelId) -> Option<Arc<dyn LanguageModel>> {
172 self.models.get(model_id).cloned()
173 }
174
175 fn map_language_model_to_info(
176 model: &Arc<dyn LanguageModel>,
177 provider: &Arc<dyn LanguageModelProvider>,
178 ) -> acp_thread::AgentModelInfo {
179 acp_thread::AgentModelInfo {
180 id: Self::model_id(model),
181 name: model.name().0,
182 description: None,
183 icon: Some(match provider.icon() {
184 IconOrSvg::Svg(path) => acp_thread::AgentModelIcon::Path(path),
185 IconOrSvg::Icon(name) => acp_thread::AgentModelIcon::Named(name),
186 }),
187 is_latest: model.is_latest(),
188 cost: model.model_cost_info().map(|cost| cost.to_shared_string()),
189 }
190 }
191
192 fn model_id(model: &Arc<dyn LanguageModel>) -> acp::ModelId {
193 acp::ModelId::new(format!("{}/{}", model.provider_id().0, model.id().0))
194 }
195
196 fn authenticate_all_language_model_providers(cx: &mut App) -> Task<()> {
197 let authenticate_all_providers = LanguageModelRegistry::global(cx)
198 .read(cx)
199 .visible_providers()
200 .iter()
201 .map(|provider| (provider.id(), provider.name(), provider.authenticate(cx)))
202 .collect::<Vec<_>>();
203
204 cx.spawn(async move |cx| {
205 for (provider_id, provider_name, authenticate_task) in authenticate_all_providers {
206 if let Err(err) = authenticate_task.await {
207 match err {
208 language_model::AuthenticateError::CredentialsNotFound => {
209 // Since we're authenticating these providers in the
210 // background for the purposes of populating the
211 // language selector, we don't care about providers
212 // where the credentials are not found.
213 }
214 language_model::AuthenticateError::ConnectionRefused => {
215 // Not logging connection refused errors as they are mostly from LM Studio's noisy auth failures.
216 // LM Studio only has one auth method (endpoint call) which fails for users who haven't enabled it.
217 // TODO: Better manage LM Studio auth logic to avoid these noisy failures.
218 }
219 _ => {
220 // Some providers have noisy failure states that we
221 // don't want to spam the logs with every time the
222 // language model selector is initialized.
223 //
224 // Ideally these should have more clear failure modes
225 // that we know are safe to ignore here, like what we do
226 // with `CredentialsNotFound` above.
227 match provider_id.0.as_ref() {
228 "lmstudio" | "ollama" => {
229 // LM Studio and Ollama both make fetch requests to the local APIs to determine if they are "authenticated".
230 //
231 // These fail noisily, so we don't log them.
232 }
233 "copilot_chat" => {
234 // Copilot Chat returns an error if Copilot is not enabled, so we don't log those errors.
235 }
236 _ => {
237 log::error!(
238 "Failed to authenticate provider: {}: {err:#}",
239 provider_name.0
240 );
241 }
242 }
243 }
244 }
245 }
246 }
247
248 cx.update(language_models::update_environment_fallback_model);
249 })
250 }
251}
252
253pub struct NativeAgent {
254 /// Session ID -> Session mapping
255 sessions: HashMap<acp::SessionId, Session>,
256 pending_sessions: HashMap<acp::SessionId, PendingSession>,
257 thread_store: Entity<ThreadStore>,
258 /// Project-specific state keyed by project EntityId
259 projects: HashMap<EntityId, ProjectState>,
260 /// Shared templates for all threads
261 templates: Arc<Templates>,
262 /// Cached model information
263 models: LanguageModels,
264 prompt_store: Option<Entity<PromptStore>>,
265 fs: Arc<dyn Fs>,
266 _subscriptions: Vec<Subscription>,
267}
268
269impl NativeAgent {
270 pub fn new(
271 thread_store: Entity<ThreadStore>,
272 templates: Arc<Templates>,
273 prompt_store: Option<Entity<PromptStore>>,
274 fs: Arc<dyn Fs>,
275 cx: &mut App,
276 ) -> Entity<NativeAgent> {
277 log::debug!("Creating new NativeAgent");
278
279 cx.new(|cx| {
280 let mut subscriptions = vec![cx.subscribe(
281 &LanguageModelRegistry::global(cx),
282 Self::handle_models_updated_event,
283 )];
284 if let Some(prompt_store) = prompt_store.as_ref() {
285 subscriptions.push(cx.subscribe(prompt_store, Self::handle_prompts_updated_event))
286 }
287
288 Self {
289 sessions: HashMap::default(),
290 pending_sessions: HashMap::default(),
291 thread_store,
292 projects: HashMap::default(),
293 templates,
294 models: LanguageModels::new(cx),
295 prompt_store,
296 fs,
297 _subscriptions: subscriptions,
298 }
299 })
300 }
301
302 fn new_session(
303 &mut self,
304 project: Entity<Project>,
305 cx: &mut Context<Self>,
306 ) -> Entity<AcpThread> {
307 let project_id = self.get_or_create_project_state(&project, cx);
308 let project_state = &self.projects[&project_id];
309
310 let registry = LanguageModelRegistry::read_global(cx);
311 let available_count = registry.available_models(cx).count();
312 log::debug!("Total available models: {}", available_count);
313
314 let default_model = registry.default_model().and_then(|default_model| {
315 self.models
316 .model_from_id(&LanguageModels::model_id(&default_model.model))
317 });
318 let thread = cx.new(|cx| {
319 Thread::new(
320 project,
321 project_state.project_context.clone(),
322 project_state.context_server_registry.clone(),
323 self.templates.clone(),
324 default_model,
325 cx,
326 )
327 });
328
329 self.register_session(thread, project_id, 1, cx)
330 }
331
332 fn register_session(
333 &mut self,
334 thread_handle: Entity<Thread>,
335 project_id: EntityId,
336 ref_count: usize,
337 cx: &mut Context<Self>,
338 ) -> Entity<AcpThread> {
339 let connection = Rc::new(NativeAgentConnection(cx.entity()));
340
341 let thread = thread_handle.read(cx);
342 let session_id = thread.id().clone();
343 let parent_session_id = thread.parent_thread_id();
344 let title = thread.title();
345 let draft_prompt = thread.draft_prompt().map(Vec::from);
346 let scroll_position = thread.ui_scroll_position();
347 let token_usage = thread.latest_token_usage();
348 let project = thread.project.clone();
349 let action_log = thread.action_log.clone();
350 let prompt_capabilities_rx = thread.prompt_capabilities_rx.clone();
351 let acp_thread = cx.new(|cx| {
352 let mut acp_thread = acp_thread::AcpThread::new(
353 parent_session_id,
354 title,
355 None,
356 connection,
357 project.clone(),
358 action_log.clone(),
359 session_id.clone(),
360 prompt_capabilities_rx,
361 cx,
362 );
363 acp_thread.set_draft_prompt(draft_prompt, cx);
364 acp_thread.set_ui_scroll_position(scroll_position);
365 acp_thread.update_token_usage(token_usage, cx);
366 acp_thread
367 });
368
369 let registry = LanguageModelRegistry::read_global(cx);
370 let summarization_model = registry.thread_summary_model(cx).map(|c| c.model);
371
372 let weak = cx.weak_entity();
373 let weak_thread = thread_handle.downgrade();
374 thread_handle.update(cx, |thread, cx| {
375 thread.set_summarization_model(summarization_model, cx);
376 thread.add_default_tools(
377 Rc::new(NativeThreadEnvironment {
378 acp_thread: acp_thread.downgrade(),
379 thread: weak_thread,
380 agent: weak,
381 }) as _,
382 cx,
383 )
384 });
385
386 let subscriptions = vec![
387 cx.subscribe(&thread_handle, Self::handle_thread_title_updated),
388 cx.subscribe(&thread_handle, Self::handle_thread_token_usage_updated),
389 cx.observe(&thread_handle, move |this, thread, cx| {
390 this.save_thread(thread, cx)
391 }),
392 ];
393
394 self.sessions.insert(
395 session_id,
396 Session {
397 thread: thread_handle,
398 acp_thread: acp_thread.clone(),
399 project_id,
400 _subscriptions: subscriptions,
401 pending_save: Task::ready(Ok(())),
402 ref_count,
403 },
404 );
405
406 self.update_available_commands_for_project(project_id, cx);
407
408 acp_thread
409 }
410
411 pub fn models(&self) -> &LanguageModels {
412 &self.models
413 }
414
415 fn get_or_create_project_state(
416 &mut self,
417 project: &Entity<Project>,
418 cx: &mut Context<Self>,
419 ) -> EntityId {
420 let project_id = project.entity_id();
421 if self.projects.contains_key(&project_id) {
422 return project_id;
423 }
424
425 let project_context = cx.new(|_| ProjectContext::new(vec![], vec![]));
426 self.register_project_with_initial_context(project.clone(), project_context, cx);
427 if let Some(state) = self.projects.get_mut(&project_id) {
428 state.project_context_needs_refresh.send(()).ok();
429 }
430 project_id
431 }
432
433 fn register_project_with_initial_context(
434 &mut self,
435 project: Entity<Project>,
436 project_context: Entity<ProjectContext>,
437 cx: &mut Context<Self>,
438 ) {
439 let project_id = project.entity_id();
440
441 let context_server_store = project.read(cx).context_server_store();
442 let context_server_registry =
443 cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
444
445 let subscriptions = vec![
446 cx.subscribe(&project, Self::handle_project_event),
447 cx.subscribe(
448 &context_server_store,
449 Self::handle_context_server_store_updated,
450 ),
451 cx.subscribe(
452 &context_server_registry,
453 Self::handle_context_server_registry_event,
454 ),
455 ];
456
457 let (project_context_needs_refresh_tx, project_context_needs_refresh_rx) =
458 watch::channel(());
459
460 self.projects.insert(
461 project_id,
462 ProjectState {
463 project,
464 project_context,
465 project_context_needs_refresh: project_context_needs_refresh_tx,
466 _maintain_project_context: cx.spawn(async move |this, cx| {
467 Self::maintain_project_context(
468 this,
469 project_id,
470 project_context_needs_refresh_rx,
471 cx,
472 )
473 .await
474 }),
475 context_server_registry,
476 _subscriptions: subscriptions,
477 },
478 );
479 }
480
481 fn session_project_state(&self, session_id: &acp::SessionId) -> Option<&ProjectState> {
482 self.sessions
483 .get(session_id)
484 .and_then(|session| self.projects.get(&session.project_id))
485 }
486
487 async fn maintain_project_context(
488 this: WeakEntity<Self>,
489 project_id: EntityId,
490 mut needs_refresh: watch::Receiver<()>,
491 cx: &mut AsyncApp,
492 ) -> Result<()> {
493 while needs_refresh.changed().await.is_ok() {
494 let project_context = this
495 .update(cx, |this, cx| {
496 let state = this
497 .projects
498 .get(&project_id)
499 .context("project state not found")?;
500 anyhow::Ok(Self::build_project_context(
501 &state.project,
502 this.prompt_store.as_ref(),
503 cx,
504 ))
505 })??
506 .await;
507 this.update(cx, |this, cx| {
508 if let Some(state) = this.projects.get(&project_id) {
509 state
510 .project_context
511 .update(cx, |current_project_context, _cx| {
512 *current_project_context = project_context;
513 });
514 }
515 })?;
516 }
517
518 Ok(())
519 }
520
521 fn build_project_context(
522 project: &Entity<Project>,
523 prompt_store: Option<&Entity<PromptStore>>,
524 cx: &mut App,
525 ) -> Task<ProjectContext> {
526 let worktrees = project.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
527 let worktree_tasks = worktrees
528 .into_iter()
529 .map(|worktree| {
530 Self::load_worktree_info_for_system_prompt(worktree, project.clone(), cx)
531 })
532 .collect::<Vec<_>>();
533 let default_user_rules_task = if let Some(prompt_store) = prompt_store.as_ref() {
534 prompt_store.read_with(cx, |prompt_store, cx| {
535 let prompts = prompt_store.default_prompt_metadata();
536 let load_tasks = prompts.into_iter().map(|prompt_metadata| {
537 let contents = prompt_store.load(prompt_metadata.id, cx);
538 async move { (contents.await, prompt_metadata) }
539 });
540 cx.background_spawn(future::join_all(load_tasks))
541 })
542 } else {
543 Task::ready(vec![])
544 };
545
546 cx.spawn(async move |_cx| {
547 let (worktrees, default_user_rules) =
548 future::join(future::join_all(worktree_tasks), default_user_rules_task).await;
549
550 let worktrees = worktrees
551 .into_iter()
552 .map(|(worktree, _rules_error)| {
553 // TODO: show error message
554 // if let Some(rules_error) = rules_error {
555 // this.update(cx, |_, cx| cx.emit(rules_error)).ok();
556 // }
557 worktree
558 })
559 .collect::<Vec<_>>();
560
561 let default_user_rules = default_user_rules
562 .into_iter()
563 .flat_map(|(contents, prompt_metadata)| match contents {
564 Ok(contents) => Some(UserRulesContext {
565 uuid: prompt_metadata.id.as_user()?,
566 title: prompt_metadata.title.map(|title| title.to_string()),
567 contents,
568 }),
569 Err(_err) => {
570 // TODO: show error message
571 // this.update(cx, |_, cx| {
572 // cx.emit(RulesLoadingError {
573 // message: format!("{err:?}").into(),
574 // });
575 // })
576 // .ok();
577 None
578 }
579 })
580 .collect::<Vec<_>>();
581
582 ProjectContext::new(worktrees, default_user_rules)
583 })
584 }
585
586 fn load_worktree_info_for_system_prompt(
587 worktree: Entity<Worktree>,
588 project: Entity<Project>,
589 cx: &mut App,
590 ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
591 let tree = worktree.read(cx);
592 let root_name = tree.root_name_str().into();
593 let abs_path = tree.abs_path();
594
595 let mut context = WorktreeContext {
596 root_name,
597 abs_path,
598 rules_file: None,
599 };
600
601 let rules_task = Self::load_worktree_rules_file(worktree, project, cx);
602 let Some(rules_task) = rules_task else {
603 return Task::ready((context, None));
604 };
605
606 cx.spawn(async move |_| {
607 let (rules_file, rules_file_error) = match rules_task.await {
608 Ok(rules_file) => (Some(rules_file), None),
609 Err(err) => (
610 None,
611 Some(RulesLoadingError {
612 message: format!("{err}").into(),
613 }),
614 ),
615 };
616 context.rules_file = rules_file;
617 (context, rules_file_error)
618 })
619 }
620
621 fn load_worktree_rules_file(
622 worktree: Entity<Worktree>,
623 project: Entity<Project>,
624 cx: &mut App,
625 ) -> Option<Task<Result<RulesFileContext>>> {
626 let worktree = worktree.read(cx);
627 let worktree_id = worktree.id();
628 let selected_rules_file = RULES_FILE_NAMES
629 .into_iter()
630 .filter_map(|name| {
631 worktree
632 .entry_for_path(RelPath::unix(name).unwrap())
633 .filter(|entry| entry.is_file())
634 .map(|entry| entry.path.clone())
635 })
636 .next();
637
638 // Note that Cline supports `.clinerules` being a directory, but that is not currently
639 // supported. This doesn't seem to occur often in GitHub repositories.
640 selected_rules_file.map(|path_in_worktree| {
641 let project_path = ProjectPath {
642 worktree_id,
643 path: path_in_worktree.clone(),
644 };
645 let buffer_task =
646 project.update(cx, |project, cx| project.open_buffer(project_path, cx));
647 let rope_task = cx.spawn(async move |cx| {
648 let buffer = buffer_task.await?;
649 let (project_entry_id, rope) = buffer.read_with(cx, |buffer, cx| {
650 let project_entry_id = buffer.entry_id(cx).context("buffer has no file")?;
651 anyhow::Ok((project_entry_id, buffer.as_rope().clone()))
652 })?;
653 anyhow::Ok((project_entry_id, rope))
654 });
655 // Build a string from the rope on a background thread.
656 cx.background_spawn(async move {
657 let (project_entry_id, rope) = rope_task.await?;
658 anyhow::Ok(RulesFileContext {
659 path_in_worktree,
660 text: rope.to_string().trim().to_string(),
661 project_entry_id: project_entry_id.to_usize(),
662 })
663 })
664 })
665 }
666
667 fn handle_thread_title_updated(
668 &mut self,
669 thread: Entity<Thread>,
670 _: &TitleUpdated,
671 cx: &mut Context<Self>,
672 ) {
673 let session_id = thread.read(cx).id();
674 let Some(session) = self.sessions.get(session_id) else {
675 return;
676 };
677
678 let thread = thread.downgrade();
679 let acp_thread = session.acp_thread.downgrade();
680 cx.spawn(async move |_, cx| {
681 let title = thread.read_with(cx, |thread, _| thread.title())?;
682 if let Some(title) = title {
683 let task =
684 acp_thread.update(cx, |acp_thread, cx| acp_thread.set_title(title, cx))?;
685 task.await?;
686 }
687 anyhow::Ok(())
688 })
689 .detach_and_log_err(cx);
690 }
691
692 fn handle_thread_token_usage_updated(
693 &mut self,
694 thread: Entity<Thread>,
695 usage: &TokenUsageUpdated,
696 cx: &mut Context<Self>,
697 ) {
698 let Some(session) = self.sessions.get(thread.read(cx).id()) else {
699 return;
700 };
701 session.acp_thread.update(cx, |acp_thread, cx| {
702 acp_thread.update_token_usage(usage.0.clone(), cx);
703 });
704 }
705
706 fn handle_project_event(
707 &mut self,
708 project: Entity<Project>,
709 event: &project::Event,
710 _cx: &mut Context<Self>,
711 ) {
712 let project_id = project.entity_id();
713 let Some(state) = self.projects.get_mut(&project_id) else {
714 return;
715 };
716 match event {
717 project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
718 state.project_context_needs_refresh.send(()).ok();
719 }
720 project::Event::WorktreeUpdatedEntries(_, items) => {
721 if items.iter().any(|(path, _, _)| {
722 RULES_FILE_NAMES
723 .iter()
724 .any(|name| path.as_ref() == RelPath::unix(name).unwrap())
725 }) {
726 state.project_context_needs_refresh.send(()).ok();
727 }
728 }
729 _ => {}
730 }
731 }
732
733 fn handle_prompts_updated_event(
734 &mut self,
735 _prompt_store: Entity<PromptStore>,
736 _event: &prompt_store::PromptsUpdatedEvent,
737 _cx: &mut Context<Self>,
738 ) {
739 for state in self.projects.values_mut() {
740 state.project_context_needs_refresh.send(()).ok();
741 }
742 }
743
744 fn handle_models_updated_event(
745 &mut self,
746 _registry: Entity<LanguageModelRegistry>,
747 event: &language_model::Event,
748 cx: &mut Context<Self>,
749 ) {
750 self.models.refresh_list(cx);
751
752 let registry = LanguageModelRegistry::read_global(cx);
753 let default_model = registry.default_model().map(|m| m.model);
754 let summarization_model = registry.thread_summary_model(cx).map(|m| m.model);
755
756 for session in self.sessions.values_mut() {
757 session.thread.update(cx, |thread, cx| {
758 let should_update_model = thread.model().is_none()
759 || (thread.is_empty()
760 && matches!(event, language_model::Event::DefaultModelChanged));
761 if should_update_model && let Some(model) = default_model.clone() {
762 thread.set_model(model, cx);
763 cx.notify();
764 }
765 if let Some(model) = summarization_model.clone() {
766 if thread.summarization_model().is_none()
767 || matches!(event, language_model::Event::ThreadSummaryModelChanged)
768 {
769 thread.set_summarization_model(Some(model), cx);
770 }
771 }
772 });
773 }
774 }
775
776 fn handle_context_server_store_updated(
777 &mut self,
778 store: Entity<project::context_server_store::ContextServerStore>,
779 _event: &project::context_server_store::ServerStatusChangedEvent,
780 cx: &mut Context<Self>,
781 ) {
782 let project_id = self.projects.iter().find_map(|(id, state)| {
783 if *state.context_server_registry.read(cx).server_store() == store {
784 Some(*id)
785 } else {
786 None
787 }
788 });
789 if let Some(project_id) = project_id {
790 self.update_available_commands_for_project(project_id, cx);
791 }
792 }
793
794 fn handle_context_server_registry_event(
795 &mut self,
796 registry: Entity<ContextServerRegistry>,
797 event: &ContextServerRegistryEvent,
798 cx: &mut Context<Self>,
799 ) {
800 match event {
801 ContextServerRegistryEvent::ToolsChanged => {}
802 ContextServerRegistryEvent::PromptsChanged => {
803 let project_id = self.projects.iter().find_map(|(id, state)| {
804 if state.context_server_registry == registry {
805 Some(*id)
806 } else {
807 None
808 }
809 });
810 if let Some(project_id) = project_id {
811 self.update_available_commands_for_project(project_id, cx);
812 }
813 }
814 }
815 }
816
817 fn update_available_commands_for_project(&self, project_id: EntityId, cx: &mut Context<Self>) {
818 let available_commands =
819 Self::build_available_commands_for_project(self.projects.get(&project_id), cx);
820 for session in self.sessions.values() {
821 if session.project_id != project_id {
822 continue;
823 }
824 session.acp_thread.update(cx, |thread, cx| {
825 thread
826 .handle_session_update(
827 acp::SessionUpdate::AvailableCommandsUpdate(
828 acp::AvailableCommandsUpdate::new(available_commands.clone()),
829 ),
830 cx,
831 )
832 .log_err();
833 });
834 }
835 }
836
837 fn build_available_commands_for_project(
838 project_state: Option<&ProjectState>,
839 cx: &App,
840 ) -> Vec<acp::AvailableCommand> {
841 let Some(state) = project_state else {
842 return vec![];
843 };
844 let registry = state.context_server_registry.read(cx);
845
846 let mut prompt_name_counts: HashMap<&str, usize> = HashMap::default();
847 for context_server_prompt in registry.prompts() {
848 *prompt_name_counts
849 .entry(context_server_prompt.prompt.name.as_str())
850 .or_insert(0) += 1;
851 }
852
853 registry
854 .prompts()
855 .flat_map(|context_server_prompt| {
856 let prompt = &context_server_prompt.prompt;
857
858 let should_prefix = prompt_name_counts
859 .get(prompt.name.as_str())
860 .copied()
861 .unwrap_or(0)
862 > 1;
863
864 let name = if should_prefix {
865 format!("{}.{}", context_server_prompt.server_id, prompt.name)
866 } else {
867 prompt.name.clone()
868 };
869
870 let mut command = acp::AvailableCommand::new(
871 name,
872 prompt.description.clone().unwrap_or_default(),
873 );
874
875 match prompt.arguments.as_deref() {
876 Some([arg]) => {
877 let hint = format!("<{}>", arg.name);
878
879 command = command.input(acp::AvailableCommandInput::Unstructured(
880 acp::UnstructuredCommandInput::new(hint),
881 ));
882 }
883 Some([]) | None => {}
884 Some(_) => {
885 // skip >1 argument commands since we don't support them yet
886 return None;
887 }
888 }
889
890 Some(command)
891 })
892 .collect()
893 }
894
895 pub fn load_thread(
896 &mut self,
897 id: acp::SessionId,
898 project: Entity<Project>,
899 cx: &mut Context<Self>,
900 ) -> Task<Result<Entity<Thread>>> {
901 let database_future = ThreadsDatabase::connect(cx);
902 cx.spawn(async move |this, cx| {
903 let database = database_future.await.map_err(|err| anyhow!(err))?;
904 let db_thread = database
905 .load_thread(id.clone())
906 .await?
907 .with_context(|| format!("no thread found with ID: {id:?}"))?;
908
909 this.update(cx, |this, cx| {
910 let project_id = this.get_or_create_project_state(&project, cx);
911 let project_state = this
912 .projects
913 .get(&project_id)
914 .context("project state not found")?;
915 let summarization_model = LanguageModelRegistry::read_global(cx)
916 .thread_summary_model(cx)
917 .map(|c| c.model);
918
919 Ok(cx.new(|cx| {
920 let mut thread = Thread::from_db(
921 id.clone(),
922 db_thread,
923 project_state.project.clone(),
924 project_state.project_context.clone(),
925 project_state.context_server_registry.clone(),
926 this.templates.clone(),
927 cx,
928 );
929 thread.set_summarization_model(summarization_model, cx);
930 thread
931 }))
932 })?
933 })
934 }
935
936 pub fn open_thread(
937 &mut self,
938 id: acp::SessionId,
939 project: Entity<Project>,
940 cx: &mut Context<Self>,
941 ) -> Task<Result<Entity<AcpThread>>> {
942 if let Some(session) = self.sessions.get_mut(&id) {
943 session.ref_count += 1;
944 return Task::ready(Ok(session.acp_thread.clone()));
945 }
946
947 if let Some(pending) = self.pending_sessions.get_mut(&id) {
948 pending.ref_count += 1;
949 let task = pending.task.clone();
950 return cx.background_spawn(async move { task.await.map_err(|err| anyhow!(err)) });
951 }
952
953 let task = self.load_thread(id.clone(), project.clone(), cx);
954 let shared_task = cx
955 .spawn({
956 let id = id.clone();
957 async move |this, cx| {
958 let thread = match task.await {
959 Ok(thread) => thread,
960 Err(err) => {
961 this.update(cx, |this, _cx| {
962 this.pending_sessions.remove(&id);
963 })
964 .ok();
965 return Err(Arc::new(err));
966 }
967 };
968 let acp_thread = this
969 .update(cx, |this, cx| {
970 let project_id = this.get_or_create_project_state(&project, cx);
971 let ref_count = this
972 .pending_sessions
973 .remove(&id)
974 .map_or(1, |pending| pending.ref_count);
975 this.register_session(thread.clone(), project_id, ref_count, cx)
976 })
977 .map_err(Arc::new)?;
978 let events = thread.update(cx, |thread, cx| thread.replay(cx));
979 cx.update(|cx| {
980 NativeAgentConnection::handle_thread_events(
981 events,
982 acp_thread.downgrade(),
983 cx,
984 )
985 })
986 .await
987 .map_err(Arc::new)?;
988 acp_thread.update(cx, |thread, cx| {
989 thread.snapshot_completed_plan(cx);
990 });
991 Ok(acp_thread)
992 }
993 })
994 .shared();
995 self.pending_sessions.insert(
996 id,
997 PendingSession {
998 task: shared_task.clone(),
999 ref_count: 1,
1000 },
1001 );
1002
1003 cx.background_spawn(async move { shared_task.await.map_err(|err| anyhow!(err)) })
1004 }
1005
1006 pub fn thread_summary(
1007 &mut self,
1008 id: acp::SessionId,
1009 project: Entity<Project>,
1010 cx: &mut Context<Self>,
1011 ) -> Task<Result<SharedString>> {
1012 let thread = self.open_thread(id.clone(), project, cx);
1013 cx.spawn(async move |this, cx| {
1014 let acp_thread = thread.await?;
1015 let result = this
1016 .update(cx, |this, cx| {
1017 this.sessions
1018 .get(&id)
1019 .unwrap()
1020 .thread
1021 .update(cx, |thread, cx| thread.summary(cx))
1022 })?
1023 .await
1024 .context("Failed to generate summary")?;
1025
1026 this.update(cx, |this, cx| this.close_session(&id, cx))?
1027 .await?;
1028 drop(acp_thread);
1029 Ok(result)
1030 })
1031 }
1032
1033 fn close_session(
1034 &mut self,
1035 session_id: &acp::SessionId,
1036 cx: &mut Context<Self>,
1037 ) -> Task<Result<()>> {
1038 let Some(session) = self.sessions.get_mut(session_id) else {
1039 return Task::ready(Ok(()));
1040 };
1041
1042 session.ref_count -= 1;
1043 if session.ref_count > 0 {
1044 return Task::ready(Ok(()));
1045 }
1046
1047 let thread = session.thread.clone();
1048 self.save_thread(thread, cx);
1049 let Some(session) = self.sessions.remove(session_id) else {
1050 return Task::ready(Ok(()));
1051 };
1052 let project_id = session.project_id;
1053
1054 let has_remaining = self.sessions.values().any(|s| s.project_id == project_id);
1055 if !has_remaining {
1056 self.projects.remove(&project_id);
1057 }
1058
1059 session.pending_save
1060 }
1061
1062 fn save_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
1063 if thread.read(cx).is_empty() {
1064 return;
1065 }
1066
1067 let id = thread.read(cx).id().clone();
1068 let Some(session) = self.sessions.get_mut(&id) else {
1069 return;
1070 };
1071
1072 let project_id = session.project_id;
1073 let Some(state) = self.projects.get(&project_id) else {
1074 return;
1075 };
1076
1077 let folder_paths = PathList::new(
1078 &state
1079 .project
1080 .read(cx)
1081 .visible_worktrees(cx)
1082 .map(|worktree| worktree.read(cx).abs_path().to_path_buf())
1083 .collect::<Vec<_>>(),
1084 );
1085
1086 let draft_prompt = session.acp_thread.read(cx).draft_prompt().map(Vec::from);
1087 let database_future = ThreadsDatabase::connect(cx);
1088 let db_thread = thread.update(cx, |thread, cx| {
1089 thread.set_draft_prompt(draft_prompt);
1090 thread.to_db(cx)
1091 });
1092 let thread_store = self.thread_store.clone();
1093 session.pending_save = cx.spawn(async move |_, cx| {
1094 let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
1095 return Ok(());
1096 };
1097 let db_thread = db_thread.await;
1098 database
1099 .save_thread(id, db_thread, folder_paths)
1100 .await
1101 .log_err();
1102 thread_store.update(cx, |store, cx| store.reload(cx));
1103 Ok(())
1104 });
1105 }
1106
1107 fn send_mcp_prompt(
1108 &self,
1109 message_id: UserMessageId,
1110 session_id: acp::SessionId,
1111 prompt_name: String,
1112 server_id: ContextServerId,
1113 arguments: HashMap<String, String>,
1114 original_content: Vec<acp::ContentBlock>,
1115 cx: &mut Context<Self>,
1116 ) -> Task<Result<acp::PromptResponse>> {
1117 let Some(state) = self.session_project_state(&session_id) else {
1118 return Task::ready(Err(anyhow!("Project state not found for session")));
1119 };
1120 let server_store = state
1121 .context_server_registry
1122 .read(cx)
1123 .server_store()
1124 .clone();
1125 let path_style = state.project.read(cx).path_style(cx);
1126
1127 cx.spawn(async move |this, cx| {
1128 let prompt =
1129 crate::get_prompt(&server_store, &server_id, &prompt_name, arguments, cx).await?;
1130
1131 let (acp_thread, thread) = this.update(cx, |this, _cx| {
1132 let session = this
1133 .sessions
1134 .get(&session_id)
1135 .context("Failed to get session")?;
1136 anyhow::Ok((session.acp_thread.clone(), session.thread.clone()))
1137 })??;
1138
1139 let mut last_is_user = true;
1140
1141 thread.update(cx, |thread, cx| {
1142 thread.push_acp_user_block(
1143 message_id,
1144 original_content.into_iter().skip(1),
1145 path_style,
1146 cx,
1147 );
1148 });
1149
1150 for message in prompt.messages {
1151 let context_server::types::PromptMessage { role, content } = message;
1152 let block = mcp_message_content_to_acp_content_block(content);
1153
1154 match role {
1155 context_server::types::Role::User => {
1156 let id = acp_thread::UserMessageId::new();
1157
1158 acp_thread.update(cx, |acp_thread, cx| {
1159 acp_thread.push_user_content_block_with_indent(
1160 Some(id.clone()),
1161 block.clone(),
1162 true,
1163 cx,
1164 );
1165 });
1166
1167 thread.update(cx, |thread, cx| {
1168 thread.push_acp_user_block(id, [block], path_style, cx);
1169 });
1170 }
1171 context_server::types::Role::Assistant => {
1172 acp_thread.update(cx, |acp_thread, cx| {
1173 acp_thread.push_assistant_content_block_with_indent(
1174 block.clone(),
1175 false,
1176 true,
1177 cx,
1178 );
1179 });
1180
1181 thread.update(cx, |thread, cx| {
1182 thread.push_acp_agent_block(block, cx);
1183 });
1184 }
1185 }
1186
1187 last_is_user = role == context_server::types::Role::User;
1188 }
1189
1190 let response_stream = thread.update(cx, |thread, cx| {
1191 if last_is_user {
1192 thread.send_existing(cx)
1193 } else {
1194 // Resume if MCP prompt did not end with a user message
1195 thread.resume(cx)
1196 }
1197 })?;
1198
1199 cx.update(|cx| {
1200 NativeAgentConnection::handle_thread_events(
1201 response_stream,
1202 acp_thread.downgrade(),
1203 cx,
1204 )
1205 })
1206 .await
1207 })
1208 }
1209}
1210
1211/// Wrapper struct that implements the AgentConnection trait
1212#[derive(Clone)]
1213pub struct NativeAgentConnection(pub Entity<NativeAgent>);
1214
1215impl NativeAgentConnection {
1216 pub fn thread(&self, session_id: &acp::SessionId, cx: &App) -> Option<Entity<Thread>> {
1217 self.0
1218 .read(cx)
1219 .sessions
1220 .get(session_id)
1221 .map(|session| session.thread.clone())
1222 }
1223
1224 pub fn load_thread(
1225 &self,
1226 id: acp::SessionId,
1227 project: Entity<Project>,
1228 cx: &mut App,
1229 ) -> Task<Result<Entity<Thread>>> {
1230 self.0
1231 .update(cx, |this, cx| this.load_thread(id, project, cx))
1232 }
1233
1234 fn run_turn(
1235 &self,
1236 session_id: acp::SessionId,
1237 cx: &mut App,
1238 f: impl 'static
1239 + FnOnce(Entity<Thread>, &mut App) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>,
1240 ) -> Task<Result<acp::PromptResponse>> {
1241 let Some((thread, acp_thread)) = self.0.update(cx, |agent, _cx| {
1242 agent
1243 .sessions
1244 .get_mut(&session_id)
1245 .map(|s| (s.thread.clone(), s.acp_thread.clone()))
1246 }) else {
1247 log::error!("Session not found in run_turn: {}", session_id);
1248 return Task::ready(Err(anyhow!("Session not found")));
1249 };
1250 log::debug!("Found session for: {}", session_id);
1251
1252 let response_stream = match f(thread, cx) {
1253 Ok(stream) => stream,
1254 Err(err) => return Task::ready(Err(err)),
1255 };
1256 Self::handle_thread_events(response_stream, acp_thread.downgrade(), cx)
1257 }
1258
1259 fn handle_thread_events(
1260 mut events: mpsc::UnboundedReceiver<Result<ThreadEvent>>,
1261 acp_thread: WeakEntity<AcpThread>,
1262 cx: &App,
1263 ) -> Task<Result<acp::PromptResponse>> {
1264 cx.spawn(async move |cx| {
1265 // Handle response stream and forward to session.acp_thread
1266 while let Some(result) = events.next().await {
1267 match result {
1268 Ok(event) => {
1269 log::trace!("Received completion event: {:?}", event);
1270
1271 match event {
1272 ThreadEvent::UserMessage(message) => {
1273 acp_thread.update(cx, |thread, cx| {
1274 for content in message.content {
1275 thread.push_user_content_block(
1276 Some(message.id.clone()),
1277 content.into(),
1278 cx,
1279 );
1280 }
1281 })?;
1282 }
1283 ThreadEvent::AgentText(text) => {
1284 acp_thread.update(cx, |thread, cx| {
1285 thread.push_assistant_content_block(text.into(), false, cx)
1286 })?;
1287 }
1288 ThreadEvent::AgentThinking(text) => {
1289 acp_thread.update(cx, |thread, cx| {
1290 thread.push_assistant_content_block(text.into(), true, cx)
1291 })?;
1292 }
1293 ThreadEvent::ToolCallAuthorization(ToolCallAuthorization {
1294 tool_call,
1295 options,
1296 response,
1297 context: _,
1298 }) => {
1299 let outcome_task = acp_thread.update(cx, |thread, cx| {
1300 thread.request_tool_call_authorization(tool_call, options, cx)
1301 })??;
1302 cx.background_spawn(async move {
1303 if let acp_thread::RequestPermissionOutcome::Selected(outcome) =
1304 outcome_task.await
1305 {
1306 response
1307 .send(outcome)
1308 .map(|_| anyhow!("authorization receiver was dropped"))
1309 .log_err();
1310 }
1311 })
1312 .detach();
1313 }
1314 ThreadEvent::ToolCall(tool_call) => {
1315 acp_thread.update(cx, |thread, cx| {
1316 thread.upsert_tool_call(tool_call, cx)
1317 })??;
1318 }
1319 ThreadEvent::ToolCallUpdate(update) => {
1320 acp_thread.update(cx, |thread, cx| {
1321 thread.update_tool_call(update, cx)
1322 })??;
1323 }
1324 ThreadEvent::Plan(plan) => {
1325 acp_thread.update(cx, |thread, cx| thread.update_plan(plan, cx))?;
1326 }
1327 ThreadEvent::SubagentSpawned(session_id) => {
1328 acp_thread.update(cx, |thread, cx| {
1329 thread.subagent_spawned(session_id, cx);
1330 })?;
1331 }
1332 ThreadEvent::Retry(status) => {
1333 acp_thread.update(cx, |thread, cx| {
1334 thread.update_retry_status(status, cx)
1335 })?;
1336 }
1337 ThreadEvent::Stop(stop_reason) => {
1338 log::debug!("Assistant message complete: {:?}", stop_reason);
1339 return Ok(acp::PromptResponse::new(stop_reason));
1340 }
1341 }
1342 }
1343 Err(e) => {
1344 log::error!("Error in model response stream: {:?}", e);
1345 return Err(e);
1346 }
1347 }
1348 }
1349
1350 log::debug!("Response stream completed");
1351 anyhow::Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
1352 })
1353 }
1354}
1355
1356struct Command<'a> {
1357 prompt_name: &'a str,
1358 arg_value: &'a str,
1359 explicit_server_id: Option<&'a str>,
1360}
1361
1362impl<'a> Command<'a> {
1363 fn parse(prompt: &'a [acp::ContentBlock]) -> Option<Self> {
1364 let acp::ContentBlock::Text(text_content) = prompt.first()? else {
1365 return None;
1366 };
1367 let text = text_content.text.trim();
1368 let command = text.strip_prefix('/')?;
1369 let (command, arg_value) = command
1370 .split_once(char::is_whitespace)
1371 .unwrap_or((command, ""));
1372
1373 if let Some((server_id, prompt_name)) = command.split_once('.') {
1374 Some(Self {
1375 prompt_name,
1376 arg_value,
1377 explicit_server_id: Some(server_id),
1378 })
1379 } else {
1380 Some(Self {
1381 prompt_name: command,
1382 arg_value,
1383 explicit_server_id: None,
1384 })
1385 }
1386 }
1387}
1388
1389struct NativeAgentModelSelector {
1390 session_id: acp::SessionId,
1391 connection: NativeAgentConnection,
1392}
1393
1394impl acp_thread::AgentModelSelector for NativeAgentModelSelector {
1395 fn list_models(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelList>> {
1396 log::debug!("NativeAgentConnection::list_models called");
1397 let list = self.connection.0.read(cx).models.model_list.clone();
1398 Task::ready(if list.is_empty() {
1399 Err(anyhow::anyhow!("No models available"))
1400 } else {
1401 Ok(list)
1402 })
1403 }
1404
1405 fn select_model(&self, model_id: acp::ModelId, cx: &mut App) -> Task<Result<()>> {
1406 log::debug!(
1407 "Setting model for session {}: {}",
1408 self.session_id,
1409 model_id
1410 );
1411 let Some(thread) = self
1412 .connection
1413 .0
1414 .read(cx)
1415 .sessions
1416 .get(&self.session_id)
1417 .map(|session| session.thread.clone())
1418 else {
1419 return Task::ready(Err(anyhow!("Session not found")));
1420 };
1421
1422 let Some(model) = self.connection.0.read(cx).models.model_from_id(&model_id) else {
1423 return Task::ready(Err(anyhow!("Invalid model ID {}", model_id)));
1424 };
1425
1426 // We want to reset the effort level when switching models, as the currently-selected effort level may
1427 // not be compatible.
1428 let effort = model
1429 .default_effort_level()
1430 .map(|effort_level| effort_level.value.to_string());
1431
1432 thread.update(cx, |thread, cx| {
1433 thread.set_model(model.clone(), cx);
1434 thread.set_thinking_effort(effort.clone(), cx);
1435 thread.set_thinking_enabled(model.supports_thinking(), cx);
1436 });
1437
1438 update_settings_file(
1439 self.connection.0.read(cx).fs.clone(),
1440 cx,
1441 move |settings, cx| {
1442 let provider = model.provider_id().0.to_string();
1443 let model = model.id().0.to_string();
1444 let enable_thinking = thread.read(cx).thinking_enabled();
1445 let speed = thread.read(cx).speed();
1446 settings
1447 .agent
1448 .get_or_insert_default()
1449 .set_model(LanguageModelSelection {
1450 provider: provider.into(),
1451 model,
1452 enable_thinking,
1453 effort,
1454 speed,
1455 });
1456 },
1457 );
1458
1459 Task::ready(Ok(()))
1460 }
1461
1462 fn selected_model(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelInfo>> {
1463 let Some(thread) = self
1464 .connection
1465 .0
1466 .read(cx)
1467 .sessions
1468 .get(&self.session_id)
1469 .map(|session| session.thread.clone())
1470 else {
1471 return Task::ready(Err(anyhow!("Session not found")));
1472 };
1473 let Some(model) = thread.read(cx).model() else {
1474 return Task::ready(Err(anyhow!("Model not found")));
1475 };
1476 let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&model.provider_id())
1477 else {
1478 return Task::ready(Err(anyhow!("Provider not found")));
1479 };
1480 Task::ready(Ok(LanguageModels::map_language_model_to_info(
1481 model, &provider,
1482 )))
1483 }
1484
1485 fn watch(&self, cx: &mut App) -> Option<watch::Receiver<()>> {
1486 Some(self.connection.0.read(cx).models.watch())
1487 }
1488
1489 fn should_render_footer(&self) -> bool {
1490 true
1491 }
1492}
1493
1494pub static ZED_AGENT_ID: LazyLock<AgentId> = LazyLock::new(|| AgentId::new("Zed Agent"));
1495
1496impl acp_thread::AgentConnection for NativeAgentConnection {
1497 fn agent_id(&self) -> AgentId {
1498 ZED_AGENT_ID.clone()
1499 }
1500
1501 fn telemetry_id(&self) -> SharedString {
1502 "zed".into()
1503 }
1504
1505 fn new_session(
1506 self: Rc<Self>,
1507 project: Entity<Project>,
1508 work_dirs: PathList,
1509 cx: &mut App,
1510 ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
1511 log::debug!("Creating new thread for project at: {work_dirs:?}");
1512 Task::ready(Ok(self
1513 .0
1514 .update(cx, |agent, cx| agent.new_session(project, cx))))
1515 }
1516
1517 fn supports_load_session(&self) -> bool {
1518 true
1519 }
1520
1521 fn load_session(
1522 self: Rc<Self>,
1523 session_id: acp::SessionId,
1524 project: Entity<Project>,
1525 _work_dirs: PathList,
1526 _title: Option<SharedString>,
1527 cx: &mut App,
1528 ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
1529 self.0
1530 .update(cx, |agent, cx| agent.open_thread(session_id, project, cx))
1531 }
1532
1533 fn supports_close_session(&self) -> bool {
1534 true
1535 }
1536
1537 fn close_session(
1538 self: Rc<Self>,
1539 session_id: &acp::SessionId,
1540 cx: &mut App,
1541 ) -> Task<Result<()>> {
1542 self.0
1543 .update(cx, |agent, cx| agent.close_session(session_id, cx))
1544 }
1545
1546 fn auth_methods(&self) -> &[acp::AuthMethod] {
1547 &[] // No auth for in-process
1548 }
1549
1550 fn authenticate(&self, _method: acp::AuthMethodId, _cx: &mut App) -> Task<Result<()>> {
1551 Task::ready(Ok(()))
1552 }
1553
1554 fn model_selector(&self, session_id: &acp::SessionId) -> Option<Rc<dyn AgentModelSelector>> {
1555 Some(Rc::new(NativeAgentModelSelector {
1556 session_id: session_id.clone(),
1557 connection: self.clone(),
1558 }) as Rc<dyn AgentModelSelector>)
1559 }
1560
1561 fn prompt(
1562 &self,
1563 id: acp_thread::UserMessageId,
1564 params: acp::PromptRequest,
1565 cx: &mut App,
1566 ) -> Task<Result<acp::PromptResponse>> {
1567 let session_id = params.session_id.clone();
1568 log::info!("Received prompt request for session: {}", session_id);
1569 log::debug!("Prompt blocks count: {}", params.prompt.len());
1570
1571 let Some(project_state) = self.0.read(cx).session_project_state(&session_id) else {
1572 log::error!("Session not found in prompt: {}", session_id);
1573 if self.0.read(cx).sessions.contains_key(&session_id) {
1574 log::error!(
1575 "Session found in sessions map, but not in project state: {}",
1576 session_id
1577 );
1578 }
1579 return Task::ready(Err(anyhow::anyhow!("Session not found")));
1580 };
1581
1582 if let Some(parsed_command) = Command::parse(¶ms.prompt) {
1583 let registry = project_state.context_server_registry.read(cx);
1584
1585 let explicit_server_id = parsed_command
1586 .explicit_server_id
1587 .map(|server_id| ContextServerId(server_id.into()));
1588
1589 if let Some(prompt) =
1590 registry.find_prompt(explicit_server_id.as_ref(), parsed_command.prompt_name)
1591 {
1592 let arguments = if !parsed_command.arg_value.is_empty()
1593 && let Some(arg_name) = prompt
1594 .prompt
1595 .arguments
1596 .as_ref()
1597 .and_then(|args| args.first())
1598 .map(|arg| arg.name.clone())
1599 {
1600 HashMap::from_iter([(arg_name, parsed_command.arg_value.to_string())])
1601 } else {
1602 Default::default()
1603 };
1604
1605 let prompt_name = prompt.prompt.name.clone();
1606 let server_id = prompt.server_id.clone();
1607
1608 return self.0.update(cx, |agent, cx| {
1609 agent.send_mcp_prompt(
1610 id,
1611 session_id.clone(),
1612 prompt_name,
1613 server_id,
1614 arguments,
1615 params.prompt,
1616 cx,
1617 )
1618 });
1619 }
1620 };
1621
1622 let path_style = project_state.project.read(cx).path_style(cx);
1623
1624 self.run_turn(session_id, cx, move |thread, cx| {
1625 let content: Vec<UserMessageContent> = params
1626 .prompt
1627 .into_iter()
1628 .map(|block| UserMessageContent::from_content_block(block, path_style))
1629 .collect::<Vec<_>>();
1630 log::debug!("Converted prompt to message: {} chars", content.len());
1631 log::debug!("Message id: {:?}", id);
1632 log::debug!("Message content: {:?}", content);
1633
1634 thread.update(cx, |thread, cx| thread.send(id, content, cx))
1635 })
1636 }
1637
1638 fn retry(
1639 &self,
1640 session_id: &acp::SessionId,
1641 _cx: &App,
1642 ) -> Option<Rc<dyn acp_thread::AgentSessionRetry>> {
1643 Some(Rc::new(NativeAgentSessionRetry {
1644 connection: self.clone(),
1645 session_id: session_id.clone(),
1646 }) as _)
1647 }
1648
1649 fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
1650 log::info!("Cancelling on session: {}", session_id);
1651 self.0.update(cx, |agent, cx| {
1652 if let Some(session) = agent.sessions.get(session_id) {
1653 session
1654 .thread
1655 .update(cx, |thread, cx| thread.cancel(cx))
1656 .detach();
1657 }
1658 });
1659 }
1660
1661 fn truncate(
1662 &self,
1663 session_id: &acp::SessionId,
1664 cx: &App,
1665 ) -> Option<Rc<dyn acp_thread::AgentSessionTruncate>> {
1666 self.0.read_with(cx, |agent, _cx| {
1667 agent.sessions.get(session_id).map(|session| {
1668 Rc::new(NativeAgentSessionTruncate {
1669 thread: session.thread.clone(),
1670 acp_thread: session.acp_thread.downgrade(),
1671 }) as _
1672 })
1673 })
1674 }
1675
1676 fn set_title(
1677 &self,
1678 session_id: &acp::SessionId,
1679 cx: &App,
1680 ) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
1681 self.0.read_with(cx, |agent, _cx| {
1682 agent
1683 .sessions
1684 .get(session_id)
1685 .filter(|s| !s.thread.read(cx).is_subagent())
1686 .map(|session| {
1687 Rc::new(NativeAgentSessionSetTitle {
1688 thread: session.thread.clone(),
1689 }) as _
1690 })
1691 })
1692 }
1693
1694 fn session_list(&self, cx: &mut App) -> Option<Rc<dyn AgentSessionList>> {
1695 let thread_store = self.0.read(cx).thread_store.clone();
1696 Some(Rc::new(NativeAgentSessionList::new(thread_store, cx)) as _)
1697 }
1698
1699 fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
1700 Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
1701 }
1702
1703 fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1704 self
1705 }
1706}
1707
1708impl acp_thread::AgentTelemetry for NativeAgentConnection {
1709 fn thread_data(
1710 &self,
1711 session_id: &acp::SessionId,
1712 cx: &mut App,
1713 ) -> Task<Result<serde_json::Value>> {
1714 let Some(session) = self.0.read(cx).sessions.get(session_id) else {
1715 return Task::ready(Err(anyhow!("Session not found")));
1716 };
1717
1718 let task = session.thread.read(cx).to_db(cx);
1719 cx.background_spawn(async move {
1720 serde_json::to_value(task.await).context("Failed to serialize thread")
1721 })
1722 }
1723}
1724
1725pub struct NativeAgentSessionList {
1726 thread_store: Entity<ThreadStore>,
1727 updates_tx: smol::channel::Sender<acp_thread::SessionListUpdate>,
1728 updates_rx: smol::channel::Receiver<acp_thread::SessionListUpdate>,
1729 _subscription: Subscription,
1730}
1731
1732impl NativeAgentSessionList {
1733 fn new(thread_store: Entity<ThreadStore>, cx: &mut App) -> Self {
1734 let (tx, rx) = smol::channel::unbounded();
1735 let this_tx = tx.clone();
1736 let subscription = cx.observe(&thread_store, move |_, _| {
1737 this_tx
1738 .try_send(acp_thread::SessionListUpdate::Refresh)
1739 .ok();
1740 });
1741 Self {
1742 thread_store,
1743 updates_tx: tx,
1744 updates_rx: rx,
1745 _subscription: subscription,
1746 }
1747 }
1748
1749 pub fn thread_store(&self) -> &Entity<ThreadStore> {
1750 &self.thread_store
1751 }
1752}
1753
1754impl AgentSessionList for NativeAgentSessionList {
1755 fn list_sessions(
1756 &self,
1757 _request: AgentSessionListRequest,
1758 cx: &mut App,
1759 ) -> Task<Result<AgentSessionListResponse>> {
1760 let sessions = self
1761 .thread_store
1762 .read(cx)
1763 .entries()
1764 .map(|entry| AgentSessionInfo::from(&entry))
1765 .collect();
1766 Task::ready(Ok(AgentSessionListResponse::new(sessions)))
1767 }
1768
1769 fn supports_delete(&self) -> bool {
1770 true
1771 }
1772
1773 fn delete_session(&self, session_id: &acp::SessionId, cx: &mut App) -> Task<Result<()>> {
1774 self.thread_store
1775 .update(cx, |store, cx| store.delete_thread(session_id.clone(), cx))
1776 }
1777
1778 fn delete_sessions(&self, cx: &mut App) -> Task<Result<()>> {
1779 self.thread_store
1780 .update(cx, |store, cx| store.delete_threads(cx))
1781 }
1782
1783 fn watch(
1784 &self,
1785 _cx: &mut App,
1786 ) -> Option<smol::channel::Receiver<acp_thread::SessionListUpdate>> {
1787 Some(self.updates_rx.clone())
1788 }
1789
1790 fn notify_refresh(&self) {
1791 self.updates_tx
1792 .try_send(acp_thread::SessionListUpdate::Refresh)
1793 .ok();
1794 }
1795
1796 fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1797 self
1798 }
1799}
1800
1801struct NativeAgentSessionTruncate {
1802 thread: Entity<Thread>,
1803 acp_thread: WeakEntity<AcpThread>,
1804}
1805
1806impl acp_thread::AgentSessionTruncate for NativeAgentSessionTruncate {
1807 fn run(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
1808 match self.thread.update(cx, |thread, cx| {
1809 thread.truncate(message_id.clone(), cx)?;
1810 Ok(thread.latest_token_usage())
1811 }) {
1812 Ok(usage) => {
1813 self.acp_thread
1814 .update(cx, |thread, cx| {
1815 thread.update_token_usage(usage, cx);
1816 })
1817 .ok();
1818 Task::ready(Ok(()))
1819 }
1820 Err(error) => Task::ready(Err(error)),
1821 }
1822 }
1823}
1824
1825struct NativeAgentSessionRetry {
1826 connection: NativeAgentConnection,
1827 session_id: acp::SessionId,
1828}
1829
1830impl acp_thread::AgentSessionRetry for NativeAgentSessionRetry {
1831 fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>> {
1832 self.connection
1833 .run_turn(self.session_id.clone(), cx, |thread, cx| {
1834 thread.update(cx, |thread, cx| thread.resume(cx))
1835 })
1836 }
1837}
1838
1839struct NativeAgentSessionSetTitle {
1840 thread: Entity<Thread>,
1841}
1842
1843impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
1844 fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
1845 self.thread
1846 .update(cx, |thread, cx| thread.set_title(title, cx));
1847 Task::ready(Ok(()))
1848 }
1849}
1850
1851pub struct NativeThreadEnvironment {
1852 agent: WeakEntity<NativeAgent>,
1853 thread: WeakEntity<Thread>,
1854 acp_thread: WeakEntity<AcpThread>,
1855}
1856
1857impl NativeThreadEnvironment {
1858 pub(crate) fn create_subagent_thread(
1859 &self,
1860 label: String,
1861 cx: &mut App,
1862 ) -> Result<Rc<dyn SubagentHandle>> {
1863 let Some(parent_thread_entity) = self.thread.upgrade() else {
1864 anyhow::bail!("Parent thread no longer exists".to_string());
1865 };
1866 let parent_thread = parent_thread_entity.read(cx);
1867 let current_depth = parent_thread.depth();
1868 let parent_session_id = parent_thread.id().clone();
1869
1870 if current_depth >= MAX_SUBAGENT_DEPTH {
1871 return Err(anyhow!(
1872 "Maximum subagent depth ({}) reached",
1873 MAX_SUBAGENT_DEPTH
1874 ));
1875 }
1876
1877 let subagent_thread: Entity<Thread> = cx.new(|cx| {
1878 let mut thread = Thread::new_subagent(&parent_thread_entity, cx);
1879 thread.set_title(label.into(), cx);
1880 thread
1881 });
1882
1883 let session_id = subagent_thread.read(cx).id().clone();
1884
1885 let acp_thread = self
1886 .agent
1887 .update(cx, |agent, cx| -> Result<Entity<AcpThread>> {
1888 let project_id = agent
1889 .sessions
1890 .get(&parent_session_id)
1891 .map(|s| s.project_id)
1892 .context("parent session not found")?;
1893 Ok(agent.register_session(subagent_thread.clone(), project_id, 1, cx))
1894 })??;
1895
1896 let depth = current_depth + 1;
1897
1898 telemetry::event!(
1899 "Subagent Started",
1900 session = parent_thread_entity.read(cx).id().to_string(),
1901 subagent_session = session_id.to_string(),
1902 depth,
1903 is_resumed = false,
1904 );
1905
1906 self.prompt_subagent(session_id, subagent_thread, acp_thread)
1907 }
1908
1909 pub(crate) fn resume_subagent_thread(
1910 &self,
1911 session_id: acp::SessionId,
1912 cx: &mut App,
1913 ) -> Result<Rc<dyn SubagentHandle>> {
1914 let (subagent_thread, acp_thread) = self.agent.update(cx, |agent, _cx| {
1915 let session = agent
1916 .sessions
1917 .get(&session_id)
1918 .ok_or_else(|| anyhow!("No subagent session found with id {session_id}"))?;
1919 anyhow::Ok((session.thread.clone(), session.acp_thread.clone()))
1920 })??;
1921
1922 let depth = subagent_thread.read(cx).depth();
1923
1924 if let Some(parent_thread_entity) = self.thread.upgrade() {
1925 telemetry::event!(
1926 "Subagent Started",
1927 session = parent_thread_entity.read(cx).id().to_string(),
1928 subagent_session = session_id.to_string(),
1929 depth,
1930 is_resumed = true,
1931 );
1932 }
1933
1934 self.prompt_subagent(session_id, subagent_thread, acp_thread)
1935 }
1936
1937 fn prompt_subagent(
1938 &self,
1939 session_id: acp::SessionId,
1940 subagent_thread: Entity<Thread>,
1941 acp_thread: Entity<acp_thread::AcpThread>,
1942 ) -> Result<Rc<dyn SubagentHandle>> {
1943 let Some(parent_thread_entity) = self.thread.upgrade() else {
1944 anyhow::bail!("Parent thread no longer exists".to_string());
1945 };
1946 Ok(Rc::new(NativeSubagentHandle::new(
1947 session_id,
1948 subagent_thread,
1949 acp_thread,
1950 parent_thread_entity,
1951 )) as _)
1952 }
1953}
1954
1955impl ThreadEnvironment for NativeThreadEnvironment {
1956 fn create_terminal(
1957 &self,
1958 command: String,
1959 cwd: Option<PathBuf>,
1960 output_byte_limit: Option<u64>,
1961 cx: &mut AsyncApp,
1962 ) -> Task<Result<Rc<dyn TerminalHandle>>> {
1963 let task = self.acp_thread.update(cx, |thread, cx| {
1964 thread.create_terminal(command, vec![], vec![], cwd, output_byte_limit, cx)
1965 });
1966
1967 let acp_thread = self.acp_thread.clone();
1968 cx.spawn(async move |cx| {
1969 let terminal = task?.await?;
1970
1971 let (drop_tx, drop_rx) = oneshot::channel();
1972 let terminal_id = terminal.read_with(cx, |terminal, _cx| terminal.id().clone());
1973
1974 cx.spawn(async move |cx| {
1975 drop_rx.await.ok();
1976 acp_thread.update(cx, |thread, cx| thread.release_terminal(terminal_id, cx))
1977 })
1978 .detach();
1979
1980 let handle = AcpTerminalHandle {
1981 terminal,
1982 _drop_tx: Some(drop_tx),
1983 };
1984
1985 Ok(Rc::new(handle) as _)
1986 })
1987 }
1988
1989 fn create_subagent(&self, label: String, cx: &mut App) -> Result<Rc<dyn SubagentHandle>> {
1990 self.create_subagent_thread(label, cx)
1991 }
1992
1993 fn resume_subagent(
1994 &self,
1995 session_id: acp::SessionId,
1996 cx: &mut App,
1997 ) -> Result<Rc<dyn SubagentHandle>> {
1998 self.resume_subagent_thread(session_id, cx)
1999 }
2000}
2001
2002#[derive(Debug, Clone)]
2003enum SubagentPromptResult {
2004 Completed,
2005 Cancelled,
2006 ContextWindowWarning,
2007 Error(String),
2008}
2009
2010pub struct NativeSubagentHandle {
2011 session_id: acp::SessionId,
2012 parent_thread: WeakEntity<Thread>,
2013 subagent_thread: Entity<Thread>,
2014 acp_thread: Entity<acp_thread::AcpThread>,
2015}
2016
2017impl NativeSubagentHandle {
2018 fn new(
2019 session_id: acp::SessionId,
2020 subagent_thread: Entity<Thread>,
2021 acp_thread: Entity<acp_thread::AcpThread>,
2022 parent_thread_entity: Entity<Thread>,
2023 ) -> Self {
2024 NativeSubagentHandle {
2025 session_id,
2026 subagent_thread,
2027 parent_thread: parent_thread_entity.downgrade(),
2028 acp_thread,
2029 }
2030 }
2031}
2032
2033impl SubagentHandle for NativeSubagentHandle {
2034 fn id(&self) -> acp::SessionId {
2035 self.session_id.clone()
2036 }
2037
2038 fn num_entries(&self, cx: &App) -> usize {
2039 self.acp_thread.read(cx).entries().len()
2040 }
2041
2042 fn send(&self, message: String, cx: &AsyncApp) -> Task<Result<String>> {
2043 let thread = self.subagent_thread.clone();
2044 let acp_thread = self.acp_thread.clone();
2045 let subagent_session_id = self.session_id.clone();
2046 let parent_thread = self.parent_thread.clone();
2047
2048 cx.spawn(async move |cx| {
2049 let (task, _subscription) = cx.update(|cx| {
2050 let ratio_before_prompt = thread
2051 .read(cx)
2052 .latest_token_usage()
2053 .map(|usage| usage.ratio());
2054
2055 parent_thread
2056 .update(cx, |parent_thread, _cx| {
2057 parent_thread.register_running_subagent(thread.downgrade())
2058 })
2059 .ok();
2060
2061 let task = acp_thread.update(cx, |acp_thread, cx| {
2062 acp_thread.send(vec![message.into()], cx)
2063 });
2064
2065 let (token_limit_tx, token_limit_rx) = oneshot::channel::<()>();
2066 let mut token_limit_tx = Some(token_limit_tx);
2067
2068 let subscription = cx.subscribe(
2069 &thread,
2070 move |_thread, event: &TokenUsageUpdated, _cx| {
2071 if let Some(usage) = &event.0 {
2072 let old_ratio = ratio_before_prompt
2073 .clone()
2074 .unwrap_or(TokenUsageRatio::Normal);
2075 let new_ratio = usage.ratio();
2076 if old_ratio == TokenUsageRatio::Normal
2077 && new_ratio == TokenUsageRatio::Warning
2078 {
2079 if let Some(tx) = token_limit_tx.take() {
2080 tx.send(()).ok();
2081 }
2082 }
2083 }
2084 },
2085 );
2086
2087 let wait_for_prompt = cx
2088 .background_spawn(async move {
2089 futures::select! {
2090 response = task.fuse() => match response {
2091 Ok(Some(response)) => {
2092 match response.stop_reason {
2093 acp::StopReason::Cancelled => SubagentPromptResult::Cancelled,
2094 acp::StopReason::MaxTokens => SubagentPromptResult::Error("The agent reached the maximum number of tokens.".into()),
2095 acp::StopReason::MaxTurnRequests => SubagentPromptResult::Error("The agent reached the maximum number of allowed requests between user turns. Try prompting again.".into()),
2096 acp::StopReason::Refusal => SubagentPromptResult::Error("The agent refused to process that prompt. Try again.".into()),
2097 acp::StopReason::EndTurn | _ => SubagentPromptResult::Completed,
2098 }
2099 }
2100 Ok(None) => SubagentPromptResult::Error("No response from the agent. You can try messaging again.".into()),
2101 Err(error) => SubagentPromptResult::Error(error.to_string()),
2102 },
2103 _ = token_limit_rx.fuse() => SubagentPromptResult::ContextWindowWarning,
2104 }
2105 });
2106
2107 (wait_for_prompt, subscription)
2108 });
2109
2110 let result = match task.await {
2111 SubagentPromptResult::Completed => thread.read_with(cx, |thread, _cx| {
2112 thread
2113 .last_message()
2114 .and_then(|message| {
2115 let content = message.as_agent_message()?
2116 .content
2117 .iter()
2118 .filter_map(|c| match c {
2119 AgentMessageContent::Text(text) => Some(text.as_str()),
2120 _ => None,
2121 })
2122 .join("\n\n");
2123 if content.is_empty() {
2124 None
2125 } else {
2126 Some( content)
2127 }
2128 })
2129 .context("No response from subagent")
2130 }),
2131 SubagentPromptResult::Cancelled => Err(anyhow!("User canceled")),
2132 SubagentPromptResult::Error(message) => Err(anyhow!("{message}")),
2133 SubagentPromptResult::ContextWindowWarning => {
2134 thread.update(cx, |thread, cx| thread.cancel(cx)).await;
2135 Err(anyhow!(
2136 "The agent is nearing the end of its context window and has been \
2137 stopped. You can prompt the thread again to have the agent wrap up \
2138 or hand off its work."
2139 ))
2140 }
2141 };
2142
2143 parent_thread
2144 .update(cx, |parent_thread, cx| {
2145 parent_thread.unregister_running_subagent(&subagent_session_id, cx)
2146 })
2147 .ok();
2148
2149 result
2150 })
2151 }
2152}
2153
2154pub struct AcpTerminalHandle {
2155 terminal: Entity<acp_thread::Terminal>,
2156 _drop_tx: Option<oneshot::Sender<()>>,
2157}
2158
2159impl TerminalHandle for AcpTerminalHandle {
2160 fn id(&self, cx: &AsyncApp) -> Result<acp::TerminalId> {
2161 Ok(self.terminal.read_with(cx, |term, _cx| term.id().clone()))
2162 }
2163
2164 fn wait_for_exit(&self, cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {
2165 Ok(self
2166 .terminal
2167 .read_with(cx, |term, _cx| term.wait_for_exit()))
2168 }
2169
2170 fn current_output(&self, cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {
2171 Ok(self
2172 .terminal
2173 .read_with(cx, |term, cx| term.current_output(cx)))
2174 }
2175
2176 fn kill(&self, cx: &AsyncApp) -> Result<()> {
2177 cx.update(|cx| {
2178 self.terminal.update(cx, |terminal, cx| {
2179 terminal.kill(cx);
2180 });
2181 });
2182 Ok(())
2183 }
2184
2185 fn was_stopped_by_user(&self, cx: &AsyncApp) -> Result<bool> {
2186 Ok(self
2187 .terminal
2188 .read_with(cx, |term, _cx| term.was_stopped_by_user()))
2189 }
2190}
2191
2192#[cfg(test)]
2193mod internal_tests {
2194 use std::path::Path;
2195
2196 use super::*;
2197 use acp_thread::{AgentConnection, AgentModelGroupName, AgentModelInfo, MentionUri};
2198 use fs::FakeFs;
2199 use gpui::TestAppContext;
2200 use indoc::formatdoc;
2201 use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider};
2202 use language_model::{
2203 LanguageModelCompletionEvent, LanguageModelProviderId, LanguageModelProviderName,
2204 };
2205 use serde_json::json;
2206 use settings::SettingsStore;
2207 use util::{path, rel_path::rel_path};
2208
2209 #[gpui::test]
2210 async fn test_maintaining_project_context(cx: &mut TestAppContext) {
2211 init_test(cx);
2212 let fs = FakeFs::new(cx.executor());
2213 fs.insert_tree(
2214 "/",
2215 json!({
2216 "a": {}
2217 }),
2218 )
2219 .await;
2220 let project = Project::test(fs.clone(), [], cx).await;
2221 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2222 let agent =
2223 cx.update(|cx| NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx));
2224
2225 // Creating a session registers the project and triggers context building.
2226 let connection = NativeAgentConnection(agent.clone());
2227 let _acp_thread = cx
2228 .update(|cx| {
2229 Rc::new(connection).new_session(
2230 project.clone(),
2231 PathList::new(&[Path::new("/")]),
2232 cx,
2233 )
2234 })
2235 .await
2236 .unwrap();
2237 cx.run_until_parked();
2238
2239 let thread = agent.read_with(cx, |agent, _cx| {
2240 agent.sessions.values().next().unwrap().thread.clone()
2241 });
2242
2243 agent.read_with(cx, |agent, cx| {
2244 let project_id = project.entity_id();
2245 let state = agent.projects.get(&project_id).unwrap();
2246 assert_eq!(state.project_context.read(cx).worktrees, vec![]);
2247 assert_eq!(thread.read(cx).project_context().read(cx).worktrees, vec![]);
2248 });
2249
2250 let worktree = project
2251 .update(cx, |project, cx| project.create_worktree("/a", true, cx))
2252 .await
2253 .unwrap();
2254 cx.run_until_parked();
2255 agent.read_with(cx, |agent, cx| {
2256 let project_id = project.entity_id();
2257 let state = agent.projects.get(&project_id).unwrap();
2258 let expected_worktrees = vec![WorktreeContext {
2259 root_name: "a".into(),
2260 abs_path: Path::new("/a").into(),
2261 rules_file: None,
2262 }];
2263 assert_eq!(state.project_context.read(cx).worktrees, expected_worktrees);
2264 assert_eq!(
2265 thread.read(cx).project_context().read(cx).worktrees,
2266 expected_worktrees
2267 );
2268 });
2269
2270 // Creating `/a/.rules` updates the project context.
2271 fs.insert_file("/a/.rules", Vec::new()).await;
2272 cx.run_until_parked();
2273 agent.read_with(cx, |agent, cx| {
2274 let project_id = project.entity_id();
2275 let state = agent.projects.get(&project_id).unwrap();
2276 let rules_entry = worktree
2277 .read(cx)
2278 .entry_for_path(rel_path(".rules"))
2279 .unwrap();
2280 let expected_worktrees = vec![WorktreeContext {
2281 root_name: "a".into(),
2282 abs_path: Path::new("/a").into(),
2283 rules_file: Some(RulesFileContext {
2284 path_in_worktree: rel_path(".rules").into(),
2285 text: "".into(),
2286 project_entry_id: rules_entry.id.to_usize(),
2287 }),
2288 }];
2289 assert_eq!(state.project_context.read(cx).worktrees, expected_worktrees);
2290 assert_eq!(
2291 thread.read(cx).project_context().read(cx).worktrees,
2292 expected_worktrees
2293 );
2294 });
2295 }
2296
2297 #[gpui::test]
2298 async fn test_listing_models(cx: &mut TestAppContext) {
2299 init_test(cx);
2300 let fs = FakeFs::new(cx.executor());
2301 fs.insert_tree("/", json!({ "a": {} })).await;
2302 let project = Project::test(fs.clone(), [], cx).await;
2303 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2304 let connection =
2305 NativeAgentConnection(cx.update(|cx| {
2306 NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx)
2307 }));
2308
2309 // Create a thread/session
2310 let acp_thread = cx
2311 .update(|cx| {
2312 Rc::new(connection.clone()).new_session(
2313 project.clone(),
2314 PathList::new(&[Path::new("/a")]),
2315 cx,
2316 )
2317 })
2318 .await
2319 .unwrap();
2320
2321 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2322
2323 let models = cx
2324 .update(|cx| {
2325 connection
2326 .model_selector(&session_id)
2327 .unwrap()
2328 .list_models(cx)
2329 })
2330 .await
2331 .unwrap();
2332
2333 let acp_thread::AgentModelList::Grouped(models) = models else {
2334 panic!("Unexpected model group");
2335 };
2336 assert_eq!(
2337 models,
2338 IndexMap::from_iter([(
2339 AgentModelGroupName("Fake".into()),
2340 vec![AgentModelInfo {
2341 id: acp::ModelId::new("fake/fake"),
2342 name: "Fake".into(),
2343 description: None,
2344 icon: Some(acp_thread::AgentModelIcon::Named(
2345 ui::IconName::ZedAssistant
2346 )),
2347 is_latest: false,
2348 cost: None,
2349 }]
2350 )])
2351 );
2352 }
2353
2354 #[gpui::test]
2355 async fn test_model_selection_persists_to_settings(cx: &mut TestAppContext) {
2356 init_test(cx);
2357 let fs = FakeFs::new(cx.executor());
2358 fs.create_dir(paths::settings_file().parent().unwrap())
2359 .await
2360 .unwrap();
2361 fs.insert_file(
2362 paths::settings_file(),
2363 json!({
2364 "agent": {
2365 "default_model": {
2366 "provider": "foo",
2367 "model": "bar"
2368 }
2369 }
2370 })
2371 .to_string()
2372 .into_bytes(),
2373 )
2374 .await;
2375 let project = Project::test(fs.clone(), [], cx).await;
2376
2377 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2378
2379 // Create the agent and connection
2380 let agent =
2381 cx.update(|cx| NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx));
2382 let connection = NativeAgentConnection(agent.clone());
2383
2384 // Create a thread/session
2385 let acp_thread = cx
2386 .update(|cx| {
2387 Rc::new(connection.clone()).new_session(
2388 project.clone(),
2389 PathList::new(&[Path::new("/a")]),
2390 cx,
2391 )
2392 })
2393 .await
2394 .unwrap();
2395
2396 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2397
2398 // Select a model
2399 let selector = connection.model_selector(&session_id).unwrap();
2400 let model_id = acp::ModelId::new("fake/fake");
2401 cx.update(|cx| selector.select_model(model_id.clone(), cx))
2402 .await
2403 .unwrap();
2404
2405 // Verify the thread has the selected model
2406 agent.read_with(cx, |agent, _| {
2407 let session = agent.sessions.get(&session_id).unwrap();
2408 session.thread.read_with(cx, |thread, _| {
2409 assert_eq!(thread.model().unwrap().id().0, "fake");
2410 });
2411 });
2412
2413 cx.run_until_parked();
2414
2415 // Verify settings file was updated
2416 let settings_content = fs.load(paths::settings_file()).await.unwrap();
2417 let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
2418
2419 // Check that the agent settings contain the selected model
2420 assert_eq!(
2421 settings_json["agent"]["default_model"]["model"],
2422 json!("fake")
2423 );
2424 assert_eq!(
2425 settings_json["agent"]["default_model"]["provider"],
2426 json!("fake")
2427 );
2428
2429 // Register a thinking model and select it.
2430 cx.update(|cx| {
2431 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2432 "fake-corp",
2433 "fake-thinking",
2434 "Fake Thinking",
2435 true,
2436 ));
2437 let thinking_provider = Arc::new(
2438 FakeLanguageModelProvider::new(
2439 LanguageModelProviderId::from("fake-corp".to_string()),
2440 LanguageModelProviderName::from("Fake Corp".to_string()),
2441 )
2442 .with_models(vec![thinking_model]),
2443 );
2444 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2445 registry.register_provider(thinking_provider, cx);
2446 });
2447 });
2448 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2449
2450 let selector = connection.model_selector(&session_id).unwrap();
2451 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2452 .await
2453 .unwrap();
2454 cx.run_until_parked();
2455
2456 // Verify enable_thinking was written to settings as true.
2457 let settings_content = fs.load(paths::settings_file()).await.unwrap();
2458 let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
2459 assert_eq!(
2460 settings_json["agent"]["default_model"]["enable_thinking"],
2461 json!(true),
2462 "selecting a thinking model should persist enable_thinking: true to settings"
2463 );
2464 }
2465
2466 #[gpui::test]
2467 async fn test_select_model_updates_thinking_enabled(cx: &mut TestAppContext) {
2468 init_test(cx);
2469 let fs = FakeFs::new(cx.executor());
2470 fs.create_dir(paths::settings_file().parent().unwrap())
2471 .await
2472 .unwrap();
2473 fs.insert_file(paths::settings_file(), b"{}".to_vec()).await;
2474 let project = Project::test(fs.clone(), [], cx).await;
2475
2476 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2477 let agent =
2478 cx.update(|cx| NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx));
2479 let connection = NativeAgentConnection(agent.clone());
2480
2481 let acp_thread = cx
2482 .update(|cx| {
2483 Rc::new(connection.clone()).new_session(
2484 project.clone(),
2485 PathList::new(&[Path::new("/a")]),
2486 cx,
2487 )
2488 })
2489 .await
2490 .unwrap();
2491 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2492
2493 // Register a second provider with a thinking model.
2494 cx.update(|cx| {
2495 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2496 "fake-corp",
2497 "fake-thinking",
2498 "Fake Thinking",
2499 true,
2500 ));
2501 let thinking_provider = Arc::new(
2502 FakeLanguageModelProvider::new(
2503 LanguageModelProviderId::from("fake-corp".to_string()),
2504 LanguageModelProviderName::from("Fake Corp".to_string()),
2505 )
2506 .with_models(vec![thinking_model]),
2507 );
2508 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2509 registry.register_provider(thinking_provider, cx);
2510 });
2511 });
2512 // Refresh the agent's model list so it picks up the new provider.
2513 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2514
2515 // Thread starts with thinking_enabled = false (the default).
2516 agent.read_with(cx, |agent, _| {
2517 let session = agent.sessions.get(&session_id).unwrap();
2518 session.thread.read_with(cx, |thread, _| {
2519 assert!(!thread.thinking_enabled(), "thinking defaults to false");
2520 });
2521 });
2522
2523 // Select the thinking model via select_model.
2524 let selector = connection.model_selector(&session_id).unwrap();
2525 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2526 .await
2527 .unwrap();
2528
2529 // select_model should have enabled thinking based on the model's supports_thinking().
2530 agent.read_with(cx, |agent, _| {
2531 let session = agent.sessions.get(&session_id).unwrap();
2532 session.thread.read_with(cx, |thread, _| {
2533 assert!(
2534 thread.thinking_enabled(),
2535 "select_model should enable thinking when model supports it"
2536 );
2537 });
2538 });
2539
2540 // Switch back to the non-thinking model.
2541 let selector = connection.model_selector(&session_id).unwrap();
2542 cx.update(|cx| selector.select_model(acp::ModelId::new("fake/fake"), cx))
2543 .await
2544 .unwrap();
2545
2546 // select_model should have disabled thinking.
2547 agent.read_with(cx, |agent, _| {
2548 let session = agent.sessions.get(&session_id).unwrap();
2549 session.thread.read_with(cx, |thread, _| {
2550 assert!(
2551 !thread.thinking_enabled(),
2552 "select_model should disable thinking when model does not support it"
2553 );
2554 });
2555 });
2556 }
2557
2558 #[gpui::test]
2559 async fn test_summarization_model_survives_transient_registry_clearing(
2560 cx: &mut TestAppContext,
2561 ) {
2562 init_test(cx);
2563 let fs = FakeFs::new(cx.executor());
2564 fs.insert_tree("/", json!({ "a": {} })).await;
2565 let project = Project::test(fs.clone(), [], cx).await;
2566
2567 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2568 let agent =
2569 cx.update(|cx| NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx));
2570 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2571
2572 let acp_thread = cx
2573 .update(|cx| {
2574 connection.clone().new_session(
2575 project.clone(),
2576 PathList::new(&[Path::new("/a")]),
2577 cx,
2578 )
2579 })
2580 .await
2581 .unwrap();
2582 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2583
2584 let thread = agent.read_with(cx, |agent, _| {
2585 agent.sessions.get(&session_id).unwrap().thread.clone()
2586 });
2587
2588 thread.read_with(cx, |thread, _| {
2589 assert!(
2590 thread.summarization_model().is_some(),
2591 "session should have a summarization model from the test registry"
2592 );
2593 });
2594
2595 // Simulate what happens during a provider blip:
2596 // update_active_language_model_from_settings calls set_default_model(None)
2597 // when it can't resolve the model, clearing all fallbacks.
2598 cx.update(|cx| {
2599 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2600 registry.set_default_model(None, cx);
2601 });
2602 });
2603 cx.run_until_parked();
2604
2605 thread.read_with(cx, |thread, _| {
2606 assert!(
2607 thread.summarization_model().is_some(),
2608 "summarization model should survive a transient default model clearing"
2609 );
2610 });
2611 }
2612
2613 #[gpui::test]
2614 async fn test_loaded_thread_preserves_thinking_enabled(cx: &mut TestAppContext) {
2615 init_test(cx);
2616 let fs = FakeFs::new(cx.executor());
2617 fs.insert_tree("/", json!({ "a": {} })).await;
2618 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2619 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2620 let agent = cx.update(|cx| {
2621 NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
2622 });
2623 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2624
2625 // Register a thinking model.
2626 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2627 "fake-corp",
2628 "fake-thinking",
2629 "Fake Thinking",
2630 true,
2631 ));
2632 let thinking_provider = Arc::new(
2633 FakeLanguageModelProvider::new(
2634 LanguageModelProviderId::from("fake-corp".to_string()),
2635 LanguageModelProviderName::from("Fake Corp".to_string()),
2636 )
2637 .with_models(vec![thinking_model.clone()]),
2638 );
2639 cx.update(|cx| {
2640 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2641 registry.register_provider(thinking_provider, cx);
2642 });
2643 });
2644 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2645
2646 // Create a thread and select the thinking model.
2647 let acp_thread = cx
2648 .update(|cx| {
2649 connection.clone().new_session(
2650 project.clone(),
2651 PathList::new(&[Path::new("/a")]),
2652 cx,
2653 )
2654 })
2655 .await
2656 .unwrap();
2657 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2658
2659 let selector = connection.model_selector(&session_id).unwrap();
2660 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2661 .await
2662 .unwrap();
2663
2664 // Verify thinking is enabled after selecting the thinking model.
2665 let thread = agent.read_with(cx, |agent, _| {
2666 agent.sessions.get(&session_id).unwrap().thread.clone()
2667 });
2668 thread.read_with(cx, |thread, _| {
2669 assert!(
2670 thread.thinking_enabled(),
2671 "thinking should be enabled after selecting thinking model"
2672 );
2673 });
2674
2675 // Send a message so the thread gets persisted.
2676 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
2677 let send = cx.foreground_executor().spawn(send);
2678 cx.run_until_parked();
2679
2680 thinking_model.send_last_completion_stream_text_chunk("Response.");
2681 thinking_model.end_last_completion_stream();
2682
2683 send.await.unwrap();
2684 cx.run_until_parked();
2685
2686 // Close the session so it can be reloaded from disk.
2687 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2688 .await
2689 .unwrap();
2690 drop(thread);
2691 drop(acp_thread);
2692 agent.read_with(cx, |agent, _| {
2693 assert!(agent.sessions.is_empty());
2694 });
2695
2696 // Reload the thread and verify thinking_enabled is still true.
2697 let reloaded_acp_thread = agent
2698 .update(cx, |agent, cx| {
2699 agent.open_thread(session_id.clone(), project.clone(), cx)
2700 })
2701 .await
2702 .unwrap();
2703 let reloaded_thread = agent.read_with(cx, |agent, _| {
2704 agent.sessions.get(&session_id).unwrap().thread.clone()
2705 });
2706 reloaded_thread.read_with(cx, |thread, _| {
2707 assert!(
2708 thread.thinking_enabled(),
2709 "thinking_enabled should be preserved when reloading a thread with a thinking model"
2710 );
2711 });
2712
2713 drop(reloaded_acp_thread);
2714 }
2715
2716 #[gpui::test]
2717 async fn test_loaded_thread_preserves_model(cx: &mut TestAppContext) {
2718 init_test(cx);
2719 let fs = FakeFs::new(cx.executor());
2720 fs.insert_tree("/", json!({ "a": {} })).await;
2721 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2722 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2723 let agent = cx.update(|cx| {
2724 NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
2725 });
2726 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2727
2728 // Register a model where id() != name(), like real Anthropic models
2729 // (e.g. id="claude-sonnet-4-5-thinking-latest", name="Claude Sonnet 4.5 Thinking").
2730 let model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2731 "fake-corp",
2732 "custom-model-id",
2733 "Custom Model Display Name",
2734 false,
2735 ));
2736 let provider = Arc::new(
2737 FakeLanguageModelProvider::new(
2738 LanguageModelProviderId::from("fake-corp".to_string()),
2739 LanguageModelProviderName::from("Fake Corp".to_string()),
2740 )
2741 .with_models(vec![model.clone()]),
2742 );
2743 cx.update(|cx| {
2744 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2745 registry.register_provider(provider, cx);
2746 });
2747 });
2748 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2749
2750 // Create a thread and select the model.
2751 let acp_thread = cx
2752 .update(|cx| {
2753 connection.clone().new_session(
2754 project.clone(),
2755 PathList::new(&[Path::new("/a")]),
2756 cx,
2757 )
2758 })
2759 .await
2760 .unwrap();
2761 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2762
2763 let selector = connection.model_selector(&session_id).unwrap();
2764 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/custom-model-id"), cx))
2765 .await
2766 .unwrap();
2767
2768 let thread = agent.read_with(cx, |agent, _| {
2769 agent.sessions.get(&session_id).unwrap().thread.clone()
2770 });
2771 thread.read_with(cx, |thread, _| {
2772 assert_eq!(
2773 thread.model().unwrap().id().0.as_ref(),
2774 "custom-model-id",
2775 "model should be set before persisting"
2776 );
2777 });
2778
2779 // Send a message so the thread gets persisted.
2780 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
2781 let send = cx.foreground_executor().spawn(send);
2782 cx.run_until_parked();
2783
2784 model.send_last_completion_stream_text_chunk("Response.");
2785 model.end_last_completion_stream();
2786
2787 send.await.unwrap();
2788 cx.run_until_parked();
2789
2790 // Close the session so it can be reloaded from disk.
2791 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2792 .await
2793 .unwrap();
2794 drop(thread);
2795 drop(acp_thread);
2796 agent.read_with(cx, |agent, _| {
2797 assert!(agent.sessions.is_empty());
2798 });
2799
2800 // Reload the thread and verify the model was preserved.
2801 let reloaded_acp_thread = agent
2802 .update(cx, |agent, cx| {
2803 agent.open_thread(session_id.clone(), project.clone(), cx)
2804 })
2805 .await
2806 .unwrap();
2807 let reloaded_thread = agent.read_with(cx, |agent, _| {
2808 agent.sessions.get(&session_id).unwrap().thread.clone()
2809 });
2810 reloaded_thread.read_with(cx, |thread, _| {
2811 let reloaded_model = thread
2812 .model()
2813 .expect("model should be present after reload");
2814 assert_eq!(
2815 reloaded_model.id().0.as_ref(),
2816 "custom-model-id",
2817 "reloaded thread should have the same model, not fall back to the default"
2818 );
2819 });
2820
2821 drop(reloaded_acp_thread);
2822 }
2823
2824 #[gpui::test]
2825 async fn test_save_load_thread(cx: &mut TestAppContext) {
2826 init_test(cx);
2827 let fs = FakeFs::new(cx.executor());
2828 fs.insert_tree(
2829 "/",
2830 json!({
2831 "a": {
2832 "b.md": "Lorem"
2833 }
2834 }),
2835 )
2836 .await;
2837 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2838 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2839 let agent = cx.update(|cx| {
2840 NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
2841 });
2842 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2843
2844 let acp_thread = cx
2845 .update(|cx| {
2846 connection
2847 .clone()
2848 .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
2849 })
2850 .await
2851 .unwrap();
2852 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2853 let thread = agent.read_with(cx, |agent, _| {
2854 agent.sessions.get(&session_id).unwrap().thread.clone()
2855 });
2856
2857 // Ensure empty threads are not saved, even if they get mutated.
2858 let model = Arc::new(FakeLanguageModel::default());
2859 let summary_model = Arc::new(FakeLanguageModel::default());
2860 thread.update(cx, |thread, cx| {
2861 thread.set_model(model.clone(), cx);
2862 thread.set_summarization_model(Some(summary_model.clone()), cx);
2863 });
2864 cx.run_until_parked();
2865 assert_eq!(thread_entries(&thread_store, cx), vec![]);
2866
2867 let send = acp_thread.update(cx, |thread, cx| {
2868 thread.send(
2869 vec![
2870 "What does ".into(),
2871 acp::ContentBlock::ResourceLink(acp::ResourceLink::new(
2872 "b.md",
2873 MentionUri::File {
2874 abs_path: path!("/a/b.md").into(),
2875 }
2876 .to_uri()
2877 .to_string(),
2878 )),
2879 " mean?".into(),
2880 ],
2881 cx,
2882 )
2883 });
2884 let send = cx.foreground_executor().spawn(send);
2885 cx.run_until_parked();
2886
2887 model.send_last_completion_stream_text_chunk("Lorem.");
2888 model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
2889 language_model::TokenUsage {
2890 input_tokens: 150,
2891 output_tokens: 75,
2892 ..Default::default()
2893 },
2894 ));
2895 model.end_last_completion_stream();
2896 cx.run_until_parked();
2897 summary_model
2898 .send_last_completion_stream_text_chunk(&format!("Explaining {}", path!("/a/b.md")));
2899 summary_model.end_last_completion_stream();
2900
2901 send.await.unwrap();
2902 let uri = MentionUri::File {
2903 abs_path: path!("/a/b.md").into(),
2904 }
2905 .to_uri();
2906 acp_thread.read_with(cx, |thread, cx| {
2907 assert_eq!(
2908 thread.to_markdown(cx),
2909 formatdoc! {"
2910 ## User
2911
2912 What does [@b.md]({uri}) mean?
2913
2914 ## Assistant
2915
2916 Lorem.
2917
2918 "}
2919 )
2920 });
2921
2922 cx.run_until_parked();
2923
2924 // Set a draft prompt with rich content blocks and scroll position
2925 // AFTER run_until_parked, so the only save that captures these
2926 // changes is the one performed by close_session itself.
2927 let draft_blocks = vec![
2928 acp::ContentBlock::Text(acp::TextContent::new("Check out ")),
2929 acp::ContentBlock::ResourceLink(acp::ResourceLink::new("b.md", uri.to_string())),
2930 acp::ContentBlock::Text(acp::TextContent::new(" please")),
2931 ];
2932 acp_thread.update(cx, |thread, cx| {
2933 thread.set_draft_prompt(Some(draft_blocks.clone()), cx);
2934 });
2935 thread.update(cx, |thread, _cx| {
2936 thread.set_ui_scroll_position(Some(gpui::ListOffset {
2937 item_ix: 5,
2938 offset_in_item: gpui::px(12.5),
2939 }));
2940 });
2941
2942 // Close the session so it can be reloaded from disk.
2943 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2944 .await
2945 .unwrap();
2946 drop(thread);
2947 drop(acp_thread);
2948 agent.read_with(cx, |agent, _| {
2949 assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
2950 });
2951
2952 // Ensure the thread can be reloaded from disk.
2953 assert_eq!(
2954 thread_entries(&thread_store, cx),
2955 vec![(
2956 session_id.clone(),
2957 format!("Explaining {}", path!("/a/b.md"))
2958 )]
2959 );
2960 let acp_thread = agent
2961 .update(cx, |agent, cx| {
2962 agent.open_thread(session_id.clone(), project.clone(), cx)
2963 })
2964 .await
2965 .unwrap();
2966 acp_thread.read_with(cx, |thread, cx| {
2967 assert_eq!(
2968 thread.to_markdown(cx),
2969 formatdoc! {"
2970 ## User
2971
2972 What does [@b.md]({uri}) mean?
2973
2974 ## Assistant
2975
2976 Lorem.
2977
2978 "}
2979 )
2980 });
2981
2982 // Ensure the draft prompt with rich content blocks survived the round-trip.
2983 acp_thread.read_with(cx, |thread, _| {
2984 assert_eq!(thread.draft_prompt(), Some(draft_blocks.as_slice()));
2985 });
2986
2987 // Ensure token usage survived the round-trip.
2988 acp_thread.read_with(cx, |thread, _| {
2989 let usage = thread
2990 .token_usage()
2991 .expect("token usage should be restored after reload");
2992 assert_eq!(usage.input_tokens, 150);
2993 assert_eq!(usage.output_tokens, 75);
2994 });
2995
2996 // Ensure scroll position survived the round-trip.
2997 acp_thread.read_with(cx, |thread, _| {
2998 let scroll = thread
2999 .ui_scroll_position()
3000 .expect("scroll position should be restored after reload");
3001 assert_eq!(scroll.item_ix, 5);
3002 assert_eq!(scroll.offset_in_item, gpui::px(12.5));
3003 });
3004 }
3005
3006 #[gpui::test]
3007 async fn test_close_session_saves_thread(cx: &mut TestAppContext) {
3008 init_test(cx);
3009 let fs = FakeFs::new(cx.executor());
3010 fs.insert_tree(
3011 "/",
3012 json!({
3013 "a": {
3014 "file.txt": "hello"
3015 }
3016 }),
3017 )
3018 .await;
3019 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
3020 let thread_store = cx.new(|cx| ThreadStore::new(cx));
3021 let agent = cx.update(|cx| {
3022 NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
3023 });
3024 let connection = Rc::new(NativeAgentConnection(agent.clone()));
3025
3026 let acp_thread = cx
3027 .update(|cx| {
3028 connection
3029 .clone()
3030 .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
3031 })
3032 .await
3033 .unwrap();
3034 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
3035 let thread = agent.read_with(cx, |agent, _| {
3036 agent.sessions.get(&session_id).unwrap().thread.clone()
3037 });
3038
3039 let model = Arc::new(FakeLanguageModel::default());
3040 thread.update(cx, |thread, cx| {
3041 thread.set_model(model.clone(), cx);
3042 });
3043
3044 // Send a message so the thread is non-empty (empty threads aren't saved).
3045 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx));
3046 let send = cx.foreground_executor().spawn(send);
3047 cx.run_until_parked();
3048
3049 model.send_last_completion_stream_text_chunk("world");
3050 model.end_last_completion_stream();
3051 send.await.unwrap();
3052 cx.run_until_parked();
3053
3054 // Set a draft prompt WITHOUT calling run_until_parked afterwards.
3055 // This means no observe-triggered save has run for this change.
3056 // The only way this data gets persisted is if close_session
3057 // itself performs the save.
3058 let draft_blocks = vec![acp::ContentBlock::Text(acp::TextContent::new(
3059 "unsaved draft",
3060 ))];
3061 acp_thread.update(cx, |thread, cx| {
3062 thread.set_draft_prompt(Some(draft_blocks.clone()), cx);
3063 });
3064
3065 // Close the session immediately — no run_until_parked in between.
3066 cx.update(|cx| connection.clone().close_session(&session_id, cx))
3067 .await
3068 .unwrap();
3069 cx.run_until_parked();
3070
3071 // Reopen and verify the draft prompt was saved.
3072 let reloaded = agent
3073 .update(cx, |agent, cx| {
3074 agent.open_thread(session_id.clone(), project.clone(), cx)
3075 })
3076 .await
3077 .unwrap();
3078 reloaded.read_with(cx, |thread, _| {
3079 assert_eq!(
3080 thread.draft_prompt(),
3081 Some(draft_blocks.as_slice()),
3082 "close_session must save the thread; draft prompt was lost"
3083 );
3084 });
3085 }
3086
3087 #[gpui::test]
3088 async fn test_thread_summary_releases_loaded_session(cx: &mut TestAppContext) {
3089 init_test(cx);
3090 let fs = FakeFs::new(cx.executor());
3091 fs.insert_tree(
3092 "/",
3093 json!({
3094 "a": {
3095 "file.txt": "hello"
3096 }
3097 }),
3098 )
3099 .await;
3100 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
3101 let thread_store = cx.new(|cx| ThreadStore::new(cx));
3102 let agent = cx.update(|cx| {
3103 NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
3104 });
3105 let connection = Rc::new(NativeAgentConnection(agent.clone()));
3106
3107 let acp_thread = cx
3108 .update(|cx| {
3109 connection
3110 .clone()
3111 .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
3112 })
3113 .await
3114 .unwrap();
3115 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
3116 let thread = agent.read_with(cx, |agent, _| {
3117 agent.sessions.get(&session_id).unwrap().thread.clone()
3118 });
3119
3120 let model = Arc::new(FakeLanguageModel::default());
3121 let summary_model = Arc::new(FakeLanguageModel::default());
3122 thread.update(cx, |thread, cx| {
3123 thread.set_model(model.clone(), cx);
3124 thread.set_summarization_model(Some(summary_model.clone()), cx);
3125 });
3126
3127 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx));
3128 let send = cx.foreground_executor().spawn(send);
3129 cx.run_until_parked();
3130
3131 model.send_last_completion_stream_text_chunk("world");
3132 model.end_last_completion_stream();
3133 send.await.unwrap();
3134 cx.run_until_parked();
3135
3136 let summary = agent.update(cx, |agent, cx| {
3137 agent.thread_summary(session_id.clone(), project.clone(), cx)
3138 });
3139 cx.run_until_parked();
3140
3141 summary_model.send_last_completion_stream_text_chunk("summary");
3142 summary_model.end_last_completion_stream();
3143
3144 assert_eq!(summary.await.unwrap(), "summary");
3145 cx.run_until_parked();
3146
3147 agent.read_with(cx, |agent, _| {
3148 let session = agent
3149 .sessions
3150 .get(&session_id)
3151 .expect("thread_summary should not close the active session");
3152 assert_eq!(
3153 session.ref_count, 1,
3154 "thread_summary should release its temporary session reference"
3155 );
3156 });
3157
3158 cx.update(|cx| connection.clone().close_session(&session_id, cx))
3159 .await
3160 .unwrap();
3161 cx.run_until_parked();
3162
3163 agent.read_with(cx, |agent, _| {
3164 assert!(
3165 agent.sessions.is_empty(),
3166 "closing the active session after thread_summary should unload it"
3167 );
3168 });
3169 }
3170
3171 #[gpui::test]
3172 async fn test_loaded_sessions_keep_state_until_last_close(cx: &mut TestAppContext) {
3173 init_test(cx);
3174 let fs = FakeFs::new(cx.executor());
3175 fs.insert_tree(
3176 "/",
3177 json!({
3178 "a": {
3179 "file.txt": "hello"
3180 }
3181 }),
3182 )
3183 .await;
3184 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
3185 let thread_store = cx.new(|cx| ThreadStore::new(cx));
3186 let agent = cx.update(|cx| {
3187 NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
3188 });
3189 let connection = Rc::new(NativeAgentConnection(agent.clone()));
3190
3191 let acp_thread = cx
3192 .update(|cx| {
3193 connection
3194 .clone()
3195 .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
3196 })
3197 .await
3198 .unwrap();
3199 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
3200 let thread = agent.read_with(cx, |agent, _| {
3201 agent.sessions.get(&session_id).unwrap().thread.clone()
3202 });
3203
3204 let model = cx.update(|cx| {
3205 LanguageModelRegistry::read_global(cx)
3206 .default_model()
3207 .map(|default_model| default_model.model)
3208 .expect("default test model should be available")
3209 });
3210 let fake_model = model.as_fake();
3211 thread.update(cx, |thread, cx| {
3212 thread.set_model(model.clone(), cx);
3213 });
3214
3215 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx));
3216 let send = cx.foreground_executor().spawn(send);
3217 cx.run_until_parked();
3218
3219 fake_model.send_last_completion_stream_text_chunk("world");
3220 fake_model.end_last_completion_stream();
3221 send.await.unwrap();
3222 cx.run_until_parked();
3223
3224 cx.update(|cx| connection.clone().close_session(&session_id, cx))
3225 .await
3226 .unwrap();
3227 drop(thread);
3228 drop(acp_thread);
3229 agent.read_with(cx, |agent, _| {
3230 assert!(agent.sessions.is_empty());
3231 });
3232
3233 let first_loaded_thread = cx.update(|cx| {
3234 connection.clone().load_session(
3235 session_id.clone(),
3236 project.clone(),
3237 PathList::new(&[Path::new("")]),
3238 None,
3239 cx,
3240 )
3241 });
3242 let second_loaded_thread = cx.update(|cx| {
3243 connection.clone().load_session(
3244 session_id.clone(),
3245 project.clone(),
3246 PathList::new(&[Path::new("")]),
3247 None,
3248 cx,
3249 )
3250 });
3251
3252 let first_loaded_thread = first_loaded_thread.await.unwrap();
3253 let second_loaded_thread = second_loaded_thread.await.unwrap();
3254
3255 cx.run_until_parked();
3256
3257 assert_eq!(
3258 first_loaded_thread.entity_id(),
3259 second_loaded_thread.entity_id(),
3260 "concurrent loads for the same session should share one AcpThread"
3261 );
3262
3263 cx.update(|cx| connection.clone().close_session(&session_id, cx))
3264 .await
3265 .unwrap();
3266
3267 agent.read_with(cx, |agent, _| {
3268 assert!(
3269 agent.sessions.contains_key(&session_id),
3270 "closing one loaded session should not drop shared session state"
3271 );
3272 });
3273
3274 let follow_up = second_loaded_thread.update(cx, |thread, cx| {
3275 thread.send(vec!["still there?".into()], cx)
3276 });
3277 let follow_up = cx.foreground_executor().spawn(follow_up);
3278 cx.run_until_parked();
3279
3280 fake_model.send_last_completion_stream_text_chunk("yes");
3281 fake_model.end_last_completion_stream();
3282 follow_up.await.unwrap();
3283 cx.run_until_parked();
3284
3285 second_loaded_thread.read_with(cx, |thread, cx| {
3286 assert_eq!(
3287 thread.to_markdown(cx),
3288 formatdoc! {"
3289 ## User
3290
3291 hello
3292
3293 ## Assistant
3294
3295 world
3296
3297 ## User
3298
3299 still there?
3300
3301 ## Assistant
3302
3303 yes
3304
3305 "}
3306 );
3307 });
3308
3309 cx.update(|cx| connection.clone().close_session(&session_id, cx))
3310 .await
3311 .unwrap();
3312
3313 cx.run_until_parked();
3314
3315 drop(first_loaded_thread);
3316 drop(second_loaded_thread);
3317 agent.read_with(cx, |agent, _| {
3318 assert!(agent.sessions.is_empty());
3319 });
3320 }
3321
3322 #[gpui::test]
3323 async fn test_rapid_title_changes_do_not_loop(cx: &mut TestAppContext) {
3324 // Regression test: rapid title changes must not cause a propagation loop
3325 // between Thread and AcpThread via handle_thread_title_updated.
3326 init_test(cx);
3327 let fs = FakeFs::new(cx.executor());
3328 fs.insert_tree("/", json!({ "a": {} })).await;
3329 let project = Project::test(fs.clone(), [], cx).await;
3330 let thread_store = cx.new(|cx| ThreadStore::new(cx));
3331 let agent = cx.update(|cx| {
3332 NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
3333 });
3334 let connection = Rc::new(NativeAgentConnection(agent.clone()));
3335
3336 let acp_thread = cx
3337 .update(|cx| {
3338 connection
3339 .clone()
3340 .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
3341 })
3342 .await
3343 .unwrap();
3344
3345 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
3346 let thread = agent.read_with(cx, |agent, _| {
3347 agent.sessions.get(&session_id).unwrap().thread.clone()
3348 });
3349
3350 let title_updated_count = Rc::new(std::cell::RefCell::new(0usize));
3351 cx.update(|cx| {
3352 let count = title_updated_count.clone();
3353 cx.subscribe(
3354 &thread,
3355 move |_entity: Entity<Thread>, _event: &TitleUpdated, _cx: &mut App| {
3356 let new_count = {
3357 let mut count = count.borrow_mut();
3358 *count += 1;
3359 *count
3360 };
3361 assert!(
3362 new_count <= 2,
3363 "TitleUpdated fired {new_count} times; \
3364 title updates are looping"
3365 );
3366 },
3367 )
3368 .detach();
3369 });
3370
3371 thread.update(cx, |thread, cx| thread.set_title("first".into(), cx));
3372 thread.update(cx, |thread, cx| thread.set_title("second".into(), cx));
3373
3374 cx.run_until_parked();
3375
3376 thread.read_with(cx, |thread, _| {
3377 assert_eq!(thread.title(), Some("second".into()));
3378 });
3379 acp_thread.read_with(cx, |acp_thread, _| {
3380 assert_eq!(acp_thread.title(), Some("second".into()));
3381 });
3382
3383 assert_eq!(*title_updated_count.borrow(), 2);
3384 }
3385
3386 fn thread_entries(
3387 thread_store: &Entity<ThreadStore>,
3388 cx: &mut TestAppContext,
3389 ) -> Vec<(acp::SessionId, String)> {
3390 thread_store.read_with(cx, |store, _| {
3391 store
3392 .entries()
3393 .map(|entry| (entry.id.clone(), entry.title.to_string()))
3394 .collect::<Vec<_>>()
3395 })
3396 }
3397
3398 fn init_test(cx: &mut TestAppContext) {
3399 env_logger::try_init().ok();
3400 cx.update(|cx| {
3401 let settings_store = SettingsStore::test(cx);
3402 cx.set_global(settings_store);
3403
3404 LanguageModelRegistry::test(cx);
3405 });
3406 }
3407}
3408
3409fn mcp_message_content_to_acp_content_block(
3410 content: context_server::types::MessageContent,
3411) -> acp::ContentBlock {
3412 match content {
3413 context_server::types::MessageContent::Text {
3414 text,
3415 annotations: _,
3416 } => text.into(),
3417 context_server::types::MessageContent::Image {
3418 data,
3419 mime_type,
3420 annotations: _,
3421 } => acp::ContentBlock::Image(acp::ImageContent::new(data, mime_type)),
3422 context_server::types::MessageContent::Audio {
3423 data,
3424 mime_type,
3425 annotations: _,
3426 } => acp::ContentBlock::Audio(acp::AudioContent::new(data, mime_type)),
3427 context_server::types::MessageContent::Resource {
3428 resource,
3429 annotations: _,
3430 } => {
3431 let mut link =
3432 acp::ResourceLink::new(resource.uri.to_string(), resource.uri.to_string());
3433 if let Some(mime_type) = resource.mime_type {
3434 link = link.mime_type(mime_type);
3435 }
3436 acp::ContentBlock::ResourceLink(link)
3437 }
3438 }
3439}