1use crate::{
2 ContextServerRegistry, Thread, ThreadEvent, ThreadsDatabase, ToolCallAuthorization,
3 UserMessageContent, templates::Templates,
4};
5use crate::{HistoryStore, TokenUsageUpdated};
6use acp_thread::{AcpThread, AgentModelSelector};
7use action_log::ActionLog;
8use agent_client_protocol as acp;
9use agent_settings::AgentSettings;
10use anyhow::{Context as _, Result, anyhow};
11use collections::{HashSet, IndexMap};
12use fs::Fs;
13use futures::channel::mpsc;
14use futures::{StreamExt, future};
15use gpui::{
16 App, AppContext, AsyncApp, Context, Entity, SharedString, Subscription, Task, WeakEntity,
17};
18use language_model::{LanguageModel, LanguageModelProvider, LanguageModelRegistry};
19use project::{Project, ProjectItem, ProjectPath, Worktree};
20use prompt_store::{
21 ProjectContext, PromptId, PromptStore, RulesFileContext, UserRulesContext, WorktreeContext,
22};
23use settings::update_settings_file;
24use std::any::Any;
25use std::collections::HashMap;
26use std::path::Path;
27use std::rc::Rc;
28use std::sync::Arc;
29use util::ResultExt;
30
31const RULES_FILE_NAMES: [&str; 9] = [
32 ".rules",
33 ".cursorrules",
34 ".windsurfrules",
35 ".clinerules",
36 ".github/copilot-instructions.md",
37 "CLAUDE.md",
38 "AGENT.md",
39 "AGENTS.md",
40 "GEMINI.md",
41];
42
43pub struct RulesLoadingError {
44 pub message: SharedString,
45}
46
47/// Holds both the internal Thread and the AcpThread for a session
48struct Session {
49 /// The internal thread that processes messages
50 thread: Entity<Thread>,
51 /// The ACP thread that handles protocol communication
52 acp_thread: WeakEntity<acp_thread::AcpThread>,
53 pending_save: Task<()>,
54 _subscriptions: Vec<Subscription>,
55}
56
57pub struct LanguageModels {
58 /// Access language model by ID
59 models: HashMap<acp_thread::AgentModelId, Arc<dyn LanguageModel>>,
60 /// Cached list for returning language model information
61 model_list: acp_thread::AgentModelList,
62 refresh_models_rx: watch::Receiver<()>,
63 refresh_models_tx: watch::Sender<()>,
64}
65
66impl LanguageModels {
67 fn new(cx: &App) -> Self {
68 let (refresh_models_tx, refresh_models_rx) = watch::channel(());
69 let mut this = Self {
70 models: HashMap::default(),
71 model_list: acp_thread::AgentModelList::Grouped(IndexMap::default()),
72 refresh_models_rx,
73 refresh_models_tx,
74 };
75 this.refresh_list(cx);
76 this
77 }
78
79 fn refresh_list(&mut self, cx: &App) {
80 let providers = LanguageModelRegistry::global(cx)
81 .read(cx)
82 .providers()
83 .into_iter()
84 .filter(|provider| provider.is_authenticated(cx))
85 .collect::<Vec<_>>();
86
87 let mut language_model_list = IndexMap::default();
88 let mut recommended_models = HashSet::default();
89
90 let mut recommended = Vec::new();
91 for provider in &providers {
92 for model in provider.recommended_models(cx) {
93 recommended_models.insert(model.id());
94 recommended.push(Self::map_language_model_to_info(&model, provider));
95 }
96 }
97 if !recommended.is_empty() {
98 language_model_list.insert(
99 acp_thread::AgentModelGroupName("Recommended".into()),
100 recommended,
101 );
102 }
103
104 let mut models = HashMap::default();
105 for provider in providers {
106 let mut provider_models = Vec::new();
107 for model in provider.provided_models(cx) {
108 let model_info = Self::map_language_model_to_info(&model, &provider);
109 let model_id = model_info.id.clone();
110 if !recommended_models.contains(&model.id()) {
111 provider_models.push(model_info);
112 }
113 models.insert(model_id, model);
114 }
115 if !provider_models.is_empty() {
116 language_model_list.insert(
117 acp_thread::AgentModelGroupName(provider.name().0.clone()),
118 provider_models,
119 );
120 }
121 }
122
123 self.models = models;
124 self.model_list = acp_thread::AgentModelList::Grouped(language_model_list);
125 self.refresh_models_tx.send(()).ok();
126 }
127
128 fn watch(&self) -> watch::Receiver<()> {
129 self.refresh_models_rx.clone()
130 }
131
132 pub fn model_from_id(
133 &self,
134 model_id: &acp_thread::AgentModelId,
135 ) -> Option<Arc<dyn LanguageModel>> {
136 self.models.get(model_id).cloned()
137 }
138
139 fn map_language_model_to_info(
140 model: &Arc<dyn LanguageModel>,
141 provider: &Arc<dyn LanguageModelProvider>,
142 ) -> acp_thread::AgentModelInfo {
143 acp_thread::AgentModelInfo {
144 id: Self::model_id(model),
145 name: model.name().0,
146 icon: Some(provider.icon()),
147 }
148 }
149
150 fn model_id(model: &Arc<dyn LanguageModel>) -> acp_thread::AgentModelId {
151 acp_thread::AgentModelId(format!("{}/{}", model.provider_id().0, model.id().0).into())
152 }
153}
154
155pub struct NativeAgent {
156 /// Session ID -> Session mapping
157 sessions: HashMap<acp::SessionId, Session>,
158 history: Entity<HistoryStore>,
159 /// Shared project context for all threads
160 project_context: Entity<ProjectContext>,
161 project_context_needs_refresh: watch::Sender<()>,
162 _maintain_project_context: Task<Result<()>>,
163 context_server_registry: Entity<ContextServerRegistry>,
164 /// Shared templates for all threads
165 templates: Arc<Templates>,
166 /// Cached model information
167 models: LanguageModels,
168 project: Entity<Project>,
169 prompt_store: Option<Entity<PromptStore>>,
170 fs: Arc<dyn Fs>,
171 _subscriptions: Vec<Subscription>,
172}
173
174impl NativeAgent {
175 pub async fn new(
176 project: Entity<Project>,
177 history: Entity<HistoryStore>,
178 templates: Arc<Templates>,
179 prompt_store: Option<Entity<PromptStore>>,
180 fs: Arc<dyn Fs>,
181 cx: &mut AsyncApp,
182 ) -> Result<Entity<NativeAgent>> {
183 log::info!("Creating new NativeAgent");
184
185 let project_context = cx
186 .update(|cx| Self::build_project_context(&project, prompt_store.as_ref(), cx))?
187 .await;
188
189 cx.new(|cx| {
190 let mut subscriptions = vec![
191 cx.subscribe(&project, Self::handle_project_event),
192 cx.subscribe(
193 &LanguageModelRegistry::global(cx),
194 Self::handle_models_updated_event,
195 ),
196 ];
197 if let Some(prompt_store) = prompt_store.as_ref() {
198 subscriptions.push(cx.subscribe(prompt_store, Self::handle_prompts_updated_event))
199 }
200
201 let (project_context_needs_refresh_tx, project_context_needs_refresh_rx) =
202 watch::channel(());
203 Self {
204 sessions: HashMap::new(),
205 history,
206 project_context: cx.new(|_| project_context),
207 project_context_needs_refresh: project_context_needs_refresh_tx,
208 _maintain_project_context: cx.spawn(async move |this, cx| {
209 Self::maintain_project_context(this, project_context_needs_refresh_rx, cx).await
210 }),
211 context_server_registry: cx.new(|cx| {
212 ContextServerRegistry::new(project.read(cx).context_server_store(), cx)
213 }),
214 templates,
215 models: LanguageModels::new(cx),
216 project,
217 prompt_store,
218 fs,
219 _subscriptions: subscriptions,
220 }
221 })
222 }
223
224 fn register_session(
225 &mut self,
226 thread_handle: Entity<Thread>,
227 cx: &mut Context<Self>,
228 ) -> Entity<AcpThread> {
229 let connection = Rc::new(NativeAgentConnection(cx.entity()));
230 let registry = LanguageModelRegistry::read_global(cx);
231 let summarization_model = registry.thread_summary_model().map(|c| c.model);
232
233 thread_handle.update(cx, |thread, cx| {
234 thread.set_summarization_model(summarization_model, cx);
235 thread.add_default_tools(cx)
236 });
237
238 let thread = thread_handle.read(cx);
239 let session_id = thread.id().clone();
240 let title = thread.title();
241 let project = thread.project.clone();
242 let action_log = thread.action_log.clone();
243 let acp_thread = cx.new(|_cx| {
244 acp_thread::AcpThread::new(
245 title,
246 connection,
247 project.clone(),
248 action_log.clone(),
249 session_id.clone(),
250 )
251 });
252 let subscriptions = vec![
253 cx.observe_release(&acp_thread, |this, acp_thread, _cx| {
254 this.sessions.remove(acp_thread.session_id());
255 }),
256 cx.subscribe(&thread_handle, Self::handle_thread_token_usage_updated),
257 cx.observe(&thread_handle, move |this, thread, cx| {
258 this.save_thread(thread, cx)
259 }),
260 ];
261
262 self.sessions.insert(
263 session_id,
264 Session {
265 thread: thread_handle,
266 acp_thread: acp_thread.downgrade(),
267 _subscriptions: subscriptions,
268 pending_save: Task::ready(()),
269 },
270 );
271 acp_thread
272 }
273
274 pub fn models(&self) -> &LanguageModels {
275 &self.models
276 }
277
278 async fn maintain_project_context(
279 this: WeakEntity<Self>,
280 mut needs_refresh: watch::Receiver<()>,
281 cx: &mut AsyncApp,
282 ) -> Result<()> {
283 while needs_refresh.changed().await.is_ok() {
284 let project_context = this
285 .update(cx, |this, cx| {
286 Self::build_project_context(&this.project, this.prompt_store.as_ref(), cx)
287 })?
288 .await;
289 this.update(cx, |this, cx| {
290 this.project_context = cx.new(|_| project_context);
291 })?;
292 }
293
294 Ok(())
295 }
296
297 fn build_project_context(
298 project: &Entity<Project>,
299 prompt_store: Option<&Entity<PromptStore>>,
300 cx: &mut App,
301 ) -> Task<ProjectContext> {
302 let worktrees = project.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
303 let worktree_tasks = worktrees
304 .into_iter()
305 .map(|worktree| {
306 Self::load_worktree_info_for_system_prompt(worktree, project.clone(), cx)
307 })
308 .collect::<Vec<_>>();
309 let default_user_rules_task = if let Some(prompt_store) = prompt_store.as_ref() {
310 prompt_store.read_with(cx, |prompt_store, cx| {
311 let prompts = prompt_store.default_prompt_metadata();
312 let load_tasks = prompts.into_iter().map(|prompt_metadata| {
313 let contents = prompt_store.load(prompt_metadata.id, cx);
314 async move { (contents.await, prompt_metadata) }
315 });
316 cx.background_spawn(future::join_all(load_tasks))
317 })
318 } else {
319 Task::ready(vec![])
320 };
321
322 cx.spawn(async move |_cx| {
323 let (worktrees, default_user_rules) =
324 future::join(future::join_all(worktree_tasks), default_user_rules_task).await;
325
326 let worktrees = worktrees
327 .into_iter()
328 .map(|(worktree, _rules_error)| {
329 // TODO: show error message
330 // if let Some(rules_error) = rules_error {
331 // this.update(cx, |_, cx| cx.emit(rules_error)).ok();
332 // }
333 worktree
334 })
335 .collect::<Vec<_>>();
336
337 let default_user_rules = default_user_rules
338 .into_iter()
339 .flat_map(|(contents, prompt_metadata)| match contents {
340 Ok(contents) => Some(UserRulesContext {
341 uuid: match prompt_metadata.id {
342 PromptId::User { uuid } => uuid,
343 PromptId::EditWorkflow => return None,
344 },
345 title: prompt_metadata.title.map(|title| title.to_string()),
346 contents,
347 }),
348 Err(_err) => {
349 // TODO: show error message
350 // this.update(cx, |_, cx| {
351 // cx.emit(RulesLoadingError {
352 // message: format!("{err:?}").into(),
353 // });
354 // })
355 // .ok();
356 None
357 }
358 })
359 .collect::<Vec<_>>();
360
361 ProjectContext::new(worktrees, default_user_rules)
362 })
363 }
364
365 fn load_worktree_info_for_system_prompt(
366 worktree: Entity<Worktree>,
367 project: Entity<Project>,
368 cx: &mut App,
369 ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
370 let tree = worktree.read(cx);
371 let root_name = tree.root_name().into();
372 let abs_path = tree.abs_path();
373
374 let mut context = WorktreeContext {
375 root_name,
376 abs_path,
377 rules_file: None,
378 };
379
380 let rules_task = Self::load_worktree_rules_file(worktree, project, cx);
381 let Some(rules_task) = rules_task else {
382 return Task::ready((context, None));
383 };
384
385 cx.spawn(async move |_| {
386 let (rules_file, rules_file_error) = match rules_task.await {
387 Ok(rules_file) => (Some(rules_file), None),
388 Err(err) => (
389 None,
390 Some(RulesLoadingError {
391 message: format!("{err}").into(),
392 }),
393 ),
394 };
395 context.rules_file = rules_file;
396 (context, rules_file_error)
397 })
398 }
399
400 fn load_worktree_rules_file(
401 worktree: Entity<Worktree>,
402 project: Entity<Project>,
403 cx: &mut App,
404 ) -> Option<Task<Result<RulesFileContext>>> {
405 let worktree = worktree.read(cx);
406 let worktree_id = worktree.id();
407 let selected_rules_file = RULES_FILE_NAMES
408 .into_iter()
409 .filter_map(|name| {
410 worktree
411 .entry_for_path(name)
412 .filter(|entry| entry.is_file())
413 .map(|entry| entry.path.clone())
414 })
415 .next();
416
417 // Note that Cline supports `.clinerules` being a directory, but that is not currently
418 // supported. This doesn't seem to occur often in GitHub repositories.
419 selected_rules_file.map(|path_in_worktree| {
420 let project_path = ProjectPath {
421 worktree_id,
422 path: path_in_worktree.clone(),
423 };
424 let buffer_task =
425 project.update(cx, |project, cx| project.open_buffer(project_path, cx));
426 let rope_task = cx.spawn(async move |cx| {
427 buffer_task.await?.read_with(cx, |buffer, cx| {
428 let project_entry_id = buffer.entry_id(cx).context("buffer has no file")?;
429 anyhow::Ok((project_entry_id, buffer.as_rope().clone()))
430 })?
431 });
432 // Build a string from the rope on a background thread.
433 cx.background_spawn(async move {
434 let (project_entry_id, rope) = rope_task.await?;
435 anyhow::Ok(RulesFileContext {
436 path_in_worktree,
437 text: rope.to_string().trim().to_string(),
438 project_entry_id: project_entry_id.to_usize(),
439 })
440 })
441 })
442 }
443
444 fn handle_thread_token_usage_updated(
445 &mut self,
446 thread: Entity<Thread>,
447 usage: &TokenUsageUpdated,
448 cx: &mut Context<Self>,
449 ) {
450 let Some(session) = self.sessions.get(thread.read(cx).id()) else {
451 return;
452 };
453 session
454 .acp_thread
455 .update(cx, |acp_thread, cx| {
456 acp_thread.update_token_usage(usage.0.clone(), cx);
457 })
458 .ok();
459 }
460
461 fn handle_project_event(
462 &mut self,
463 _project: Entity<Project>,
464 event: &project::Event,
465 _cx: &mut Context<Self>,
466 ) {
467 match event {
468 project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
469 self.project_context_needs_refresh.send(()).ok();
470 }
471 project::Event::WorktreeUpdatedEntries(_, items) => {
472 if items.iter().any(|(path, _, _)| {
473 RULES_FILE_NAMES
474 .iter()
475 .any(|name| path.as_ref() == Path::new(name))
476 }) {
477 self.project_context_needs_refresh.send(()).ok();
478 }
479 }
480 _ => {}
481 }
482 }
483
484 fn handle_prompts_updated_event(
485 &mut self,
486 _prompt_store: Entity<PromptStore>,
487 _event: &prompt_store::PromptsUpdatedEvent,
488 _cx: &mut Context<Self>,
489 ) {
490 self.project_context_needs_refresh.send(()).ok();
491 }
492
493 fn handle_models_updated_event(
494 &mut self,
495 _registry: Entity<LanguageModelRegistry>,
496 _event: &language_model::Event,
497 cx: &mut Context<Self>,
498 ) {
499 self.models.refresh_list(cx);
500
501 let registry = LanguageModelRegistry::read_global(cx);
502 let default_model = registry.default_model().map(|m| m.model);
503 let summarization_model = registry.thread_summary_model().map(|m| m.model);
504
505 for session in self.sessions.values_mut() {
506 session.thread.update(cx, |thread, cx| {
507 if thread.model().is_none()
508 && let Some(model) = default_model.clone()
509 {
510 thread.set_model(model, cx);
511 cx.notify();
512 }
513 thread.set_summarization_model(summarization_model.clone(), cx);
514 });
515 }
516 }
517
518 pub fn open_thread(
519 &mut self,
520 id: acp::SessionId,
521 cx: &mut Context<Self>,
522 ) -> Task<Result<Entity<AcpThread>>> {
523 let database_future = ThreadsDatabase::connect(cx);
524 cx.spawn(async move |this, cx| {
525 let database = database_future.await.map_err(|err| anyhow!(err))?;
526 let db_thread = database
527 .load_thread(id.clone())
528 .await?
529 .with_context(|| format!("no thread found with ID: {id:?}"))?;
530
531 let thread = this.update(cx, |this, cx| {
532 let action_log = cx.new(|_cx| ActionLog::new(this.project.clone()));
533 cx.new(|cx| {
534 Thread::from_db(
535 id.clone(),
536 db_thread,
537 this.project.clone(),
538 this.project_context.clone(),
539 this.context_server_registry.clone(),
540 action_log.clone(),
541 this.templates.clone(),
542 cx,
543 )
544 })
545 })?;
546 let acp_thread =
547 this.update(cx, |this, cx| this.register_session(thread.clone(), cx))?;
548 let events = thread.update(cx, |thread, cx| thread.replay(cx))?;
549 cx.update(|cx| {
550 NativeAgentConnection::handle_thread_events(events, acp_thread.downgrade(), cx)
551 })?
552 .await?;
553 Ok(acp_thread)
554 })
555 }
556
557 pub fn thread_summary(
558 &mut self,
559 id: acp::SessionId,
560 cx: &mut Context<Self>,
561 ) -> Task<Result<SharedString>> {
562 let thread = self.open_thread(id.clone(), cx);
563 cx.spawn(async move |this, cx| {
564 let acp_thread = thread.await?;
565 let result = this
566 .update(cx, |this, cx| {
567 this.sessions
568 .get(&id)
569 .unwrap()
570 .thread
571 .update(cx, |thread, cx| thread.summary(cx))
572 })?
573 .await?;
574 drop(acp_thread);
575 Ok(result)
576 })
577 }
578
579 fn save_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
580 if thread.read(cx).is_empty() {
581 return;
582 }
583
584 let database_future = ThreadsDatabase::connect(cx);
585 let (id, db_thread) =
586 thread.update(cx, |thread, cx| (thread.id().clone(), thread.to_db(cx)));
587 let Some(session) = self.sessions.get_mut(&id) else {
588 return;
589 };
590 let history = self.history.clone();
591 session.pending_save = cx.spawn(async move |_, cx| {
592 let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
593 return;
594 };
595 let db_thread = db_thread.await;
596 database.save_thread(id, db_thread).await.log_err();
597 history.update(cx, |history, cx| history.reload(cx)).ok();
598 });
599 }
600}
601
602/// Wrapper struct that implements the AgentConnection trait
603#[derive(Clone)]
604pub struct NativeAgentConnection(pub Entity<NativeAgent>);
605
606impl NativeAgentConnection {
607 pub fn thread(&self, session_id: &acp::SessionId, cx: &App) -> Option<Entity<Thread>> {
608 self.0
609 .read(cx)
610 .sessions
611 .get(session_id)
612 .map(|session| session.thread.clone())
613 }
614
615 fn run_turn(
616 &self,
617 session_id: acp::SessionId,
618 cx: &mut App,
619 f: impl 'static
620 + FnOnce(Entity<Thread>, &mut App) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>,
621 ) -> Task<Result<acp::PromptResponse>> {
622 let Some((thread, acp_thread)) = self.0.update(cx, |agent, _cx| {
623 agent
624 .sessions
625 .get_mut(&session_id)
626 .map(|s| (s.thread.clone(), s.acp_thread.clone()))
627 }) else {
628 return Task::ready(Err(anyhow!("Session not found")));
629 };
630 log::debug!("Found session for: {}", session_id);
631
632 let response_stream = match f(thread, cx) {
633 Ok(stream) => stream,
634 Err(err) => return Task::ready(Err(err)),
635 };
636 Self::handle_thread_events(response_stream, acp_thread, cx)
637 }
638
639 fn handle_thread_events(
640 mut events: mpsc::UnboundedReceiver<Result<ThreadEvent>>,
641 acp_thread: WeakEntity<AcpThread>,
642 cx: &App,
643 ) -> Task<Result<acp::PromptResponse>> {
644 cx.spawn(async move |cx| {
645 // Handle response stream and forward to session.acp_thread
646 while let Some(result) = events.next().await {
647 match result {
648 Ok(event) => {
649 log::trace!("Received completion event: {:?}", event);
650
651 match event {
652 ThreadEvent::UserMessage(message) => {
653 acp_thread.update(cx, |thread, cx| {
654 for content in message.content {
655 thread.push_user_content_block(
656 Some(message.id.clone()),
657 content.into(),
658 cx,
659 );
660 }
661 })?;
662 }
663 ThreadEvent::AgentText(text) => {
664 acp_thread.update(cx, |thread, cx| {
665 thread.push_assistant_content_block(
666 acp::ContentBlock::Text(acp::TextContent {
667 text,
668 annotations: None,
669 }),
670 false,
671 cx,
672 )
673 })?;
674 }
675 ThreadEvent::AgentThinking(text) => {
676 acp_thread.update(cx, |thread, cx| {
677 thread.push_assistant_content_block(
678 acp::ContentBlock::Text(acp::TextContent {
679 text,
680 annotations: None,
681 }),
682 true,
683 cx,
684 )
685 })?;
686 }
687 ThreadEvent::ToolCallAuthorization(ToolCallAuthorization {
688 tool_call,
689 options,
690 response,
691 }) => {
692 let recv = acp_thread.update(cx, |thread, cx| {
693 thread.request_tool_call_authorization(tool_call, options, cx)
694 })?;
695 cx.background_spawn(async move {
696 if let Some(recv) = recv.log_err()
697 && let Some(option) = recv
698 .await
699 .context("authorization sender was dropped")
700 .log_err()
701 {
702 response
703 .send(option)
704 .map(|_| anyhow!("authorization receiver was dropped"))
705 .log_err();
706 }
707 })
708 .detach();
709 }
710 ThreadEvent::ToolCall(tool_call) => {
711 acp_thread.update(cx, |thread, cx| {
712 thread.upsert_tool_call(tool_call, cx)
713 })??;
714 }
715 ThreadEvent::ToolCallUpdate(update) => {
716 acp_thread.update(cx, |thread, cx| {
717 thread.update_tool_call(update, cx)
718 })??;
719 }
720 ThreadEvent::TitleUpdate(title) => {
721 acp_thread
722 .update(cx, |thread, cx| thread.update_title(title, cx))??;
723 }
724 ThreadEvent::Retry(status) => {
725 acp_thread.update(cx, |thread, cx| {
726 thread.update_retry_status(status, cx)
727 })?;
728 }
729 ThreadEvent::Stop(stop_reason) => {
730 log::debug!("Assistant message complete: {:?}", stop_reason);
731 return Ok(acp::PromptResponse { stop_reason });
732 }
733 }
734 }
735 Err(e) => {
736 log::error!("Error in model response stream: {:?}", e);
737 return Err(e);
738 }
739 }
740 }
741
742 log::info!("Response stream completed");
743 anyhow::Ok(acp::PromptResponse {
744 stop_reason: acp::StopReason::EndTurn,
745 })
746 })
747 }
748}
749
750impl AgentModelSelector for NativeAgentConnection {
751 fn list_models(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelList>> {
752 log::debug!("NativeAgentConnection::list_models called");
753 let list = self.0.read(cx).models.model_list.clone();
754 Task::ready(if list.is_empty() {
755 Err(anyhow::anyhow!("No models available"))
756 } else {
757 Ok(list)
758 })
759 }
760
761 fn select_model(
762 &self,
763 session_id: acp::SessionId,
764 model_id: acp_thread::AgentModelId,
765 cx: &mut App,
766 ) -> Task<Result<()>> {
767 log::info!("Setting model for session {}: {}", session_id, model_id);
768 let Some(thread) = self
769 .0
770 .read(cx)
771 .sessions
772 .get(&session_id)
773 .map(|session| session.thread.clone())
774 else {
775 return Task::ready(Err(anyhow!("Session not found")));
776 };
777
778 let Some(model) = self.0.read(cx).models.model_from_id(&model_id) else {
779 return Task::ready(Err(anyhow!("Invalid model ID {}", model_id)));
780 };
781
782 thread.update(cx, |thread, cx| {
783 thread.set_model(model.clone(), cx);
784 });
785
786 update_settings_file::<AgentSettings>(
787 self.0.read(cx).fs.clone(),
788 cx,
789 move |settings, _cx| {
790 settings.set_model(model);
791 },
792 );
793
794 Task::ready(Ok(()))
795 }
796
797 fn selected_model(
798 &self,
799 session_id: &acp::SessionId,
800 cx: &mut App,
801 ) -> Task<Result<acp_thread::AgentModelInfo>> {
802 let session_id = session_id.clone();
803
804 let Some(thread) = self
805 .0
806 .read(cx)
807 .sessions
808 .get(&session_id)
809 .map(|session| session.thread.clone())
810 else {
811 return Task::ready(Err(anyhow!("Session not found")));
812 };
813 let Some(model) = thread.read(cx).model() else {
814 return Task::ready(Err(anyhow!("Model not found")));
815 };
816 let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&model.provider_id())
817 else {
818 return Task::ready(Err(anyhow!("Provider not found")));
819 };
820 Task::ready(Ok(LanguageModels::map_language_model_to_info(
821 model, &provider,
822 )))
823 }
824
825 fn watch(&self, cx: &mut App) -> watch::Receiver<()> {
826 self.0.read(cx).models.watch()
827 }
828}
829
830impl acp_thread::AgentConnection for NativeAgentConnection {
831 fn new_thread(
832 self: Rc<Self>,
833 project: Entity<Project>,
834 cwd: &Path,
835 cx: &mut App,
836 ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
837 let agent = self.0.clone();
838 log::info!("Creating new thread for project at: {:?}", cwd);
839
840 cx.spawn(async move |cx| {
841 log::debug!("Starting thread creation in async context");
842
843 let action_log = cx.new(|_cx| ActionLog::new(project.clone()))?;
844 // Create Thread
845 let thread = agent.update(
846 cx,
847 |agent, cx: &mut gpui::Context<NativeAgent>| -> Result<_> {
848 // Fetch default model from registry settings
849 let registry = LanguageModelRegistry::read_global(cx);
850 // Log available models for debugging
851 let available_count = registry.available_models(cx).count();
852 log::debug!("Total available models: {}", available_count);
853
854 let default_model = registry.default_model().and_then(|default_model| {
855 agent
856 .models
857 .model_from_id(&LanguageModels::model_id(&default_model.model))
858 });
859
860 let thread = cx.new(|cx| {
861 Thread::new(
862 project.clone(),
863 agent.project_context.clone(),
864 agent.context_server_registry.clone(),
865 action_log.clone(),
866 agent.templates.clone(),
867 default_model,
868 cx,
869 )
870 });
871
872 Ok(thread)
873 },
874 )??;
875 agent.update(cx, |agent, cx| agent.register_session(thread, cx))
876 })
877 }
878
879 fn auth_methods(&self) -> &[acp::AuthMethod] {
880 &[] // No auth for in-process
881 }
882
883 fn authenticate(&self, _method: acp::AuthMethodId, _cx: &mut App) -> Task<Result<()>> {
884 Task::ready(Ok(()))
885 }
886
887 fn model_selector(&self) -> Option<Rc<dyn AgentModelSelector>> {
888 Some(Rc::new(self.clone()) as Rc<dyn AgentModelSelector>)
889 }
890
891 fn prompt(
892 &self,
893 id: Option<acp_thread::UserMessageId>,
894 params: acp::PromptRequest,
895 cx: &mut App,
896 ) -> Task<Result<acp::PromptResponse>> {
897 let id = id.expect("UserMessageId is required");
898 let session_id = params.session_id.clone();
899 log::info!("Received prompt request for session: {}", session_id);
900 log::debug!("Prompt blocks count: {}", params.prompt.len());
901
902 self.run_turn(session_id, cx, |thread, cx| {
903 let content: Vec<UserMessageContent> = params
904 .prompt
905 .into_iter()
906 .map(Into::into)
907 .collect::<Vec<_>>();
908 log::info!("Converted prompt to message: {} chars", content.len());
909 log::debug!("Message id: {:?}", id);
910 log::debug!("Message content: {:?}", content);
911
912 thread.update(cx, |thread, cx| thread.send(id, content, cx))
913 })
914 }
915
916 fn prompt_capabilities(&self) -> acp::PromptCapabilities {
917 acp::PromptCapabilities {
918 image: true,
919 audio: false,
920 embedded_context: true,
921 }
922 }
923
924 fn resume(
925 &self,
926 session_id: &acp::SessionId,
927 _cx: &mut App,
928 ) -> Option<Rc<dyn acp_thread::AgentSessionResume>> {
929 Some(Rc::new(NativeAgentSessionResume {
930 connection: self.clone(),
931 session_id: session_id.clone(),
932 }) as _)
933 }
934
935 fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
936 log::info!("Cancelling on session: {}", session_id);
937 self.0.update(cx, |agent, cx| {
938 if let Some(agent) = agent.sessions.get(session_id) {
939 agent.thread.update(cx, |thread, cx| thread.cancel(cx));
940 }
941 });
942 }
943
944 fn session_editor(
945 &self,
946 session_id: &agent_client_protocol::SessionId,
947 cx: &mut App,
948 ) -> Option<Rc<dyn acp_thread::AgentSessionEditor>> {
949 self.0.update(cx, |agent, _cx| {
950 agent.sessions.get(session_id).map(|session| {
951 Rc::new(NativeAgentSessionEditor {
952 thread: session.thread.clone(),
953 acp_thread: session.acp_thread.clone(),
954 }) as _
955 })
956 })
957 }
958
959 fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
960 Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
961 }
962
963 fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
964 self
965 }
966}
967
968impl acp_thread::AgentTelemetry for NativeAgentConnection {
969 fn agent_name(&self) -> String {
970 "Zed".into()
971 }
972
973 fn thread_data(
974 &self,
975 session_id: &acp::SessionId,
976 cx: &mut App,
977 ) -> Task<Result<serde_json::Value>> {
978 let Some(session) = self.0.read(cx).sessions.get(session_id) else {
979 return Task::ready(Err(anyhow!("Session not found")));
980 };
981
982 let task = session.thread.read(cx).to_db(cx);
983 cx.background_spawn(async move {
984 serde_json::to_value(task.await).context("Failed to serialize thread")
985 })
986 }
987}
988
989struct NativeAgentSessionEditor {
990 thread: Entity<Thread>,
991 acp_thread: WeakEntity<AcpThread>,
992}
993
994impl acp_thread::AgentSessionEditor for NativeAgentSessionEditor {
995 fn truncate(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
996 match self.thread.update(cx, |thread, cx| {
997 thread.truncate(message_id.clone(), cx)?;
998 Ok(thread.latest_token_usage())
999 }) {
1000 Ok(usage) => {
1001 self.acp_thread
1002 .update(cx, |thread, cx| {
1003 thread.update_token_usage(usage, cx);
1004 })
1005 .ok();
1006 Task::ready(Ok(()))
1007 }
1008 Err(error) => Task::ready(Err(error)),
1009 }
1010 }
1011}
1012
1013struct NativeAgentSessionResume {
1014 connection: NativeAgentConnection,
1015 session_id: acp::SessionId,
1016}
1017
1018impl acp_thread::AgentSessionResume for NativeAgentSessionResume {
1019 fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>> {
1020 self.connection
1021 .run_turn(self.session_id.clone(), cx, |thread, cx| {
1022 thread.update(cx, |thread, cx| thread.resume(cx))
1023 })
1024 }
1025}
1026
1027#[cfg(test)]
1028mod tests {
1029 use crate::HistoryEntryId;
1030
1031 use super::*;
1032 use acp_thread::{
1033 AgentConnection, AgentModelGroupName, AgentModelId, AgentModelInfo, MentionUri,
1034 };
1035 use fs::FakeFs;
1036 use gpui::TestAppContext;
1037 use indoc::indoc;
1038 use language_model::fake_provider::FakeLanguageModel;
1039 use serde_json::json;
1040 use settings::SettingsStore;
1041 use util::path;
1042
1043 #[gpui::test]
1044 async fn test_maintaining_project_context(cx: &mut TestAppContext) {
1045 init_test(cx);
1046 let fs = FakeFs::new(cx.executor());
1047 fs.insert_tree(
1048 "/",
1049 json!({
1050 "a": {}
1051 }),
1052 )
1053 .await;
1054 let project = Project::test(fs.clone(), [], cx).await;
1055 let context_store = cx.new(|cx| assistant_context::ContextStore::fake(project.clone(), cx));
1056 let history_store = cx.new(|cx| HistoryStore::new(context_store, cx));
1057 let agent = NativeAgent::new(
1058 project.clone(),
1059 history_store,
1060 Templates::new(),
1061 None,
1062 fs.clone(),
1063 &mut cx.to_async(),
1064 )
1065 .await
1066 .unwrap();
1067 agent.read_with(cx, |agent, cx| {
1068 assert_eq!(agent.project_context.read(cx).worktrees, vec![])
1069 });
1070
1071 let worktree = project
1072 .update(cx, |project, cx| project.create_worktree("/a", true, cx))
1073 .await
1074 .unwrap();
1075 cx.run_until_parked();
1076 agent.read_with(cx, |agent, cx| {
1077 assert_eq!(
1078 agent.project_context.read(cx).worktrees,
1079 vec![WorktreeContext {
1080 root_name: "a".into(),
1081 abs_path: Path::new("/a").into(),
1082 rules_file: None
1083 }]
1084 )
1085 });
1086
1087 // Creating `/a/.rules` updates the project context.
1088 fs.insert_file("/a/.rules", Vec::new()).await;
1089 cx.run_until_parked();
1090 agent.read_with(cx, |agent, cx| {
1091 let rules_entry = worktree.read(cx).entry_for_path(".rules").unwrap();
1092 assert_eq!(
1093 agent.project_context.read(cx).worktrees,
1094 vec![WorktreeContext {
1095 root_name: "a".into(),
1096 abs_path: Path::new("/a").into(),
1097 rules_file: Some(RulesFileContext {
1098 path_in_worktree: Path::new(".rules").into(),
1099 text: "".into(),
1100 project_entry_id: rules_entry.id.to_usize()
1101 })
1102 }]
1103 )
1104 });
1105 }
1106
1107 #[gpui::test]
1108 async fn test_listing_models(cx: &mut TestAppContext) {
1109 init_test(cx);
1110 let fs = FakeFs::new(cx.executor());
1111 fs.insert_tree("/", json!({ "a": {} })).await;
1112 let project = Project::test(fs.clone(), [], cx).await;
1113 let context_store = cx.new(|cx| assistant_context::ContextStore::fake(project.clone(), cx));
1114 let history_store = cx.new(|cx| HistoryStore::new(context_store, cx));
1115 let connection = NativeAgentConnection(
1116 NativeAgent::new(
1117 project.clone(),
1118 history_store,
1119 Templates::new(),
1120 None,
1121 fs.clone(),
1122 &mut cx.to_async(),
1123 )
1124 .await
1125 .unwrap(),
1126 );
1127
1128 let models = cx.update(|cx| connection.list_models(cx)).await.unwrap();
1129
1130 let acp_thread::AgentModelList::Grouped(models) = models else {
1131 panic!("Unexpected model group");
1132 };
1133 assert_eq!(
1134 models,
1135 IndexMap::from_iter([(
1136 AgentModelGroupName("Fake".into()),
1137 vec![AgentModelInfo {
1138 id: AgentModelId("fake/fake".into()),
1139 name: "Fake".into(),
1140 icon: Some(ui::IconName::ZedAssistant),
1141 }]
1142 )])
1143 );
1144 }
1145
1146 #[gpui::test]
1147 async fn test_model_selection_persists_to_settings(cx: &mut TestAppContext) {
1148 init_test(cx);
1149 let fs = FakeFs::new(cx.executor());
1150 fs.create_dir(paths::settings_file().parent().unwrap())
1151 .await
1152 .unwrap();
1153 fs.insert_file(
1154 paths::settings_file(),
1155 json!({
1156 "agent": {
1157 "default_model": {
1158 "provider": "foo",
1159 "model": "bar"
1160 }
1161 }
1162 })
1163 .to_string()
1164 .into_bytes(),
1165 )
1166 .await;
1167 let project = Project::test(fs.clone(), [], cx).await;
1168
1169 let context_store = cx.new(|cx| assistant_context::ContextStore::fake(project.clone(), cx));
1170 let history_store = cx.new(|cx| HistoryStore::new(context_store, cx));
1171
1172 // Create the agent and connection
1173 let agent = NativeAgent::new(
1174 project.clone(),
1175 history_store,
1176 Templates::new(),
1177 None,
1178 fs.clone(),
1179 &mut cx.to_async(),
1180 )
1181 .await
1182 .unwrap();
1183 let connection = NativeAgentConnection(agent.clone());
1184
1185 // Create a thread/session
1186 let acp_thread = cx
1187 .update(|cx| {
1188 Rc::new(connection.clone()).new_thread(project.clone(), Path::new("/a"), cx)
1189 })
1190 .await
1191 .unwrap();
1192
1193 let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1194
1195 // Select a model
1196 let model_id = AgentModelId("fake/fake".into());
1197 cx.update(|cx| connection.select_model(session_id.clone(), model_id.clone(), cx))
1198 .await
1199 .unwrap();
1200
1201 // Verify the thread has the selected model
1202 agent.read_with(cx, |agent, _| {
1203 let session = agent.sessions.get(&session_id).unwrap();
1204 session.thread.read_with(cx, |thread, _| {
1205 assert_eq!(thread.model().unwrap().id().0, "fake");
1206 });
1207 });
1208
1209 cx.run_until_parked();
1210
1211 // Verify settings file was updated
1212 let settings_content = fs.load(paths::settings_file()).await.unwrap();
1213 let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
1214
1215 // Check that the agent settings contain the selected model
1216 assert_eq!(
1217 settings_json["agent"]["default_model"]["model"],
1218 json!("fake")
1219 );
1220 assert_eq!(
1221 settings_json["agent"]["default_model"]["provider"],
1222 json!("fake")
1223 );
1224 }
1225
1226 #[gpui::test]
1227 #[cfg_attr(target_os = "windows", ignore)] // TODO: Fix this test on Windows
1228 async fn test_save_load_thread(cx: &mut TestAppContext) {
1229 init_test(cx);
1230 let fs = FakeFs::new(cx.executor());
1231 fs.insert_tree(
1232 "/",
1233 json!({
1234 "a": {
1235 "b.md": "Lorem"
1236 }
1237 }),
1238 )
1239 .await;
1240 let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
1241 let context_store = cx.new(|cx| assistant_context::ContextStore::fake(project.clone(), cx));
1242 let history_store = cx.new(|cx| HistoryStore::new(context_store, cx));
1243 let agent = NativeAgent::new(
1244 project.clone(),
1245 history_store.clone(),
1246 Templates::new(),
1247 None,
1248 fs.clone(),
1249 &mut cx.to_async(),
1250 )
1251 .await
1252 .unwrap();
1253 let connection = Rc::new(NativeAgentConnection(agent.clone()));
1254
1255 let acp_thread = cx
1256 .update(|cx| {
1257 connection
1258 .clone()
1259 .new_thread(project.clone(), Path::new(""), cx)
1260 })
1261 .await
1262 .unwrap();
1263 let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
1264 let thread = agent.read_with(cx, |agent, _| {
1265 agent.sessions.get(&session_id).unwrap().thread.clone()
1266 });
1267
1268 // Ensure empty threads are not saved, even if they get mutated.
1269 let model = Arc::new(FakeLanguageModel::default());
1270 let summary_model = Arc::new(FakeLanguageModel::default());
1271 thread.update(cx, |thread, cx| {
1272 thread.set_model(model.clone(), cx);
1273 thread.set_summarization_model(Some(summary_model.clone()), cx);
1274 });
1275 cx.run_until_parked();
1276 assert_eq!(history_entries(&history_store, cx), vec![]);
1277
1278 let send = acp_thread.update(cx, |thread, cx| {
1279 thread.send(
1280 vec![
1281 "What does ".into(),
1282 acp::ContentBlock::ResourceLink(acp::ResourceLink {
1283 name: "b.md".into(),
1284 uri: MentionUri::File {
1285 abs_path: path!("/a/b.md").into(),
1286 }
1287 .to_uri()
1288 .to_string(),
1289 annotations: None,
1290 description: None,
1291 mime_type: None,
1292 size: None,
1293 title: None,
1294 }),
1295 " mean?".into(),
1296 ],
1297 cx,
1298 )
1299 });
1300 let send = cx.foreground_executor().spawn(send);
1301 cx.run_until_parked();
1302
1303 model.send_last_completion_stream_text_chunk("Lorem.");
1304 model.end_last_completion_stream();
1305 cx.run_until_parked();
1306 summary_model.send_last_completion_stream_text_chunk("Explaining /a/b.md");
1307 summary_model.end_last_completion_stream();
1308
1309 send.await.unwrap();
1310 acp_thread.read_with(cx, |thread, cx| {
1311 assert_eq!(
1312 thread.to_markdown(cx),
1313 indoc! {"
1314 ## User
1315
1316 What does [@b.md](file:///a/b.md) mean?
1317
1318 ## Assistant
1319
1320 Lorem.
1321
1322 "}
1323 )
1324 });
1325
1326 // Drop the ACP thread, which should cause the session to be dropped as well.
1327 cx.update(|_| {
1328 drop(thread);
1329 drop(acp_thread);
1330 });
1331 agent.read_with(cx, |agent, _| {
1332 assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
1333 });
1334
1335 // Ensure the thread can be reloaded from disk.
1336 assert_eq!(
1337 history_entries(&history_store, cx),
1338 vec![(
1339 HistoryEntryId::AcpThread(session_id.clone()),
1340 "Explaining /a/b.md".into()
1341 )]
1342 );
1343 let acp_thread = agent
1344 .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
1345 .await
1346 .unwrap();
1347 acp_thread.read_with(cx, |thread, cx| {
1348 assert_eq!(
1349 thread.to_markdown(cx),
1350 indoc! {"
1351 ## User
1352
1353 What does [@b.md](file:///a/b.md) mean?
1354
1355 ## Assistant
1356
1357 Lorem.
1358
1359 "}
1360 )
1361 });
1362 }
1363
1364 fn history_entries(
1365 history: &Entity<HistoryStore>,
1366 cx: &mut TestAppContext,
1367 ) -> Vec<(HistoryEntryId, String)> {
1368 history.read_with(cx, |history, cx| {
1369 history
1370 .entries(cx)
1371 .iter()
1372 .map(|e| (e.id(), e.title().to_string()))
1373 .collect::<Vec<_>>()
1374 })
1375 }
1376
1377 fn init_test(cx: &mut TestAppContext) {
1378 env_logger::try_init().ok();
1379 cx.update(|cx| {
1380 let settings_store = SettingsStore::test(cx);
1381 cx.set_global(settings_store);
1382 Project::init_settings(cx);
1383 agent_settings::init(cx);
1384 language::init(cx);
1385 LanguageModelRegistry::test(cx);
1386 });
1387 }
1388}