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