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