agent.rs

   1mod db;
   2mod edit_agent;
   3mod legacy_thread;
   4mod native_agent_server;
   5pub mod outline;
   6mod pattern_extraction;
   7mod shell_parser;
   8mod templates;
   9#[cfg(test)]
  10mod tests;
  11mod thread;
  12mod thread_store;
  13mod tool_permissions;
  14mod tools;
  15
  16use context_server::ContextServerId;
  17pub use db::*;
  18pub use native_agent_server::NativeAgentServer;
  19pub use pattern_extraction::*;
  20pub use templates::*;
  21pub use thread::*;
  22pub use thread_store::*;
  23pub use tool_permissions::*;
  24pub use tools::*;
  25
  26use acp_thread::{
  27    AcpThread, AgentModelSelector, AgentSessionInfo, AgentSessionList, AgentSessionListRequest,
  28    AgentSessionListResponse, UserMessageId,
  29};
  30use agent_client_protocol as acp;
  31use anyhow::{Context as _, Result, anyhow};
  32use chrono::{DateTime, Utc};
  33use collections::{HashMap, HashSet, IndexMap};
  34use fs::Fs;
  35use futures::channel::{mpsc, oneshot};
  36use futures::future::Shared;
  37use futures::{StreamExt, future};
  38use gpui::{
  39    App, AppContext, AsyncApp, Context, Entity, SharedString, Subscription, Task, WeakEntity,
  40};
  41use language_model::{IconOrSvg, LanguageModel, LanguageModelProvider, LanguageModelRegistry};
  42use project::{Project, ProjectItem, ProjectPath, Worktree};
  43use prompt_store::{
  44    ProjectContext, PromptStore, RULES_FILE_NAMES, RulesFileContext, UserRulesContext,
  45    WorktreeContext,
  46};
  47use serde::{Deserialize, Serialize};
  48use settings::{LanguageModelSelection, update_settings_file};
  49use std::any::Any;
  50use std::path::{Path, PathBuf};
  51use std::rc::Rc;
  52use std::sync::Arc;
  53use util::ResultExt;
  54use util::rel_path::RelPath;
  55
  56#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
  57pub struct ProjectSnapshot {
  58    pub worktree_snapshots: Vec<project::telemetry_snapshot::TelemetryWorktreeSnapshot>,
  59    pub timestamp: DateTime<Utc>,
  60}
  61
  62pub struct RulesLoadingError {
  63    pub message: SharedString,
  64}
  65
  66/// Holds both the internal Thread and the AcpThread for a session
  67struct Session {
  68    /// The internal thread that processes messages
  69    thread: Entity<Thread>,
  70    /// The ACP thread that handles protocol communication
  71    acp_thread: WeakEntity<acp_thread::AcpThread>,
  72    pending_save: Task<()>,
  73    _subscriptions: Vec<Subscription>,
  74}
  75
  76pub struct LanguageModels {
  77    /// Access language model by ID
  78    models: HashMap<acp::ModelId, Arc<dyn LanguageModel>>,
  79    /// Cached list for returning language model information
  80    model_list: acp_thread::AgentModelList,
  81    refresh_models_rx: watch::Receiver<()>,
  82    refresh_models_tx: watch::Sender<()>,
  83    _authenticate_all_providers_task: Task<()>,
  84}
  85
  86impl LanguageModels {
  87    fn new(cx: &mut App) -> Self {
  88        let (refresh_models_tx, refresh_models_rx) = watch::channel(());
  89
  90        let mut this = Self {
  91            models: HashMap::default(),
  92            model_list: acp_thread::AgentModelList::Grouped(IndexMap::default()),
  93            refresh_models_rx,
  94            refresh_models_tx,
  95            _authenticate_all_providers_task: Self::authenticate_all_language_model_providers(cx),
  96        };
  97        this.refresh_list(cx);
  98        this
  99    }
 100
 101    fn refresh_list(&mut self, cx: &App) {
 102        let providers = LanguageModelRegistry::global(cx)
 103            .read(cx)
 104            .visible_providers()
 105            .into_iter()
 106            .filter(|provider| provider.is_authenticated(cx))
 107            .collect::<Vec<_>>();
 108
 109        let mut language_model_list = IndexMap::default();
 110        let mut recommended_models = HashSet::default();
 111
 112        let mut recommended = Vec::new();
 113        for provider in &providers {
 114            for model in provider.recommended_models(cx) {
 115                recommended_models.insert((model.provider_id(), model.id()));
 116                recommended.push(Self::map_language_model_to_info(&model, provider));
 117            }
 118        }
 119        if !recommended.is_empty() {
 120            language_model_list.insert(
 121                acp_thread::AgentModelGroupName("Recommended".into()),
 122                recommended,
 123            );
 124        }
 125
 126        let mut models = HashMap::default();
 127        for provider in providers {
 128            let mut provider_models = Vec::new();
 129            for model in provider.provided_models(cx) {
 130                let model_info = Self::map_language_model_to_info(&model, &provider);
 131                let model_id = model_info.id.clone();
 132                provider_models.push(model_info);
 133                models.insert(model_id, model);
 134            }
 135            if !provider_models.is_empty() {
 136                language_model_list.insert(
 137                    acp_thread::AgentModelGroupName(provider.name().0.clone()),
 138                    provider_models,
 139                );
 140            }
 141        }
 142
 143        self.models = models;
 144        self.model_list = acp_thread::AgentModelList::Grouped(language_model_list);
 145        self.refresh_models_tx.send(()).ok();
 146    }
 147
 148    fn watch(&self) -> watch::Receiver<()> {
 149        self.refresh_models_rx.clone()
 150    }
 151
 152    pub fn model_from_id(&self, model_id: &acp::ModelId) -> Option<Arc<dyn LanguageModel>> {
 153        self.models.get(model_id).cloned()
 154    }
 155
 156    fn map_language_model_to_info(
 157        model: &Arc<dyn LanguageModel>,
 158        provider: &Arc<dyn LanguageModelProvider>,
 159    ) -> acp_thread::AgentModelInfo {
 160        acp_thread::AgentModelInfo {
 161            id: Self::model_id(model),
 162            name: model.name().0,
 163            description: None,
 164            icon: Some(match provider.icon() {
 165                IconOrSvg::Svg(path) => acp_thread::AgentModelIcon::Path(path),
 166                IconOrSvg::Icon(name) => acp_thread::AgentModelIcon::Named(name),
 167            }),
 168        }
 169    }
 170
 171    fn model_id(model: &Arc<dyn LanguageModel>) -> acp::ModelId {
 172        acp::ModelId::new(format!("{}/{}", model.provider_id().0, model.id().0))
 173    }
 174
 175    fn authenticate_all_language_model_providers(cx: &mut App) -> Task<()> {
 176        let authenticate_all_providers = LanguageModelRegistry::global(cx)
 177            .read(cx)
 178            .visible_providers()
 179            .iter()
 180            .map(|provider| (provider.id(), provider.name(), provider.authenticate(cx)))
 181            .collect::<Vec<_>>();
 182
 183        cx.background_spawn(async move {
 184            for (provider_id, provider_name, authenticate_task) in authenticate_all_providers {
 185                if let Err(err) = authenticate_task.await {
 186                    match err {
 187                        language_model::AuthenticateError::CredentialsNotFound => {
 188                            // Since we're authenticating these providers in the
 189                            // background for the purposes of populating the
 190                            // language selector, we don't care about providers
 191                            // where the credentials are not found.
 192                        }
 193                        language_model::AuthenticateError::ConnectionRefused => {
 194                            // Not logging connection refused errors as they are mostly from LM Studio's noisy auth failures.
 195                            // LM Studio only has one auth method (endpoint call) which fails for users who haven't enabled it.
 196                            // TODO: Better manage LM Studio auth logic to avoid these noisy failures.
 197                        }
 198                        _ => {
 199                            // Some providers have noisy failure states that we
 200                            // don't want to spam the logs with every time the
 201                            // language model selector is initialized.
 202                            //
 203                            // Ideally these should have more clear failure modes
 204                            // that we know are safe to ignore here, like what we do
 205                            // with `CredentialsNotFound` above.
 206                            match provider_id.0.as_ref() {
 207                                "lmstudio" | "ollama" => {
 208                                    // LM Studio and Ollama both make fetch requests to the local APIs to determine if they are "authenticated".
 209                                    //
 210                                    // These fail noisily, so we don't log them.
 211                                }
 212                                "copilot_chat" => {
 213                                    // Copilot Chat returns an error if Copilot is not enabled, so we don't log those errors.
 214                                }
 215                                _ => {
 216                                    log::error!(
 217                                        "Failed to authenticate provider: {}: {err:#}",
 218                                        provider_name.0
 219                                    );
 220                                }
 221                            }
 222                        }
 223                    }
 224                }
 225            }
 226        })
 227    }
 228}
 229
 230pub struct NativeAgent {
 231    /// Session ID -> Session mapping
 232    sessions: HashMap<acp::SessionId, Session>,
 233    thread_store: Entity<ThreadStore>,
 234    /// Shared project context for all threads
 235    project_context: Entity<ProjectContext>,
 236    project_context_needs_refresh: watch::Sender<()>,
 237    _maintain_project_context: Task<Result<()>>,
 238    context_server_registry: Entity<ContextServerRegistry>,
 239    /// Shared templates for all threads
 240    templates: Arc<Templates>,
 241    /// Cached model information
 242    models: LanguageModels,
 243    project: Entity<Project>,
 244    prompt_store: Option<Entity<PromptStore>>,
 245    fs: Arc<dyn Fs>,
 246    _subscriptions: Vec<Subscription>,
 247}
 248
 249impl NativeAgent {
 250    pub async fn new(
 251        project: Entity<Project>,
 252        thread_store: Entity<ThreadStore>,
 253        templates: Arc<Templates>,
 254        prompt_store: Option<Entity<PromptStore>>,
 255        fs: Arc<dyn Fs>,
 256        cx: &mut AsyncApp,
 257    ) -> Result<Entity<NativeAgent>> {
 258        log::debug!("Creating new NativeAgent");
 259
 260        let project_context = cx
 261            .update(|cx| Self::build_project_context(&project, prompt_store.as_ref(), cx))
 262            .await;
 263
 264        Ok(cx.new(|cx| {
 265            let context_server_store = project.read(cx).context_server_store();
 266            let context_server_registry =
 267                cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
 268
 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                cx.subscribe(
 276                    &context_server_store,
 277                    Self::handle_context_server_store_updated,
 278                ),
 279                cx.subscribe(
 280                    &context_server_registry,
 281                    Self::handle_context_server_registry_event,
 282                ),
 283            ];
 284            if let Some(prompt_store) = prompt_store.as_ref() {
 285                subscriptions.push(cx.subscribe(prompt_store, Self::handle_prompts_updated_event))
 286            }
 287
 288            let (project_context_needs_refresh_tx, project_context_needs_refresh_rx) =
 289                watch::channel(());
 290            Self {
 291                sessions: HashMap::default(),
 292                thread_store,
 293                project_context: cx.new(|_| project_context),
 294                project_context_needs_refresh: project_context_needs_refresh_tx,
 295                _maintain_project_context: cx.spawn(async move |this, cx| {
 296                    Self::maintain_project_context(this, project_context_needs_refresh_rx, cx).await
 297                }),
 298                context_server_registry,
 299                templates,
 300                models: LanguageModels::new(cx),
 301                project,
 302                prompt_store,
 303                fs,
 304                _subscriptions: subscriptions,
 305            }
 306        }))
 307    }
 308
 309    fn new_session(
 310        &mut self,
 311        project: Entity<Project>,
 312        cx: &mut Context<Self>,
 313    ) -> Entity<AcpThread> {
 314        // Create Thread
 315        // Fetch default model from registry settings
 316        let registry = LanguageModelRegistry::read_global(cx);
 317        // Log available models for debugging
 318        let available_count = registry.available_models(cx).count();
 319        log::debug!("Total available models: {}", available_count);
 320
 321        let default_model = registry.default_model().and_then(|default_model| {
 322            self.models
 323                .model_from_id(&LanguageModels::model_id(&default_model.model))
 324        });
 325        let thread = cx.new(|cx| {
 326            Thread::new(
 327                project.clone(),
 328                self.project_context.clone(),
 329                self.context_server_registry.clone(),
 330                self.templates.clone(),
 331                default_model,
 332                cx,
 333            )
 334        });
 335
 336        self.register_session(thread, cx)
 337    }
 338
 339    fn register_session(
 340        &mut self,
 341        thread_handle: Entity<Thread>,
 342        cx: &mut Context<Self>,
 343    ) -> Entity<AcpThread> {
 344        let connection = Rc::new(NativeAgentConnection(cx.entity()));
 345
 346        let thread = thread_handle.read(cx);
 347        let session_id = thread.id().clone();
 348        let title = thread.title();
 349        let project = thread.project.clone();
 350        let action_log = thread.action_log.clone();
 351        let prompt_capabilities_rx = thread.prompt_capabilities_rx.clone();
 352        let acp_thread = cx.new(|cx| {
 353            acp_thread::AcpThread::new(
 354                title,
 355                connection,
 356                project.clone(),
 357                action_log.clone(),
 358                session_id.clone(),
 359                prompt_capabilities_rx,
 360                cx,
 361            )
 362        });
 363
 364        let registry = LanguageModelRegistry::read_global(cx);
 365        let summarization_model = registry.thread_summary_model().map(|c| c.model);
 366
 367        thread_handle.update(cx, |thread, cx| {
 368            thread.set_summarization_model(summarization_model, cx);
 369            thread.add_default_tools(
 370                Rc::new(AcpThreadEnvironment {
 371                    acp_thread: acp_thread.downgrade(),
 372                }) as _,
 373                cx,
 374            )
 375        });
 376
 377        let subscriptions = vec![
 378            cx.observe_release(&acp_thread, |this, acp_thread, _cx| {
 379                this.sessions.remove(acp_thread.session_id());
 380            }),
 381            cx.subscribe(&thread_handle, Self::handle_thread_title_updated),
 382            cx.subscribe(&thread_handle, Self::handle_thread_token_usage_updated),
 383            cx.observe(&thread_handle, move |this, thread, cx| {
 384                this.save_thread(thread, cx)
 385            }),
 386        ];
 387
 388        self.sessions.insert(
 389            session_id,
 390            Session {
 391                thread: thread_handle,
 392                acp_thread: acp_thread.downgrade(),
 393                _subscriptions: subscriptions,
 394                pending_save: Task::ready(()),
 395            },
 396        );
 397
 398        self.update_available_commands(cx);
 399
 400        acp_thread
 401    }
 402
 403    pub fn models(&self) -> &LanguageModels {
 404        &self.models
 405    }
 406
 407    async fn maintain_project_context(
 408        this: WeakEntity<Self>,
 409        mut needs_refresh: watch::Receiver<()>,
 410        cx: &mut AsyncApp,
 411    ) -> Result<()> {
 412        while needs_refresh.changed().await.is_ok() {
 413            let project_context = this
 414                .update(cx, |this, cx| {
 415                    Self::build_project_context(&this.project, this.prompt_store.as_ref(), cx)
 416                })?
 417                .await;
 418            this.update(cx, |this, cx| {
 419                this.project_context = cx.new(|_| project_context);
 420            })?;
 421        }
 422
 423        Ok(())
 424    }
 425
 426    fn build_project_context(
 427        project: &Entity<Project>,
 428        prompt_store: Option<&Entity<PromptStore>>,
 429        cx: &mut App,
 430    ) -> Task<ProjectContext> {
 431        let worktrees = project.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
 432        let worktree_tasks = worktrees
 433            .into_iter()
 434            .map(|worktree| {
 435                Self::load_worktree_info_for_system_prompt(worktree, project.clone(), cx)
 436            })
 437            .collect::<Vec<_>>();
 438        let default_user_rules_task = if let Some(prompt_store) = prompt_store.as_ref() {
 439            prompt_store.read_with(cx, |prompt_store, cx| {
 440                let prompts = prompt_store.default_prompt_metadata();
 441                let load_tasks = prompts.into_iter().map(|prompt_metadata| {
 442                    let contents = prompt_store.load(prompt_metadata.id, cx);
 443                    async move { (contents.await, prompt_metadata) }
 444                });
 445                cx.background_spawn(future::join_all(load_tasks))
 446            })
 447        } else {
 448            Task::ready(vec![])
 449        };
 450
 451        cx.spawn(async move |_cx| {
 452            let (worktrees, default_user_rules) =
 453                future::join(future::join_all(worktree_tasks), default_user_rules_task).await;
 454
 455            let worktrees = worktrees
 456                .into_iter()
 457                .map(|(worktree, _rules_error)| {
 458                    // TODO: show error message
 459                    // if let Some(rules_error) = rules_error {
 460                    //     this.update(cx, |_, cx| cx.emit(rules_error)).ok();
 461                    // }
 462                    worktree
 463                })
 464                .collect::<Vec<_>>();
 465
 466            let default_user_rules = default_user_rules
 467                .into_iter()
 468                .flat_map(|(contents, prompt_metadata)| match contents {
 469                    Ok(contents) => Some(UserRulesContext {
 470                        uuid: prompt_metadata.id.as_user()?,
 471                        title: prompt_metadata.title.map(|title| title.to_string()),
 472                        contents,
 473                    }),
 474                    Err(_err) => {
 475                        // TODO: show error message
 476                        // this.update(cx, |_, cx| {
 477                        //     cx.emit(RulesLoadingError {
 478                        //         message: format!("{err:?}").into(),
 479                        //     });
 480                        // })
 481                        // .ok();
 482                        None
 483                    }
 484                })
 485                .collect::<Vec<_>>();
 486
 487            ProjectContext::new(worktrees, default_user_rules)
 488        })
 489    }
 490
 491    fn load_worktree_info_for_system_prompt(
 492        worktree: Entity<Worktree>,
 493        project: Entity<Project>,
 494        cx: &mut App,
 495    ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
 496        let tree = worktree.read(cx);
 497        let root_name = tree.root_name_str().into();
 498        let abs_path = tree.abs_path();
 499
 500        let mut context = WorktreeContext {
 501            root_name,
 502            abs_path,
 503            rules_file: None,
 504        };
 505
 506        let rules_task = Self::load_worktree_rules_file(worktree, project, cx);
 507        let Some(rules_task) = rules_task else {
 508            return Task::ready((context, None));
 509        };
 510
 511        cx.spawn(async move |_| {
 512            let (rules_file, rules_file_error) = match rules_task.await {
 513                Ok(rules_file) => (Some(rules_file), None),
 514                Err(err) => (
 515                    None,
 516                    Some(RulesLoadingError {
 517                        message: format!("{err}").into(),
 518                    }),
 519                ),
 520            };
 521            context.rules_file = rules_file;
 522            (context, rules_file_error)
 523        })
 524    }
 525
 526    fn load_worktree_rules_file(
 527        worktree: Entity<Worktree>,
 528        project: Entity<Project>,
 529        cx: &mut App,
 530    ) -> Option<Task<Result<RulesFileContext>>> {
 531        let worktree = worktree.read(cx);
 532        let worktree_id = worktree.id();
 533        let selected_rules_file = RULES_FILE_NAMES
 534            .into_iter()
 535            .filter_map(|name| {
 536                worktree
 537                    .entry_for_path(RelPath::unix(name).unwrap())
 538                    .filter(|entry| entry.is_file())
 539                    .map(|entry| entry.path.clone())
 540            })
 541            .next();
 542
 543        // Note that Cline supports `.clinerules` being a directory, but that is not currently
 544        // supported. This doesn't seem to occur often in GitHub repositories.
 545        selected_rules_file.map(|path_in_worktree| {
 546            let project_path = ProjectPath {
 547                worktree_id,
 548                path: path_in_worktree.clone(),
 549            };
 550            let buffer_task =
 551                project.update(cx, |project, cx| project.open_buffer(project_path, cx));
 552            let rope_task = cx.spawn(async move |cx| {
 553                let buffer = buffer_task.await?;
 554                let (project_entry_id, rope) = buffer.read_with(cx, |buffer, cx| {
 555                    let project_entry_id = buffer.entry_id(cx).context("buffer has no file")?;
 556                    anyhow::Ok((project_entry_id, buffer.as_rope().clone()))
 557                })?;
 558                anyhow::Ok((project_entry_id, rope))
 559            });
 560            // Build a string from the rope on a background thread.
 561            cx.background_spawn(async move {
 562                let (project_entry_id, rope) = rope_task.await?;
 563                anyhow::Ok(RulesFileContext {
 564                    path_in_worktree,
 565                    text: rope.to_string().trim().to_string(),
 566                    project_entry_id: project_entry_id.to_usize(),
 567                })
 568            })
 569        })
 570    }
 571
 572    fn handle_thread_title_updated(
 573        &mut self,
 574        thread: Entity<Thread>,
 575        _: &TitleUpdated,
 576        cx: &mut Context<Self>,
 577    ) {
 578        let session_id = thread.read(cx).id();
 579        let Some(session) = self.sessions.get(session_id) else {
 580            return;
 581        };
 582        let thread = thread.downgrade();
 583        let acp_thread = session.acp_thread.clone();
 584        cx.spawn(async move |_, cx| {
 585            let title = thread.read_with(cx, |thread, _| thread.title())?;
 586            let task = acp_thread.update(cx, |acp_thread, cx| acp_thread.set_title(title, cx))?;
 587            task.await
 588        })
 589        .detach_and_log_err(cx);
 590    }
 591
 592    fn handle_thread_token_usage_updated(
 593        &mut self,
 594        thread: Entity<Thread>,
 595        usage: &TokenUsageUpdated,
 596        cx: &mut Context<Self>,
 597    ) {
 598        let Some(session) = self.sessions.get(thread.read(cx).id()) else {
 599            return;
 600        };
 601        session
 602            .acp_thread
 603            .update(cx, |acp_thread, cx| {
 604                acp_thread.update_token_usage(usage.0.clone(), cx);
 605            })
 606            .ok();
 607    }
 608
 609    fn handle_project_event(
 610        &mut self,
 611        _project: Entity<Project>,
 612        event: &project::Event,
 613        _cx: &mut Context<Self>,
 614    ) {
 615        match event {
 616            project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
 617                self.project_context_needs_refresh.send(()).ok();
 618            }
 619            project::Event::WorktreeUpdatedEntries(_, items) => {
 620                if items.iter().any(|(path, _, _)| {
 621                    RULES_FILE_NAMES
 622                        .iter()
 623                        .any(|name| path.as_ref() == RelPath::unix(name).unwrap())
 624                }) {
 625                    self.project_context_needs_refresh.send(()).ok();
 626                }
 627            }
 628            _ => {}
 629        }
 630    }
 631
 632    fn handle_prompts_updated_event(
 633        &mut self,
 634        _prompt_store: Entity<PromptStore>,
 635        _event: &prompt_store::PromptsUpdatedEvent,
 636        _cx: &mut Context<Self>,
 637    ) {
 638        self.project_context_needs_refresh.send(()).ok();
 639    }
 640
 641    fn handle_models_updated_event(
 642        &mut self,
 643        _registry: Entity<LanguageModelRegistry>,
 644        _event: &language_model::Event,
 645        cx: &mut Context<Self>,
 646    ) {
 647        self.models.refresh_list(cx);
 648
 649        let registry = LanguageModelRegistry::read_global(cx);
 650        let default_model = registry.default_model().map(|m| m.model);
 651        let summarization_model = registry.thread_summary_model().map(|m| m.model);
 652
 653        for session in self.sessions.values_mut() {
 654            session.thread.update(cx, |thread, cx| {
 655                if thread.model().is_none()
 656                    && let Some(model) = default_model.clone()
 657                {
 658                    thread.set_model(model, cx);
 659                    cx.notify();
 660                }
 661                thread.set_summarization_model(summarization_model.clone(), cx);
 662            });
 663        }
 664    }
 665
 666    fn handle_context_server_store_updated(
 667        &mut self,
 668        _store: Entity<project::context_server_store::ContextServerStore>,
 669        _event: &project::context_server_store::Event,
 670        cx: &mut Context<Self>,
 671    ) {
 672        self.update_available_commands(cx);
 673    }
 674
 675    fn handle_context_server_registry_event(
 676        &mut self,
 677        _registry: Entity<ContextServerRegistry>,
 678        event: &ContextServerRegistryEvent,
 679        cx: &mut Context<Self>,
 680    ) {
 681        match event {
 682            ContextServerRegistryEvent::ToolsChanged => {}
 683            ContextServerRegistryEvent::PromptsChanged => {
 684                self.update_available_commands(cx);
 685            }
 686        }
 687    }
 688
 689    fn update_available_commands(&self, cx: &mut Context<Self>) {
 690        let available_commands = self.build_available_commands(cx);
 691        for session in self.sessions.values() {
 692            if let Some(acp_thread) = session.acp_thread.upgrade() {
 693                acp_thread.update(cx, |thread, cx| {
 694                    thread
 695                        .handle_session_update(
 696                            acp::SessionUpdate::AvailableCommandsUpdate(
 697                                acp::AvailableCommandsUpdate::new(available_commands.clone()),
 698                            ),
 699                            cx,
 700                        )
 701                        .log_err();
 702                });
 703            }
 704        }
 705    }
 706
 707    fn build_available_commands(&self, cx: &App) -> Vec<acp::AvailableCommand> {
 708        let registry = self.context_server_registry.read(cx);
 709
 710        let mut prompt_name_counts: HashMap<&str, usize> = HashMap::default();
 711        for context_server_prompt in registry.prompts() {
 712            *prompt_name_counts
 713                .entry(context_server_prompt.prompt.name.as_str())
 714                .or_insert(0) += 1;
 715        }
 716
 717        registry
 718            .prompts()
 719            .flat_map(|context_server_prompt| {
 720                let prompt = &context_server_prompt.prompt;
 721
 722                let should_prefix = prompt_name_counts
 723                    .get(prompt.name.as_str())
 724                    .copied()
 725                    .unwrap_or(0)
 726                    > 1;
 727
 728                let name = if should_prefix {
 729                    format!("{}.{}", context_server_prompt.server_id, prompt.name)
 730                } else {
 731                    prompt.name.clone()
 732                };
 733
 734                let mut command = acp::AvailableCommand::new(
 735                    name,
 736                    prompt.description.clone().unwrap_or_default(),
 737                );
 738
 739                match prompt.arguments.as_deref() {
 740                    Some([arg]) => {
 741                        let hint = format!("<{}>", arg.name);
 742
 743                        command = command.input(acp::AvailableCommandInput::Unstructured(
 744                            acp::UnstructuredCommandInput::new(hint),
 745                        ));
 746                    }
 747                    Some([]) | None => {}
 748                    Some(_) => {
 749                        // skip >1 argument commands since we don't support them yet
 750                        return None;
 751                    }
 752                }
 753
 754                Some(command)
 755            })
 756            .collect()
 757    }
 758
 759    pub fn load_thread(
 760        &mut self,
 761        id: acp::SessionId,
 762        cx: &mut Context<Self>,
 763    ) -> Task<Result<Entity<Thread>>> {
 764        let database_future = ThreadsDatabase::connect(cx);
 765        cx.spawn(async move |this, cx| {
 766            let database = database_future.await.map_err(|err| anyhow!(err))?;
 767            let db_thread = database
 768                .load_thread(id.clone())
 769                .await?
 770                .with_context(|| format!("no thread found with ID: {id:?}"))?;
 771
 772            this.update(cx, |this, cx| {
 773                let summarization_model = LanguageModelRegistry::read_global(cx)
 774                    .thread_summary_model()
 775                    .map(|c| c.model);
 776
 777                cx.new(|cx| {
 778                    let mut thread = Thread::from_db(
 779                        id.clone(),
 780                        db_thread,
 781                        this.project.clone(),
 782                        this.project_context.clone(),
 783                        this.context_server_registry.clone(),
 784                        this.templates.clone(),
 785                        cx,
 786                    );
 787                    thread.set_summarization_model(summarization_model, cx);
 788                    thread
 789                })
 790            })
 791        })
 792    }
 793
 794    pub fn open_thread(
 795        &mut self,
 796        id: acp::SessionId,
 797        cx: &mut Context<Self>,
 798    ) -> Task<Result<Entity<AcpThread>>> {
 799        let task = self.load_thread(id, cx);
 800        cx.spawn(async move |this, cx| {
 801            let thread = task.await?;
 802            let acp_thread =
 803                this.update(cx, |this, cx| this.register_session(thread.clone(), cx))?;
 804            let events = thread.update(cx, |thread, cx| thread.replay(cx));
 805            cx.update(|cx| {
 806                NativeAgentConnection::handle_thread_events(events, acp_thread.downgrade(), cx)
 807            })
 808            .await?;
 809            Ok(acp_thread)
 810        })
 811    }
 812
 813    pub fn thread_summary(
 814        &mut self,
 815        id: acp::SessionId,
 816        cx: &mut Context<Self>,
 817    ) -> Task<Result<SharedString>> {
 818        let thread = self.open_thread(id.clone(), cx);
 819        cx.spawn(async move |this, cx| {
 820            let acp_thread = thread.await?;
 821            let result = this
 822                .update(cx, |this, cx| {
 823                    this.sessions
 824                        .get(&id)
 825                        .unwrap()
 826                        .thread
 827                        .update(cx, |thread, cx| thread.summary(cx))
 828                })?
 829                .await
 830                .context("Failed to generate summary")?;
 831            drop(acp_thread);
 832            Ok(result)
 833        })
 834    }
 835
 836    fn save_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
 837        if thread.read(cx).is_empty() {
 838            return;
 839        }
 840
 841        let database_future = ThreadsDatabase::connect(cx);
 842        let (id, db_thread) =
 843            thread.update(cx, |thread, cx| (thread.id().clone(), thread.to_db(cx)));
 844        let Some(session) = self.sessions.get_mut(&id) else {
 845            return;
 846        };
 847        let thread_store = self.thread_store.clone();
 848        session.pending_save = cx.spawn(async move |_, cx| {
 849            let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
 850                return;
 851            };
 852            let db_thread = db_thread.await;
 853            database.save_thread(id, db_thread).await.log_err();
 854            thread_store.update(cx, |store, cx| store.reload(cx));
 855        });
 856    }
 857
 858    fn send_mcp_prompt(
 859        &self,
 860        message_id: UserMessageId,
 861        session_id: agent_client_protocol::SessionId,
 862        prompt_name: String,
 863        server_id: ContextServerId,
 864        arguments: HashMap<String, String>,
 865        original_content: Vec<acp::ContentBlock>,
 866        cx: &mut Context<Self>,
 867    ) -> Task<Result<acp::PromptResponse>> {
 868        let server_store = self.context_server_registry.read(cx).server_store().clone();
 869        let path_style = self.project.read(cx).path_style(cx);
 870
 871        cx.spawn(async move |this, cx| {
 872            let prompt =
 873                crate::get_prompt(&server_store, &server_id, &prompt_name, arguments, cx).await?;
 874
 875            let (acp_thread, thread) = this.update(cx, |this, _cx| {
 876                let session = this
 877                    .sessions
 878                    .get(&session_id)
 879                    .context("Failed to get session")?;
 880                anyhow::Ok((session.acp_thread.clone(), session.thread.clone()))
 881            })??;
 882
 883            let mut last_is_user = true;
 884
 885            thread.update(cx, |thread, cx| {
 886                thread.push_acp_user_block(
 887                    message_id,
 888                    original_content.into_iter().skip(1),
 889                    path_style,
 890                    cx,
 891                );
 892            });
 893
 894            for message in prompt.messages {
 895                let context_server::types::PromptMessage { role, content } = message;
 896                let block = mcp_message_content_to_acp_content_block(content);
 897
 898                match role {
 899                    context_server::types::Role::User => {
 900                        let id = acp_thread::UserMessageId::new();
 901
 902                        acp_thread.update(cx, |acp_thread, cx| {
 903                            acp_thread.push_user_content_block_with_indent(
 904                                Some(id.clone()),
 905                                block.clone(),
 906                                true,
 907                                cx,
 908                            );
 909                        })?;
 910
 911                        thread.update(cx, |thread, cx| {
 912                            thread.push_acp_user_block(id, [block], path_style, cx);
 913                        });
 914                    }
 915                    context_server::types::Role::Assistant => {
 916                        acp_thread.update(cx, |acp_thread, cx| {
 917                            acp_thread.push_assistant_content_block_with_indent(
 918                                block.clone(),
 919                                false,
 920                                true,
 921                                cx,
 922                            );
 923                        })?;
 924
 925                        thread.update(cx, |thread, cx| {
 926                            thread.push_acp_agent_block(block, cx);
 927                        });
 928                    }
 929                }
 930
 931                last_is_user = role == context_server::types::Role::User;
 932            }
 933
 934            let response_stream = thread.update(cx, |thread, cx| {
 935                if last_is_user {
 936                    thread.send_existing(cx)
 937                } else {
 938                    // Resume if MCP prompt did not end with a user message
 939                    thread.resume(cx)
 940                }
 941            })?;
 942
 943            cx.update(|cx| {
 944                NativeAgentConnection::handle_thread_events(response_stream, acp_thread, cx)
 945            })
 946            .await
 947        })
 948    }
 949}
 950
 951/// Wrapper struct that implements the AgentConnection trait
 952#[derive(Clone)]
 953pub struct NativeAgentConnection(pub Entity<NativeAgent>);
 954
 955impl NativeAgentConnection {
 956    pub fn thread(&self, session_id: &acp::SessionId, cx: &App) -> Option<Entity<Thread>> {
 957        self.0
 958            .read(cx)
 959            .sessions
 960            .get(session_id)
 961            .map(|session| session.thread.clone())
 962    }
 963
 964    pub fn load_thread(&self, id: acp::SessionId, cx: &mut App) -> Task<Result<Entity<Thread>>> {
 965        self.0.update(cx, |this, cx| this.load_thread(id, cx))
 966    }
 967
 968    fn run_turn(
 969        &self,
 970        session_id: acp::SessionId,
 971        cx: &mut App,
 972        f: impl 'static
 973        + FnOnce(Entity<Thread>, &mut App) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>,
 974    ) -> Task<Result<acp::PromptResponse>> {
 975        let Some((thread, acp_thread)) = self.0.update(cx, |agent, _cx| {
 976            agent
 977                .sessions
 978                .get_mut(&session_id)
 979                .map(|s| (s.thread.clone(), s.acp_thread.clone()))
 980        }) else {
 981            return Task::ready(Err(anyhow!("Session not found")));
 982        };
 983        log::debug!("Found session for: {}", session_id);
 984
 985        let response_stream = match f(thread, cx) {
 986            Ok(stream) => stream,
 987            Err(err) => return Task::ready(Err(err)),
 988        };
 989        Self::handle_thread_events(response_stream, acp_thread, cx)
 990    }
 991
 992    fn handle_thread_events(
 993        mut events: mpsc::UnboundedReceiver<Result<ThreadEvent>>,
 994        acp_thread: WeakEntity<AcpThread>,
 995        cx: &App,
 996    ) -> Task<Result<acp::PromptResponse>> {
 997        cx.spawn(async move |cx| {
 998            // Handle response stream and forward to session.acp_thread
 999            while let Some(result) = events.next().await {
1000                match result {
1001                    Ok(event) => {
1002                        log::trace!("Received completion event: {:?}", event);
1003
1004                        match event {
1005                            ThreadEvent::UserMessage(message) => {
1006                                acp_thread.update(cx, |thread, cx| {
1007                                    for content in message.content {
1008                                        thread.push_user_content_block(
1009                                            Some(message.id.clone()),
1010                                            content.into(),
1011                                            cx,
1012                                        );
1013                                    }
1014                                })?;
1015                            }
1016                            ThreadEvent::AgentText(text) => {
1017                                acp_thread.update(cx, |thread, cx| {
1018                                    thread.push_assistant_content_block(text.into(), false, cx)
1019                                })?;
1020                            }
1021                            ThreadEvent::AgentThinking(text) => {
1022                                acp_thread.update(cx, |thread, cx| {
1023                                    thread.push_assistant_content_block(text.into(), true, cx)
1024                                })?;
1025                            }
1026                            ThreadEvent::ToolCallAuthorization(ToolCallAuthorization {
1027                                tool_call,
1028                                options,
1029                                response,
1030                                context: _,
1031                            }) => {
1032                                let outcome_task = acp_thread.update(cx, |thread, cx| {
1033                                    thread.request_tool_call_authorization(
1034                                        tool_call, options, true, cx,
1035                                    )
1036                                })??;
1037                                cx.background_spawn(async move {
1038                                    if let acp::RequestPermissionOutcome::Selected(
1039                                        acp::SelectedPermissionOutcome { option_id, .. },
1040                                    ) = outcome_task.await
1041                                    {
1042                                        response
1043                                            .send(option_id)
1044                                            .map(|_| anyhow!("authorization receiver was dropped"))
1045                                            .log_err();
1046                                    }
1047                                })
1048                                .detach();
1049                            }
1050                            ThreadEvent::ToolCall(tool_call) => {
1051                                acp_thread.update(cx, |thread, cx| {
1052                                    thread.upsert_tool_call(tool_call, cx)
1053                                })??;
1054                            }
1055                            ThreadEvent::ToolCallUpdate(update) => {
1056                                acp_thread.update(cx, |thread, cx| {
1057                                    thread.update_tool_call(update, cx)
1058                                })??;
1059                            }
1060                            ThreadEvent::Retry(status) => {
1061                                acp_thread.update(cx, |thread, cx| {
1062                                    thread.update_retry_status(status, cx)
1063                                })?;
1064                            }
1065                            ThreadEvent::Stop(stop_reason) => {
1066                                log::debug!("Assistant message complete: {:?}", stop_reason);
1067                                return Ok(acp::PromptResponse::new(stop_reason));
1068                            }
1069                        }
1070                    }
1071                    Err(e) => {
1072                        log::error!("Error in model response stream: {:?}", e);
1073                        return Err(e);
1074                    }
1075                }
1076            }
1077
1078            log::debug!("Response stream completed");
1079            anyhow::Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
1080        })
1081    }
1082}
1083
1084struct Command<'a> {
1085    prompt_name: &'a str,
1086    arg_value: &'a str,
1087    explicit_server_id: Option<&'a str>,
1088}
1089
1090impl<'a> Command<'a> {
1091    fn parse(prompt: &'a [acp::ContentBlock]) -> Option<Self> {
1092        let acp::ContentBlock::Text(text_content) = prompt.first()? else {
1093            return None;
1094        };
1095        let text = text_content.text.trim();
1096        let command = text.strip_prefix('/')?;
1097        let (command, arg_value) = command
1098            .split_once(char::is_whitespace)
1099            .unwrap_or((command, ""));
1100
1101        if let Some((server_id, prompt_name)) = command.split_once('.') {
1102            Some(Self {
1103                prompt_name,
1104                arg_value,
1105                explicit_server_id: Some(server_id),
1106            })
1107        } else {
1108            Some(Self {
1109                prompt_name: command,
1110                arg_value,
1111                explicit_server_id: None,
1112            })
1113        }
1114    }
1115}
1116
1117struct NativeAgentModelSelector {
1118    session_id: acp::SessionId,
1119    connection: NativeAgentConnection,
1120}
1121
1122impl acp_thread::AgentModelSelector for NativeAgentModelSelector {
1123    fn list_models(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelList>> {
1124        log::debug!("NativeAgentConnection::list_models called");
1125        let list = self.connection.0.read(cx).models.model_list.clone();
1126        Task::ready(if list.is_empty() {
1127            Err(anyhow::anyhow!("No models available"))
1128        } else {
1129            Ok(list)
1130        })
1131    }
1132
1133    fn select_model(&self, model_id: acp::ModelId, cx: &mut App) -> Task<Result<()>> {
1134        log::debug!(
1135            "Setting model for session {}: {}",
1136            self.session_id,
1137            model_id
1138        );
1139        let Some(thread) = self
1140            .connection
1141            .0
1142            .read(cx)
1143            .sessions
1144            .get(&self.session_id)
1145            .map(|session| session.thread.clone())
1146        else {
1147            return Task::ready(Err(anyhow!("Session not found")));
1148        };
1149
1150        let Some(model) = self.connection.0.read(cx).models.model_from_id(&model_id) else {
1151            return Task::ready(Err(anyhow!("Invalid model ID {}", model_id)));
1152        };
1153
1154        // We want to reset the effort level when switching models, as the currently-selected effort level may
1155        // not be compatible.
1156        let effort = model
1157            .default_effort_level()
1158            .map(|effort_level| effort_level.value.to_string());
1159
1160        thread.update(cx, |thread, cx| {
1161            thread.set_model(model.clone(), cx);
1162            thread.set_thinking_effort(effort.clone(), cx);
1163        });
1164
1165        update_settings_file(
1166            self.connection.0.read(cx).fs.clone(),
1167            cx,
1168            move |settings, cx| {
1169                let provider = model.provider_id().0.to_string();
1170                let model = model.id().0.to_string();
1171                let enable_thinking = settings
1172                    .agent
1173                    .as_ref()
1174                    .and_then(|agent| {
1175                        agent
1176                            .default_model
1177                            .as_ref()
1178                            .map(|default_model| default_model.enable_thinking)
1179                    })
1180                    .unwrap_or_else(|| thread.read(cx).thinking_enabled());
1181                settings
1182                    .agent
1183                    .get_or_insert_default()
1184                    .set_model(LanguageModelSelection {
1185                        provider: provider.into(),
1186                        model,
1187                        enable_thinking,
1188                        effort,
1189                    });
1190            },
1191        );
1192
1193        Task::ready(Ok(()))
1194    }
1195
1196    fn selected_model(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelInfo>> {
1197        let Some(thread) = self
1198            .connection
1199            .0
1200            .read(cx)
1201            .sessions
1202            .get(&self.session_id)
1203            .map(|session| session.thread.clone())
1204        else {
1205            return Task::ready(Err(anyhow!("Session not found")));
1206        };
1207        let Some(model) = thread.read(cx).model() else {
1208            return Task::ready(Err(anyhow!("Model not found")));
1209        };
1210        let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&model.provider_id())
1211        else {
1212            return Task::ready(Err(anyhow!("Provider not found")));
1213        };
1214        Task::ready(Ok(LanguageModels::map_language_model_to_info(
1215            model, &provider,
1216        )))
1217    }
1218
1219    fn watch(&self, cx: &mut App) -> Option<watch::Receiver<()>> {
1220        Some(self.connection.0.read(cx).models.watch())
1221    }
1222
1223    fn should_render_footer(&self) -> bool {
1224        true
1225    }
1226}
1227
1228impl acp_thread::AgentConnection for NativeAgentConnection {
1229    fn telemetry_id(&self) -> SharedString {
1230        "zed".into()
1231    }
1232
1233    fn new_thread(
1234        self: Rc<Self>,
1235        project: Entity<Project>,
1236        cwd: &Path,
1237        cx: &mut App,
1238    ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
1239        log::debug!("Creating new thread for project at: {cwd:?}");
1240        Task::ready(Ok(self
1241            .0
1242            .update(cx, |agent, cx| agent.new_session(project, cx))))
1243    }
1244
1245    fn supports_load_session(&self, _cx: &App) -> bool {
1246        true
1247    }
1248
1249    fn load_session(
1250        self: Rc<Self>,
1251        session: AgentSessionInfo,
1252        _project: Entity<Project>,
1253        _cwd: &Path,
1254        cx: &mut App,
1255    ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
1256        self.0
1257            .update(cx, |agent, cx| agent.open_thread(session.session_id, cx))
1258    }
1259
1260    fn auth_methods(&self) -> &[acp::AuthMethod] {
1261        &[] // No auth for in-process
1262    }
1263
1264    fn authenticate(&self, _method: acp::AuthMethodId, _cx: &mut App) -> Task<Result<()>> {
1265        Task::ready(Ok(()))
1266    }
1267
1268    fn model_selector(&self, session_id: &acp::SessionId) -> Option<Rc<dyn AgentModelSelector>> {
1269        Some(Rc::new(NativeAgentModelSelector {
1270            session_id: session_id.clone(),
1271            connection: self.clone(),
1272        }) as Rc<dyn AgentModelSelector>)
1273    }
1274
1275    fn prompt(
1276        &self,
1277        id: Option<acp_thread::UserMessageId>,
1278        params: acp::PromptRequest,
1279        cx: &mut App,
1280    ) -> Task<Result<acp::PromptResponse>> {
1281        let id = id.expect("UserMessageId is required");
1282        let session_id = params.session_id.clone();
1283        log::info!("Received prompt request for session: {}", session_id);
1284        log::debug!("Prompt blocks count: {}", params.prompt.len());
1285
1286        if let Some(parsed_command) = Command::parse(&params.prompt) {
1287            let registry = self.0.read(cx).context_server_registry.read(cx);
1288
1289            let explicit_server_id = parsed_command
1290                .explicit_server_id
1291                .map(|server_id| ContextServerId(server_id.into()));
1292
1293            if let Some(prompt) =
1294                registry.find_prompt(explicit_server_id.as_ref(), parsed_command.prompt_name)
1295            {
1296                let arguments = if !parsed_command.arg_value.is_empty()
1297                    && let Some(arg_name) = prompt
1298                        .prompt
1299                        .arguments
1300                        .as_ref()
1301                        .and_then(|args| args.first())
1302                        .map(|arg| arg.name.clone())
1303                {
1304                    HashMap::from_iter([(arg_name, parsed_command.arg_value.to_string())])
1305                } else {
1306                    Default::default()
1307                };
1308
1309                let prompt_name = prompt.prompt.name.clone();
1310                let server_id = prompt.server_id.clone();
1311
1312                return self.0.update(cx, |agent, cx| {
1313                    agent.send_mcp_prompt(
1314                        id,
1315                        session_id.clone(),
1316                        prompt_name,
1317                        server_id,
1318                        arguments,
1319                        params.prompt,
1320                        cx,
1321                    )
1322                });
1323            };
1324        };
1325
1326        let path_style = self.0.read(cx).project.read(cx).path_style(cx);
1327
1328        self.run_turn(session_id, cx, move |thread, cx| {
1329            let content: Vec<UserMessageContent> = params
1330                .prompt
1331                .into_iter()
1332                .map(|block| UserMessageContent::from_content_block(block, path_style))
1333                .collect::<Vec<_>>();
1334            log::debug!("Converted prompt to message: {} chars", content.len());
1335            log::debug!("Message id: {:?}", id);
1336            log::debug!("Message content: {:?}", content);
1337
1338            thread.update(cx, |thread, cx| thread.send(id, content, cx))
1339        })
1340    }
1341
1342    fn retry(
1343        &self,
1344        session_id: &acp::SessionId,
1345        _cx: &App,
1346    ) -> Option<Rc<dyn acp_thread::AgentSessionRetry>> {
1347        Some(Rc::new(NativeAgentSessionRetry {
1348            connection: self.clone(),
1349            session_id: session_id.clone(),
1350        }) as _)
1351    }
1352
1353    fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
1354        log::info!("Cancelling on session: {}", session_id);
1355        self.0.update(cx, |agent, cx| {
1356            if let Some(agent) = agent.sessions.get(session_id) {
1357                agent
1358                    .thread
1359                    .update(cx, |thread, cx| thread.cancel(cx))
1360                    .detach();
1361            }
1362        });
1363    }
1364
1365    fn truncate(
1366        &self,
1367        session_id: &agent_client_protocol::SessionId,
1368        cx: &App,
1369    ) -> Option<Rc<dyn acp_thread::AgentSessionTruncate>> {
1370        self.0.read_with(cx, |agent, _cx| {
1371            agent.sessions.get(session_id).map(|session| {
1372                Rc::new(NativeAgentSessionTruncate {
1373                    thread: session.thread.clone(),
1374                    acp_thread: session.acp_thread.clone(),
1375                }) as _
1376            })
1377        })
1378    }
1379
1380    fn set_title(
1381        &self,
1382        session_id: &acp::SessionId,
1383        _cx: &App,
1384    ) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
1385        Some(Rc::new(NativeAgentSessionSetTitle {
1386            connection: self.clone(),
1387            session_id: session_id.clone(),
1388        }) as _)
1389    }
1390
1391    fn session_list(&self, cx: &mut App) -> Option<Rc<dyn AgentSessionList>> {
1392        let thread_store = self.0.read(cx).thread_store.clone();
1393        Some(Rc::new(NativeAgentSessionList::new(thread_store, cx)) as _)
1394    }
1395
1396    fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
1397        Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
1398    }
1399
1400    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1401        self
1402    }
1403}
1404
1405impl acp_thread::AgentTelemetry for NativeAgentConnection {
1406    fn thread_data(
1407        &self,
1408        session_id: &acp::SessionId,
1409        cx: &mut App,
1410    ) -> Task<Result<serde_json::Value>> {
1411        let Some(session) = self.0.read(cx).sessions.get(session_id) else {
1412            return Task::ready(Err(anyhow!("Session not found")));
1413        };
1414
1415        let task = session.thread.read(cx).to_db(cx);
1416        cx.background_spawn(async move {
1417            serde_json::to_value(task.await).context("Failed to serialize thread")
1418        })
1419    }
1420}
1421
1422pub struct NativeAgentSessionList {
1423    thread_store: Entity<ThreadStore>,
1424    updates_tx: smol::channel::Sender<acp_thread::SessionListUpdate>,
1425    updates_rx: smol::channel::Receiver<acp_thread::SessionListUpdate>,
1426    _subscription: Subscription,
1427}
1428
1429impl NativeAgentSessionList {
1430    fn new(thread_store: Entity<ThreadStore>, cx: &mut App) -> Self {
1431        let (tx, rx) = smol::channel::unbounded();
1432        let this_tx = tx.clone();
1433        let subscription = cx.observe(&thread_store, move |_, _| {
1434            this_tx
1435                .try_send(acp_thread::SessionListUpdate::Refresh)
1436                .ok();
1437        });
1438        Self {
1439            thread_store,
1440            updates_tx: tx,
1441            updates_rx: rx,
1442            _subscription: subscription,
1443        }
1444    }
1445
1446    fn to_session_info(entry: DbThreadMetadata) -> AgentSessionInfo {
1447        AgentSessionInfo {
1448            session_id: entry.id,
1449            cwd: None,
1450            title: Some(entry.title),
1451            updated_at: Some(entry.updated_at),
1452            meta: None,
1453        }
1454    }
1455
1456    pub fn thread_store(&self) -> &Entity<ThreadStore> {
1457        &self.thread_store
1458    }
1459}
1460
1461impl AgentSessionList for NativeAgentSessionList {
1462    fn list_sessions(
1463        &self,
1464        _request: AgentSessionListRequest,
1465        cx: &mut App,
1466    ) -> Task<Result<AgentSessionListResponse>> {
1467        let sessions = self
1468            .thread_store
1469            .read(cx)
1470            .entries()
1471            .map(Self::to_session_info)
1472            .collect();
1473        Task::ready(Ok(AgentSessionListResponse::new(sessions)))
1474    }
1475
1476    fn supports_delete(&self) -> bool {
1477        true
1478    }
1479
1480    fn delete_session(&self, session_id: &acp::SessionId, cx: &mut App) -> Task<Result<()>> {
1481        self.thread_store
1482            .update(cx, |store, cx| store.delete_thread(session_id.clone(), cx))
1483    }
1484
1485    fn delete_sessions(&self, cx: &mut App) -> Task<Result<()>> {
1486        self.thread_store
1487            .update(cx, |store, cx| store.delete_threads(cx))
1488    }
1489
1490    fn watch(
1491        &self,
1492        _cx: &mut App,
1493    ) -> Option<smol::channel::Receiver<acp_thread::SessionListUpdate>> {
1494        Some(self.updates_rx.clone())
1495    }
1496
1497    fn notify_refresh(&self) {
1498        self.updates_tx
1499            .try_send(acp_thread::SessionListUpdate::Refresh)
1500            .ok();
1501    }
1502
1503    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1504        self
1505    }
1506}
1507
1508struct NativeAgentSessionTruncate {
1509    thread: Entity<Thread>,
1510    acp_thread: WeakEntity<AcpThread>,
1511}
1512
1513impl acp_thread::AgentSessionTruncate for NativeAgentSessionTruncate {
1514    fn run(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
1515        match self.thread.update(cx, |thread, cx| {
1516            thread.truncate(message_id.clone(), cx)?;
1517            Ok(thread.latest_token_usage())
1518        }) {
1519            Ok(usage) => {
1520                self.acp_thread
1521                    .update(cx, |thread, cx| {
1522                        thread.update_token_usage(usage, cx);
1523                    })
1524                    .ok();
1525                Task::ready(Ok(()))
1526            }
1527            Err(error) => Task::ready(Err(error)),
1528        }
1529    }
1530}
1531
1532struct NativeAgentSessionRetry {
1533    connection: NativeAgentConnection,
1534    session_id: acp::SessionId,
1535}
1536
1537impl acp_thread::AgentSessionRetry for NativeAgentSessionRetry {
1538    fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>> {
1539        self.connection
1540            .run_turn(self.session_id.clone(), cx, |thread, cx| {
1541                thread.update(cx, |thread, cx| thread.resume(cx))
1542            })
1543    }
1544}
1545
1546struct NativeAgentSessionSetTitle {
1547    connection: NativeAgentConnection,
1548    session_id: acp::SessionId,
1549}
1550
1551impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
1552    fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
1553        let Some(session) = self.connection.0.read(cx).sessions.get(&self.session_id) else {
1554            return Task::ready(Err(anyhow!("session not found")));
1555        };
1556        let thread = session.thread.clone();
1557        thread.update(cx, |thread, cx| thread.set_title(title, cx));
1558        Task::ready(Ok(()))
1559    }
1560}
1561
1562pub struct AcpThreadEnvironment {
1563    acp_thread: WeakEntity<AcpThread>,
1564}
1565
1566impl ThreadEnvironment for AcpThreadEnvironment {
1567    fn create_terminal(
1568        &self,
1569        command: String,
1570        cwd: Option<PathBuf>,
1571        output_byte_limit: Option<u64>,
1572        cx: &mut AsyncApp,
1573    ) -> Task<Result<Rc<dyn TerminalHandle>>> {
1574        let task = self.acp_thread.update(cx, |thread, cx| {
1575            thread.create_terminal(command, vec![], vec![], cwd, output_byte_limit, cx)
1576        });
1577
1578        let acp_thread = self.acp_thread.clone();
1579        cx.spawn(async move |cx| {
1580            let terminal = task?.await?;
1581
1582            let (drop_tx, drop_rx) = oneshot::channel();
1583            let terminal_id = terminal.read_with(cx, |terminal, _cx| terminal.id().clone());
1584
1585            cx.spawn(async move |cx| {
1586                drop_rx.await.ok();
1587                acp_thread.update(cx, |thread, cx| thread.release_terminal(terminal_id, cx))
1588            })
1589            .detach();
1590
1591            let handle = AcpTerminalHandle {
1592                terminal,
1593                _drop_tx: Some(drop_tx),
1594            };
1595
1596            Ok(Rc::new(handle) as _)
1597        })
1598    }
1599}
1600
1601pub struct AcpTerminalHandle {
1602    terminal: Entity<acp_thread::Terminal>,
1603    _drop_tx: Option<oneshot::Sender<()>>,
1604}
1605
1606impl TerminalHandle for AcpTerminalHandle {
1607    fn id(&self, cx: &AsyncApp) -> Result<acp::TerminalId> {
1608        Ok(self.terminal.read_with(cx, |term, _cx| term.id().clone()))
1609    }
1610
1611    fn wait_for_exit(&self, cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {
1612        Ok(self
1613            .terminal
1614            .read_with(cx, |term, _cx| term.wait_for_exit()))
1615    }
1616
1617    fn current_output(&self, cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {
1618        Ok(self
1619            .terminal
1620            .read_with(cx, |term, cx| term.current_output(cx)))
1621    }
1622
1623    fn kill(&self, cx: &AsyncApp) -> Result<()> {
1624        cx.update(|cx| {
1625            self.terminal.update(cx, |terminal, cx| {
1626                terminal.kill(cx);
1627            });
1628        });
1629        Ok(())
1630    }
1631
1632    fn was_stopped_by_user(&self, cx: &AsyncApp) -> Result<bool> {
1633        Ok(self
1634            .terminal
1635            .read_with(cx, |term, _cx| term.was_stopped_by_user()))
1636    }
1637}
1638
1639#[cfg(test)]
1640mod internal_tests {
1641    use super::*;
1642    use acp_thread::{AgentConnection, AgentModelGroupName, AgentModelInfo, MentionUri};
1643    use fs::FakeFs;
1644    use gpui::TestAppContext;
1645    use indoc::formatdoc;
1646    use language_model::fake_provider::FakeLanguageModel;
1647    use serde_json::json;
1648    use settings::SettingsStore;
1649    use util::{path, rel_path::rel_path};
1650
1651    #[gpui::test]
1652    async fn test_maintaining_project_context(cx: &mut TestAppContext) {
1653        init_test(cx);
1654        let fs = FakeFs::new(cx.executor());
1655        fs.insert_tree(
1656            "/",
1657            json!({
1658                "a": {}
1659            }),
1660        )
1661        .await;
1662        let project = Project::test(fs.clone(), [], cx).await;
1663        let thread_store = cx.new(|cx| ThreadStore::new(cx));
1664        let agent = NativeAgent::new(
1665            project.clone(),
1666            thread_store,
1667            Templates::new(),
1668            None,
1669            fs.clone(),
1670            &mut cx.to_async(),
1671        )
1672        .await
1673        .unwrap();
1674        agent.read_with(cx, |agent, cx| {
1675            assert_eq!(agent.project_context.read(cx).worktrees, vec![])
1676        });
1677
1678        let worktree = project
1679            .update(cx, |project, cx| project.create_worktree("/a", true, cx))
1680            .await
1681            .unwrap();
1682        cx.run_until_parked();
1683        agent.read_with(cx, |agent, cx| {
1684            assert_eq!(
1685                agent.project_context.read(cx).worktrees,
1686                vec![WorktreeContext {
1687                    root_name: "a".into(),
1688                    abs_path: Path::new("/a").into(),
1689                    rules_file: None
1690                }]
1691            )
1692        });
1693
1694        // Creating `/a/.rules` updates the project context.
1695        fs.insert_file("/a/.rules", Vec::new()).await;
1696        cx.run_until_parked();
1697        agent.read_with(cx, |agent, cx| {
1698            let rules_entry = worktree
1699                .read(cx)
1700                .entry_for_path(rel_path(".rules"))
1701                .unwrap();
1702            assert_eq!(
1703                agent.project_context.read(cx).worktrees,
1704                vec![WorktreeContext {
1705                    root_name: "a".into(),
1706                    abs_path: Path::new("/a").into(),
1707                    rules_file: Some(RulesFileContext {
1708                        path_in_worktree: rel_path(".rules").into(),
1709                        text: "".into(),
1710                        project_entry_id: rules_entry.id.to_usize()
1711                    })
1712                }]
1713            )
1714        });
1715    }
1716
1717    #[gpui::test]
1718    async fn test_listing_models(cx: &mut TestAppContext) {
1719        init_test(cx);
1720        let fs = FakeFs::new(cx.executor());
1721        fs.insert_tree("/", json!({ "a": {}  })).await;
1722        let project = Project::test(fs.clone(), [], cx).await;
1723        let thread_store = cx.new(|cx| ThreadStore::new(cx));
1724        let connection = NativeAgentConnection(
1725            NativeAgent::new(
1726                project.clone(),
1727                thread_store,
1728                Templates::new(),
1729                None,
1730                fs.clone(),
1731                &mut cx.to_async(),
1732            )
1733            .await
1734            .unwrap(),
1735        );
1736
1737        // Create a thread/session
1738        let acp_thread = cx
1739            .update(|cx| {
1740                Rc::new(connection.clone()).new_thread(project.clone(), Path::new("/a"), cx)
1741            })
1742            .await
1743            .unwrap();
1744
1745        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1746
1747        let models = cx
1748            .update(|cx| {
1749                connection
1750                    .model_selector(&session_id)
1751                    .unwrap()
1752                    .list_models(cx)
1753            })
1754            .await
1755            .unwrap();
1756
1757        let acp_thread::AgentModelList::Grouped(models) = models else {
1758            panic!("Unexpected model group");
1759        };
1760        assert_eq!(
1761            models,
1762            IndexMap::from_iter([(
1763                AgentModelGroupName("Fake".into()),
1764                vec![AgentModelInfo {
1765                    id: acp::ModelId::new("fake/fake"),
1766                    name: "Fake".into(),
1767                    description: None,
1768                    icon: Some(acp_thread::AgentModelIcon::Named(
1769                        ui::IconName::ZedAssistant
1770                    )),
1771                }]
1772            )])
1773        );
1774    }
1775
1776    #[gpui::test]
1777    async fn test_model_selection_persists_to_settings(cx: &mut TestAppContext) {
1778        init_test(cx);
1779        let fs = FakeFs::new(cx.executor());
1780        fs.create_dir(paths::settings_file().parent().unwrap())
1781            .await
1782            .unwrap();
1783        fs.insert_file(
1784            paths::settings_file(),
1785            json!({
1786                "agent": {
1787                    "default_model": {
1788                        "provider": "foo",
1789                        "model": "bar"
1790                    }
1791                }
1792            })
1793            .to_string()
1794            .into_bytes(),
1795        )
1796        .await;
1797        let project = Project::test(fs.clone(), [], cx).await;
1798
1799        let thread_store = cx.new(|cx| ThreadStore::new(cx));
1800
1801        // Create the agent and connection
1802        let agent = NativeAgent::new(
1803            project.clone(),
1804            thread_store,
1805            Templates::new(),
1806            None,
1807            fs.clone(),
1808            &mut cx.to_async(),
1809        )
1810        .await
1811        .unwrap();
1812        let connection = NativeAgentConnection(agent.clone());
1813
1814        // Create a thread/session
1815        let acp_thread = cx
1816            .update(|cx| {
1817                Rc::new(connection.clone()).new_thread(project.clone(), Path::new("/a"), cx)
1818            })
1819            .await
1820            .unwrap();
1821
1822        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1823
1824        // Select a model
1825        let selector = connection.model_selector(&session_id).unwrap();
1826        let model_id = acp::ModelId::new("fake/fake");
1827        cx.update(|cx| selector.select_model(model_id.clone(), cx))
1828            .await
1829            .unwrap();
1830
1831        // Verify the thread has the selected model
1832        agent.read_with(cx, |agent, _| {
1833            let session = agent.sessions.get(&session_id).unwrap();
1834            session.thread.read_with(cx, |thread, _| {
1835                assert_eq!(thread.model().unwrap().id().0, "fake");
1836            });
1837        });
1838
1839        cx.run_until_parked();
1840
1841        // Verify settings file was updated
1842        let settings_content = fs.load(paths::settings_file()).await.unwrap();
1843        let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
1844
1845        // Check that the agent settings contain the selected model
1846        assert_eq!(
1847            settings_json["agent"]["default_model"]["model"],
1848            json!("fake")
1849        );
1850        assert_eq!(
1851            settings_json["agent"]["default_model"]["provider"],
1852            json!("fake")
1853        );
1854    }
1855
1856    #[gpui::test]
1857    async fn test_save_load_thread(cx: &mut TestAppContext) {
1858        init_test(cx);
1859        let fs = FakeFs::new(cx.executor());
1860        fs.insert_tree(
1861            "/",
1862            json!({
1863                "a": {
1864                    "b.md": "Lorem"
1865                }
1866            }),
1867        )
1868        .await;
1869        let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
1870        let thread_store = cx.new(|cx| ThreadStore::new(cx));
1871        let agent = NativeAgent::new(
1872            project.clone(),
1873            thread_store.clone(),
1874            Templates::new(),
1875            None,
1876            fs.clone(),
1877            &mut cx.to_async(),
1878        )
1879        .await
1880        .unwrap();
1881        let connection = Rc::new(NativeAgentConnection(agent.clone()));
1882
1883        let acp_thread = cx
1884            .update(|cx| {
1885                connection
1886                    .clone()
1887                    .new_thread(project.clone(), Path::new(""), cx)
1888            })
1889            .await
1890            .unwrap();
1891        let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
1892        let thread = agent.read_with(cx, |agent, _| {
1893            agent.sessions.get(&session_id).unwrap().thread.clone()
1894        });
1895
1896        // Ensure empty threads are not saved, even if they get mutated.
1897        let model = Arc::new(FakeLanguageModel::default());
1898        let summary_model = Arc::new(FakeLanguageModel::default());
1899        thread.update(cx, |thread, cx| {
1900            thread.set_model(model.clone(), cx);
1901            thread.set_summarization_model(Some(summary_model.clone()), cx);
1902        });
1903        cx.run_until_parked();
1904        assert_eq!(thread_entries(&thread_store, cx), vec![]);
1905
1906        let send = acp_thread.update(cx, |thread, cx| {
1907            thread.send(
1908                vec![
1909                    "What does ".into(),
1910                    acp::ContentBlock::ResourceLink(acp::ResourceLink::new(
1911                        "b.md",
1912                        MentionUri::File {
1913                            abs_path: path!("/a/b.md").into(),
1914                        }
1915                        .to_uri()
1916                        .to_string(),
1917                    )),
1918                    " mean?".into(),
1919                ],
1920                cx,
1921            )
1922        });
1923        let send = cx.foreground_executor().spawn(send);
1924        cx.run_until_parked();
1925
1926        model.send_last_completion_stream_text_chunk("Lorem.");
1927        model.end_last_completion_stream();
1928        cx.run_until_parked();
1929        summary_model
1930            .send_last_completion_stream_text_chunk(&format!("Explaining {}", path!("/a/b.md")));
1931        summary_model.end_last_completion_stream();
1932
1933        send.await.unwrap();
1934        let uri = MentionUri::File {
1935            abs_path: path!("/a/b.md").into(),
1936        }
1937        .to_uri();
1938        acp_thread.read_with(cx, |thread, cx| {
1939            assert_eq!(
1940                thread.to_markdown(cx),
1941                formatdoc! {"
1942                    ## User
1943
1944                    What does [@b.md]({uri}) mean?
1945
1946                    ## Assistant
1947
1948                    Lorem.
1949
1950                "}
1951            )
1952        });
1953
1954        cx.run_until_parked();
1955
1956        // Drop the ACP thread, which should cause the session to be dropped as well.
1957        cx.update(|_| {
1958            drop(thread);
1959            drop(acp_thread);
1960        });
1961        agent.read_with(cx, |agent, _| {
1962            assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
1963        });
1964
1965        // Ensure the thread can be reloaded from disk.
1966        assert_eq!(
1967            thread_entries(&thread_store, cx),
1968            vec![(
1969                session_id.clone(),
1970                format!("Explaining {}", path!("/a/b.md"))
1971            )]
1972        );
1973        let acp_thread = agent
1974            .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
1975            .await
1976            .unwrap();
1977        acp_thread.read_with(cx, |thread, cx| {
1978            assert_eq!(
1979                thread.to_markdown(cx),
1980                formatdoc! {"
1981                    ## User
1982
1983                    What does [@b.md]({uri}) mean?
1984
1985                    ## Assistant
1986
1987                    Lorem.
1988
1989                "}
1990            )
1991        });
1992    }
1993
1994    fn thread_entries(
1995        thread_store: &Entity<ThreadStore>,
1996        cx: &mut TestAppContext,
1997    ) -> Vec<(acp::SessionId, String)> {
1998        thread_store.read_with(cx, |store, _| {
1999            store
2000                .entries()
2001                .map(|entry| (entry.id.clone(), entry.title.to_string()))
2002                .collect::<Vec<_>>()
2003        })
2004    }
2005
2006    fn init_test(cx: &mut TestAppContext) {
2007        env_logger::try_init().ok();
2008        cx.update(|cx| {
2009            let settings_store = SettingsStore::test(cx);
2010            cx.set_global(settings_store);
2011
2012            LanguageModelRegistry::test(cx);
2013        });
2014    }
2015}
2016
2017fn mcp_message_content_to_acp_content_block(
2018    content: context_server::types::MessageContent,
2019) -> acp::ContentBlock {
2020    match content {
2021        context_server::types::MessageContent::Text {
2022            text,
2023            annotations: _,
2024        } => text.into(),
2025        context_server::types::MessageContent::Image {
2026            data,
2027            mime_type,
2028            annotations: _,
2029        } => acp::ContentBlock::Image(acp::ImageContent::new(data, mime_type)),
2030        context_server::types::MessageContent::Audio {
2031            data,
2032            mime_type,
2033            annotations: _,
2034        } => acp::ContentBlock::Audio(acp::AudioContent::new(data, mime_type)),
2035        context_server::types::MessageContent::Resource {
2036            resource,
2037            annotations: _,
2038        } => {
2039            let mut link =
2040                acp::ResourceLink::new(resource.uri.to_string(), resource.uri.to_string());
2041            if let Some(mime_type) = resource.mime_type {
2042                link = link.mime_type(mime_type);
2043            }
2044            acp::ContentBlock::ResourceLink(link)
2045        }
2046    }
2047}