agent.rs

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