agent.rs

   1mod db;
   2mod edit_agent;
   3mod history_store;
   4mod legacy_thread;
   5mod native_agent_server;
   6pub mod outline;
   7mod templates;
   8mod thread;
   9mod tool_schema;
  10mod tools;
  11
  12#[cfg(test)]
  13mod tests;
  14
  15pub use db::*;
  16pub use history_store::*;
  17pub use native_agent_server::NativeAgentServer;
  18pub use templates::*;
  19pub use thread::*;
  20pub use tools::*;
  21
  22use acp_thread::{AcpThread, AgentModelSelector};
  23use agent_client_protocol as acp;
  24use anyhow::{Context as _, Result, anyhow};
  25use chrono::{DateTime, Utc};
  26use collections::{HashSet, IndexMap};
  27use fs::Fs;
  28use futures::channel::{mpsc, oneshot};
  29use futures::future::Shared;
  30use futures::{StreamExt, future};
  31use gpui::{
  32    App, AppContext, AsyncApp, Context, Entity, SharedString, Subscription, Task, WeakEntity,
  33};
  34use language_model::{LanguageModel, LanguageModelProvider, LanguageModelRegistry};
  35use project::{Project, ProjectItem, ProjectPath, Worktree};
  36use prompt_store::{
  37    ProjectContext, PromptStore, RulesFileContext, UserRulesContext, WorktreeContext,
  38};
  39use serde::{Deserialize, Serialize};
  40use settings::{LanguageModelSelection, update_settings_file};
  41use std::any::Any;
  42use std::collections::HashMap;
  43use std::path::{Path, PathBuf};
  44use std::rc::Rc;
  45use std::sync::Arc;
  46use util::ResultExt;
  47use util::rel_path::RelPath;
  48
  49#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
  50pub struct ProjectSnapshot {
  51    pub worktree_snapshots: Vec<project::telemetry_snapshot::TelemetryWorktreeSnapshot>,
  52    pub timestamp: DateTime<Utc>,
  53}
  54
  55const RULES_FILE_NAMES: [&str; 9] = [
  56    ".rules",
  57    ".cursorrules",
  58    ".windsurfrules",
  59    ".clinerules",
  60    ".github/copilot-instructions.md",
  61    "CLAUDE.md",
  62    "AGENT.md",
  63    "AGENTS.md",
  64    "GEMINI.md",
  65];
  66
  67pub struct RulesLoadingError {
  68    pub message: SharedString,
  69}
  70
  71/// Holds both the internal Thread and the AcpThread for a session
  72struct Session {
  73    /// The internal thread that processes messages
  74    thread: Entity<Thread>,
  75    /// The ACP thread that handles protocol communication
  76    acp_thread: WeakEntity<acp_thread::AcpThread>,
  77    pending_save: Task<()>,
  78    _subscriptions: Vec<Subscription>,
  79}
  80
  81pub struct LanguageModels {
  82    /// Access language model by ID
  83    models: HashMap<acp::ModelId, Arc<dyn LanguageModel>>,
  84    /// Cached list for returning language model information
  85    model_list: acp_thread::AgentModelList,
  86    refresh_models_rx: watch::Receiver<()>,
  87    refresh_models_tx: watch::Sender<()>,
  88    _authenticate_all_providers_task: Task<()>,
  89}
  90
  91impl LanguageModels {
  92    fn new(cx: &mut App) -> Self {
  93        let (refresh_models_tx, refresh_models_rx) = watch::channel(());
  94
  95        let mut this = Self {
  96            models: HashMap::default(),
  97            model_list: acp_thread::AgentModelList::Grouped(IndexMap::default()),
  98            refresh_models_rx,
  99            refresh_models_tx,
 100            _authenticate_all_providers_task: Self::authenticate_all_language_model_providers(cx),
 101        };
 102        this.refresh_list(cx);
 103        this
 104    }
 105
 106    fn refresh_list(&mut self, cx: &App) {
 107        let providers = LanguageModelRegistry::global(cx)
 108            .read(cx)
 109            .providers()
 110            .into_iter()
 111            .filter(|provider| provider.is_authenticated(cx))
 112            .collect::<Vec<_>>();
 113
 114        let mut language_model_list = IndexMap::default();
 115        let mut recommended_models = HashSet::default();
 116
 117        let mut recommended = Vec::new();
 118        for provider in &providers {
 119            for model in provider.recommended_models(cx) {
 120                recommended_models.insert((model.provider_id(), model.id()));
 121                recommended.push(Self::map_language_model_to_info(&model, provider));
 122            }
 123        }
 124        if !recommended.is_empty() {
 125            language_model_list.insert(
 126                acp_thread::AgentModelGroupName("Recommended".into()),
 127                recommended,
 128            );
 129        }
 130
 131        let mut models = HashMap::default();
 132        for provider in providers {
 133            let mut provider_models = Vec::new();
 134            for model in provider.provided_models(cx) {
 135                let model_info = Self::map_language_model_to_info(&model, &provider);
 136                let model_id = model_info.id.clone();
 137                if !recommended_models.contains(&(model.provider_id(), model.id())) {
 138                    provider_models.push(model_info);
 139                }
 140                models.insert(model_id, model);
 141            }
 142            if !provider_models.is_empty() {
 143                language_model_list.insert(
 144                    acp_thread::AgentModelGroupName(provider.name().0.clone()),
 145                    provider_models,
 146                );
 147            }
 148        }
 149
 150        self.models = models;
 151        self.model_list = acp_thread::AgentModelList::Grouped(language_model_list);
 152        self.refresh_models_tx.send(()).ok();
 153    }
 154
 155    fn watch(&self) -> watch::Receiver<()> {
 156        self.refresh_models_rx.clone()
 157    }
 158
 159    pub fn model_from_id(&self, model_id: &acp::ModelId) -> Option<Arc<dyn LanguageModel>> {
 160        self.models.get(model_id).cloned()
 161    }
 162
 163    fn map_language_model_to_info(
 164        model: &Arc<dyn LanguageModel>,
 165        provider: &Arc<dyn LanguageModelProvider>,
 166    ) -> acp_thread::AgentModelInfo {
 167        acp_thread::AgentModelInfo {
 168            id: Self::model_id(model),
 169            name: model.name().0,
 170            description: None,
 171            icon: Some(provider.icon()),
 172        }
 173    }
 174
 175    fn model_id(model: &Arc<dyn LanguageModel>) -> acp::ModelId {
 176        acp::ModelId(format!("{}/{}", model.provider_id().0, model.id().0).into())
 177    }
 178
 179    fn authenticate_all_language_model_providers(cx: &mut App) -> Task<()> {
 180        let authenticate_all_providers = LanguageModelRegistry::global(cx)
 181            .read(cx)
 182            .providers()
 183            .iter()
 184            .map(|provider| (provider.id(), provider.name(), provider.authenticate(cx)))
 185            .collect::<Vec<_>>();
 186
 187        cx.background_spawn(async move {
 188            for (provider_id, provider_name, authenticate_task) in authenticate_all_providers {
 189                if let Err(err) = authenticate_task.await {
 190                    match err {
 191                        language_model::AuthenticateError::CredentialsNotFound => {
 192                            // Since we're authenticating these providers in the
 193                            // background for the purposes of populating the
 194                            // language selector, we don't care about providers
 195                            // where the credentials are not found.
 196                        }
 197                        language_model::AuthenticateError::ConnectionRefused => {
 198                            // Not logging connection refused errors as they are mostly from LM Studio's noisy auth failures.
 199                            // LM Studio only has one auth method (endpoint call) which fails for users who haven't enabled it.
 200                            // TODO: Better manage LM Studio auth logic to avoid these noisy failures.
 201                        }
 202                        _ => {
 203                            // Some providers have noisy failure states that we
 204                            // don't want to spam the logs with every time the
 205                            // language model selector is initialized.
 206                            //
 207                            // Ideally these should have more clear failure modes
 208                            // that we know are safe to ignore here, like what we do
 209                            // with `CredentialsNotFound` above.
 210                            match provider_id.0.as_ref() {
 211                                "lmstudio" | "ollama" => {
 212                                    // LM Studio and Ollama both make fetch requests to the local APIs to determine if they are "authenticated".
 213                                    //
 214                                    // These fail noisily, so we don't log them.
 215                                }
 216                                "copilot_chat" => {
 217                                    // Copilot Chat returns an error if Copilot is not enabled, so we don't log those errors.
 218                                }
 219                                _ => {
 220                                    log::error!(
 221                                        "Failed to authenticate provider: {}: {err}",
 222                                        provider_name.0
 223                                    );
 224                                }
 225                            }
 226                        }
 227                    }
 228                }
 229            }
 230        })
 231    }
 232}
 233
 234pub struct NativeAgent {
 235    /// Session ID -> Session mapping
 236    sessions: HashMap<acp::SessionId, Session>,
 237    history: Entity<HistoryStore>,
 238    /// Shared project context for all threads
 239    project_context: Entity<ProjectContext>,
 240    project_context_needs_refresh: watch::Sender<()>,
 241    _maintain_project_context: Task<Result<()>>,
 242    context_server_registry: Entity<ContextServerRegistry>,
 243    /// Shared templates for all threads
 244    templates: Arc<Templates>,
 245    /// Cached model information
 246    models: LanguageModels,
 247    project: Entity<Project>,
 248    prompt_store: Option<Entity<PromptStore>>,
 249    fs: Arc<dyn Fs>,
 250    _subscriptions: Vec<Subscription>,
 251}
 252
 253impl NativeAgent {
 254    pub async fn new(
 255        project: Entity<Project>,
 256        history: Entity<HistoryStore>,
 257        templates: Arc<Templates>,
 258        prompt_store: Option<Entity<PromptStore>>,
 259        fs: Arc<dyn Fs>,
 260        cx: &mut AsyncApp,
 261    ) -> Result<Entity<NativeAgent>> {
 262        log::debug!("Creating new NativeAgent");
 263
 264        let project_context = cx
 265            .update(|cx| Self::build_project_context(&project, prompt_store.as_ref(), cx))?
 266            .await;
 267
 268        cx.new(|cx| {
 269            let mut subscriptions = vec![
 270                cx.subscribe(&project, Self::handle_project_event),
 271                cx.subscribe(
 272                    &LanguageModelRegistry::global(cx),
 273                    Self::handle_models_updated_event,
 274                ),
 275            ];
 276            if let Some(prompt_store) = prompt_store.as_ref() {
 277                subscriptions.push(cx.subscribe(prompt_store, Self::handle_prompts_updated_event))
 278            }
 279
 280            let (project_context_needs_refresh_tx, project_context_needs_refresh_rx) =
 281                watch::channel(());
 282            Self {
 283                sessions: HashMap::new(),
 284                history,
 285                project_context: cx.new(|_| project_context),
 286                project_context_needs_refresh: project_context_needs_refresh_tx,
 287                _maintain_project_context: cx.spawn(async move |this, cx| {
 288                    Self::maintain_project_context(this, project_context_needs_refresh_rx, cx).await
 289                }),
 290                context_server_registry: cx.new(|cx| {
 291                    ContextServerRegistry::new(project.read(cx).context_server_store(), cx)
 292                }),
 293                templates,
 294                models: LanguageModels::new(cx),
 295                project,
 296                prompt_store,
 297                fs,
 298                _subscriptions: subscriptions,
 299            }
 300        })
 301    }
 302
 303    fn register_session(
 304        &mut self,
 305        thread_handle: Entity<Thread>,
 306        cx: &mut Context<Self>,
 307    ) -> Entity<AcpThread> {
 308        let connection = Rc::new(NativeAgentConnection(cx.entity()));
 309
 310        let thread = thread_handle.read(cx);
 311        let session_id = thread.id().clone();
 312        let title = thread.title();
 313        let project = thread.project.clone();
 314        let action_log = thread.action_log.clone();
 315        let prompt_capabilities_rx = thread.prompt_capabilities_rx.clone();
 316        let acp_thread = cx.new(|cx| {
 317            acp_thread::AcpThread::new(
 318                title,
 319                connection,
 320                project.clone(),
 321                action_log.clone(),
 322                session_id.clone(),
 323                prompt_capabilities_rx,
 324                cx,
 325            )
 326        });
 327
 328        let registry = LanguageModelRegistry::read_global(cx);
 329        let summarization_model = registry.thread_summary_model().map(|c| c.model);
 330
 331        thread_handle.update(cx, |thread, cx| {
 332            thread.set_summarization_model(summarization_model, cx);
 333            thread.add_default_tools(
 334                Rc::new(AcpThreadEnvironment {
 335                    acp_thread: acp_thread.downgrade(),
 336                }) as _,
 337                cx,
 338            )
 339        });
 340
 341        let subscriptions = vec![
 342            cx.observe_release(&acp_thread, |this, acp_thread, _cx| {
 343                this.sessions.remove(acp_thread.session_id());
 344            }),
 345            cx.subscribe(&thread_handle, Self::handle_thread_title_updated),
 346            cx.subscribe(&thread_handle, Self::handle_thread_token_usage_updated),
 347            cx.observe(&thread_handle, move |this, thread, cx| {
 348                this.save_thread(thread, cx)
 349            }),
 350        ];
 351
 352        self.sessions.insert(
 353            session_id,
 354            Session {
 355                thread: thread_handle,
 356                acp_thread: acp_thread.downgrade(),
 357                _subscriptions: subscriptions,
 358                pending_save: Task::ready(()),
 359            },
 360        );
 361        acp_thread
 362    }
 363
 364    pub fn models(&self) -> &LanguageModels {
 365        &self.models
 366    }
 367
 368    async fn maintain_project_context(
 369        this: WeakEntity<Self>,
 370        mut needs_refresh: watch::Receiver<()>,
 371        cx: &mut AsyncApp,
 372    ) -> Result<()> {
 373        while needs_refresh.changed().await.is_ok() {
 374            let project_context = this
 375                .update(cx, |this, cx| {
 376                    Self::build_project_context(&this.project, this.prompt_store.as_ref(), cx)
 377                })?
 378                .await;
 379            this.update(cx, |this, cx| {
 380                this.project_context = cx.new(|_| project_context);
 381            })?;
 382        }
 383
 384        Ok(())
 385    }
 386
 387    fn build_project_context(
 388        project: &Entity<Project>,
 389        prompt_store: Option<&Entity<PromptStore>>,
 390        cx: &mut App,
 391    ) -> Task<ProjectContext> {
 392        let worktrees = project.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
 393        let worktree_tasks = worktrees
 394            .into_iter()
 395            .map(|worktree| {
 396                Self::load_worktree_info_for_system_prompt(worktree, project.clone(), cx)
 397            })
 398            .collect::<Vec<_>>();
 399        let default_user_rules_task = if let Some(prompt_store) = prompt_store.as_ref() {
 400            prompt_store.read_with(cx, |prompt_store, cx| {
 401                let prompts = prompt_store.default_prompt_metadata();
 402                let load_tasks = prompts.into_iter().map(|prompt_metadata| {
 403                    let contents = prompt_store.load(prompt_metadata.id, cx);
 404                    async move { (contents.await, prompt_metadata) }
 405                });
 406                cx.background_spawn(future::join_all(load_tasks))
 407            })
 408        } else {
 409            Task::ready(vec![])
 410        };
 411
 412        cx.spawn(async move |_cx| {
 413            let (worktrees, default_user_rules) =
 414                future::join(future::join_all(worktree_tasks), default_user_rules_task).await;
 415
 416            let worktrees = worktrees
 417                .into_iter()
 418                .map(|(worktree, _rules_error)| {
 419                    // TODO: show error message
 420                    // if let Some(rules_error) = rules_error {
 421                    //     this.update(cx, |_, cx| cx.emit(rules_error)).ok();
 422                    // }
 423                    worktree
 424                })
 425                .collect::<Vec<_>>();
 426
 427            let default_user_rules = default_user_rules
 428                .into_iter()
 429                .flat_map(|(contents, prompt_metadata)| match contents {
 430                    Ok(contents) => Some(UserRulesContext {
 431                        uuid: match prompt_metadata.id {
 432                            prompt_store::PromptId::User { uuid } => uuid,
 433                            prompt_store::PromptId::EditWorkflow => return None,
 434                        },
 435                        title: prompt_metadata.title.map(|title| title.to_string()),
 436                        contents,
 437                    }),
 438                    Err(_err) => {
 439                        // TODO: show error message
 440                        // this.update(cx, |_, cx| {
 441                        //     cx.emit(RulesLoadingError {
 442                        //         message: format!("{err:?}").into(),
 443                        //     });
 444                        // })
 445                        // .ok();
 446                        None
 447                    }
 448                })
 449                .collect::<Vec<_>>();
 450
 451            ProjectContext::new(worktrees, default_user_rules)
 452        })
 453    }
 454
 455    fn load_worktree_info_for_system_prompt(
 456        worktree: Entity<Worktree>,
 457        project: Entity<Project>,
 458        cx: &mut App,
 459    ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
 460        let tree = worktree.read(cx);
 461        let root_name = tree.root_name_str().into();
 462        let abs_path = tree.abs_path();
 463
 464        let mut context = WorktreeContext {
 465            root_name,
 466            abs_path,
 467            rules_file: None,
 468        };
 469
 470        let rules_task = Self::load_worktree_rules_file(worktree, project, cx);
 471        let Some(rules_task) = rules_task else {
 472            return Task::ready((context, None));
 473        };
 474
 475        cx.spawn(async move |_| {
 476            let (rules_file, rules_file_error) = match rules_task.await {
 477                Ok(rules_file) => (Some(rules_file), None),
 478                Err(err) => (
 479                    None,
 480                    Some(RulesLoadingError {
 481                        message: format!("{err}").into(),
 482                    }),
 483                ),
 484            };
 485            context.rules_file = rules_file;
 486            (context, rules_file_error)
 487        })
 488    }
 489
 490    fn load_worktree_rules_file(
 491        worktree: Entity<Worktree>,
 492        project: Entity<Project>,
 493        cx: &mut App,
 494    ) -> Option<Task<Result<RulesFileContext>>> {
 495        let worktree = worktree.read(cx);
 496        let worktree_id = worktree.id();
 497        let selected_rules_file = RULES_FILE_NAMES
 498            .into_iter()
 499            .filter_map(|name| {
 500                worktree
 501                    .entry_for_path(RelPath::unix(name).unwrap())
 502                    .filter(|entry| entry.is_file())
 503                    .map(|entry| entry.path.clone())
 504            })
 505            .next();
 506
 507        // Note that Cline supports `.clinerules` being a directory, but that is not currently
 508        // supported. This doesn't seem to occur often in GitHub repositories.
 509        selected_rules_file.map(|path_in_worktree| {
 510            let project_path = ProjectPath {
 511                worktree_id,
 512                path: path_in_worktree.clone(),
 513            };
 514            let buffer_task =
 515                project.update(cx, |project, cx| project.open_buffer(project_path, cx));
 516            let rope_task = cx.spawn(async move |cx| {
 517                buffer_task.await?.read_with(cx, |buffer, cx| {
 518                    let project_entry_id = buffer.entry_id(cx).context("buffer has no file")?;
 519                    anyhow::Ok((project_entry_id, buffer.as_rope().clone()))
 520                })?
 521            });
 522            // Build a string from the rope on a background thread.
 523            cx.background_spawn(async move {
 524                let (project_entry_id, rope) = rope_task.await?;
 525                anyhow::Ok(RulesFileContext {
 526                    path_in_worktree,
 527                    text: rope.to_string().trim().to_string(),
 528                    project_entry_id: project_entry_id.to_usize(),
 529                })
 530            })
 531        })
 532    }
 533
 534    fn handle_thread_title_updated(
 535        &mut self,
 536        thread: Entity<Thread>,
 537        _: &TitleUpdated,
 538        cx: &mut Context<Self>,
 539    ) {
 540        let session_id = thread.read(cx).id();
 541        let Some(session) = self.sessions.get(session_id) else {
 542            return;
 543        };
 544        let thread = thread.downgrade();
 545        let acp_thread = session.acp_thread.clone();
 546        cx.spawn(async move |_, cx| {
 547            let title = thread.read_with(cx, |thread, _| thread.title())?;
 548            let task = acp_thread.update(cx, |acp_thread, cx| acp_thread.set_title(title, cx))?;
 549            task.await
 550        })
 551        .detach_and_log_err(cx);
 552    }
 553
 554    fn handle_thread_token_usage_updated(
 555        &mut self,
 556        thread: Entity<Thread>,
 557        usage: &TokenUsageUpdated,
 558        cx: &mut Context<Self>,
 559    ) {
 560        let Some(session) = self.sessions.get(thread.read(cx).id()) else {
 561            return;
 562        };
 563        session
 564            .acp_thread
 565            .update(cx, |acp_thread, cx| {
 566                acp_thread.update_token_usage(usage.0.clone(), cx);
 567            })
 568            .ok();
 569    }
 570
 571    fn handle_project_event(
 572        &mut self,
 573        _project: Entity<Project>,
 574        event: &project::Event,
 575        _cx: &mut Context<Self>,
 576    ) {
 577        match event {
 578            project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
 579                self.project_context_needs_refresh.send(()).ok();
 580            }
 581            project::Event::WorktreeUpdatedEntries(_, items) => {
 582                if items.iter().any(|(path, _, _)| {
 583                    RULES_FILE_NAMES
 584                        .iter()
 585                        .any(|name| path.as_ref() == RelPath::unix(name).unwrap())
 586                }) {
 587                    self.project_context_needs_refresh.send(()).ok();
 588                }
 589            }
 590            _ => {}
 591        }
 592    }
 593
 594    fn handle_prompts_updated_event(
 595        &mut self,
 596        _prompt_store: Entity<PromptStore>,
 597        _event: &prompt_store::PromptsUpdatedEvent,
 598        _cx: &mut Context<Self>,
 599    ) {
 600        self.project_context_needs_refresh.send(()).ok();
 601    }
 602
 603    fn handle_models_updated_event(
 604        &mut self,
 605        _registry: Entity<LanguageModelRegistry>,
 606        _event: &language_model::Event,
 607        cx: &mut Context<Self>,
 608    ) {
 609        self.models.refresh_list(cx);
 610
 611        let registry = LanguageModelRegistry::read_global(cx);
 612        let default_model = registry.default_model().map(|m| m.model);
 613        let summarization_model = registry.thread_summary_model().map(|m| m.model);
 614
 615        for session in self.sessions.values_mut() {
 616            session.thread.update(cx, |thread, cx| {
 617                if thread.model().is_none()
 618                    && let Some(model) = default_model.clone()
 619                {
 620                    thread.set_model(model, cx);
 621                    cx.notify();
 622                }
 623                thread.set_summarization_model(summarization_model.clone(), cx);
 624            });
 625        }
 626    }
 627
 628    pub fn load_thread(
 629        &mut self,
 630        id: acp::SessionId,
 631        cx: &mut Context<Self>,
 632    ) -> Task<Result<Entity<Thread>>> {
 633        let database_future = ThreadsDatabase::connect(cx);
 634        cx.spawn(async move |this, cx| {
 635            let database = database_future.await.map_err(|err| anyhow!(err))?;
 636            let db_thread = database
 637                .load_thread(id.clone())
 638                .await?
 639                .with_context(|| format!("no thread found with ID: {id:?}"))?;
 640
 641            this.update(cx, |this, cx| {
 642                let summarization_model = LanguageModelRegistry::read_global(cx)
 643                    .thread_summary_model()
 644                    .map(|c| c.model);
 645
 646                cx.new(|cx| {
 647                    let mut thread = Thread::from_db(
 648                        id.clone(),
 649                        db_thread,
 650                        this.project.clone(),
 651                        this.project_context.clone(),
 652                        this.context_server_registry.clone(),
 653                        this.templates.clone(),
 654                        cx,
 655                    );
 656                    thread.set_summarization_model(summarization_model, cx);
 657                    thread
 658                })
 659            })
 660        })
 661    }
 662
 663    pub fn open_thread(
 664        &mut self,
 665        id: acp::SessionId,
 666        cx: &mut Context<Self>,
 667    ) -> Task<Result<Entity<AcpThread>>> {
 668        let task = self.load_thread(id, cx);
 669        cx.spawn(async move |this, cx| {
 670            let thread = task.await?;
 671            let acp_thread =
 672                this.update(cx, |this, cx| this.register_session(thread.clone(), cx))?;
 673            let events = thread.update(cx, |thread, cx| thread.replay(cx))?;
 674            cx.update(|cx| {
 675                NativeAgentConnection::handle_thread_events(events, acp_thread.downgrade(), cx)
 676            })?
 677            .await?;
 678            Ok(acp_thread)
 679        })
 680    }
 681
 682    pub fn thread_summary(
 683        &mut self,
 684        id: acp::SessionId,
 685        cx: &mut Context<Self>,
 686    ) -> Task<Result<SharedString>> {
 687        let thread = self.open_thread(id.clone(), cx);
 688        cx.spawn(async move |this, cx| {
 689            let acp_thread = thread.await?;
 690            let result = this
 691                .update(cx, |this, cx| {
 692                    this.sessions
 693                        .get(&id)
 694                        .unwrap()
 695                        .thread
 696                        .update(cx, |thread, cx| thread.summary(cx))
 697                })?
 698                .await
 699                .context("Failed to generate summary")?;
 700            drop(acp_thread);
 701            Ok(result)
 702        })
 703    }
 704
 705    fn save_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
 706        if thread.read(cx).is_empty() {
 707            return;
 708        }
 709
 710        let database_future = ThreadsDatabase::connect(cx);
 711        let (id, db_thread) =
 712            thread.update(cx, |thread, cx| (thread.id().clone(), thread.to_db(cx)));
 713        let Some(session) = self.sessions.get_mut(&id) else {
 714            return;
 715        };
 716        let history = self.history.clone();
 717        session.pending_save = cx.spawn(async move |_, cx| {
 718            let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
 719                return;
 720            };
 721            let db_thread = db_thread.await;
 722            database.save_thread(id, db_thread).await.log_err();
 723            history.update(cx, |history, cx| history.reload(cx)).ok();
 724        });
 725    }
 726}
 727
 728/// Wrapper struct that implements the AgentConnection trait
 729#[derive(Clone)]
 730pub struct NativeAgentConnection(pub Entity<NativeAgent>);
 731
 732impl NativeAgentConnection {
 733    pub fn thread(&self, session_id: &acp::SessionId, cx: &App) -> Option<Entity<Thread>> {
 734        self.0
 735            .read(cx)
 736            .sessions
 737            .get(session_id)
 738            .map(|session| session.thread.clone())
 739    }
 740
 741    pub fn load_thread(&self, id: acp::SessionId, cx: &mut App) -> Task<Result<Entity<Thread>>> {
 742        self.0.update(cx, |this, cx| this.load_thread(id, cx))
 743    }
 744
 745    fn run_turn(
 746        &self,
 747        session_id: acp::SessionId,
 748        cx: &mut App,
 749        f: impl 'static
 750        + FnOnce(Entity<Thread>, &mut App) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>,
 751    ) -> Task<Result<acp::PromptResponse>> {
 752        let Some((thread, acp_thread)) = self.0.update(cx, |agent, _cx| {
 753            agent
 754                .sessions
 755                .get_mut(&session_id)
 756                .map(|s| (s.thread.clone(), s.acp_thread.clone()))
 757        }) else {
 758            return Task::ready(Err(anyhow!("Session not found")));
 759        };
 760        log::debug!("Found session for: {}", session_id);
 761
 762        let response_stream = match f(thread, cx) {
 763            Ok(stream) => stream,
 764            Err(err) => return Task::ready(Err(err)),
 765        };
 766        Self::handle_thread_events(response_stream, acp_thread, cx)
 767    }
 768
 769    fn handle_thread_events(
 770        mut events: mpsc::UnboundedReceiver<Result<ThreadEvent>>,
 771        acp_thread: WeakEntity<AcpThread>,
 772        cx: &App,
 773    ) -> Task<Result<acp::PromptResponse>> {
 774        cx.spawn(async move |cx| {
 775            // Handle response stream and forward to session.acp_thread
 776            while let Some(result) = events.next().await {
 777                match result {
 778                    Ok(event) => {
 779                        log::trace!("Received completion event: {:?}", event);
 780
 781                        match event {
 782                            ThreadEvent::UserMessage(message) => {
 783                                acp_thread.update(cx, |thread, cx| {
 784                                    for content in message.content {
 785                                        thread.push_user_content_block(
 786                                            Some(message.id.clone()),
 787                                            content.into(),
 788                                            cx,
 789                                        );
 790                                    }
 791                                })?;
 792                            }
 793                            ThreadEvent::AgentText(text) => {
 794                                acp_thread.update(cx, |thread, cx| {
 795                                    thread.push_assistant_content_block(
 796                                        acp::ContentBlock::Text(acp::TextContent {
 797                                            text,
 798                                            annotations: None,
 799                                            meta: None,
 800                                        }),
 801                                        false,
 802                                        cx,
 803                                    )
 804                                })?;
 805                            }
 806                            ThreadEvent::AgentThinking(text) => {
 807                                acp_thread.update(cx, |thread, cx| {
 808                                    thread.push_assistant_content_block(
 809                                        acp::ContentBlock::Text(acp::TextContent {
 810                                            text,
 811                                            annotations: None,
 812                                            meta: None,
 813                                        }),
 814                                        true,
 815                                        cx,
 816                                    )
 817                                })?;
 818                            }
 819                            ThreadEvent::ToolCallAuthorization(ToolCallAuthorization {
 820                                tool_call,
 821                                options,
 822                                response,
 823                            }) => {
 824                                let outcome_task = acp_thread.update(cx, |thread, cx| {
 825                                    thread.request_tool_call_authorization(
 826                                        tool_call, options, true, cx,
 827                                    )
 828                                })??;
 829                                cx.background_spawn(async move {
 830                                    if let acp::RequestPermissionOutcome::Selected { option_id } =
 831                                        outcome_task.await
 832                                    {
 833                                        response
 834                                            .send(option_id)
 835                                            .map(|_| anyhow!("authorization receiver was dropped"))
 836                                            .log_err();
 837                                    }
 838                                })
 839                                .detach();
 840                            }
 841                            ThreadEvent::ToolCall(tool_call) => {
 842                                acp_thread.update(cx, |thread, cx| {
 843                                    thread.upsert_tool_call(tool_call, cx)
 844                                })??;
 845                            }
 846                            ThreadEvent::ToolCallUpdate(update) => {
 847                                acp_thread.update(cx, |thread, cx| {
 848                                    thread.update_tool_call(update, cx)
 849                                })??;
 850                            }
 851                            ThreadEvent::Retry(status) => {
 852                                acp_thread.update(cx, |thread, cx| {
 853                                    thread.update_retry_status(status, cx)
 854                                })?;
 855                            }
 856                            ThreadEvent::Stop(stop_reason) => {
 857                                log::debug!("Assistant message complete: {:?}", stop_reason);
 858                                return Ok(acp::PromptResponse {
 859                                    stop_reason,
 860                                    meta: None,
 861                                });
 862                            }
 863                        }
 864                    }
 865                    Err(e) => {
 866                        log::error!("Error in model response stream: {:?}", e);
 867                        return Err(e);
 868                    }
 869                }
 870            }
 871
 872            log::debug!("Response stream completed");
 873            anyhow::Ok(acp::PromptResponse {
 874                stop_reason: acp::StopReason::EndTurn,
 875                meta: None,
 876            })
 877        })
 878    }
 879}
 880
 881struct NativeAgentModelSelector {
 882    session_id: acp::SessionId,
 883    connection: NativeAgentConnection,
 884}
 885
 886impl acp_thread::AgentModelSelector for NativeAgentModelSelector {
 887    fn list_models(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelList>> {
 888        log::debug!("NativeAgentConnection::list_models called");
 889        let list = self.connection.0.read(cx).models.model_list.clone();
 890        Task::ready(if list.is_empty() {
 891            Err(anyhow::anyhow!("No models available"))
 892        } else {
 893            Ok(list)
 894        })
 895    }
 896
 897    fn select_model(&self, model_id: acp::ModelId, cx: &mut App) -> Task<Result<()>> {
 898        log::debug!(
 899            "Setting model for session {}: {}",
 900            self.session_id,
 901            model_id
 902        );
 903        let Some(thread) = self
 904            .connection
 905            .0
 906            .read(cx)
 907            .sessions
 908            .get(&self.session_id)
 909            .map(|session| session.thread.clone())
 910        else {
 911            return Task::ready(Err(anyhow!("Session not found")));
 912        };
 913
 914        let Some(model) = self.connection.0.read(cx).models.model_from_id(&model_id) else {
 915            return Task::ready(Err(anyhow!("Invalid model ID {}", model_id)));
 916        };
 917
 918        thread.update(cx, |thread, cx| {
 919            thread.set_model(model.clone(), cx);
 920        });
 921
 922        update_settings_file(
 923            self.connection.0.read(cx).fs.clone(),
 924            cx,
 925            move |settings, _cx| {
 926                let provider = model.provider_id().0.to_string();
 927                let model = model.id().0.to_string();
 928                settings
 929                    .agent
 930                    .get_or_insert_default()
 931                    .set_model(LanguageModelSelection {
 932                        provider: provider.into(),
 933                        model,
 934                    });
 935            },
 936        );
 937
 938        Task::ready(Ok(()))
 939    }
 940
 941    fn selected_model(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelInfo>> {
 942        let Some(thread) = self
 943            .connection
 944            .0
 945            .read(cx)
 946            .sessions
 947            .get(&self.session_id)
 948            .map(|session| session.thread.clone())
 949        else {
 950            return Task::ready(Err(anyhow!("Session not found")));
 951        };
 952        let Some(model) = thread.read(cx).model() else {
 953            return Task::ready(Err(anyhow!("Model not found")));
 954        };
 955        let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&model.provider_id())
 956        else {
 957            return Task::ready(Err(anyhow!("Provider not found")));
 958        };
 959        Task::ready(Ok(LanguageModels::map_language_model_to_info(
 960            model, &provider,
 961        )))
 962    }
 963
 964    fn watch(&self, cx: &mut App) -> Option<watch::Receiver<()>> {
 965        Some(self.connection.0.read(cx).models.watch())
 966    }
 967}
 968
 969impl acp_thread::AgentConnection for NativeAgentConnection {
 970    fn new_thread(
 971        self: Rc<Self>,
 972        project: Entity<Project>,
 973        cwd: &Path,
 974        cx: &mut App,
 975    ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
 976        let agent = self.0.clone();
 977        log::debug!("Creating new thread for project at: {:?}", cwd);
 978
 979        cx.spawn(async move |cx| {
 980            log::debug!("Starting thread creation in async context");
 981
 982            // Create Thread
 983            let thread = agent.update(
 984                cx,
 985                |agent, cx: &mut gpui::Context<NativeAgent>| -> Result<_> {
 986                    // Fetch default model from registry settings
 987                    let registry = LanguageModelRegistry::read_global(cx);
 988                    // Log available models for debugging
 989                    let available_count = registry.available_models(cx).count();
 990                    log::debug!("Total available models: {}", available_count);
 991
 992                    let default_model = registry.default_model().and_then(|default_model| {
 993                        agent
 994                            .models
 995                            .model_from_id(&LanguageModels::model_id(&default_model.model))
 996                    });
 997                    Ok(cx.new(|cx| {
 998                        Thread::new(
 999                            project.clone(),
1000                            agent.project_context.clone(),
1001                            agent.context_server_registry.clone(),
1002                            agent.templates.clone(),
1003                            default_model,
1004                            cx,
1005                        )
1006                    }))
1007                },
1008            )??;
1009            agent.update(cx, |agent, cx| agent.register_session(thread, cx))
1010        })
1011    }
1012
1013    fn auth_methods(&self) -> &[acp::AuthMethod] {
1014        &[] // No auth for in-process
1015    }
1016
1017    fn authenticate(&self, _method: acp::AuthMethodId, _cx: &mut App) -> Task<Result<()>> {
1018        Task::ready(Ok(()))
1019    }
1020
1021    fn model_selector(&self, session_id: &acp::SessionId) -> Option<Rc<dyn AgentModelSelector>> {
1022        Some(Rc::new(NativeAgentModelSelector {
1023            session_id: session_id.clone(),
1024            connection: self.clone(),
1025        }) as Rc<dyn AgentModelSelector>)
1026    }
1027
1028    fn prompt(
1029        &self,
1030        id: Option<acp_thread::UserMessageId>,
1031        params: acp::PromptRequest,
1032        cx: &mut App,
1033    ) -> Task<Result<acp::PromptResponse>> {
1034        let id = id.expect("UserMessageId is required");
1035        let session_id = params.session_id.clone();
1036        log::info!("Received prompt request for session: {}", session_id);
1037        log::debug!("Prompt blocks count: {}", params.prompt.len());
1038
1039        self.run_turn(session_id, cx, |thread, cx| {
1040            let content: Vec<UserMessageContent> = params
1041                .prompt
1042                .into_iter()
1043                .map(Into::into)
1044                .collect::<Vec<_>>();
1045            log::debug!("Converted prompt to message: {} chars", content.len());
1046            log::debug!("Message id: {:?}", id);
1047            log::debug!("Message content: {:?}", content);
1048
1049            thread.update(cx, |thread, cx| thread.send(id, content, cx))
1050        })
1051    }
1052
1053    fn resume(
1054        &self,
1055        session_id: &acp::SessionId,
1056        _cx: &App,
1057    ) -> Option<Rc<dyn acp_thread::AgentSessionResume>> {
1058        Some(Rc::new(NativeAgentSessionResume {
1059            connection: self.clone(),
1060            session_id: session_id.clone(),
1061        }) as _)
1062    }
1063
1064    fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
1065        log::info!("Cancelling on session: {}", session_id);
1066        self.0.update(cx, |agent, cx| {
1067            if let Some(agent) = agent.sessions.get(session_id) {
1068                agent.thread.update(cx, |thread, cx| thread.cancel(cx));
1069            }
1070        });
1071    }
1072
1073    fn truncate(
1074        &self,
1075        session_id: &agent_client_protocol::SessionId,
1076        cx: &App,
1077    ) -> Option<Rc<dyn acp_thread::AgentSessionTruncate>> {
1078        self.0.read_with(cx, |agent, _cx| {
1079            agent.sessions.get(session_id).map(|session| {
1080                Rc::new(NativeAgentSessionTruncate {
1081                    thread: session.thread.clone(),
1082                    acp_thread: session.acp_thread.clone(),
1083                }) as _
1084            })
1085        })
1086    }
1087
1088    fn set_title(
1089        &self,
1090        session_id: &acp::SessionId,
1091        _cx: &App,
1092    ) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
1093        Some(Rc::new(NativeAgentSessionSetTitle {
1094            connection: self.clone(),
1095            session_id: session_id.clone(),
1096        }) as _)
1097    }
1098
1099    fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
1100        Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
1101    }
1102
1103    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1104        self
1105    }
1106}
1107
1108impl acp_thread::AgentTelemetry for NativeAgentConnection {
1109    fn agent_name(&self) -> String {
1110        "Zed".into()
1111    }
1112
1113    fn thread_data(
1114        &self,
1115        session_id: &acp::SessionId,
1116        cx: &mut App,
1117    ) -> Task<Result<serde_json::Value>> {
1118        let Some(session) = self.0.read(cx).sessions.get(session_id) else {
1119            return Task::ready(Err(anyhow!("Session not found")));
1120        };
1121
1122        let task = session.thread.read(cx).to_db(cx);
1123        cx.background_spawn(async move {
1124            serde_json::to_value(task.await).context("Failed to serialize thread")
1125        })
1126    }
1127}
1128
1129struct NativeAgentSessionTruncate {
1130    thread: Entity<Thread>,
1131    acp_thread: WeakEntity<AcpThread>,
1132}
1133
1134impl acp_thread::AgentSessionTruncate for NativeAgentSessionTruncate {
1135    fn run(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
1136        match self.thread.update(cx, |thread, cx| {
1137            thread.truncate(message_id.clone(), cx)?;
1138            Ok(thread.latest_token_usage())
1139        }) {
1140            Ok(usage) => {
1141                self.acp_thread
1142                    .update(cx, |thread, cx| {
1143                        thread.update_token_usage(usage, cx);
1144                    })
1145                    .ok();
1146                Task::ready(Ok(()))
1147            }
1148            Err(error) => Task::ready(Err(error)),
1149        }
1150    }
1151}
1152
1153struct NativeAgentSessionResume {
1154    connection: NativeAgentConnection,
1155    session_id: acp::SessionId,
1156}
1157
1158impl acp_thread::AgentSessionResume for NativeAgentSessionResume {
1159    fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>> {
1160        self.connection
1161            .run_turn(self.session_id.clone(), cx, |thread, cx| {
1162                thread.update(cx, |thread, cx| thread.resume(cx))
1163            })
1164    }
1165}
1166
1167struct NativeAgentSessionSetTitle {
1168    connection: NativeAgentConnection,
1169    session_id: acp::SessionId,
1170}
1171
1172impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
1173    fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
1174        let Some(session) = self.connection.0.read(cx).sessions.get(&self.session_id) else {
1175            return Task::ready(Err(anyhow!("session not found")));
1176        };
1177        let thread = session.thread.clone();
1178        thread.update(cx, |thread, cx| thread.set_title(title, cx));
1179        Task::ready(Ok(()))
1180    }
1181}
1182
1183pub struct AcpThreadEnvironment {
1184    acp_thread: WeakEntity<AcpThread>,
1185}
1186
1187impl ThreadEnvironment for AcpThreadEnvironment {
1188    fn create_terminal(
1189        &self,
1190        command: String,
1191        cwd: Option<PathBuf>,
1192        output_byte_limit: Option<u64>,
1193        cx: &mut AsyncApp,
1194    ) -> Task<Result<Rc<dyn TerminalHandle>>> {
1195        let task = self.acp_thread.update(cx, |thread, cx| {
1196            thread.create_terminal(command, vec![], vec![], cwd, output_byte_limit, cx)
1197        });
1198
1199        let acp_thread = self.acp_thread.clone();
1200        cx.spawn(async move |cx| {
1201            let terminal = task?.await?;
1202
1203            let (drop_tx, drop_rx) = oneshot::channel();
1204            let terminal_id = terminal.read_with(cx, |terminal, _cx| terminal.id().clone())?;
1205
1206            cx.spawn(async move |cx| {
1207                drop_rx.await.ok();
1208                acp_thread.update(cx, |thread, cx| thread.release_terminal(terminal_id, cx))
1209            })
1210            .detach();
1211
1212            let handle = AcpTerminalHandle {
1213                terminal,
1214                _drop_tx: Some(drop_tx),
1215            };
1216
1217            Ok(Rc::new(handle) as _)
1218        })
1219    }
1220}
1221
1222pub struct AcpTerminalHandle {
1223    terminal: Entity<acp_thread::Terminal>,
1224    _drop_tx: Option<oneshot::Sender<()>>,
1225}
1226
1227impl TerminalHandle for AcpTerminalHandle {
1228    fn id(&self, cx: &AsyncApp) -> Result<acp::TerminalId> {
1229        self.terminal.read_with(cx, |term, _cx| term.id().clone())
1230    }
1231
1232    fn wait_for_exit(&self, cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {
1233        self.terminal
1234            .read_with(cx, |term, _cx| term.wait_for_exit())
1235    }
1236
1237    fn current_output(&self, cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {
1238        self.terminal
1239            .read_with(cx, |term, cx| term.current_output(cx))
1240    }
1241}
1242
1243#[cfg(test)]
1244mod internal_tests {
1245    use crate::HistoryEntryId;
1246
1247    use super::*;
1248    use acp_thread::{AgentConnection, AgentModelGroupName, AgentModelInfo, MentionUri};
1249    use fs::FakeFs;
1250    use gpui::TestAppContext;
1251    use indoc::formatdoc;
1252    use language_model::fake_provider::FakeLanguageModel;
1253    use serde_json::json;
1254    use settings::SettingsStore;
1255    use util::{path, rel_path::rel_path};
1256
1257    #[gpui::test]
1258    async fn test_maintaining_project_context(cx: &mut TestAppContext) {
1259        init_test(cx);
1260        let fs = FakeFs::new(cx.executor());
1261        fs.insert_tree(
1262            "/",
1263            json!({
1264                "a": {}
1265            }),
1266        )
1267        .await;
1268        let project = Project::test(fs.clone(), [], cx).await;
1269        let context_store = cx.new(|cx| assistant_context::ContextStore::fake(project.clone(), cx));
1270        let history_store = cx.new(|cx| HistoryStore::new(context_store, cx));
1271        let agent = NativeAgent::new(
1272            project.clone(),
1273            history_store,
1274            Templates::new(),
1275            None,
1276            fs.clone(),
1277            &mut cx.to_async(),
1278        )
1279        .await
1280        .unwrap();
1281        agent.read_with(cx, |agent, cx| {
1282            assert_eq!(agent.project_context.read(cx).worktrees, vec![])
1283        });
1284
1285        let worktree = project
1286            .update(cx, |project, cx| project.create_worktree("/a", true, cx))
1287            .await
1288            .unwrap();
1289        cx.run_until_parked();
1290        agent.read_with(cx, |agent, cx| {
1291            assert_eq!(
1292                agent.project_context.read(cx).worktrees,
1293                vec![WorktreeContext {
1294                    root_name: "a".into(),
1295                    abs_path: Path::new("/a").into(),
1296                    rules_file: None
1297                }]
1298            )
1299        });
1300
1301        // Creating `/a/.rules` updates the project context.
1302        fs.insert_file("/a/.rules", Vec::new()).await;
1303        cx.run_until_parked();
1304        agent.read_with(cx, |agent, cx| {
1305            let rules_entry = worktree
1306                .read(cx)
1307                .entry_for_path(rel_path(".rules"))
1308                .unwrap();
1309            assert_eq!(
1310                agent.project_context.read(cx).worktrees,
1311                vec![WorktreeContext {
1312                    root_name: "a".into(),
1313                    abs_path: Path::new("/a").into(),
1314                    rules_file: Some(RulesFileContext {
1315                        path_in_worktree: rel_path(".rules").into(),
1316                        text: "".into(),
1317                        project_entry_id: rules_entry.id.to_usize()
1318                    })
1319                }]
1320            )
1321        });
1322    }
1323
1324    #[gpui::test]
1325    async fn test_listing_models(cx: &mut TestAppContext) {
1326        init_test(cx);
1327        let fs = FakeFs::new(cx.executor());
1328        fs.insert_tree("/", json!({ "a": {}  })).await;
1329        let project = Project::test(fs.clone(), [], cx).await;
1330        let context_store = cx.new(|cx| assistant_context::ContextStore::fake(project.clone(), cx));
1331        let history_store = cx.new(|cx| HistoryStore::new(context_store, cx));
1332        let connection = NativeAgentConnection(
1333            NativeAgent::new(
1334                project.clone(),
1335                history_store,
1336                Templates::new(),
1337                None,
1338                fs.clone(),
1339                &mut cx.to_async(),
1340            )
1341            .await
1342            .unwrap(),
1343        );
1344
1345        // Create a thread/session
1346        let acp_thread = cx
1347            .update(|cx| {
1348                Rc::new(connection.clone()).new_thread(project.clone(), Path::new("/a"), cx)
1349            })
1350            .await
1351            .unwrap();
1352
1353        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1354
1355        let models = cx
1356            .update(|cx| {
1357                connection
1358                    .model_selector(&session_id)
1359                    .unwrap()
1360                    .list_models(cx)
1361            })
1362            .await
1363            .unwrap();
1364
1365        let acp_thread::AgentModelList::Grouped(models) = models else {
1366            panic!("Unexpected model group");
1367        };
1368        assert_eq!(
1369            models,
1370            IndexMap::from_iter([(
1371                AgentModelGroupName("Fake".into()),
1372                vec![AgentModelInfo {
1373                    id: acp::ModelId("fake/fake".into()),
1374                    name: "Fake".into(),
1375                    description: None,
1376                    icon: Some(ui::IconName::ZedAssistant),
1377                }]
1378            )])
1379        );
1380    }
1381
1382    #[gpui::test]
1383    async fn test_model_selection_persists_to_settings(cx: &mut TestAppContext) {
1384        init_test(cx);
1385        let fs = FakeFs::new(cx.executor());
1386        fs.create_dir(paths::settings_file().parent().unwrap())
1387            .await
1388            .unwrap();
1389        fs.insert_file(
1390            paths::settings_file(),
1391            json!({
1392                "agent": {
1393                    "default_model": {
1394                        "provider": "foo",
1395                        "model": "bar"
1396                    }
1397                }
1398            })
1399            .to_string()
1400            .into_bytes(),
1401        )
1402        .await;
1403        let project = Project::test(fs.clone(), [], cx).await;
1404
1405        let context_store = cx.new(|cx| assistant_context::ContextStore::fake(project.clone(), cx));
1406        let history_store = cx.new(|cx| HistoryStore::new(context_store, cx));
1407
1408        // Create the agent and connection
1409        let agent = NativeAgent::new(
1410            project.clone(),
1411            history_store,
1412            Templates::new(),
1413            None,
1414            fs.clone(),
1415            &mut cx.to_async(),
1416        )
1417        .await
1418        .unwrap();
1419        let connection = NativeAgentConnection(agent.clone());
1420
1421        // Create a thread/session
1422        let acp_thread = cx
1423            .update(|cx| {
1424                Rc::new(connection.clone()).new_thread(project.clone(), Path::new("/a"), cx)
1425            })
1426            .await
1427            .unwrap();
1428
1429        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1430
1431        // Select a model
1432        let selector = connection.model_selector(&session_id).unwrap();
1433        let model_id = acp::ModelId("fake/fake".into());
1434        cx.update(|cx| selector.select_model(model_id.clone(), cx))
1435            .await
1436            .unwrap();
1437
1438        // Verify the thread has the selected model
1439        agent.read_with(cx, |agent, _| {
1440            let session = agent.sessions.get(&session_id).unwrap();
1441            session.thread.read_with(cx, |thread, _| {
1442                assert_eq!(thread.model().unwrap().id().0, "fake");
1443            });
1444        });
1445
1446        cx.run_until_parked();
1447
1448        // Verify settings file was updated
1449        let settings_content = fs.load(paths::settings_file()).await.unwrap();
1450        let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
1451
1452        // Check that the agent settings contain the selected model
1453        assert_eq!(
1454            settings_json["agent"]["default_model"]["model"],
1455            json!("fake")
1456        );
1457        assert_eq!(
1458            settings_json["agent"]["default_model"]["provider"],
1459            json!("fake")
1460        );
1461    }
1462
1463    #[gpui::test]
1464    async fn test_save_load_thread(cx: &mut TestAppContext) {
1465        init_test(cx);
1466        let fs = FakeFs::new(cx.executor());
1467        fs.insert_tree(
1468            "/",
1469            json!({
1470                "a": {
1471                    "b.md": "Lorem"
1472                }
1473            }),
1474        )
1475        .await;
1476        let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
1477        let context_store = cx.new(|cx| assistant_context::ContextStore::fake(project.clone(), cx));
1478        let history_store = cx.new(|cx| HistoryStore::new(context_store, cx));
1479        let agent = NativeAgent::new(
1480            project.clone(),
1481            history_store.clone(),
1482            Templates::new(),
1483            None,
1484            fs.clone(),
1485            &mut cx.to_async(),
1486        )
1487        .await
1488        .unwrap();
1489        let connection = Rc::new(NativeAgentConnection(agent.clone()));
1490
1491        let acp_thread = cx
1492            .update(|cx| {
1493                connection
1494                    .clone()
1495                    .new_thread(project.clone(), Path::new(""), cx)
1496            })
1497            .await
1498            .unwrap();
1499        let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
1500        let thread = agent.read_with(cx, |agent, _| {
1501            agent.sessions.get(&session_id).unwrap().thread.clone()
1502        });
1503
1504        // Ensure empty threads are not saved, even if they get mutated.
1505        let model = Arc::new(FakeLanguageModel::default());
1506        let summary_model = Arc::new(FakeLanguageModel::default());
1507        thread.update(cx, |thread, cx| {
1508            thread.set_model(model.clone(), cx);
1509            thread.set_summarization_model(Some(summary_model.clone()), cx);
1510        });
1511        cx.run_until_parked();
1512        assert_eq!(history_entries(&history_store, cx), vec![]);
1513
1514        let send = acp_thread.update(cx, |thread, cx| {
1515            thread.send(
1516                vec![
1517                    "What does ".into(),
1518                    acp::ContentBlock::ResourceLink(acp::ResourceLink {
1519                        name: "b.md".into(),
1520                        uri: MentionUri::File {
1521                            abs_path: path!("/a/b.md").into(),
1522                        }
1523                        .to_uri()
1524                        .to_string(),
1525                        annotations: None,
1526                        description: None,
1527                        mime_type: None,
1528                        size: None,
1529                        title: None,
1530                        meta: None,
1531                    }),
1532                    " mean?".into(),
1533                ],
1534                cx,
1535            )
1536        });
1537        let send = cx.foreground_executor().spawn(send);
1538        cx.run_until_parked();
1539
1540        model.send_last_completion_stream_text_chunk("Lorem.");
1541        model.end_last_completion_stream();
1542        cx.run_until_parked();
1543        summary_model
1544            .send_last_completion_stream_text_chunk(&format!("Explaining {}", path!("/a/b.md")));
1545        summary_model.end_last_completion_stream();
1546
1547        send.await.unwrap();
1548        let uri = MentionUri::File {
1549            abs_path: path!("/a/b.md").into(),
1550        }
1551        .to_uri();
1552        acp_thread.read_with(cx, |thread, cx| {
1553            assert_eq!(
1554                thread.to_markdown(cx),
1555                formatdoc! {"
1556                    ## User
1557
1558                    What does [@b.md]({uri}) mean?
1559
1560                    ## Assistant
1561
1562                    Lorem.
1563
1564                "}
1565            )
1566        });
1567
1568        cx.run_until_parked();
1569
1570        // Drop the ACP thread, which should cause the session to be dropped as well.
1571        cx.update(|_| {
1572            drop(thread);
1573            drop(acp_thread);
1574        });
1575        agent.read_with(cx, |agent, _| {
1576            assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
1577        });
1578
1579        // Ensure the thread can be reloaded from disk.
1580        assert_eq!(
1581            history_entries(&history_store, cx),
1582            vec![(
1583                HistoryEntryId::AcpThread(session_id.clone()),
1584                format!("Explaining {}", path!("/a/b.md"))
1585            )]
1586        );
1587        let acp_thread = agent
1588            .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
1589            .await
1590            .unwrap();
1591        acp_thread.read_with(cx, |thread, cx| {
1592            assert_eq!(
1593                thread.to_markdown(cx),
1594                formatdoc! {"
1595                    ## User
1596
1597                    What does [@b.md]({uri}) mean?
1598
1599                    ## Assistant
1600
1601                    Lorem.
1602
1603                "}
1604            )
1605        });
1606    }
1607
1608    fn history_entries(
1609        history: &Entity<HistoryStore>,
1610        cx: &mut TestAppContext,
1611    ) -> Vec<(HistoryEntryId, String)> {
1612        history.read_with(cx, |history, _| {
1613            history
1614                .entries()
1615                .map(|e| (e.id(), e.title().to_string()))
1616                .collect::<Vec<_>>()
1617        })
1618    }
1619
1620    fn init_test(cx: &mut TestAppContext) {
1621        env_logger::try_init().ok();
1622        cx.update(|cx| {
1623            let settings_store = SettingsStore::test(cx);
1624            cx.set_global(settings_store);
1625            Project::init_settings(cx);
1626            agent_settings::init(cx);
1627            language::init(cx);
1628            LanguageModelRegistry::test(cx);
1629        });
1630    }
1631}