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::*;
17pub use native_agent_server::NativeAgentServer;
18pub use pattern_extraction::*;
19pub use shell_command_parser::extract_commands;
20pub use templates::*;
21pub use thread::*;
22pub use thread_store::*;
23pub use tool_permissions::*;
24pub use tools::*;
25
26use acp_thread::{
27 AcpThread, AgentModelSelector, AgentSessionInfo, AgentSessionList, AgentSessionListRequest,
28 AgentSessionListResponse, UserMessageId,
29};
30use agent_client_protocol as acp;
31use anyhow::{Context as _, Result, anyhow};
32use chrono::{DateTime, Utc};
33use collections::{HashMap, HashSet, IndexMap};
34use fs::Fs;
35use futures::channel::{mpsc, oneshot};
36use futures::future::Shared;
37use futures::{FutureExt as _, StreamExt as _, future};
38use gpui::{
39 App, AppContext, AsyncApp, Context, Entity, SharedString, Subscription, Task, WeakEntity,
40};
41use language_model::{IconOrSvg, LanguageModel, LanguageModelProvider, LanguageModelRegistry};
42use project::{Project, ProjectItem, ProjectPath, Worktree};
43use prompt_store::{
44 ProjectContext, PromptStore, RULES_FILE_NAMES, RulesFileContext, UserRulesContext,
45 WorktreeContext,
46};
47use serde::{Deserialize, Serialize};
48use settings::{LanguageModelSelection, update_settings_file};
49use std::any::Any;
50use std::path::{Path, PathBuf};
51use std::rc::Rc;
52use std::sync::Arc;
53use std::time::Duration;
54use util::ResultExt;
55use util::rel_path::RelPath;
56
57#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
58pub struct ProjectSnapshot {
59 pub worktree_snapshots: Vec<project::telemetry_snapshot::TelemetryWorktreeSnapshot>,
60 pub timestamp: DateTime<Utc>,
61}
62
63pub struct RulesLoadingError {
64 pub message: SharedString,
65}
66
67/// Holds both the internal Thread and the AcpThread for a session
68struct Session {
69 /// The internal thread that processes messages
70 thread: Entity<Thread>,
71 /// The ACP thread that handles protocol communication
72 acp_thread: Entity<acp_thread::AcpThread>,
73 pending_save: Task<()>,
74 _subscriptions: Vec<Subscription>,
75}
76
77pub struct LanguageModels {
78 /// Access language model by ID
79 models: HashMap<acp::ModelId, Arc<dyn LanguageModel>>,
80 /// Cached list for returning language model information
81 model_list: acp_thread::AgentModelList,
82 refresh_models_rx: watch::Receiver<()>,
83 refresh_models_tx: watch::Sender<()>,
84 _authenticate_all_providers_task: Task<()>,
85}
86
87impl LanguageModels {
88 fn new(cx: &mut App) -> Self {
89 let (refresh_models_tx, refresh_models_rx) = watch::channel(());
90
91 let mut this = Self {
92 models: HashMap::default(),
93 model_list: acp_thread::AgentModelList::Grouped(IndexMap::default()),
94 refresh_models_rx,
95 refresh_models_tx,
96 _authenticate_all_providers_task: Self::authenticate_all_language_model_providers(cx),
97 };
98 this.refresh_list(cx);
99 this
100 }
101
102 fn refresh_list(&mut self, cx: &App) {
103 let providers = LanguageModelRegistry::global(cx)
104 .read(cx)
105 .visible_providers()
106 .into_iter()
107 .filter(|provider| provider.is_authenticated(cx))
108 .collect::<Vec<_>>();
109
110 let mut language_model_list = IndexMap::default();
111 let mut recommended_models = HashSet::default();
112
113 let mut recommended = Vec::new();
114 for provider in &providers {
115 for model in provider.recommended_models(cx) {
116 recommended_models.insert((model.provider_id(), model.id()));
117 recommended.push(Self::map_language_model_to_info(&model, provider));
118 }
119 }
120 if !recommended.is_empty() {
121 language_model_list.insert(
122 acp_thread::AgentModelGroupName("Recommended".into()),
123 recommended,
124 );
125 }
126
127 let mut models = HashMap::default();
128 for provider in providers {
129 let mut provider_models = Vec::new();
130 for model in provider.provided_models(cx) {
131 let model_info = Self::map_language_model_to_info(&model, &provider);
132 let model_id = model_info.id.clone();
133 provider_models.push(model_info);
134 models.insert(model_id, model);
135 }
136 if !provider_models.is_empty() {
137 language_model_list.insert(
138 acp_thread::AgentModelGroupName(provider.name().0.clone()),
139 provider_models,
140 );
141 }
142 }
143
144 self.models = models;
145 self.model_list = acp_thread::AgentModelList::Grouped(language_model_list);
146 self.refresh_models_tx.send(()).ok();
147 }
148
149 fn watch(&self) -> watch::Receiver<()> {
150 self.refresh_models_rx.clone()
151 }
152
153 pub fn model_from_id(&self, model_id: &acp::ModelId) -> Option<Arc<dyn LanguageModel>> {
154 self.models.get(model_id).cloned()
155 }
156
157 fn map_language_model_to_info(
158 model: &Arc<dyn LanguageModel>,
159 provider: &Arc<dyn LanguageModelProvider>,
160 ) -> acp_thread::AgentModelInfo {
161 acp_thread::AgentModelInfo {
162 id: Self::model_id(model),
163 name: model.name().0,
164 description: None,
165 icon: Some(match provider.icon() {
166 IconOrSvg::Svg(path) => acp_thread::AgentModelIcon::Path(path),
167 IconOrSvg::Icon(name) => acp_thread::AgentModelIcon::Named(name),
168 }),
169 is_latest: model.is_latest(),
170 }
171 }
172
173 fn model_id(model: &Arc<dyn LanguageModel>) -> acp::ModelId {
174 acp::ModelId::new(format!("{}/{}", model.provider_id().0, model.id().0))
175 }
176
177 fn authenticate_all_language_model_providers(cx: &mut App) -> Task<()> {
178 let authenticate_all_providers = LanguageModelRegistry::global(cx)
179 .read(cx)
180 .visible_providers()
181 .iter()
182 .map(|provider| (provider.id(), provider.name(), provider.authenticate(cx)))
183 .collect::<Vec<_>>();
184
185 cx.background_spawn(async move {
186 for (provider_id, provider_name, authenticate_task) in authenticate_all_providers {
187 if let Err(err) = authenticate_task.await {
188 match err {
189 language_model::AuthenticateError::CredentialsNotFound => {
190 // Since we're authenticating these providers in the
191 // background for the purposes of populating the
192 // language selector, we don't care about providers
193 // where the credentials are not found.
194 }
195 language_model::AuthenticateError::ConnectionRefused => {
196 // Not logging connection refused errors as they are mostly from LM Studio's noisy auth failures.
197 // LM Studio only has one auth method (endpoint call) which fails for users who haven't enabled it.
198 // TODO: Better manage LM Studio auth logic to avoid these noisy failures.
199 }
200 _ => {
201 // Some providers have noisy failure states that we
202 // don't want to spam the logs with every time the
203 // language model selector is initialized.
204 //
205 // Ideally these should have more clear failure modes
206 // that we know are safe to ignore here, like what we do
207 // with `CredentialsNotFound` above.
208 match provider_id.0.as_ref() {
209 "lmstudio" | "ollama" => {
210 // LM Studio and Ollama both make fetch requests to the local APIs to determine if they are "authenticated".
211 //
212 // These fail noisily, so we don't log them.
213 }
214 "copilot_chat" => {
215 // Copilot Chat returns an error if Copilot is not enabled, so we don't log those errors.
216 }
217 _ => {
218 log::error!(
219 "Failed to authenticate provider: {}: {err:#}",
220 provider_name.0
221 );
222 }
223 }
224 }
225 }
226 }
227 }
228 })
229 }
230}
231
232pub struct NativeAgent {
233 /// Session ID -> Session mapping
234 sessions: HashMap<acp::SessionId, Session>,
235 thread_store: Entity<ThreadStore>,
236 /// Shared project context for all threads
237 project_context: Entity<ProjectContext>,
238 project_context_needs_refresh: watch::Sender<()>,
239 _maintain_project_context: Task<Result<()>>,
240 context_server_registry: Entity<ContextServerRegistry>,
241 /// Shared templates for all threads
242 templates: Arc<Templates>,
243 /// Cached model information
244 models: LanguageModels,
245 project: Entity<Project>,
246 prompt_store: Option<Entity<PromptStore>>,
247 fs: Arc<dyn Fs>,
248 _subscriptions: Vec<Subscription>,
249}
250
251impl NativeAgent {
252 pub async fn new(
253 project: Entity<Project>,
254 thread_store: Entity<ThreadStore>,
255 templates: Arc<Templates>,
256 prompt_store: Option<Entity<PromptStore>>,
257 fs: Arc<dyn Fs>,
258 cx: &mut AsyncApp,
259 ) -> Result<Entity<NativeAgent>> {
260 log::debug!("Creating new NativeAgent");
261
262 let project_context = cx
263 .update(|cx| Self::build_project_context(&project, prompt_store.as_ref(), cx))
264 .await;
265
266 Ok(cx.new(|cx| {
267 let context_server_store = project.read(cx).context_server_store();
268 let context_server_registry =
269 cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
270
271 let mut subscriptions = vec![
272 cx.subscribe(&project, Self::handle_project_event),
273 cx.subscribe(
274 &LanguageModelRegistry::global(cx),
275 Self::handle_models_updated_event,
276 ),
277 cx.subscribe(
278 &context_server_store,
279 Self::handle_context_server_store_updated,
280 ),
281 cx.subscribe(
282 &context_server_registry,
283 Self::handle_context_server_registry_event,
284 ),
285 ];
286 if let Some(prompt_store) = prompt_store.as_ref() {
287 subscriptions.push(cx.subscribe(prompt_store, Self::handle_prompts_updated_event))
288 }
289
290 let (project_context_needs_refresh_tx, project_context_needs_refresh_rx) =
291 watch::channel(());
292 Self {
293 sessions: HashMap::default(),
294 thread_store,
295 project_context: cx.new(|_| project_context),
296 project_context_needs_refresh: project_context_needs_refresh_tx,
297 _maintain_project_context: cx.spawn(async move |this, cx| {
298 Self::maintain_project_context(this, project_context_needs_refresh_rx, cx).await
299 }),
300 context_server_registry,
301 templates,
302 models: LanguageModels::new(cx),
303 project,
304 prompt_store,
305 fs,
306 _subscriptions: subscriptions,
307 }
308 }))
309 }
310
311 fn new_session(
312 &mut self,
313 project: Entity<Project>,
314 cx: &mut Context<Self>,
315 ) -> Entity<AcpThread> {
316 // Create Thread
317 // Fetch default model from registry settings
318 let registry = LanguageModelRegistry::read_global(cx);
319 // Log available models for debugging
320 let available_count = registry.available_models(cx).count();
321 log::debug!("Total available models: {}", available_count);
322
323 let default_model = registry.default_model().and_then(|default_model| {
324 self.models
325 .model_from_id(&LanguageModels::model_id(&default_model.model))
326 });
327 let thread = cx.new(|cx| {
328 Thread::new(
329 project.clone(),
330 self.project_context.clone(),
331 self.context_server_registry.clone(),
332 self.templates.clone(),
333 default_model,
334 cx,
335 )
336 });
337
338 self.register_session(thread, None, cx)
339 }
340
341 fn register_session(
342 &mut self,
343 thread_handle: Entity<Thread>,
344 allowed_tool_names: Option<Vec<&str>>,
345 cx: &mut Context<Self>,
346 ) -> Entity<AcpThread> {
347 let connection = Rc::new(NativeAgentConnection(cx.entity()));
348
349 let thread = thread_handle.read(cx);
350 let session_id = thread.id().clone();
351 let parent_session_id = thread.parent_thread_id();
352 let title = thread.title();
353 let project = thread.project.clone();
354 let action_log = thread.action_log.clone();
355 let prompt_capabilities_rx = thread.prompt_capabilities_rx.clone();
356 let acp_thread = cx.new(|cx| {
357 acp_thread::AcpThread::new(
358 parent_session_id,
359 title,
360 connection,
361 project.clone(),
362 action_log.clone(),
363 session_id.clone(),
364 prompt_capabilities_rx,
365 cx,
366 )
367 });
368
369 let registry = LanguageModelRegistry::read_global(cx);
370 let summarization_model = registry.thread_summary_model().map(|c| c.model);
371
372 let weak = cx.weak_entity();
373 thread_handle.update(cx, |thread, cx| {
374 thread.set_summarization_model(summarization_model, cx);
375 thread.add_default_tools(
376 allowed_tool_names,
377 Rc::new(NativeThreadEnvironment {
378 acp_thread: acp_thread.downgrade(),
379 agent: weak,
380 }) as _,
381 cx,
382 )
383 });
384
385 let subscriptions = vec![
386 cx.subscribe(&thread_handle, Self::handle_thread_title_updated),
387 cx.subscribe(&thread_handle, Self::handle_thread_token_usage_updated),
388 cx.observe(&thread_handle, move |this, thread, cx| {
389 this.save_thread(thread, cx)
390 }),
391 ];
392
393 self.sessions.insert(
394 session_id,
395 Session {
396 thread: thread_handle,
397 acp_thread: acp_thread.clone(),
398 _subscriptions: subscriptions,
399 pending_save: Task::ready(()),
400 },
401 );
402
403 self.update_available_commands(cx);
404
405 acp_thread
406 }
407
408 pub fn models(&self) -> &LanguageModels {
409 &self.models
410 }
411
412 async fn maintain_project_context(
413 this: WeakEntity<Self>,
414 mut needs_refresh: watch::Receiver<()>,
415 cx: &mut AsyncApp,
416 ) -> Result<()> {
417 while needs_refresh.changed().await.is_ok() {
418 let project_context = this
419 .update(cx, |this, cx| {
420 Self::build_project_context(&this.project, this.prompt_store.as_ref(), cx)
421 })?
422 .await;
423 this.update(cx, |this, cx| {
424 this.project_context = cx.new(|_| project_context);
425 })?;
426 }
427
428 Ok(())
429 }
430
431 fn build_project_context(
432 project: &Entity<Project>,
433 prompt_store: Option<&Entity<PromptStore>>,
434 cx: &mut App,
435 ) -> Task<ProjectContext> {
436 let worktrees = project.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
437 let worktree_tasks = worktrees
438 .into_iter()
439 .map(|worktree| {
440 Self::load_worktree_info_for_system_prompt(worktree, project.clone(), cx)
441 })
442 .collect::<Vec<_>>();
443 let default_user_rules_task = if let Some(prompt_store) = prompt_store.as_ref() {
444 prompt_store.read_with(cx, |prompt_store, cx| {
445 let prompts = prompt_store.default_prompt_metadata();
446 let load_tasks = prompts.into_iter().map(|prompt_metadata| {
447 let contents = prompt_store.load(prompt_metadata.id, cx);
448 async move { (contents.await, prompt_metadata) }
449 });
450 cx.background_spawn(future::join_all(load_tasks))
451 })
452 } else {
453 Task::ready(vec![])
454 };
455
456 cx.spawn(async move |_cx| {
457 let (worktrees, default_user_rules) =
458 future::join(future::join_all(worktree_tasks), default_user_rules_task).await;
459
460 let worktrees = worktrees
461 .into_iter()
462 .map(|(worktree, _rules_error)| {
463 // TODO: show error message
464 // if let Some(rules_error) = rules_error {
465 // this.update(cx, |_, cx| cx.emit(rules_error)).ok();
466 // }
467 worktree
468 })
469 .collect::<Vec<_>>();
470
471 let default_user_rules = default_user_rules
472 .into_iter()
473 .flat_map(|(contents, prompt_metadata)| match contents {
474 Ok(contents) => Some(UserRulesContext {
475 uuid: prompt_metadata.id.as_user()?,
476 title: prompt_metadata.title.map(|title| title.to_string()),
477 contents,
478 }),
479 Err(_err) => {
480 // TODO: show error message
481 // this.update(cx, |_, cx| {
482 // cx.emit(RulesLoadingError {
483 // message: format!("{err:?}").into(),
484 // });
485 // })
486 // .ok();
487 None
488 }
489 })
490 .collect::<Vec<_>>();
491
492 ProjectContext::new(worktrees, default_user_rules)
493 })
494 }
495
496 fn load_worktree_info_for_system_prompt(
497 worktree: Entity<Worktree>,
498 project: Entity<Project>,
499 cx: &mut App,
500 ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
501 let tree = worktree.read(cx);
502 let root_name = tree.root_name_str().into();
503 let abs_path = tree.abs_path();
504
505 let mut context = WorktreeContext {
506 root_name,
507 abs_path,
508 rules_file: None,
509 };
510
511 let rules_task = Self::load_worktree_rules_file(worktree, project, cx);
512 let Some(rules_task) = rules_task else {
513 return Task::ready((context, None));
514 };
515
516 cx.spawn(async move |_| {
517 let (rules_file, rules_file_error) = match rules_task.await {
518 Ok(rules_file) => (Some(rules_file), None),
519 Err(err) => (
520 None,
521 Some(RulesLoadingError {
522 message: format!("{err}").into(),
523 }),
524 ),
525 };
526 context.rules_file = rules_file;
527 (context, rules_file_error)
528 })
529 }
530
531 fn load_worktree_rules_file(
532 worktree: Entity<Worktree>,
533 project: Entity<Project>,
534 cx: &mut App,
535 ) -> Option<Task<Result<RulesFileContext>>> {
536 let worktree = worktree.read(cx);
537 let worktree_id = worktree.id();
538 let selected_rules_file = RULES_FILE_NAMES
539 .into_iter()
540 .filter_map(|name| {
541 worktree
542 .entry_for_path(RelPath::unix(name).unwrap())
543 .filter(|entry| entry.is_file())
544 .map(|entry| entry.path.clone())
545 })
546 .next();
547
548 // Note that Cline supports `.clinerules` being a directory, but that is not currently
549 // supported. This doesn't seem to occur often in GitHub repositories.
550 selected_rules_file.map(|path_in_worktree| {
551 let project_path = ProjectPath {
552 worktree_id,
553 path: path_in_worktree.clone(),
554 };
555 let buffer_task =
556 project.update(cx, |project, cx| project.open_buffer(project_path, cx));
557 let rope_task = cx.spawn(async move |cx| {
558 let buffer = buffer_task.await?;
559 let (project_entry_id, rope) = buffer.read_with(cx, |buffer, cx| {
560 let project_entry_id = buffer.entry_id(cx).context("buffer has no file")?;
561 anyhow::Ok((project_entry_id, buffer.as_rope().clone()))
562 })?;
563 anyhow::Ok((project_entry_id, rope))
564 });
565 // Build a string from the rope on a background thread.
566 cx.background_spawn(async move {
567 let (project_entry_id, rope) = rope_task.await?;
568 anyhow::Ok(RulesFileContext {
569 path_in_worktree,
570 text: rope.to_string().trim().to_string(),
571 project_entry_id: project_entry_id.to_usize(),
572 })
573 })
574 })
575 }
576
577 fn handle_thread_title_updated(
578 &mut self,
579 thread: Entity<Thread>,
580 _: &TitleUpdated,
581 cx: &mut Context<Self>,
582 ) {
583 let session_id = thread.read(cx).id();
584 let Some(session) = self.sessions.get(session_id) else {
585 return;
586 };
587 let thread = thread.downgrade();
588 let acp_thread = session.acp_thread.downgrade();
589 cx.spawn(async move |_, cx| {
590 let title = thread.read_with(cx, |thread, _| thread.title())?;
591 let task = acp_thread.update(cx, |acp_thread, cx| acp_thread.set_title(title, cx))?;
592 task.await
593 })
594 .detach_and_log_err(cx);
595 }
596
597 fn handle_thread_token_usage_updated(
598 &mut self,
599 thread: Entity<Thread>,
600 usage: &TokenUsageUpdated,
601 cx: &mut Context<Self>,
602 ) {
603 let Some(session) = self.sessions.get(thread.read(cx).id()) else {
604 return;
605 };
606 session.acp_thread.update(cx, |acp_thread, cx| {
607 acp_thread.update_token_usage(usage.0.clone(), cx);
608 });
609 }
610
611 fn handle_project_event(
612 &mut self,
613 _project: Entity<Project>,
614 event: &project::Event,
615 _cx: &mut Context<Self>,
616 ) {
617 match event {
618 project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
619 self.project_context_needs_refresh.send(()).ok();
620 }
621 project::Event::WorktreeUpdatedEntries(_, items) => {
622 if items.iter().any(|(path, _, _)| {
623 RULES_FILE_NAMES
624 .iter()
625 .any(|name| path.as_ref() == RelPath::unix(name).unwrap())
626 }) {
627 self.project_context_needs_refresh.send(()).ok();
628 }
629 }
630 _ => {}
631 }
632 }
633
634 fn handle_prompts_updated_event(
635 &mut self,
636 _prompt_store: Entity<PromptStore>,
637 _event: &prompt_store::PromptsUpdatedEvent,
638 _cx: &mut Context<Self>,
639 ) {
640 self.project_context_needs_refresh.send(()).ok();
641 }
642
643 fn handle_models_updated_event(
644 &mut self,
645 _registry: Entity<LanguageModelRegistry>,
646 _event: &language_model::Event,
647 cx: &mut Context<Self>,
648 ) {
649 self.models.refresh_list(cx);
650
651 let registry = LanguageModelRegistry::read_global(cx);
652 let default_model = registry.default_model().map(|m| m.model);
653 let summarization_model = registry.thread_summary_model().map(|m| m.model);
654
655 for session in self.sessions.values_mut() {
656 session.thread.update(cx, |thread, cx| {
657 if thread.model().is_none()
658 && let Some(model) = default_model.clone()
659 {
660 thread.set_model(model, cx);
661 cx.notify();
662 }
663 thread.set_summarization_model(summarization_model.clone(), cx);
664 });
665 }
666 }
667
668 fn handle_context_server_store_updated(
669 &mut self,
670 _store: Entity<project::context_server_store::ContextServerStore>,
671 _event: &project::context_server_store::ServerStatusChangedEvent,
672 cx: &mut Context<Self>,
673 ) {
674 self.update_available_commands(cx);
675 }
676
677 fn handle_context_server_registry_event(
678 &mut self,
679 _registry: Entity<ContextServerRegistry>,
680 event: &ContextServerRegistryEvent,
681 cx: &mut Context<Self>,
682 ) {
683 match event {
684 ContextServerRegistryEvent::ToolsChanged => {}
685 ContextServerRegistryEvent::PromptsChanged => {
686 self.update_available_commands(cx);
687 }
688 }
689 }
690
691 fn update_available_commands(&self, cx: &mut Context<Self>) {
692 let available_commands = self.build_available_commands(cx);
693 for session in self.sessions.values() {
694 session.acp_thread.update(cx, |thread, cx| {
695 thread
696 .handle_session_update(
697 acp::SessionUpdate::AvailableCommandsUpdate(
698 acp::AvailableCommandsUpdate::new(available_commands.clone()),
699 ),
700 cx,
701 )
702 .log_err();
703 });
704 }
705 }
706
707 fn build_available_commands(&self, cx: &App) -> Vec<acp::AvailableCommand> {
708 let registry = self.context_server_registry.read(cx);
709
710 let mut prompt_name_counts: HashMap<&str, usize> = HashMap::default();
711 for context_server_prompt in registry.prompts() {
712 *prompt_name_counts
713 .entry(context_server_prompt.prompt.name.as_str())
714 .or_insert(0) += 1;
715 }
716
717 registry
718 .prompts()
719 .flat_map(|context_server_prompt| {
720 let prompt = &context_server_prompt.prompt;
721
722 let should_prefix = prompt_name_counts
723 .get(prompt.name.as_str())
724 .copied()
725 .unwrap_or(0)
726 > 1;
727
728 let name = if should_prefix {
729 format!("{}.{}", context_server_prompt.server_id, prompt.name)
730 } else {
731 prompt.name.clone()
732 };
733
734 let mut command = acp::AvailableCommand::new(
735 name,
736 prompt.description.clone().unwrap_or_default(),
737 );
738
739 match prompt.arguments.as_deref() {
740 Some([arg]) => {
741 let hint = format!("<{}>", arg.name);
742
743 command = command.input(acp::AvailableCommandInput::Unstructured(
744 acp::UnstructuredCommandInput::new(hint),
745 ));
746 }
747 Some([]) | None => {}
748 Some(_) => {
749 // skip >1 argument commands since we don't support them yet
750 return None;
751 }
752 }
753
754 Some(command)
755 })
756 .collect()
757 }
758
759 pub fn load_thread(
760 &mut self,
761 id: acp::SessionId,
762 cx: &mut Context<Self>,
763 ) -> Task<Result<Entity<Thread>>> {
764 let database_future = ThreadsDatabase::connect(cx);
765 cx.spawn(async move |this, cx| {
766 let database = database_future.await.map_err(|err| anyhow!(err))?;
767 let db_thread = database
768 .load_thread(id.clone())
769 .await?
770 .with_context(|| format!("no thread found with ID: {id:?}"))?;
771
772 this.update(cx, |this, cx| {
773 let summarization_model = LanguageModelRegistry::read_global(cx)
774 .thread_summary_model()
775 .map(|c| c.model);
776
777 cx.new(|cx| {
778 let mut thread = Thread::from_db(
779 id.clone(),
780 db_thread,
781 this.project.clone(),
782 this.project_context.clone(),
783 this.context_server_registry.clone(),
784 this.templates.clone(),
785 cx,
786 );
787 thread.set_summarization_model(summarization_model, cx);
788 thread
789 })
790 })
791 })
792 }
793
794 pub fn open_thread(
795 &mut self,
796 id: acp::SessionId,
797 cx: &mut Context<Self>,
798 ) -> Task<Result<Entity<AcpThread>>> {
799 if let Some(session) = self.sessions.get(&id) {
800 return Task::ready(Ok(session.acp_thread.clone()));
801 }
802
803 let task = self.load_thread(id, cx);
804 cx.spawn(async move |this, cx| {
805 let thread = task.await?;
806 let acp_thread = this.update(cx, |this, cx| {
807 this.register_session(thread.clone(), None, cx)
808 })?;
809 let events = thread.update(cx, |thread, cx| thread.replay(cx));
810 cx.update(|cx| {
811 NativeAgentConnection::handle_thread_events(events, acp_thread.downgrade(), cx)
812 })
813 .await?;
814 Ok(acp_thread)
815 })
816 }
817
818 pub fn thread_summary(
819 &mut self,
820 id: acp::SessionId,
821 cx: &mut Context<Self>,
822 ) -> Task<Result<SharedString>> {
823 let thread = self.open_thread(id.clone(), cx);
824 cx.spawn(async move |this, cx| {
825 let acp_thread = thread.await?;
826 let result = this
827 .update(cx, |this, cx| {
828 this.sessions
829 .get(&id)
830 .unwrap()
831 .thread
832 .update(cx, |thread, cx| thread.summary(cx))
833 })?
834 .await
835 .context("Failed to generate summary")?;
836 drop(acp_thread);
837 Ok(result)
838 })
839 }
840
841 fn save_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
842 if thread.read(cx).is_empty() {
843 return;
844 }
845
846 let database_future = ThreadsDatabase::connect(cx);
847 let (id, db_thread) =
848 thread.update(cx, |thread, cx| (thread.id().clone(), thread.to_db(cx)));
849 let Some(session) = self.sessions.get_mut(&id) else {
850 return;
851 };
852 let thread_store = self.thread_store.clone();
853 session.pending_save = cx.spawn(async move |_, cx| {
854 let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
855 return;
856 };
857 let db_thread = db_thread.await;
858 database.save_thread(id, db_thread).await.log_err();
859 thread_store.update(cx, |store, cx| store.reload(cx));
860 });
861 }
862
863 fn send_mcp_prompt(
864 &self,
865 message_id: UserMessageId,
866 session_id: agent_client_protocol::SessionId,
867 prompt_name: String,
868 server_id: ContextServerId,
869 arguments: HashMap<String, String>,
870 original_content: Vec<acp::ContentBlock>,
871 cx: &mut Context<Self>,
872 ) -> Task<Result<acp::PromptResponse>> {
873 let server_store = self.context_server_registry.read(cx).server_store().clone();
874 let path_style = self.project.read(cx).path_style(cx);
875
876 cx.spawn(async move |this, cx| {
877 let prompt =
878 crate::get_prompt(&server_store, &server_id, &prompt_name, arguments, cx).await?;
879
880 let (acp_thread, thread) = this.update(cx, |this, _cx| {
881 let session = this
882 .sessions
883 .get(&session_id)
884 .context("Failed to get session")?;
885 anyhow::Ok((session.acp_thread.clone(), session.thread.clone()))
886 })??;
887
888 let mut last_is_user = true;
889
890 thread.update(cx, |thread, cx| {
891 thread.push_acp_user_block(
892 message_id,
893 original_content.into_iter().skip(1),
894 path_style,
895 cx,
896 );
897 });
898
899 for message in prompt.messages {
900 let context_server::types::PromptMessage { role, content } = message;
901 let block = mcp_message_content_to_acp_content_block(content);
902
903 match role {
904 context_server::types::Role::User => {
905 let id = acp_thread::UserMessageId::new();
906
907 acp_thread.update(cx, |acp_thread, cx| {
908 acp_thread.push_user_content_block_with_indent(
909 Some(id.clone()),
910 block.clone(),
911 true,
912 cx,
913 );
914 });
915
916 thread.update(cx, |thread, cx| {
917 thread.push_acp_user_block(id, [block], path_style, cx);
918 });
919 }
920 context_server::types::Role::Assistant => {
921 acp_thread.update(cx, |acp_thread, cx| {
922 acp_thread.push_assistant_content_block_with_indent(
923 block.clone(),
924 false,
925 true,
926 cx,
927 );
928 });
929
930 thread.update(cx, |thread, cx| {
931 thread.push_acp_agent_block(block, cx);
932 });
933 }
934 }
935
936 last_is_user = role == context_server::types::Role::User;
937 }
938
939 let response_stream = thread.update(cx, |thread, cx| {
940 if last_is_user {
941 thread.send_existing(cx)
942 } else {
943 // Resume if MCP prompt did not end with a user message
944 thread.resume(cx)
945 }
946 })?;
947
948 cx.update(|cx| {
949 NativeAgentConnection::handle_thread_events(
950 response_stream,
951 acp_thread.downgrade(),
952 cx,
953 )
954 })
955 .await
956 })
957 }
958}
959
960/// Wrapper struct that implements the AgentConnection trait
961#[derive(Clone)]
962pub struct NativeAgentConnection(pub Entity<NativeAgent>);
963
964impl NativeAgentConnection {
965 pub fn thread(&self, session_id: &acp::SessionId, cx: &App) -> Option<Entity<Thread>> {
966 self.0
967 .read(cx)
968 .sessions
969 .get(session_id)
970 .map(|session| session.thread.clone())
971 }
972
973 pub fn load_thread(&self, id: acp::SessionId, cx: &mut App) -> Task<Result<Entity<Thread>>> {
974 self.0.update(cx, |this, cx| this.load_thread(id, cx))
975 }
976
977 fn run_turn(
978 &self,
979 session_id: acp::SessionId,
980 cx: &mut App,
981 f: impl 'static
982 + FnOnce(Entity<Thread>, &mut App) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>,
983 ) -> Task<Result<acp::PromptResponse>> {
984 let Some((thread, acp_thread)) = self.0.update(cx, |agent, _cx| {
985 agent
986 .sessions
987 .get_mut(&session_id)
988 .map(|s| (s.thread.clone(), s.acp_thread.clone()))
989 }) else {
990 return Task::ready(Err(anyhow!("Session not found")));
991 };
992 log::debug!("Found session for: {}", session_id);
993
994 let response_stream = match f(thread, cx) {
995 Ok(stream) => stream,
996 Err(err) => return Task::ready(Err(err)),
997 };
998 Self::handle_thread_events(response_stream, acp_thread.downgrade(), cx)
999 }
1000
1001 fn handle_thread_events(
1002 mut events: mpsc::UnboundedReceiver<Result<ThreadEvent>>,
1003 acp_thread: WeakEntity<AcpThread>,
1004 cx: &App,
1005 ) -> Task<Result<acp::PromptResponse>> {
1006 cx.spawn(async move |cx| {
1007 // Handle response stream and forward to session.acp_thread
1008 while let Some(result) = events.next().await {
1009 match result {
1010 Ok(event) => {
1011 log::trace!("Received completion event: {:?}", event);
1012
1013 match event {
1014 ThreadEvent::UserMessage(message) => {
1015 acp_thread.update(cx, |thread, cx| {
1016 for content in message.content {
1017 thread.push_user_content_block(
1018 Some(message.id.clone()),
1019 content.into(),
1020 cx,
1021 );
1022 }
1023 })?;
1024 }
1025 ThreadEvent::AgentText(text) => {
1026 acp_thread.update(cx, |thread, cx| {
1027 thread.push_assistant_content_block(text.into(), false, cx)
1028 })?;
1029 }
1030 ThreadEvent::AgentThinking(text) => {
1031 acp_thread.update(cx, |thread, cx| {
1032 thread.push_assistant_content_block(text.into(), true, cx)
1033 })?;
1034 }
1035 ThreadEvent::ToolCallAuthorization(ToolCallAuthorization {
1036 tool_call,
1037 options,
1038 response,
1039 context: _,
1040 }) => {
1041 let outcome_task = acp_thread.update(cx, |thread, cx| {
1042 thread.request_tool_call_authorization(tool_call, options, cx)
1043 })??;
1044 cx.background_spawn(async move {
1045 if let acp::RequestPermissionOutcome::Selected(
1046 acp::SelectedPermissionOutcome { option_id, .. },
1047 ) = outcome_task.await
1048 {
1049 response
1050 .send(option_id)
1051 .map(|_| anyhow!("authorization receiver was dropped"))
1052 .log_err();
1053 }
1054 })
1055 .detach();
1056 }
1057 ThreadEvent::ToolCall(tool_call) => {
1058 acp_thread.update(cx, |thread, cx| {
1059 thread.upsert_tool_call(tool_call, cx)
1060 })??;
1061 }
1062 ThreadEvent::ToolCallUpdate(update) => {
1063 acp_thread.update(cx, |thread, cx| {
1064 thread.update_tool_call(update, cx)
1065 })??;
1066 }
1067 ThreadEvent::SubagentSpawned(session_id) => {
1068 acp_thread.update(cx, |thread, cx| {
1069 thread.subagent_spawned(session_id, cx);
1070 })?;
1071 }
1072 ThreadEvent::Retry(status) => {
1073 acp_thread.update(cx, |thread, cx| {
1074 thread.update_retry_status(status, cx)
1075 })?;
1076 }
1077 ThreadEvent::Stop(stop_reason) => {
1078 log::debug!("Assistant message complete: {:?}", stop_reason);
1079 return Ok(acp::PromptResponse::new(stop_reason));
1080 }
1081 }
1082 }
1083 Err(e) => {
1084 log::error!("Error in model response stream: {:?}", e);
1085 return Err(e);
1086 }
1087 }
1088 }
1089
1090 log::debug!("Response stream completed");
1091 anyhow::Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
1092 })
1093 }
1094}
1095
1096struct Command<'a> {
1097 prompt_name: &'a str,
1098 arg_value: &'a str,
1099 explicit_server_id: Option<&'a str>,
1100}
1101
1102impl<'a> Command<'a> {
1103 fn parse(prompt: &'a [acp::ContentBlock]) -> Option<Self> {
1104 let acp::ContentBlock::Text(text_content) = prompt.first()? else {
1105 return None;
1106 };
1107 let text = text_content.text.trim();
1108 let command = text.strip_prefix('/')?;
1109 let (command, arg_value) = command
1110 .split_once(char::is_whitespace)
1111 .unwrap_or((command, ""));
1112
1113 if let Some((server_id, prompt_name)) = command.split_once('.') {
1114 Some(Self {
1115 prompt_name,
1116 arg_value,
1117 explicit_server_id: Some(server_id),
1118 })
1119 } else {
1120 Some(Self {
1121 prompt_name: command,
1122 arg_value,
1123 explicit_server_id: None,
1124 })
1125 }
1126 }
1127}
1128
1129struct NativeAgentModelSelector {
1130 session_id: acp::SessionId,
1131 connection: NativeAgentConnection,
1132}
1133
1134impl acp_thread::AgentModelSelector for NativeAgentModelSelector {
1135 fn list_models(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelList>> {
1136 log::debug!("NativeAgentConnection::list_models called");
1137 let list = self.connection.0.read(cx).models.model_list.clone();
1138 Task::ready(if list.is_empty() {
1139 Err(anyhow::anyhow!("No models available"))
1140 } else {
1141 Ok(list)
1142 })
1143 }
1144
1145 fn select_model(&self, model_id: acp::ModelId, cx: &mut App) -> Task<Result<()>> {
1146 log::debug!(
1147 "Setting model for session {}: {}",
1148 self.session_id,
1149 model_id
1150 );
1151 let Some(thread) = self
1152 .connection
1153 .0
1154 .read(cx)
1155 .sessions
1156 .get(&self.session_id)
1157 .map(|session| session.thread.clone())
1158 else {
1159 return Task::ready(Err(anyhow!("Session not found")));
1160 };
1161
1162 let Some(model) = self.connection.0.read(cx).models.model_from_id(&model_id) else {
1163 return Task::ready(Err(anyhow!("Invalid model ID {}", model_id)));
1164 };
1165
1166 // We want to reset the effort level when switching models, as the currently-selected effort level may
1167 // not be compatible.
1168 let effort = model
1169 .default_effort_level()
1170 .map(|effort_level| effort_level.value.to_string());
1171
1172 thread.update(cx, |thread, cx| {
1173 thread.set_model(model.clone(), cx);
1174 thread.set_thinking_effort(effort.clone(), cx);
1175 thread.set_thinking_enabled(model.supports_thinking(), cx);
1176 });
1177
1178 update_settings_file(
1179 self.connection.0.read(cx).fs.clone(),
1180 cx,
1181 move |settings, cx| {
1182 let provider = model.provider_id().0.to_string();
1183 let model = model.id().0.to_string();
1184 let enable_thinking = thread.read(cx).thinking_enabled();
1185 settings
1186 .agent
1187 .get_or_insert_default()
1188 .set_model(LanguageModelSelection {
1189 provider: provider.into(),
1190 model,
1191 enable_thinking,
1192 effort,
1193 });
1194 },
1195 );
1196
1197 Task::ready(Ok(()))
1198 }
1199
1200 fn selected_model(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelInfo>> {
1201 let Some(thread) = self
1202 .connection
1203 .0
1204 .read(cx)
1205 .sessions
1206 .get(&self.session_id)
1207 .map(|session| session.thread.clone())
1208 else {
1209 return Task::ready(Err(anyhow!("Session not found")));
1210 };
1211 let Some(model) = thread.read(cx).model() else {
1212 return Task::ready(Err(anyhow!("Model not found")));
1213 };
1214 let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&model.provider_id())
1215 else {
1216 return Task::ready(Err(anyhow!("Provider not found")));
1217 };
1218 Task::ready(Ok(LanguageModels::map_language_model_to_info(
1219 model, &provider,
1220 )))
1221 }
1222
1223 fn watch(&self, cx: &mut App) -> Option<watch::Receiver<()>> {
1224 Some(self.connection.0.read(cx).models.watch())
1225 }
1226
1227 fn should_render_footer(&self) -> bool {
1228 true
1229 }
1230}
1231
1232impl acp_thread::AgentConnection for NativeAgentConnection {
1233 fn telemetry_id(&self) -> SharedString {
1234 "zed".into()
1235 }
1236
1237 fn new_session(
1238 self: Rc<Self>,
1239 project: Entity<Project>,
1240 cwd: &Path,
1241 cx: &mut App,
1242 ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
1243 log::debug!("Creating new thread for project at: {cwd:?}");
1244 Task::ready(Ok(self
1245 .0
1246 .update(cx, |agent, cx| agent.new_session(project, cx))))
1247 }
1248
1249 fn supports_load_session(&self, _cx: &App) -> bool {
1250 true
1251 }
1252
1253 fn load_session(
1254 self: Rc<Self>,
1255 session: AgentSessionInfo,
1256 _project: Entity<Project>,
1257 _cwd: &Path,
1258 cx: &mut App,
1259 ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
1260 self.0
1261 .update(cx, |agent, cx| agent.open_thread(session.session_id, cx))
1262 }
1263
1264 fn supports_close_session(&self, _cx: &App) -> bool {
1265 true
1266 }
1267
1268 fn close_session(&self, session_id: &acp::SessionId, cx: &mut App) -> Task<Result<()>> {
1269 self.0.update(cx, |agent, _cx| {
1270 agent.sessions.remove(session_id);
1271 });
1272 Task::ready(Ok(()))
1273 }
1274
1275 fn auth_methods(&self) -> &[acp::AuthMethod] {
1276 &[] // No auth for in-process
1277 }
1278
1279 fn authenticate(&self, _method: acp::AuthMethodId, _cx: &mut App) -> Task<Result<()>> {
1280 Task::ready(Ok(()))
1281 }
1282
1283 fn model_selector(&self, session_id: &acp::SessionId) -> Option<Rc<dyn AgentModelSelector>> {
1284 Some(Rc::new(NativeAgentModelSelector {
1285 session_id: session_id.clone(),
1286 connection: self.clone(),
1287 }) as Rc<dyn AgentModelSelector>)
1288 }
1289
1290 fn prompt(
1291 &self,
1292 id: Option<acp_thread::UserMessageId>,
1293 params: acp::PromptRequest,
1294 cx: &mut App,
1295 ) -> Task<Result<acp::PromptResponse>> {
1296 let id = id.expect("UserMessageId is required");
1297 let session_id = params.session_id.clone();
1298 log::info!("Received prompt request for session: {}", session_id);
1299 log::debug!("Prompt blocks count: {}", params.prompt.len());
1300
1301 if let Some(parsed_command) = Command::parse(¶ms.prompt) {
1302 let registry = self.0.read(cx).context_server_registry.read(cx);
1303
1304 let explicit_server_id = parsed_command
1305 .explicit_server_id
1306 .map(|server_id| ContextServerId(server_id.into()));
1307
1308 if let Some(prompt) =
1309 registry.find_prompt(explicit_server_id.as_ref(), parsed_command.prompt_name)
1310 {
1311 let arguments = if !parsed_command.arg_value.is_empty()
1312 && let Some(arg_name) = prompt
1313 .prompt
1314 .arguments
1315 .as_ref()
1316 .and_then(|args| args.first())
1317 .map(|arg| arg.name.clone())
1318 {
1319 HashMap::from_iter([(arg_name, parsed_command.arg_value.to_string())])
1320 } else {
1321 Default::default()
1322 };
1323
1324 let prompt_name = prompt.prompt.name.clone();
1325 let server_id = prompt.server_id.clone();
1326
1327 return self.0.update(cx, |agent, cx| {
1328 agent.send_mcp_prompt(
1329 id,
1330 session_id.clone(),
1331 prompt_name,
1332 server_id,
1333 arguments,
1334 params.prompt,
1335 cx,
1336 )
1337 });
1338 };
1339 };
1340
1341 let path_style = self.0.read(cx).project.read(cx).path_style(cx);
1342
1343 self.run_turn(session_id, cx, move |thread, cx| {
1344 let content: Vec<UserMessageContent> = params
1345 .prompt
1346 .into_iter()
1347 .map(|block| UserMessageContent::from_content_block(block, path_style))
1348 .collect::<Vec<_>>();
1349 log::debug!("Converted prompt to message: {} chars", content.len());
1350 log::debug!("Message id: {:?}", id);
1351 log::debug!("Message content: {:?}", content);
1352
1353 thread.update(cx, |thread, cx| thread.send(id, content, cx))
1354 })
1355 }
1356
1357 fn retry(
1358 &self,
1359 session_id: &acp::SessionId,
1360 _cx: &App,
1361 ) -> Option<Rc<dyn acp_thread::AgentSessionRetry>> {
1362 Some(Rc::new(NativeAgentSessionRetry {
1363 connection: self.clone(),
1364 session_id: session_id.clone(),
1365 }) as _)
1366 }
1367
1368 fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
1369 log::info!("Cancelling on session: {}", session_id);
1370 self.0.update(cx, |agent, cx| {
1371 if let Some(agent) = agent.sessions.get(session_id) {
1372 agent
1373 .thread
1374 .update(cx, |thread, cx| thread.cancel(cx))
1375 .detach();
1376 }
1377 });
1378 }
1379
1380 fn truncate(
1381 &self,
1382 session_id: &agent_client_protocol::SessionId,
1383 cx: &App,
1384 ) -> Option<Rc<dyn acp_thread::AgentSessionTruncate>> {
1385 self.0.read_with(cx, |agent, _cx| {
1386 agent.sessions.get(session_id).map(|session| {
1387 Rc::new(NativeAgentSessionTruncate {
1388 thread: session.thread.clone(),
1389 acp_thread: session.acp_thread.downgrade(),
1390 }) as _
1391 })
1392 })
1393 }
1394
1395 fn set_title(
1396 &self,
1397 session_id: &acp::SessionId,
1398 cx: &App,
1399 ) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
1400 self.0.read_with(cx, |agent, _cx| {
1401 agent
1402 .sessions
1403 .get(session_id)
1404 .filter(|s| !s.thread.read(cx).is_subagent())
1405 .map(|session| {
1406 Rc::new(NativeAgentSessionSetTitle {
1407 thread: session.thread.clone(),
1408 }) as _
1409 })
1410 })
1411 }
1412
1413 fn session_list(&self, cx: &mut App) -> Option<Rc<dyn AgentSessionList>> {
1414 let thread_store = self.0.read(cx).thread_store.clone();
1415 Some(Rc::new(NativeAgentSessionList::new(thread_store, cx)) as _)
1416 }
1417
1418 fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
1419 Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
1420 }
1421
1422 fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1423 self
1424 }
1425}
1426
1427impl acp_thread::AgentTelemetry for NativeAgentConnection {
1428 fn thread_data(
1429 &self,
1430 session_id: &acp::SessionId,
1431 cx: &mut App,
1432 ) -> Task<Result<serde_json::Value>> {
1433 let Some(session) = self.0.read(cx).sessions.get(session_id) else {
1434 return Task::ready(Err(anyhow!("Session not found")));
1435 };
1436
1437 let task = session.thread.read(cx).to_db(cx);
1438 cx.background_spawn(async move {
1439 serde_json::to_value(task.await).context("Failed to serialize thread")
1440 })
1441 }
1442}
1443
1444pub struct NativeAgentSessionList {
1445 thread_store: Entity<ThreadStore>,
1446 updates_tx: smol::channel::Sender<acp_thread::SessionListUpdate>,
1447 updates_rx: smol::channel::Receiver<acp_thread::SessionListUpdate>,
1448 _subscription: Subscription,
1449}
1450
1451impl NativeAgentSessionList {
1452 fn new(thread_store: Entity<ThreadStore>, cx: &mut App) -> Self {
1453 let (tx, rx) = smol::channel::unbounded();
1454 let this_tx = tx.clone();
1455 let subscription = cx.observe(&thread_store, move |_, _| {
1456 this_tx
1457 .try_send(acp_thread::SessionListUpdate::Refresh)
1458 .ok();
1459 });
1460 Self {
1461 thread_store,
1462 updates_tx: tx,
1463 updates_rx: rx,
1464 _subscription: subscription,
1465 }
1466 }
1467
1468 fn to_session_info(entry: DbThreadMetadata) -> AgentSessionInfo {
1469 AgentSessionInfo {
1470 session_id: entry.id,
1471 cwd: None,
1472 title: Some(entry.title),
1473 updated_at: Some(entry.updated_at),
1474 meta: None,
1475 }
1476 }
1477
1478 pub fn thread_store(&self) -> &Entity<ThreadStore> {
1479 &self.thread_store
1480 }
1481}
1482
1483impl AgentSessionList for NativeAgentSessionList {
1484 fn list_sessions(
1485 &self,
1486 _request: AgentSessionListRequest,
1487 cx: &mut App,
1488 ) -> Task<Result<AgentSessionListResponse>> {
1489 let sessions = self
1490 .thread_store
1491 .read(cx)
1492 .entries()
1493 .map(Self::to_session_info)
1494 .collect();
1495 Task::ready(Ok(AgentSessionListResponse::new(sessions)))
1496 }
1497
1498 fn supports_delete(&self) -> bool {
1499 true
1500 }
1501
1502 fn delete_session(&self, session_id: &acp::SessionId, cx: &mut App) -> Task<Result<()>> {
1503 self.thread_store
1504 .update(cx, |store, cx| store.delete_thread(session_id.clone(), cx))
1505 }
1506
1507 fn delete_sessions(&self, cx: &mut App) -> Task<Result<()>> {
1508 self.thread_store
1509 .update(cx, |store, cx| store.delete_threads(cx))
1510 }
1511
1512 fn watch(
1513 &self,
1514 _cx: &mut App,
1515 ) -> Option<smol::channel::Receiver<acp_thread::SessionListUpdate>> {
1516 Some(self.updates_rx.clone())
1517 }
1518
1519 fn notify_refresh(&self) {
1520 self.updates_tx
1521 .try_send(acp_thread::SessionListUpdate::Refresh)
1522 .ok();
1523 }
1524
1525 fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1526 self
1527 }
1528}
1529
1530struct NativeAgentSessionTruncate {
1531 thread: Entity<Thread>,
1532 acp_thread: WeakEntity<AcpThread>,
1533}
1534
1535impl acp_thread::AgentSessionTruncate for NativeAgentSessionTruncate {
1536 fn run(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
1537 match self.thread.update(cx, |thread, cx| {
1538 thread.truncate(message_id.clone(), cx)?;
1539 Ok(thread.latest_token_usage())
1540 }) {
1541 Ok(usage) => {
1542 self.acp_thread
1543 .update(cx, |thread, cx| {
1544 thread.update_token_usage(usage, cx);
1545 })
1546 .ok();
1547 Task::ready(Ok(()))
1548 }
1549 Err(error) => Task::ready(Err(error)),
1550 }
1551 }
1552}
1553
1554struct NativeAgentSessionRetry {
1555 connection: NativeAgentConnection,
1556 session_id: acp::SessionId,
1557}
1558
1559impl acp_thread::AgentSessionRetry for NativeAgentSessionRetry {
1560 fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>> {
1561 self.connection
1562 .run_turn(self.session_id.clone(), cx, |thread, cx| {
1563 thread.update(cx, |thread, cx| thread.resume(cx))
1564 })
1565 }
1566}
1567
1568struct NativeAgentSessionSetTitle {
1569 thread: Entity<Thread>,
1570}
1571
1572impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
1573 fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
1574 self.thread
1575 .update(cx, |thread, cx| thread.set_title(title, cx));
1576 Task::ready(Ok(()))
1577 }
1578}
1579
1580pub struct NativeThreadEnvironment {
1581 agent: WeakEntity<NativeAgent>,
1582 acp_thread: WeakEntity<AcpThread>,
1583}
1584
1585impl NativeThreadEnvironment {
1586 pub(crate) fn create_subagent_thread(
1587 agent: WeakEntity<NativeAgent>,
1588 parent_thread_entity: Entity<Thread>,
1589 label: String,
1590 initial_prompt: String,
1591 timeout: Option<Duration>,
1592 allowed_tools: Option<Vec<String>>,
1593 cx: &mut App,
1594 ) -> Result<Rc<dyn SubagentHandle>> {
1595 let parent_thread = parent_thread_entity.read(cx);
1596 let current_depth = parent_thread.depth();
1597
1598 if current_depth >= MAX_SUBAGENT_DEPTH {
1599 return Err(anyhow!(
1600 "Maximum subagent depth ({}) reached",
1601 MAX_SUBAGENT_DEPTH
1602 ));
1603 }
1604
1605 let running_count = parent_thread.running_subagent_count();
1606 if running_count >= MAX_PARALLEL_SUBAGENTS {
1607 return Err(anyhow!(
1608 "Maximum parallel subagents ({}) reached. Wait for existing subagents to complete.",
1609 MAX_PARALLEL_SUBAGENTS
1610 ));
1611 }
1612
1613 let allowed_tools = match allowed_tools {
1614 Some(tools) => {
1615 let parent_tool_names: std::collections::HashSet<&str> =
1616 parent_thread.tools.keys().map(|s| s.as_str()).collect();
1617 Some(
1618 tools
1619 .into_iter()
1620 .filter(|t| parent_tool_names.contains(t.as_str()))
1621 .collect::<Vec<_>>(),
1622 )
1623 }
1624 None => Some(parent_thread.tools.keys().map(|s| s.to_string()).collect()),
1625 };
1626
1627 let subagent_thread: Entity<Thread> = cx.new(|cx| {
1628 let mut thread = Thread::new_subagent(&parent_thread_entity, cx);
1629 thread.set_title(label.into(), cx);
1630 thread
1631 });
1632
1633 let session_id = subagent_thread.read(cx).id().clone();
1634
1635 let acp_thread = agent.update(cx, |agent, cx| {
1636 agent.register_session(
1637 subagent_thread.clone(),
1638 allowed_tools
1639 .as_ref()
1640 .map(|v| v.iter().map(|s| s.as_str()).collect()),
1641 cx,
1642 )
1643 })?;
1644
1645 parent_thread_entity.update(cx, |parent_thread, _cx| {
1646 parent_thread.register_running_subagent(subagent_thread.downgrade())
1647 });
1648
1649 let task = acp_thread.update(cx, |agent, cx| agent.send(vec![initial_prompt.into()], cx));
1650
1651 let timeout_timer = timeout.map(|d| cx.background_executor().timer(d));
1652 let wait_for_prompt_to_complete = cx
1653 .background_spawn(async move {
1654 if let Some(timer) = timeout_timer {
1655 futures::select! {
1656 _ = timer.fuse() => SubagentInitialPromptResult::Timeout,
1657 _ = task.fuse() => SubagentInitialPromptResult::Completed,
1658 }
1659 } else {
1660 task.await.log_err();
1661 SubagentInitialPromptResult::Completed
1662 }
1663 })
1664 .shared();
1665
1666 let mut user_stop_rx: watch::Receiver<bool> =
1667 acp_thread.update(cx, |thread, _| thread.user_stop_receiver());
1668
1669 let user_cancelled = cx
1670 .background_spawn(async move {
1671 loop {
1672 if *user_stop_rx.borrow() {
1673 return;
1674 }
1675 if user_stop_rx.changed().await.is_err() {
1676 std::future::pending::<()>().await;
1677 }
1678 }
1679 })
1680 .shared();
1681
1682 Ok(Rc::new(NativeSubagentHandle {
1683 session_id,
1684 subagent_thread,
1685 parent_thread: parent_thread_entity.downgrade(),
1686 acp_thread,
1687 wait_for_prompt_to_complete,
1688 user_cancelled,
1689 }) as _)
1690 }
1691}
1692
1693impl ThreadEnvironment for NativeThreadEnvironment {
1694 fn create_terminal(
1695 &self,
1696 command: String,
1697 cwd: Option<PathBuf>,
1698 output_byte_limit: Option<u64>,
1699 cx: &mut AsyncApp,
1700 ) -> Task<Result<Rc<dyn TerminalHandle>>> {
1701 let task = self.acp_thread.update(cx, |thread, cx| {
1702 thread.create_terminal(command, vec![], vec![], cwd, output_byte_limit, cx)
1703 });
1704
1705 let acp_thread = self.acp_thread.clone();
1706 cx.spawn(async move |cx| {
1707 let terminal = task?.await?;
1708
1709 let (drop_tx, drop_rx) = oneshot::channel();
1710 let terminal_id = terminal.read_with(cx, |terminal, _cx| terminal.id().clone());
1711
1712 cx.spawn(async move |cx| {
1713 drop_rx.await.ok();
1714 acp_thread.update(cx, |thread, cx| thread.release_terminal(terminal_id, cx))
1715 })
1716 .detach();
1717
1718 let handle = AcpTerminalHandle {
1719 terminal,
1720 _drop_tx: Some(drop_tx),
1721 };
1722
1723 Ok(Rc::new(handle) as _)
1724 })
1725 }
1726
1727 fn create_subagent(
1728 &self,
1729 parent_thread_entity: Entity<Thread>,
1730 label: String,
1731 initial_prompt: String,
1732 timeout: Option<Duration>,
1733 allowed_tools: Option<Vec<String>>,
1734 cx: &mut App,
1735 ) -> Result<Rc<dyn SubagentHandle>> {
1736 Self::create_subagent_thread(
1737 self.agent.clone(),
1738 parent_thread_entity,
1739 label,
1740 initial_prompt,
1741 timeout,
1742 allowed_tools,
1743 cx,
1744 )
1745 }
1746}
1747
1748#[derive(Debug, Clone, Copy)]
1749enum SubagentInitialPromptResult {
1750 Completed,
1751 Timeout,
1752}
1753
1754pub struct NativeSubagentHandle {
1755 session_id: acp::SessionId,
1756 parent_thread: WeakEntity<Thread>,
1757 subagent_thread: Entity<Thread>,
1758 acp_thread: Entity<AcpThread>,
1759 wait_for_prompt_to_complete: Shared<Task<SubagentInitialPromptResult>>,
1760 user_cancelled: Shared<Task<()>>,
1761}
1762
1763impl SubagentHandle for NativeSubagentHandle {
1764 fn id(&self) -> acp::SessionId {
1765 self.session_id.clone()
1766 }
1767
1768 fn wait_for_summary(&self, summary_prompt: String, cx: &AsyncApp) -> Task<Result<String>> {
1769 let thread = self.subagent_thread.clone();
1770 let acp_thread = self.acp_thread.clone();
1771 let wait_for_prompt = self.wait_for_prompt_to_complete.clone();
1772
1773 let wait_for_summary_task = cx.spawn(async move |cx| {
1774 let timed_out = match wait_for_prompt.await {
1775 SubagentInitialPromptResult::Completed => false,
1776 SubagentInitialPromptResult::Timeout => true,
1777 };
1778
1779 let summary_prompt = if timed_out {
1780 thread.update(cx, |thread, cx| thread.cancel(cx)).await;
1781 format!("{}\n{}", "The time to complete the task was exceeded. Stop with the task and follow the directions below:", summary_prompt)
1782 } else {
1783 summary_prompt
1784 };
1785
1786 acp_thread
1787 .update(cx, |thread, cx| thread.send(vec![summary_prompt.into()], cx))
1788 .await?;
1789
1790 thread.read_with(cx, |thread, _cx| {
1791 thread
1792 .last_message()
1793 .map(|m| m.to_markdown())
1794 .context("No response from subagent")
1795 })
1796 });
1797
1798 let user_cancelled = self.user_cancelled.clone();
1799 let thread = self.subagent_thread.clone();
1800 let subagent_session_id = self.session_id.clone();
1801 let parent_thread = self.parent_thread.clone();
1802 cx.spawn(async move |cx| {
1803 let result = futures::select! {
1804 result = wait_for_summary_task.fuse() => result,
1805 _ = user_cancelled.fuse() => {
1806 thread.update(cx, |thread, cx| thread.cancel(cx).detach());
1807 Err(anyhow!("User cancelled"))
1808 },
1809 };
1810 parent_thread
1811 .update(cx, |parent_thread, cx| {
1812 parent_thread.unregister_running_subagent(&subagent_session_id, cx)
1813 })
1814 .ok();
1815 result
1816 })
1817 }
1818}
1819
1820pub struct AcpTerminalHandle {
1821 terminal: Entity<acp_thread::Terminal>,
1822 _drop_tx: Option<oneshot::Sender<()>>,
1823}
1824
1825impl TerminalHandle for AcpTerminalHandle {
1826 fn id(&self, cx: &AsyncApp) -> Result<acp::TerminalId> {
1827 Ok(self.terminal.read_with(cx, |term, _cx| term.id().clone()))
1828 }
1829
1830 fn wait_for_exit(&self, cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {
1831 Ok(self
1832 .terminal
1833 .read_with(cx, |term, _cx| term.wait_for_exit()))
1834 }
1835
1836 fn current_output(&self, cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {
1837 Ok(self
1838 .terminal
1839 .read_with(cx, |term, cx| term.current_output(cx)))
1840 }
1841
1842 fn kill(&self, cx: &AsyncApp) -> Result<()> {
1843 cx.update(|cx| {
1844 self.terminal.update(cx, |terminal, cx| {
1845 terminal.kill(cx);
1846 });
1847 });
1848 Ok(())
1849 }
1850
1851 fn was_stopped_by_user(&self, cx: &AsyncApp) -> Result<bool> {
1852 Ok(self
1853 .terminal
1854 .read_with(cx, |term, _cx| term.was_stopped_by_user()))
1855 }
1856}
1857
1858#[cfg(test)]
1859mod internal_tests {
1860 use super::*;
1861 use acp_thread::{AgentConnection, AgentModelGroupName, AgentModelInfo, MentionUri};
1862 use fs::FakeFs;
1863 use gpui::TestAppContext;
1864 use indoc::formatdoc;
1865 use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider};
1866 use language_model::{LanguageModelProviderId, LanguageModelProviderName};
1867 use serde_json::json;
1868 use settings::SettingsStore;
1869 use util::{path, rel_path::rel_path};
1870
1871 #[gpui::test]
1872 async fn test_maintaining_project_context(cx: &mut TestAppContext) {
1873 init_test(cx);
1874 let fs = FakeFs::new(cx.executor());
1875 fs.insert_tree(
1876 "/",
1877 json!({
1878 "a": {}
1879 }),
1880 )
1881 .await;
1882 let project = Project::test(fs.clone(), [], cx).await;
1883 let thread_store = cx.new(|cx| ThreadStore::new(cx));
1884 let agent = NativeAgent::new(
1885 project.clone(),
1886 thread_store,
1887 Templates::new(),
1888 None,
1889 fs.clone(),
1890 &mut cx.to_async(),
1891 )
1892 .await
1893 .unwrap();
1894 agent.read_with(cx, |agent, cx| {
1895 assert_eq!(agent.project_context.read(cx).worktrees, vec![])
1896 });
1897
1898 let worktree = project
1899 .update(cx, |project, cx| project.create_worktree("/a", true, cx))
1900 .await
1901 .unwrap();
1902 cx.run_until_parked();
1903 agent.read_with(cx, |agent, cx| {
1904 assert_eq!(
1905 agent.project_context.read(cx).worktrees,
1906 vec![WorktreeContext {
1907 root_name: "a".into(),
1908 abs_path: Path::new("/a").into(),
1909 rules_file: None
1910 }]
1911 )
1912 });
1913
1914 // Creating `/a/.rules` updates the project context.
1915 fs.insert_file("/a/.rules", Vec::new()).await;
1916 cx.run_until_parked();
1917 agent.read_with(cx, |agent, cx| {
1918 let rules_entry = worktree
1919 .read(cx)
1920 .entry_for_path(rel_path(".rules"))
1921 .unwrap();
1922 assert_eq!(
1923 agent.project_context.read(cx).worktrees,
1924 vec![WorktreeContext {
1925 root_name: "a".into(),
1926 abs_path: Path::new("/a").into(),
1927 rules_file: Some(RulesFileContext {
1928 path_in_worktree: rel_path(".rules").into(),
1929 text: "".into(),
1930 project_entry_id: rules_entry.id.to_usize()
1931 })
1932 }]
1933 )
1934 });
1935 }
1936
1937 #[gpui::test]
1938 async fn test_listing_models(cx: &mut TestAppContext) {
1939 init_test(cx);
1940 let fs = FakeFs::new(cx.executor());
1941 fs.insert_tree("/", json!({ "a": {} })).await;
1942 let project = Project::test(fs.clone(), [], cx).await;
1943 let thread_store = cx.new(|cx| ThreadStore::new(cx));
1944 let connection = NativeAgentConnection(
1945 NativeAgent::new(
1946 project.clone(),
1947 thread_store,
1948 Templates::new(),
1949 None,
1950 fs.clone(),
1951 &mut cx.to_async(),
1952 )
1953 .await
1954 .unwrap(),
1955 );
1956
1957 // Create a thread/session
1958 let acp_thread = cx
1959 .update(|cx| {
1960 Rc::new(connection.clone()).new_session(project.clone(), Path::new("/a"), cx)
1961 })
1962 .await
1963 .unwrap();
1964
1965 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1966
1967 let models = cx
1968 .update(|cx| {
1969 connection
1970 .model_selector(&session_id)
1971 .unwrap()
1972 .list_models(cx)
1973 })
1974 .await
1975 .unwrap();
1976
1977 let acp_thread::AgentModelList::Grouped(models) = models else {
1978 panic!("Unexpected model group");
1979 };
1980 assert_eq!(
1981 models,
1982 IndexMap::from_iter([(
1983 AgentModelGroupName("Fake".into()),
1984 vec![AgentModelInfo {
1985 id: acp::ModelId::new("fake/fake"),
1986 name: "Fake".into(),
1987 description: None,
1988 icon: Some(acp_thread::AgentModelIcon::Named(
1989 ui::IconName::ZedAssistant
1990 )),
1991 is_latest: false,
1992 }]
1993 )])
1994 );
1995 }
1996
1997 #[gpui::test]
1998 async fn test_model_selection_persists_to_settings(cx: &mut TestAppContext) {
1999 init_test(cx);
2000 let fs = FakeFs::new(cx.executor());
2001 fs.create_dir(paths::settings_file().parent().unwrap())
2002 .await
2003 .unwrap();
2004 fs.insert_file(
2005 paths::settings_file(),
2006 json!({
2007 "agent": {
2008 "default_model": {
2009 "provider": "foo",
2010 "model": "bar"
2011 }
2012 }
2013 })
2014 .to_string()
2015 .into_bytes(),
2016 )
2017 .await;
2018 let project = Project::test(fs.clone(), [], cx).await;
2019
2020 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2021
2022 // Create the agent and connection
2023 let agent = NativeAgent::new(
2024 project.clone(),
2025 thread_store,
2026 Templates::new(),
2027 None,
2028 fs.clone(),
2029 &mut cx.to_async(),
2030 )
2031 .await
2032 .unwrap();
2033 let connection = NativeAgentConnection(agent.clone());
2034
2035 // Create a thread/session
2036 let acp_thread = cx
2037 .update(|cx| {
2038 Rc::new(connection.clone()).new_session(project.clone(), Path::new("/a"), cx)
2039 })
2040 .await
2041 .unwrap();
2042
2043 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2044
2045 // Select a model
2046 let selector = connection.model_selector(&session_id).unwrap();
2047 let model_id = acp::ModelId::new("fake/fake");
2048 cx.update(|cx| selector.select_model(model_id.clone(), cx))
2049 .await
2050 .unwrap();
2051
2052 // Verify the thread has the selected model
2053 agent.read_with(cx, |agent, _| {
2054 let session = agent.sessions.get(&session_id).unwrap();
2055 session.thread.read_with(cx, |thread, _| {
2056 assert_eq!(thread.model().unwrap().id().0, "fake");
2057 });
2058 });
2059
2060 cx.run_until_parked();
2061
2062 // Verify settings file was updated
2063 let settings_content = fs.load(paths::settings_file()).await.unwrap();
2064 let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
2065
2066 // Check that the agent settings contain the selected model
2067 assert_eq!(
2068 settings_json["agent"]["default_model"]["model"],
2069 json!("fake")
2070 );
2071 assert_eq!(
2072 settings_json["agent"]["default_model"]["provider"],
2073 json!("fake")
2074 );
2075
2076 // Register a thinking model and select it.
2077 cx.update(|cx| {
2078 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2079 "fake-corp",
2080 "fake-thinking",
2081 "Fake Thinking",
2082 true,
2083 ));
2084 let thinking_provider = Arc::new(
2085 FakeLanguageModelProvider::new(
2086 LanguageModelProviderId::from("fake-corp".to_string()),
2087 LanguageModelProviderName::from("Fake Corp".to_string()),
2088 )
2089 .with_models(vec![thinking_model]),
2090 );
2091 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2092 registry.register_provider(thinking_provider, cx);
2093 });
2094 });
2095 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2096
2097 let selector = connection.model_selector(&session_id).unwrap();
2098 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2099 .await
2100 .unwrap();
2101 cx.run_until_parked();
2102
2103 // Verify enable_thinking was written to settings as true.
2104 let settings_content = fs.load(paths::settings_file()).await.unwrap();
2105 let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
2106 assert_eq!(
2107 settings_json["agent"]["default_model"]["enable_thinking"],
2108 json!(true),
2109 "selecting a thinking model should persist enable_thinking: true to settings"
2110 );
2111 }
2112
2113 #[gpui::test]
2114 async fn test_select_model_updates_thinking_enabled(cx: &mut TestAppContext) {
2115 init_test(cx);
2116 let fs = FakeFs::new(cx.executor());
2117 fs.create_dir(paths::settings_file().parent().unwrap())
2118 .await
2119 .unwrap();
2120 fs.insert_file(paths::settings_file(), b"{}".to_vec()).await;
2121 let project = Project::test(fs.clone(), [], cx).await;
2122
2123 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2124 let agent = NativeAgent::new(
2125 project.clone(),
2126 thread_store,
2127 Templates::new(),
2128 None,
2129 fs.clone(),
2130 &mut cx.to_async(),
2131 )
2132 .await
2133 .unwrap();
2134 let connection = NativeAgentConnection(agent.clone());
2135
2136 let acp_thread = cx
2137 .update(|cx| {
2138 Rc::new(connection.clone()).new_session(project.clone(), Path::new("/a"), cx)
2139 })
2140 .await
2141 .unwrap();
2142 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2143
2144 // Register a second provider with a thinking model.
2145 cx.update(|cx| {
2146 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2147 "fake-corp",
2148 "fake-thinking",
2149 "Fake Thinking",
2150 true,
2151 ));
2152 let thinking_provider = Arc::new(
2153 FakeLanguageModelProvider::new(
2154 LanguageModelProviderId::from("fake-corp".to_string()),
2155 LanguageModelProviderName::from("Fake Corp".to_string()),
2156 )
2157 .with_models(vec![thinking_model]),
2158 );
2159 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2160 registry.register_provider(thinking_provider, cx);
2161 });
2162 });
2163 // Refresh the agent's model list so it picks up the new provider.
2164 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2165
2166 // Thread starts with thinking_enabled = false (the default).
2167 agent.read_with(cx, |agent, _| {
2168 let session = agent.sessions.get(&session_id).unwrap();
2169 session.thread.read_with(cx, |thread, _| {
2170 assert!(!thread.thinking_enabled(), "thinking defaults to false");
2171 });
2172 });
2173
2174 // Select the thinking model via select_model.
2175 let selector = connection.model_selector(&session_id).unwrap();
2176 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2177 .await
2178 .unwrap();
2179
2180 // select_model should have enabled thinking based on the model's supports_thinking().
2181 agent.read_with(cx, |agent, _| {
2182 let session = agent.sessions.get(&session_id).unwrap();
2183 session.thread.read_with(cx, |thread, _| {
2184 assert!(
2185 thread.thinking_enabled(),
2186 "select_model should enable thinking when model supports it"
2187 );
2188 });
2189 });
2190
2191 // Switch back to the non-thinking model.
2192 let selector = connection.model_selector(&session_id).unwrap();
2193 cx.update(|cx| selector.select_model(acp::ModelId::new("fake/fake"), cx))
2194 .await
2195 .unwrap();
2196
2197 // select_model should have disabled thinking.
2198 agent.read_with(cx, |agent, _| {
2199 let session = agent.sessions.get(&session_id).unwrap();
2200 session.thread.read_with(cx, |thread, _| {
2201 assert!(
2202 !thread.thinking_enabled(),
2203 "select_model should disable thinking when model does not support it"
2204 );
2205 });
2206 });
2207 }
2208
2209 #[gpui::test]
2210 async fn test_loaded_thread_preserves_thinking_enabled(cx: &mut TestAppContext) {
2211 init_test(cx);
2212 let fs = FakeFs::new(cx.executor());
2213 fs.insert_tree("/", json!({ "a": {} })).await;
2214 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2215 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2216 let agent = NativeAgent::new(
2217 project.clone(),
2218 thread_store.clone(),
2219 Templates::new(),
2220 None,
2221 fs.clone(),
2222 &mut cx.to_async(),
2223 )
2224 .await
2225 .unwrap();
2226 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2227
2228 // Register a thinking model.
2229 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2230 "fake-corp",
2231 "fake-thinking",
2232 "Fake Thinking",
2233 true,
2234 ));
2235 let thinking_provider = Arc::new(
2236 FakeLanguageModelProvider::new(
2237 LanguageModelProviderId::from("fake-corp".to_string()),
2238 LanguageModelProviderName::from("Fake Corp".to_string()),
2239 )
2240 .with_models(vec![thinking_model.clone()]),
2241 );
2242 cx.update(|cx| {
2243 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2244 registry.register_provider(thinking_provider, cx);
2245 });
2246 });
2247 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2248
2249 // Create a thread and select the thinking model.
2250 let acp_thread = cx
2251 .update(|cx| {
2252 connection
2253 .clone()
2254 .new_session(project.clone(), Path::new("/a"), cx)
2255 })
2256 .await
2257 .unwrap();
2258 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2259
2260 let selector = connection.model_selector(&session_id).unwrap();
2261 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2262 .await
2263 .unwrap();
2264
2265 // Verify thinking is enabled after selecting the thinking model.
2266 let thread = agent.read_with(cx, |agent, _| {
2267 agent.sessions.get(&session_id).unwrap().thread.clone()
2268 });
2269 thread.read_with(cx, |thread, _| {
2270 assert!(
2271 thread.thinking_enabled(),
2272 "thinking should be enabled after selecting thinking model"
2273 );
2274 });
2275
2276 // Send a message so the thread gets persisted.
2277 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
2278 let send = cx.foreground_executor().spawn(send);
2279 cx.run_until_parked();
2280
2281 thinking_model.send_last_completion_stream_text_chunk("Response.");
2282 thinking_model.end_last_completion_stream();
2283
2284 send.await.unwrap();
2285 cx.run_until_parked();
2286
2287 // Close the session so it can be reloaded from disk.
2288 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2289 .await
2290 .unwrap();
2291 drop(thread);
2292 drop(acp_thread);
2293 agent.read_with(cx, |agent, _| {
2294 assert!(agent.sessions.is_empty());
2295 });
2296
2297 // Reload the thread and verify thinking_enabled is still true.
2298 let reloaded_acp_thread = agent
2299 .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
2300 .await
2301 .unwrap();
2302 let reloaded_thread = agent.read_with(cx, |agent, _| {
2303 agent.sessions.get(&session_id).unwrap().thread.clone()
2304 });
2305 reloaded_thread.read_with(cx, |thread, _| {
2306 assert!(
2307 thread.thinking_enabled(),
2308 "thinking_enabled should be preserved when reloading a thread with a thinking model"
2309 );
2310 });
2311
2312 drop(reloaded_acp_thread);
2313 }
2314
2315 #[gpui::test]
2316 async fn test_loaded_thread_preserves_model(cx: &mut TestAppContext) {
2317 init_test(cx);
2318 let fs = FakeFs::new(cx.executor());
2319 fs.insert_tree("/", json!({ "a": {} })).await;
2320 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2321 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2322 let agent = NativeAgent::new(
2323 project.clone(),
2324 thread_store.clone(),
2325 Templates::new(),
2326 None,
2327 fs.clone(),
2328 &mut cx.to_async(),
2329 )
2330 .await
2331 .unwrap();
2332 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2333
2334 // Register a model where id() != name(), like real Anthropic models
2335 // (e.g. id="claude-sonnet-4-5-thinking-latest", name="Claude Sonnet 4.5 Thinking").
2336 let model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2337 "fake-corp",
2338 "custom-model-id",
2339 "Custom Model Display Name",
2340 false,
2341 ));
2342 let provider = Arc::new(
2343 FakeLanguageModelProvider::new(
2344 LanguageModelProviderId::from("fake-corp".to_string()),
2345 LanguageModelProviderName::from("Fake Corp".to_string()),
2346 )
2347 .with_models(vec![model.clone()]),
2348 );
2349 cx.update(|cx| {
2350 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2351 registry.register_provider(provider, cx);
2352 });
2353 });
2354 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2355
2356 // Create a thread and select the model.
2357 let acp_thread = cx
2358 .update(|cx| {
2359 connection
2360 .clone()
2361 .new_session(project.clone(), Path::new("/a"), cx)
2362 })
2363 .await
2364 .unwrap();
2365 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2366
2367 let selector = connection.model_selector(&session_id).unwrap();
2368 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/custom-model-id"), cx))
2369 .await
2370 .unwrap();
2371
2372 let thread = agent.read_with(cx, |agent, _| {
2373 agent.sessions.get(&session_id).unwrap().thread.clone()
2374 });
2375 thread.read_with(cx, |thread, _| {
2376 assert_eq!(
2377 thread.model().unwrap().id().0.as_ref(),
2378 "custom-model-id",
2379 "model should be set before persisting"
2380 );
2381 });
2382
2383 // Send a message so the thread gets persisted.
2384 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
2385 let send = cx.foreground_executor().spawn(send);
2386 cx.run_until_parked();
2387
2388 model.send_last_completion_stream_text_chunk("Response.");
2389 model.end_last_completion_stream();
2390
2391 send.await.unwrap();
2392 cx.run_until_parked();
2393
2394 // Close the session so it can be reloaded from disk.
2395 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2396 .await
2397 .unwrap();
2398 drop(thread);
2399 drop(acp_thread);
2400 agent.read_with(cx, |agent, _| {
2401 assert!(agent.sessions.is_empty());
2402 });
2403
2404 // Reload the thread and verify the model was preserved.
2405 let reloaded_acp_thread = agent
2406 .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
2407 .await
2408 .unwrap();
2409 let reloaded_thread = agent.read_with(cx, |agent, _| {
2410 agent.sessions.get(&session_id).unwrap().thread.clone()
2411 });
2412 reloaded_thread.read_with(cx, |thread, _| {
2413 let reloaded_model = thread
2414 .model()
2415 .expect("model should be present after reload");
2416 assert_eq!(
2417 reloaded_model.id().0.as_ref(),
2418 "custom-model-id",
2419 "reloaded thread should have the same model, not fall back to the default"
2420 );
2421 });
2422
2423 drop(reloaded_acp_thread);
2424 }
2425
2426 #[gpui::test]
2427 async fn test_save_load_thread(cx: &mut TestAppContext) {
2428 init_test(cx);
2429 let fs = FakeFs::new(cx.executor());
2430 fs.insert_tree(
2431 "/",
2432 json!({
2433 "a": {
2434 "b.md": "Lorem"
2435 }
2436 }),
2437 )
2438 .await;
2439 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2440 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2441 let agent = NativeAgent::new(
2442 project.clone(),
2443 thread_store.clone(),
2444 Templates::new(),
2445 None,
2446 fs.clone(),
2447 &mut cx.to_async(),
2448 )
2449 .await
2450 .unwrap();
2451 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2452
2453 let acp_thread = cx
2454 .update(|cx| {
2455 connection
2456 .clone()
2457 .new_session(project.clone(), Path::new(""), cx)
2458 })
2459 .await
2460 .unwrap();
2461 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2462 let thread = agent.read_with(cx, |agent, _| {
2463 agent.sessions.get(&session_id).unwrap().thread.clone()
2464 });
2465
2466 // Ensure empty threads are not saved, even if they get mutated.
2467 let model = Arc::new(FakeLanguageModel::default());
2468 let summary_model = Arc::new(FakeLanguageModel::default());
2469 thread.update(cx, |thread, cx| {
2470 thread.set_model(model.clone(), cx);
2471 thread.set_summarization_model(Some(summary_model.clone()), cx);
2472 });
2473 cx.run_until_parked();
2474 assert_eq!(thread_entries(&thread_store, cx), vec![]);
2475
2476 let send = acp_thread.update(cx, |thread, cx| {
2477 thread.send(
2478 vec![
2479 "What does ".into(),
2480 acp::ContentBlock::ResourceLink(acp::ResourceLink::new(
2481 "b.md",
2482 MentionUri::File {
2483 abs_path: path!("/a/b.md").into(),
2484 }
2485 .to_uri()
2486 .to_string(),
2487 )),
2488 " mean?".into(),
2489 ],
2490 cx,
2491 )
2492 });
2493 let send = cx.foreground_executor().spawn(send);
2494 cx.run_until_parked();
2495
2496 model.send_last_completion_stream_text_chunk("Lorem.");
2497 model.end_last_completion_stream();
2498 cx.run_until_parked();
2499 summary_model
2500 .send_last_completion_stream_text_chunk(&format!("Explaining {}", path!("/a/b.md")));
2501 summary_model.end_last_completion_stream();
2502
2503 send.await.unwrap();
2504 let uri = MentionUri::File {
2505 abs_path: path!("/a/b.md").into(),
2506 }
2507 .to_uri();
2508 acp_thread.read_with(cx, |thread, cx| {
2509 assert_eq!(
2510 thread.to_markdown(cx),
2511 formatdoc! {"
2512 ## User
2513
2514 What does [@b.md]({uri}) mean?
2515
2516 ## Assistant
2517
2518 Lorem.
2519
2520 "}
2521 )
2522 });
2523
2524 cx.run_until_parked();
2525
2526 // Close the session so it can be reloaded from disk.
2527 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2528 .await
2529 .unwrap();
2530 drop(thread);
2531 drop(acp_thread);
2532 agent.read_with(cx, |agent, _| {
2533 assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
2534 });
2535
2536 // Ensure the thread can be reloaded from disk.
2537 assert_eq!(
2538 thread_entries(&thread_store, cx),
2539 vec![(
2540 session_id.clone(),
2541 format!("Explaining {}", path!("/a/b.md"))
2542 )]
2543 );
2544 let acp_thread = agent
2545 .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
2546 .await
2547 .unwrap();
2548 acp_thread.read_with(cx, |thread, cx| {
2549 assert_eq!(
2550 thread.to_markdown(cx),
2551 formatdoc! {"
2552 ## User
2553
2554 What does [@b.md]({uri}) mean?
2555
2556 ## Assistant
2557
2558 Lorem.
2559
2560 "}
2561 )
2562 });
2563 }
2564
2565 fn thread_entries(
2566 thread_store: &Entity<ThreadStore>,
2567 cx: &mut TestAppContext,
2568 ) -> Vec<(acp::SessionId, String)> {
2569 thread_store.read_with(cx, |store, _| {
2570 store
2571 .entries()
2572 .map(|entry| (entry.id.clone(), entry.title.to_string()))
2573 .collect::<Vec<_>>()
2574 })
2575 }
2576
2577 fn init_test(cx: &mut TestAppContext) {
2578 env_logger::try_init().ok();
2579 cx.update(|cx| {
2580 let settings_store = SettingsStore::test(cx);
2581 cx.set_global(settings_store);
2582
2583 LanguageModelRegistry::test(cx);
2584 });
2585 }
2586}
2587
2588fn mcp_message_content_to_acp_content_block(
2589 content: context_server::types::MessageContent,
2590) -> acp::ContentBlock {
2591 match content {
2592 context_server::types::MessageContent::Text {
2593 text,
2594 annotations: _,
2595 } => text.into(),
2596 context_server::types::MessageContent::Image {
2597 data,
2598 mime_type,
2599 annotations: _,
2600 } => acp::ContentBlock::Image(acp::ImageContent::new(data, mime_type)),
2601 context_server::types::MessageContent::Audio {
2602 data,
2603 mime_type,
2604 annotations: _,
2605 } => acp::ContentBlock::Audio(acp::AudioContent::new(data, mime_type)),
2606 context_server::types::MessageContent::Resource {
2607 resource,
2608 annotations: _,
2609 } => {
2610 let mut link =
2611 acp::ResourceLink::new(resource.uri.to_string(), resource.uri.to_string());
2612 if let Some(mime_type) = resource.mime_type {
2613 link = link.mime_type(mime_type);
2614 }
2615 acp::ContentBlock::ResourceLink(link)
2616 }
2617 }
2618}