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