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