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 Some(Rc::new(NativeAgentSessionSetTitle {
1401 connection: self.clone(),
1402 session_id: session_id.clone(),
1403 }) as _)
1404 }
1405
1406 fn session_list(&self, cx: &mut App) -> Option<Rc<dyn AgentSessionList>> {
1407 let thread_store = self.0.read(cx).thread_store.clone();
1408 Some(Rc::new(NativeAgentSessionList::new(thread_store, cx)) as _)
1409 }
1410
1411 fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
1412 Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
1413 }
1414
1415 fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1416 self
1417 }
1418}
1419
1420impl acp_thread::AgentTelemetry for NativeAgentConnection {
1421 fn thread_data(
1422 &self,
1423 session_id: &acp::SessionId,
1424 cx: &mut App,
1425 ) -> Task<Result<serde_json::Value>> {
1426 let Some(session) = self.0.read(cx).sessions.get(session_id) else {
1427 return Task::ready(Err(anyhow!("Session not found")));
1428 };
1429
1430 let task = session.thread.read(cx).to_db(cx);
1431 cx.background_spawn(async move {
1432 serde_json::to_value(task.await).context("Failed to serialize thread")
1433 })
1434 }
1435}
1436
1437pub struct NativeAgentSessionList {
1438 thread_store: Entity<ThreadStore>,
1439 updates_tx: smol::channel::Sender<acp_thread::SessionListUpdate>,
1440 updates_rx: smol::channel::Receiver<acp_thread::SessionListUpdate>,
1441 _subscription: Subscription,
1442}
1443
1444impl NativeAgentSessionList {
1445 fn new(thread_store: Entity<ThreadStore>, cx: &mut App) -> Self {
1446 let (tx, rx) = smol::channel::unbounded();
1447 let this_tx = tx.clone();
1448 let subscription = cx.observe(&thread_store, move |_, _| {
1449 this_tx
1450 .try_send(acp_thread::SessionListUpdate::Refresh)
1451 .ok();
1452 });
1453 Self {
1454 thread_store,
1455 updates_tx: tx,
1456 updates_rx: rx,
1457 _subscription: subscription,
1458 }
1459 }
1460
1461 fn to_session_info(entry: DbThreadMetadata) -> AgentSessionInfo {
1462 AgentSessionInfo {
1463 session_id: entry.id,
1464 cwd: None,
1465 title: Some(entry.title),
1466 updated_at: Some(entry.updated_at),
1467 meta: None,
1468 }
1469 }
1470
1471 pub fn thread_store(&self) -> &Entity<ThreadStore> {
1472 &self.thread_store
1473 }
1474}
1475
1476impl AgentSessionList for NativeAgentSessionList {
1477 fn list_sessions(
1478 &self,
1479 _request: AgentSessionListRequest,
1480 cx: &mut App,
1481 ) -> Task<Result<AgentSessionListResponse>> {
1482 let sessions = self
1483 .thread_store
1484 .read(cx)
1485 .entries()
1486 .map(Self::to_session_info)
1487 .collect();
1488 Task::ready(Ok(AgentSessionListResponse::new(sessions)))
1489 }
1490
1491 fn supports_delete(&self) -> bool {
1492 true
1493 }
1494
1495 fn delete_session(&self, session_id: &acp::SessionId, cx: &mut App) -> Task<Result<()>> {
1496 self.thread_store
1497 .update(cx, |store, cx| store.delete_thread(session_id.clone(), cx))
1498 }
1499
1500 fn delete_sessions(&self, cx: &mut App) -> Task<Result<()>> {
1501 self.thread_store
1502 .update(cx, |store, cx| store.delete_threads(cx))
1503 }
1504
1505 fn watch(
1506 &self,
1507 _cx: &mut App,
1508 ) -> Option<smol::channel::Receiver<acp_thread::SessionListUpdate>> {
1509 Some(self.updates_rx.clone())
1510 }
1511
1512 fn notify_refresh(&self) {
1513 self.updates_tx
1514 .try_send(acp_thread::SessionListUpdate::Refresh)
1515 .ok();
1516 }
1517
1518 fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1519 self
1520 }
1521}
1522
1523struct NativeAgentSessionTruncate {
1524 thread: Entity<Thread>,
1525 acp_thread: WeakEntity<AcpThread>,
1526}
1527
1528impl acp_thread::AgentSessionTruncate for NativeAgentSessionTruncate {
1529 fn run(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
1530 match self.thread.update(cx, |thread, cx| {
1531 thread.truncate(message_id.clone(), cx)?;
1532 Ok(thread.latest_token_usage())
1533 }) {
1534 Ok(usage) => {
1535 self.acp_thread
1536 .update(cx, |thread, cx| {
1537 thread.update_token_usage(usage, cx);
1538 })
1539 .ok();
1540 Task::ready(Ok(()))
1541 }
1542 Err(error) => Task::ready(Err(error)),
1543 }
1544 }
1545}
1546
1547struct NativeAgentSessionRetry {
1548 connection: NativeAgentConnection,
1549 session_id: acp::SessionId,
1550}
1551
1552impl acp_thread::AgentSessionRetry for NativeAgentSessionRetry {
1553 fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>> {
1554 self.connection
1555 .run_turn(self.session_id.clone(), cx, |thread, cx| {
1556 thread.update(cx, |thread, cx| thread.resume(cx))
1557 })
1558 }
1559}
1560
1561struct NativeAgentSessionSetTitle {
1562 connection: NativeAgentConnection,
1563 session_id: acp::SessionId,
1564}
1565
1566impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
1567 fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
1568 let Some(session) = self.connection.0.read(cx).sessions.get(&self.session_id) else {
1569 return Task::ready(Err(anyhow!("session not found")));
1570 };
1571 let thread = session.thread.clone();
1572 thread.update(cx, |thread, cx| thread.set_title(title, cx));
1573 Task::ready(Ok(()))
1574 }
1575}
1576
1577pub struct NativeThreadEnvironment {
1578 agent: WeakEntity<NativeAgent>,
1579 acp_thread: WeakEntity<AcpThread>,
1580}
1581
1582impl NativeThreadEnvironment {
1583 pub(crate) fn create_subagent_thread(
1584 agent: WeakEntity<NativeAgent>,
1585 parent_thread_entity: Entity<Thread>,
1586 label: String,
1587 initial_prompt: String,
1588 timeout: Option<Duration>,
1589 allowed_tools: Option<Vec<String>>,
1590 cx: &mut App,
1591 ) -> Result<Rc<dyn SubagentHandle>> {
1592 let parent_thread = parent_thread_entity.read(cx);
1593 let current_depth = parent_thread.depth();
1594
1595 if current_depth >= MAX_SUBAGENT_DEPTH {
1596 return Err(anyhow!(
1597 "Maximum subagent depth ({}) reached",
1598 MAX_SUBAGENT_DEPTH
1599 ));
1600 }
1601
1602 let running_count = parent_thread.running_subagent_count();
1603 if running_count >= MAX_PARALLEL_SUBAGENTS {
1604 return Err(anyhow!(
1605 "Maximum parallel subagents ({}) reached. Wait for existing subagents to complete.",
1606 MAX_PARALLEL_SUBAGENTS
1607 ));
1608 }
1609
1610 let allowed_tools = match allowed_tools {
1611 Some(tools) => {
1612 let parent_tool_names: std::collections::HashSet<&str> =
1613 parent_thread.tools.keys().map(|s| s.as_str()).collect();
1614 Some(
1615 tools
1616 .into_iter()
1617 .filter(|t| parent_tool_names.contains(t.as_str()))
1618 .collect::<Vec<_>>(),
1619 )
1620 }
1621 None => Some(parent_thread.tools.keys().map(|s| s.to_string()).collect()),
1622 };
1623
1624 let subagent_thread: Entity<Thread> = cx.new(|cx| {
1625 let mut thread = Thread::new_subagent(&parent_thread_entity, cx);
1626 thread.set_title(label.into(), cx);
1627 thread
1628 });
1629
1630 let session_id = subagent_thread.read(cx).id().clone();
1631
1632 let acp_thread = agent.update(cx, |agent, cx| {
1633 agent.register_session(
1634 subagent_thread.clone(),
1635 allowed_tools
1636 .as_ref()
1637 .map(|v| v.iter().map(|s| s.as_str()).collect()),
1638 cx,
1639 )
1640 })?;
1641
1642 parent_thread_entity.update(cx, |parent_thread, _cx| {
1643 parent_thread.register_running_subagent(subagent_thread.downgrade())
1644 });
1645
1646 let task = acp_thread.update(cx, |agent, cx| agent.send(vec![initial_prompt.into()], cx));
1647
1648 let timeout_timer = timeout.map(|d| cx.background_executor().timer(d));
1649 let wait_for_prompt_to_complete = cx
1650 .background_spawn(async move {
1651 if let Some(timer) = timeout_timer {
1652 futures::select! {
1653 _ = timer.fuse() => SubagentInitialPromptResult::Timeout,
1654 _ = task.fuse() => SubagentInitialPromptResult::Completed,
1655 }
1656 } else {
1657 task.await.log_err();
1658 SubagentInitialPromptResult::Completed
1659 }
1660 })
1661 .shared();
1662
1663 let mut user_stop_rx: watch::Receiver<bool> =
1664 acp_thread.update(cx, |thread, _| thread.user_stop_receiver());
1665
1666 let user_cancelled = cx
1667 .background_spawn(async move {
1668 loop {
1669 if *user_stop_rx.borrow() {
1670 return;
1671 }
1672 if user_stop_rx.changed().await.is_err() {
1673 std::future::pending::<()>().await;
1674 }
1675 }
1676 })
1677 .shared();
1678
1679 Ok(Rc::new(NativeSubagentHandle {
1680 session_id,
1681 subagent_thread,
1682 parent_thread: parent_thread_entity.downgrade(),
1683 acp_thread,
1684 wait_for_prompt_to_complete,
1685 user_cancelled,
1686 }) as _)
1687 }
1688}
1689
1690impl ThreadEnvironment for NativeThreadEnvironment {
1691 fn create_terminal(
1692 &self,
1693 command: String,
1694 cwd: Option<PathBuf>,
1695 output_byte_limit: Option<u64>,
1696 cx: &mut AsyncApp,
1697 ) -> Task<Result<Rc<dyn TerminalHandle>>> {
1698 let task = self.acp_thread.update(cx, |thread, cx| {
1699 thread.create_terminal(command, vec![], vec![], cwd, output_byte_limit, cx)
1700 });
1701
1702 let acp_thread = self.acp_thread.clone();
1703 cx.spawn(async move |cx| {
1704 let terminal = task?.await?;
1705
1706 let (drop_tx, drop_rx) = oneshot::channel();
1707 let terminal_id = terminal.read_with(cx, |terminal, _cx| terminal.id().clone());
1708
1709 cx.spawn(async move |cx| {
1710 drop_rx.await.ok();
1711 acp_thread.update(cx, |thread, cx| thread.release_terminal(terminal_id, cx))
1712 })
1713 .detach();
1714
1715 let handle = AcpTerminalHandle {
1716 terminal,
1717 _drop_tx: Some(drop_tx),
1718 };
1719
1720 Ok(Rc::new(handle) as _)
1721 })
1722 }
1723
1724 fn create_subagent(
1725 &self,
1726 parent_thread_entity: Entity<Thread>,
1727 label: String,
1728 initial_prompt: String,
1729 timeout: Option<Duration>,
1730 allowed_tools: Option<Vec<String>>,
1731 cx: &mut App,
1732 ) -> Result<Rc<dyn SubagentHandle>> {
1733 Self::create_subagent_thread(
1734 self.agent.clone(),
1735 parent_thread_entity,
1736 label,
1737 initial_prompt,
1738 timeout,
1739 allowed_tools,
1740 cx,
1741 )
1742 }
1743}
1744
1745#[derive(Debug, Clone, Copy)]
1746enum SubagentInitialPromptResult {
1747 Completed,
1748 Timeout,
1749}
1750
1751pub struct NativeSubagentHandle {
1752 session_id: acp::SessionId,
1753 parent_thread: WeakEntity<Thread>,
1754 subagent_thread: Entity<Thread>,
1755 acp_thread: Entity<AcpThread>,
1756 wait_for_prompt_to_complete: Shared<Task<SubagentInitialPromptResult>>,
1757 user_cancelled: Shared<Task<()>>,
1758}
1759
1760impl SubagentHandle for NativeSubagentHandle {
1761 fn id(&self) -> acp::SessionId {
1762 self.session_id.clone()
1763 }
1764
1765 fn wait_for_summary(&self, summary_prompt: String, cx: &AsyncApp) -> Task<Result<String>> {
1766 let thread = self.subagent_thread.clone();
1767 let acp_thread = self.acp_thread.clone();
1768 let wait_for_prompt = self.wait_for_prompt_to_complete.clone();
1769
1770 let wait_for_summary_task = cx.spawn(async move |cx| {
1771 let timed_out = match wait_for_prompt.await {
1772 SubagentInitialPromptResult::Completed => false,
1773 SubagentInitialPromptResult::Timeout => true,
1774 };
1775
1776 let summary_prompt = if timed_out {
1777 thread.update(cx, |thread, cx| thread.cancel(cx)).await;
1778 format!("{}\n{}", "The time to complete the task was exceeded. Stop with the task and follow the directions below:", summary_prompt)
1779 } else {
1780 summary_prompt
1781 };
1782
1783 acp_thread
1784 .update(cx, |thread, cx| thread.send(vec![summary_prompt.into()], cx))
1785 .await?;
1786
1787 thread.read_with(cx, |thread, _cx| {
1788 thread
1789 .last_message()
1790 .map(|m| m.to_markdown())
1791 .context("No response from subagent")
1792 })
1793 });
1794
1795 let user_cancelled = self.user_cancelled.clone();
1796 let thread = self.subagent_thread.clone();
1797 let subagent_session_id = self.session_id.clone();
1798 let parent_thread = self.parent_thread.clone();
1799 cx.spawn(async move |cx| {
1800 let result = futures::select! {
1801 result = wait_for_summary_task.fuse() => result,
1802 _ = user_cancelled.fuse() => {
1803 thread.update(cx, |thread, cx| thread.cancel(cx).detach());
1804 Err(anyhow!("User cancelled"))
1805 },
1806 };
1807 parent_thread
1808 .update(cx, |parent_thread, cx| {
1809 parent_thread.unregister_running_subagent(&subagent_session_id, cx)
1810 })
1811 .ok();
1812 result
1813 })
1814 }
1815}
1816
1817pub struct AcpTerminalHandle {
1818 terminal: Entity<acp_thread::Terminal>,
1819 _drop_tx: Option<oneshot::Sender<()>>,
1820}
1821
1822impl TerminalHandle for AcpTerminalHandle {
1823 fn id(&self, cx: &AsyncApp) -> Result<acp::TerminalId> {
1824 Ok(self.terminal.read_with(cx, |term, _cx| term.id().clone()))
1825 }
1826
1827 fn wait_for_exit(&self, cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {
1828 Ok(self
1829 .terminal
1830 .read_with(cx, |term, _cx| term.wait_for_exit()))
1831 }
1832
1833 fn current_output(&self, cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {
1834 Ok(self
1835 .terminal
1836 .read_with(cx, |term, cx| term.current_output(cx)))
1837 }
1838
1839 fn kill(&self, cx: &AsyncApp) -> Result<()> {
1840 cx.update(|cx| {
1841 self.terminal.update(cx, |terminal, cx| {
1842 terminal.kill(cx);
1843 });
1844 });
1845 Ok(())
1846 }
1847
1848 fn was_stopped_by_user(&self, cx: &AsyncApp) -> Result<bool> {
1849 Ok(self
1850 .terminal
1851 .read_with(cx, |term, _cx| term.was_stopped_by_user()))
1852 }
1853}
1854
1855#[cfg(test)]
1856mod internal_tests {
1857 use super::*;
1858 use acp_thread::{AgentConnection, AgentModelGroupName, AgentModelInfo, MentionUri};
1859 use fs::FakeFs;
1860 use gpui::TestAppContext;
1861 use indoc::formatdoc;
1862 use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider};
1863 use language_model::{LanguageModelProviderId, LanguageModelProviderName};
1864 use serde_json::json;
1865 use settings::SettingsStore;
1866 use util::{path, rel_path::rel_path};
1867
1868 #[gpui::test]
1869 async fn test_maintaining_project_context(cx: &mut TestAppContext) {
1870 init_test(cx);
1871 let fs = FakeFs::new(cx.executor());
1872 fs.insert_tree(
1873 "/",
1874 json!({
1875 "a": {}
1876 }),
1877 )
1878 .await;
1879 let project = Project::test(fs.clone(), [], cx).await;
1880 let thread_store = cx.new(|cx| ThreadStore::new(cx));
1881 let agent = NativeAgent::new(
1882 project.clone(),
1883 thread_store,
1884 Templates::new(),
1885 None,
1886 fs.clone(),
1887 &mut cx.to_async(),
1888 )
1889 .await
1890 .unwrap();
1891 agent.read_with(cx, |agent, cx| {
1892 assert_eq!(agent.project_context.read(cx).worktrees, vec![])
1893 });
1894
1895 let worktree = project
1896 .update(cx, |project, cx| project.create_worktree("/a", true, cx))
1897 .await
1898 .unwrap();
1899 cx.run_until_parked();
1900 agent.read_with(cx, |agent, cx| {
1901 assert_eq!(
1902 agent.project_context.read(cx).worktrees,
1903 vec![WorktreeContext {
1904 root_name: "a".into(),
1905 abs_path: Path::new("/a").into(),
1906 rules_file: None
1907 }]
1908 )
1909 });
1910
1911 // Creating `/a/.rules` updates the project context.
1912 fs.insert_file("/a/.rules", Vec::new()).await;
1913 cx.run_until_parked();
1914 agent.read_with(cx, |agent, cx| {
1915 let rules_entry = worktree
1916 .read(cx)
1917 .entry_for_path(rel_path(".rules"))
1918 .unwrap();
1919 assert_eq!(
1920 agent.project_context.read(cx).worktrees,
1921 vec![WorktreeContext {
1922 root_name: "a".into(),
1923 abs_path: Path::new("/a").into(),
1924 rules_file: Some(RulesFileContext {
1925 path_in_worktree: rel_path(".rules").into(),
1926 text: "".into(),
1927 project_entry_id: rules_entry.id.to_usize()
1928 })
1929 }]
1930 )
1931 });
1932 }
1933
1934 #[gpui::test]
1935 async fn test_listing_models(cx: &mut TestAppContext) {
1936 init_test(cx);
1937 let fs = FakeFs::new(cx.executor());
1938 fs.insert_tree("/", json!({ "a": {} })).await;
1939 let project = Project::test(fs.clone(), [], cx).await;
1940 let thread_store = cx.new(|cx| ThreadStore::new(cx));
1941 let connection = NativeAgentConnection(
1942 NativeAgent::new(
1943 project.clone(),
1944 thread_store,
1945 Templates::new(),
1946 None,
1947 fs.clone(),
1948 &mut cx.to_async(),
1949 )
1950 .await
1951 .unwrap(),
1952 );
1953
1954 // Create a thread/session
1955 let acp_thread = cx
1956 .update(|cx| {
1957 Rc::new(connection.clone()).new_session(project.clone(), Path::new("/a"), cx)
1958 })
1959 .await
1960 .unwrap();
1961
1962 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1963
1964 let models = cx
1965 .update(|cx| {
1966 connection
1967 .model_selector(&session_id)
1968 .unwrap()
1969 .list_models(cx)
1970 })
1971 .await
1972 .unwrap();
1973
1974 let acp_thread::AgentModelList::Grouped(models) = models else {
1975 panic!("Unexpected model group");
1976 };
1977 assert_eq!(
1978 models,
1979 IndexMap::from_iter([(
1980 AgentModelGroupName("Fake".into()),
1981 vec![AgentModelInfo {
1982 id: acp::ModelId::new("fake/fake"),
1983 name: "Fake".into(),
1984 description: None,
1985 icon: Some(acp_thread::AgentModelIcon::Named(
1986 ui::IconName::ZedAssistant
1987 )),
1988 is_latest: false,
1989 }]
1990 )])
1991 );
1992 }
1993
1994 #[gpui::test]
1995 async fn test_model_selection_persists_to_settings(cx: &mut TestAppContext) {
1996 init_test(cx);
1997 let fs = FakeFs::new(cx.executor());
1998 fs.create_dir(paths::settings_file().parent().unwrap())
1999 .await
2000 .unwrap();
2001 fs.insert_file(
2002 paths::settings_file(),
2003 json!({
2004 "agent": {
2005 "default_model": {
2006 "provider": "foo",
2007 "model": "bar"
2008 }
2009 }
2010 })
2011 .to_string()
2012 .into_bytes(),
2013 )
2014 .await;
2015 let project = Project::test(fs.clone(), [], cx).await;
2016
2017 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2018
2019 // Create the agent and connection
2020 let agent = NativeAgent::new(
2021 project.clone(),
2022 thread_store,
2023 Templates::new(),
2024 None,
2025 fs.clone(),
2026 &mut cx.to_async(),
2027 )
2028 .await
2029 .unwrap();
2030 let connection = NativeAgentConnection(agent.clone());
2031
2032 // Create a thread/session
2033 let acp_thread = cx
2034 .update(|cx| {
2035 Rc::new(connection.clone()).new_session(project.clone(), Path::new("/a"), cx)
2036 })
2037 .await
2038 .unwrap();
2039
2040 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2041
2042 // Select a model
2043 let selector = connection.model_selector(&session_id).unwrap();
2044 let model_id = acp::ModelId::new("fake/fake");
2045 cx.update(|cx| selector.select_model(model_id.clone(), cx))
2046 .await
2047 .unwrap();
2048
2049 // Verify the thread has the selected model
2050 agent.read_with(cx, |agent, _| {
2051 let session = agent.sessions.get(&session_id).unwrap();
2052 session.thread.read_with(cx, |thread, _| {
2053 assert_eq!(thread.model().unwrap().id().0, "fake");
2054 });
2055 });
2056
2057 cx.run_until_parked();
2058
2059 // Verify settings file was updated
2060 let settings_content = fs.load(paths::settings_file()).await.unwrap();
2061 let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
2062
2063 // Check that the agent settings contain the selected model
2064 assert_eq!(
2065 settings_json["agent"]["default_model"]["model"],
2066 json!("fake")
2067 );
2068 assert_eq!(
2069 settings_json["agent"]["default_model"]["provider"],
2070 json!("fake")
2071 );
2072
2073 // Register a thinking model and select it.
2074 cx.update(|cx| {
2075 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2076 "fake-corp",
2077 "fake-thinking",
2078 "Fake Thinking",
2079 true,
2080 ));
2081 let thinking_provider = Arc::new(
2082 FakeLanguageModelProvider::new(
2083 LanguageModelProviderId::from("fake-corp".to_string()),
2084 LanguageModelProviderName::from("Fake Corp".to_string()),
2085 )
2086 .with_models(vec![thinking_model]),
2087 );
2088 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2089 registry.register_provider(thinking_provider, cx);
2090 });
2091 });
2092 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2093
2094 let selector = connection.model_selector(&session_id).unwrap();
2095 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2096 .await
2097 .unwrap();
2098 cx.run_until_parked();
2099
2100 // Verify enable_thinking was written to settings as true.
2101 let settings_content = fs.load(paths::settings_file()).await.unwrap();
2102 let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
2103 assert_eq!(
2104 settings_json["agent"]["default_model"]["enable_thinking"],
2105 json!(true),
2106 "selecting a thinking model should persist enable_thinking: true to settings"
2107 );
2108 }
2109
2110 #[gpui::test]
2111 async fn test_select_model_updates_thinking_enabled(cx: &mut TestAppContext) {
2112 init_test(cx);
2113 let fs = FakeFs::new(cx.executor());
2114 fs.create_dir(paths::settings_file().parent().unwrap())
2115 .await
2116 .unwrap();
2117 fs.insert_file(paths::settings_file(), b"{}".to_vec()).await;
2118 let project = Project::test(fs.clone(), [], cx).await;
2119
2120 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2121 let agent = NativeAgent::new(
2122 project.clone(),
2123 thread_store,
2124 Templates::new(),
2125 None,
2126 fs.clone(),
2127 &mut cx.to_async(),
2128 )
2129 .await
2130 .unwrap();
2131 let connection = NativeAgentConnection(agent.clone());
2132
2133 let acp_thread = cx
2134 .update(|cx| {
2135 Rc::new(connection.clone()).new_session(project.clone(), Path::new("/a"), cx)
2136 })
2137 .await
2138 .unwrap();
2139 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2140
2141 // Register a second provider with a thinking model.
2142 cx.update(|cx| {
2143 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2144 "fake-corp",
2145 "fake-thinking",
2146 "Fake Thinking",
2147 true,
2148 ));
2149 let thinking_provider = Arc::new(
2150 FakeLanguageModelProvider::new(
2151 LanguageModelProviderId::from("fake-corp".to_string()),
2152 LanguageModelProviderName::from("Fake Corp".to_string()),
2153 )
2154 .with_models(vec![thinking_model]),
2155 );
2156 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2157 registry.register_provider(thinking_provider, cx);
2158 });
2159 });
2160 // Refresh the agent's model list so it picks up the new provider.
2161 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2162
2163 // Thread starts with thinking_enabled = false (the default).
2164 agent.read_with(cx, |agent, _| {
2165 let session = agent.sessions.get(&session_id).unwrap();
2166 session.thread.read_with(cx, |thread, _| {
2167 assert!(!thread.thinking_enabled(), "thinking defaults to false");
2168 });
2169 });
2170
2171 // Select the thinking model via select_model.
2172 let selector = connection.model_selector(&session_id).unwrap();
2173 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2174 .await
2175 .unwrap();
2176
2177 // select_model should have enabled thinking based on the model's supports_thinking().
2178 agent.read_with(cx, |agent, _| {
2179 let session = agent.sessions.get(&session_id).unwrap();
2180 session.thread.read_with(cx, |thread, _| {
2181 assert!(
2182 thread.thinking_enabled(),
2183 "select_model should enable thinking when model supports it"
2184 );
2185 });
2186 });
2187
2188 // Switch back to the non-thinking model.
2189 let selector = connection.model_selector(&session_id).unwrap();
2190 cx.update(|cx| selector.select_model(acp::ModelId::new("fake/fake"), cx))
2191 .await
2192 .unwrap();
2193
2194 // select_model should have disabled thinking.
2195 agent.read_with(cx, |agent, _| {
2196 let session = agent.sessions.get(&session_id).unwrap();
2197 session.thread.read_with(cx, |thread, _| {
2198 assert!(
2199 !thread.thinking_enabled(),
2200 "select_model should disable thinking when model does not support it"
2201 );
2202 });
2203 });
2204 }
2205
2206 #[gpui::test]
2207 async fn test_loaded_thread_preserves_thinking_enabled(cx: &mut TestAppContext) {
2208 init_test(cx);
2209 let fs = FakeFs::new(cx.executor());
2210 fs.insert_tree("/", json!({ "a": {} })).await;
2211 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2212 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2213 let agent = NativeAgent::new(
2214 project.clone(),
2215 thread_store.clone(),
2216 Templates::new(),
2217 None,
2218 fs.clone(),
2219 &mut cx.to_async(),
2220 )
2221 .await
2222 .unwrap();
2223 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2224
2225 // Register a thinking model.
2226 let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2227 "fake-corp",
2228 "fake-thinking",
2229 "Fake Thinking",
2230 true,
2231 ));
2232 let thinking_provider = Arc::new(
2233 FakeLanguageModelProvider::new(
2234 LanguageModelProviderId::from("fake-corp".to_string()),
2235 LanguageModelProviderName::from("Fake Corp".to_string()),
2236 )
2237 .with_models(vec![thinking_model.clone()]),
2238 );
2239 cx.update(|cx| {
2240 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2241 registry.register_provider(thinking_provider, cx);
2242 });
2243 });
2244 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2245
2246 // Create a thread and select the thinking model.
2247 let acp_thread = cx
2248 .update(|cx| {
2249 connection
2250 .clone()
2251 .new_session(project.clone(), Path::new("/a"), cx)
2252 })
2253 .await
2254 .unwrap();
2255 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2256
2257 let selector = connection.model_selector(&session_id).unwrap();
2258 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2259 .await
2260 .unwrap();
2261
2262 // Verify thinking is enabled after selecting the thinking model.
2263 let thread = agent.read_with(cx, |agent, _| {
2264 agent.sessions.get(&session_id).unwrap().thread.clone()
2265 });
2266 thread.read_with(cx, |thread, _| {
2267 assert!(
2268 thread.thinking_enabled(),
2269 "thinking should be enabled after selecting thinking model"
2270 );
2271 });
2272
2273 // Send a message so the thread gets persisted.
2274 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
2275 let send = cx.foreground_executor().spawn(send);
2276 cx.run_until_parked();
2277
2278 thinking_model.send_last_completion_stream_text_chunk("Response.");
2279 thinking_model.end_last_completion_stream();
2280
2281 send.await.unwrap();
2282 cx.run_until_parked();
2283
2284 // Close the session so it can be reloaded from disk.
2285 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2286 .await
2287 .unwrap();
2288 drop(thread);
2289 drop(acp_thread);
2290 agent.read_with(cx, |agent, _| {
2291 assert!(agent.sessions.is_empty());
2292 });
2293
2294 // Reload the thread and verify thinking_enabled is still true.
2295 let reloaded_acp_thread = agent
2296 .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
2297 .await
2298 .unwrap();
2299 let reloaded_thread = agent.read_with(cx, |agent, _| {
2300 agent.sessions.get(&session_id).unwrap().thread.clone()
2301 });
2302 reloaded_thread.read_with(cx, |thread, _| {
2303 assert!(
2304 thread.thinking_enabled(),
2305 "thinking_enabled should be preserved when reloading a thread with a thinking model"
2306 );
2307 });
2308
2309 drop(reloaded_acp_thread);
2310 }
2311
2312 #[gpui::test]
2313 async fn test_loaded_thread_preserves_model(cx: &mut TestAppContext) {
2314 init_test(cx);
2315 let fs = FakeFs::new(cx.executor());
2316 fs.insert_tree("/", json!({ "a": {} })).await;
2317 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2318 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2319 let agent = NativeAgent::new(
2320 project.clone(),
2321 thread_store.clone(),
2322 Templates::new(),
2323 None,
2324 fs.clone(),
2325 &mut cx.to_async(),
2326 )
2327 .await
2328 .unwrap();
2329 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2330
2331 // Register a model where id() != name(), like real Anthropic models
2332 // (e.g. id="claude-sonnet-4-5-thinking-latest", name="Claude Sonnet 4.5 Thinking").
2333 let model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2334 "fake-corp",
2335 "custom-model-id",
2336 "Custom Model Display Name",
2337 false,
2338 ));
2339 let provider = Arc::new(
2340 FakeLanguageModelProvider::new(
2341 LanguageModelProviderId::from("fake-corp".to_string()),
2342 LanguageModelProviderName::from("Fake Corp".to_string()),
2343 )
2344 .with_models(vec![model.clone()]),
2345 );
2346 cx.update(|cx| {
2347 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2348 registry.register_provider(provider, cx);
2349 });
2350 });
2351 agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2352
2353 // Create a thread and select the model.
2354 let acp_thread = cx
2355 .update(|cx| {
2356 connection
2357 .clone()
2358 .new_session(project.clone(), Path::new("/a"), cx)
2359 })
2360 .await
2361 .unwrap();
2362 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2363
2364 let selector = connection.model_selector(&session_id).unwrap();
2365 cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/custom-model-id"), cx))
2366 .await
2367 .unwrap();
2368
2369 let thread = agent.read_with(cx, |agent, _| {
2370 agent.sessions.get(&session_id).unwrap().thread.clone()
2371 });
2372 thread.read_with(cx, |thread, _| {
2373 assert_eq!(
2374 thread.model().unwrap().id().0.as_ref(),
2375 "custom-model-id",
2376 "model should be set before persisting"
2377 );
2378 });
2379
2380 // Send a message so the thread gets persisted.
2381 let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
2382 let send = cx.foreground_executor().spawn(send);
2383 cx.run_until_parked();
2384
2385 model.send_last_completion_stream_text_chunk("Response.");
2386 model.end_last_completion_stream();
2387
2388 send.await.unwrap();
2389 cx.run_until_parked();
2390
2391 // Close the session so it can be reloaded from disk.
2392 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2393 .await
2394 .unwrap();
2395 drop(thread);
2396 drop(acp_thread);
2397 agent.read_with(cx, |agent, _| {
2398 assert!(agent.sessions.is_empty());
2399 });
2400
2401 // Reload the thread and verify the model was preserved.
2402 let reloaded_acp_thread = agent
2403 .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
2404 .await
2405 .unwrap();
2406 let reloaded_thread = agent.read_with(cx, |agent, _| {
2407 agent.sessions.get(&session_id).unwrap().thread.clone()
2408 });
2409 reloaded_thread.read_with(cx, |thread, _| {
2410 let reloaded_model = thread
2411 .model()
2412 .expect("model should be present after reload");
2413 assert_eq!(
2414 reloaded_model.id().0.as_ref(),
2415 "custom-model-id",
2416 "reloaded thread should have the same model, not fall back to the default"
2417 );
2418 });
2419
2420 drop(reloaded_acp_thread);
2421 }
2422
2423 #[gpui::test]
2424 async fn test_save_load_thread(cx: &mut TestAppContext) {
2425 init_test(cx);
2426 let fs = FakeFs::new(cx.executor());
2427 fs.insert_tree(
2428 "/",
2429 json!({
2430 "a": {
2431 "b.md": "Lorem"
2432 }
2433 }),
2434 )
2435 .await;
2436 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2437 let thread_store = cx.new(|cx| ThreadStore::new(cx));
2438 let agent = NativeAgent::new(
2439 project.clone(),
2440 thread_store.clone(),
2441 Templates::new(),
2442 None,
2443 fs.clone(),
2444 &mut cx.to_async(),
2445 )
2446 .await
2447 .unwrap();
2448 let connection = Rc::new(NativeAgentConnection(agent.clone()));
2449
2450 let acp_thread = cx
2451 .update(|cx| {
2452 connection
2453 .clone()
2454 .new_session(project.clone(), Path::new(""), cx)
2455 })
2456 .await
2457 .unwrap();
2458 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2459 let thread = agent.read_with(cx, |agent, _| {
2460 agent.sessions.get(&session_id).unwrap().thread.clone()
2461 });
2462
2463 // Ensure empty threads are not saved, even if they get mutated.
2464 let model = Arc::new(FakeLanguageModel::default());
2465 let summary_model = Arc::new(FakeLanguageModel::default());
2466 thread.update(cx, |thread, cx| {
2467 thread.set_model(model.clone(), cx);
2468 thread.set_summarization_model(Some(summary_model.clone()), cx);
2469 });
2470 cx.run_until_parked();
2471 assert_eq!(thread_entries(&thread_store, cx), vec![]);
2472
2473 let send = acp_thread.update(cx, |thread, cx| {
2474 thread.send(
2475 vec![
2476 "What does ".into(),
2477 acp::ContentBlock::ResourceLink(acp::ResourceLink::new(
2478 "b.md",
2479 MentionUri::File {
2480 abs_path: path!("/a/b.md").into(),
2481 }
2482 .to_uri()
2483 .to_string(),
2484 )),
2485 " mean?".into(),
2486 ],
2487 cx,
2488 )
2489 });
2490 let send = cx.foreground_executor().spawn(send);
2491 cx.run_until_parked();
2492
2493 model.send_last_completion_stream_text_chunk("Lorem.");
2494 model.end_last_completion_stream();
2495 cx.run_until_parked();
2496 summary_model
2497 .send_last_completion_stream_text_chunk(&format!("Explaining {}", path!("/a/b.md")));
2498 summary_model.end_last_completion_stream();
2499
2500 send.await.unwrap();
2501 let uri = MentionUri::File {
2502 abs_path: path!("/a/b.md").into(),
2503 }
2504 .to_uri();
2505 acp_thread.read_with(cx, |thread, cx| {
2506 assert_eq!(
2507 thread.to_markdown(cx),
2508 formatdoc! {"
2509 ## User
2510
2511 What does [@b.md]({uri}) mean?
2512
2513 ## Assistant
2514
2515 Lorem.
2516
2517 "}
2518 )
2519 });
2520
2521 cx.run_until_parked();
2522
2523 // Close the session so it can be reloaded from disk.
2524 cx.update(|cx| connection.clone().close_session(&session_id, cx))
2525 .await
2526 .unwrap();
2527 drop(thread);
2528 drop(acp_thread);
2529 agent.read_with(cx, |agent, _| {
2530 assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
2531 });
2532
2533 // Ensure the thread can be reloaded from disk.
2534 assert_eq!(
2535 thread_entries(&thread_store, cx),
2536 vec![(
2537 session_id.clone(),
2538 format!("Explaining {}", path!("/a/b.md"))
2539 )]
2540 );
2541 let acp_thread = agent
2542 .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
2543 .await
2544 .unwrap();
2545 acp_thread.read_with(cx, |thread, cx| {
2546 assert_eq!(
2547 thread.to_markdown(cx),
2548 formatdoc! {"
2549 ## User
2550
2551 What does [@b.md]({uri}) mean?
2552
2553 ## Assistant
2554
2555 Lorem.
2556
2557 "}
2558 )
2559 });
2560 }
2561
2562 fn thread_entries(
2563 thread_store: &Entity<ThreadStore>,
2564 cx: &mut TestAppContext,
2565 ) -> Vec<(acp::SessionId, String)> {
2566 thread_store.read_with(cx, |store, _| {
2567 store
2568 .entries()
2569 .map(|entry| (entry.id.clone(), entry.title.to_string()))
2570 .collect::<Vec<_>>()
2571 })
2572 }
2573
2574 fn init_test(cx: &mut TestAppContext) {
2575 env_logger::try_init().ok();
2576 cx.update(|cx| {
2577 let settings_store = SettingsStore::test(cx);
2578 cx.set_global(settings_store);
2579
2580 LanguageModelRegistry::test(cx);
2581 });
2582 }
2583}
2584
2585fn mcp_message_content_to_acp_content_block(
2586 content: context_server::types::MessageContent,
2587) -> acp::ContentBlock {
2588 match content {
2589 context_server::types::MessageContent::Text {
2590 text,
2591 annotations: _,
2592 } => text.into(),
2593 context_server::types::MessageContent::Image {
2594 data,
2595 mime_type,
2596 annotations: _,
2597 } => acp::ContentBlock::Image(acp::ImageContent::new(data, mime_type)),
2598 context_server::types::MessageContent::Audio {
2599 data,
2600 mime_type,
2601 annotations: _,
2602 } => acp::ContentBlock::Audio(acp::AudioContent::new(data, mime_type)),
2603 context_server::types::MessageContent::Resource {
2604 resource,
2605 annotations: _,
2606 } => {
2607 let mut link =
2608 acp::ResourceLink::new(resource.uri.to_string(), resource.uri.to_string());
2609 if let Some(mime_type) = resource.mime_type {
2610 link = link.mime_type(mime_type);
2611 }
2612 acp::ContentBlock::ResourceLink(link)
2613 }
2614 }
2615}