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