1use std::ops::Range;
   2use std::path::Path;
   3use std::sync::Arc;
   4
   5use anyhow::{Context as _, Result, anyhow};
   6use collections::{BTreeMap, HashMap, HashSet};
   7use futures::future::join_all;
   8use futures::{self, Future, FutureExt, future};
   9use gpui::{App, AppContext as _, Context, Entity, Image, SharedString, Task, WeakEntity};
  10use language::Buffer;
  11use language_model::LanguageModelImage;
  12use project::{Project, ProjectEntryId, ProjectItem, ProjectPath, Worktree};
  13use prompt_store::UserPromptId;
  14use rope::{Point, Rope};
  15use text::{Anchor, BufferId, OffsetRangeExt};
  16use util::{ResultExt as _, maybe};
  17
  18use crate::ThreadStore;
  19use crate::context::{
  20    AssistantContext, ContextBuffer, ContextId, ContextSymbol, ContextSymbolId, DirectoryContext,
  21    FetchedUrlContext, FileContext, ImageContext, RulesContext, SelectionContext, SymbolContext,
  22    ThreadContext,
  23};
  24use crate::context_strip::SuggestedContext;
  25use crate::thread::{Thread, ThreadId};
  26
  27pub struct ContextStore {
  28    project: WeakEntity<Project>,
  29    context: Vec<AssistantContext>,
  30    thread_store: Option<WeakEntity<ThreadStore>>,
  31    next_context_id: ContextId,
  32    files: BTreeMap<BufferId, ContextId>,
  33    directories: HashMap<ProjectPath, ContextId>,
  34    symbols: HashMap<ContextSymbolId, ContextId>,
  35    symbol_buffers: HashMap<ContextSymbolId, Entity<Buffer>>,
  36    symbols_by_path: HashMap<ProjectPath, Vec<ContextSymbolId>>,
  37    threads: HashMap<ThreadId, ContextId>,
  38    thread_summary_tasks: Vec<Task<()>>,
  39    fetched_urls: HashMap<String, ContextId>,
  40    user_rules: HashMap<UserPromptId, ContextId>,
  41}
  42
  43impl ContextStore {
  44    pub fn new(
  45        project: WeakEntity<Project>,
  46        thread_store: Option<WeakEntity<ThreadStore>>,
  47    ) -> Self {
  48        Self {
  49            project,
  50            thread_store,
  51            context: Vec::new(),
  52            next_context_id: ContextId(0),
  53            files: BTreeMap::default(),
  54            directories: HashMap::default(),
  55            symbols: HashMap::default(),
  56            symbol_buffers: HashMap::default(),
  57            symbols_by_path: HashMap::default(),
  58            threads: HashMap::default(),
  59            thread_summary_tasks: Vec::new(),
  60            fetched_urls: HashMap::default(),
  61            user_rules: HashMap::default(),
  62        }
  63    }
  64
  65    pub fn context(&self) -> &Vec<AssistantContext> {
  66        &self.context
  67    }
  68
  69    pub fn context_for_id(&self, id: ContextId) -> Option<&AssistantContext> {
  70        self.context().iter().find(|context| context.id() == id)
  71    }
  72
  73    pub fn clear(&mut self) {
  74        self.context.clear();
  75        self.files.clear();
  76        self.directories.clear();
  77        self.threads.clear();
  78        self.fetched_urls.clear();
  79        self.user_rules.clear();
  80    }
  81
  82    pub fn add_file_from_path(
  83        &mut self,
  84        project_path: ProjectPath,
  85        remove_if_exists: bool,
  86        cx: &mut Context<Self>,
  87    ) -> Task<Result<()>> {
  88        let Some(project) = self.project.upgrade() else {
  89            return Task::ready(Err(anyhow!("failed to read project")));
  90        };
  91
  92        cx.spawn(async move |this, cx| {
  93            let open_buffer_task = project.update(cx, |project, cx| {
  94                project.open_buffer(project_path.clone(), cx)
  95            })?;
  96
  97            let buffer = open_buffer_task.await?;
  98            let buffer_id = this.update(cx, |_, cx| buffer.read(cx).remote_id())?;
  99
 100            let already_included = this.update(cx, |this, cx| {
 101                match this.will_include_buffer(buffer_id, &project_path) {
 102                    Some(FileInclusion::Direct(context_id)) => {
 103                        if remove_if_exists {
 104                            this.remove_context(context_id, cx);
 105                        }
 106                        true
 107                    }
 108                    Some(FileInclusion::InDirectory(_)) => true,
 109                    None => false,
 110                }
 111            })?;
 112
 113            if already_included {
 114                return anyhow::Ok(());
 115            }
 116
 117            let context_buffer = this
 118                .update(cx, |_, cx| load_context_buffer(buffer, cx))??
 119                .await;
 120
 121            this.update(cx, |this, cx| {
 122                this.insert_file(context_buffer, cx);
 123            })?;
 124
 125            anyhow::Ok(())
 126        })
 127    }
 128
 129    pub fn add_file_from_buffer(
 130        &mut self,
 131        buffer: Entity<Buffer>,
 132        cx: &mut Context<Self>,
 133    ) -> Task<Result<()>> {
 134        cx.spawn(async move |this, cx| {
 135            let context_buffer = this
 136                .update(cx, |_, cx| load_context_buffer(buffer, cx))??
 137                .await;
 138
 139            this.update(cx, |this, cx| this.insert_file(context_buffer, cx))?;
 140
 141            anyhow::Ok(())
 142        })
 143    }
 144
 145    fn insert_file(&mut self, context_buffer: ContextBuffer, cx: &mut Context<Self>) {
 146        let id = self.next_context_id.post_inc();
 147        self.files.insert(context_buffer.id, id);
 148        self.context
 149            .push(AssistantContext::File(FileContext { id, context_buffer }));
 150        cx.notify();
 151    }
 152
 153    pub fn add_directory(
 154        &mut self,
 155        project_path: ProjectPath,
 156        remove_if_exists: bool,
 157        cx: &mut Context<Self>,
 158    ) -> Task<Result<()>> {
 159        let Some(project) = self.project.upgrade() else {
 160            return Task::ready(Err(anyhow!("failed to read project")));
 161        };
 162
 163        let Some(entry_id) = project
 164            .read(cx)
 165            .entry_for_path(&project_path, cx)
 166            .map(|entry| entry.id)
 167        else {
 168            return Task::ready(Err(anyhow!("no entry found for directory context")));
 169        };
 170
 171        let already_included = match self.includes_directory(&project_path) {
 172            Some(FileInclusion::Direct(context_id)) => {
 173                if remove_if_exists {
 174                    self.remove_context(context_id, cx);
 175                }
 176                true
 177            }
 178            Some(FileInclusion::InDirectory(_)) => true,
 179            None => false,
 180        };
 181        if already_included {
 182            return Task::ready(Ok(()));
 183        }
 184
 185        let worktree_id = project_path.worktree_id;
 186        cx.spawn(async move |this, cx| {
 187            let worktree = project.update(cx, |project, cx| {
 188                project
 189                    .worktree_for_id(worktree_id, cx)
 190                    .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))
 191            })??;
 192
 193            let files = worktree.update(cx, |worktree, _cx| {
 194                collect_files_in_path(worktree, &project_path.path)
 195            })?;
 196
 197            let open_buffers_task = project.update(cx, |project, cx| {
 198                let tasks = files.iter().map(|file_path| {
 199                    project.open_buffer(
 200                        ProjectPath {
 201                            worktree_id,
 202                            path: file_path.clone(),
 203                        },
 204                        cx,
 205                    )
 206                });
 207                future::join_all(tasks)
 208            })?;
 209
 210            let buffers = open_buffers_task.await;
 211
 212            let context_buffer_tasks = this.update(cx, |_, cx| {
 213                buffers
 214                    .into_iter()
 215                    .flatten()
 216                    .flat_map(move |buffer| load_context_buffer(buffer, cx).log_err())
 217                    .collect::<Vec<_>>()
 218            })?;
 219
 220            let context_buffers = future::join_all(context_buffer_tasks).await;
 221
 222            if context_buffers.is_empty() {
 223                let full_path = cx.update(|cx| worktree.read(cx).full_path(&project_path.path))?;
 224                return Err(anyhow!("No text files found in {}", &full_path.display()));
 225            }
 226
 227            this.update(cx, |this, cx| {
 228                this.insert_directory(worktree, entry_id, project_path, context_buffers, cx);
 229            })?;
 230
 231            anyhow::Ok(())
 232        })
 233    }
 234
 235    fn insert_directory(
 236        &mut self,
 237        worktree: Entity<Worktree>,
 238        entry_id: ProjectEntryId,
 239        project_path: ProjectPath,
 240        context_buffers: Vec<ContextBuffer>,
 241        cx: &mut Context<Self>,
 242    ) {
 243        let id = self.next_context_id.post_inc();
 244        let last_path = project_path.path.clone();
 245        self.directories.insert(project_path, id);
 246
 247        self.context
 248            .push(AssistantContext::Directory(DirectoryContext {
 249                id,
 250                worktree,
 251                entry_id,
 252                last_path,
 253                context_buffers,
 254            }));
 255        cx.notify();
 256    }
 257
 258    pub fn add_symbol(
 259        &mut self,
 260        buffer: Entity<Buffer>,
 261        symbol_name: SharedString,
 262        symbol_range: Range<Anchor>,
 263        symbol_enclosing_range: Range<Anchor>,
 264        remove_if_exists: bool,
 265        cx: &mut Context<Self>,
 266    ) -> Task<Result<bool>> {
 267        let buffer_ref = buffer.read(cx);
 268        let Some(project_path) = buffer_ref.project_path(cx) else {
 269            return Task::ready(Err(anyhow!("buffer has no path")));
 270        };
 271
 272        if let Some(symbols_for_path) = self.symbols_by_path.get(&project_path) {
 273            let mut matching_symbol_id = None;
 274            for symbol in symbols_for_path {
 275                if &symbol.name == &symbol_name {
 276                    let snapshot = buffer_ref.snapshot();
 277                    if symbol.range.to_offset(&snapshot) == symbol_range.to_offset(&snapshot) {
 278                        matching_symbol_id = self.symbols.get(symbol).cloned();
 279                        break;
 280                    }
 281                }
 282            }
 283
 284            if let Some(id) = matching_symbol_id {
 285                if remove_if_exists {
 286                    self.remove_context(id, cx);
 287                }
 288                return Task::ready(Ok(false));
 289            }
 290        }
 291
 292        let context_buffer_task =
 293            match load_context_buffer_range(buffer, symbol_enclosing_range.clone(), cx) {
 294                Ok((_line_range, context_buffer_task)) => context_buffer_task,
 295                Err(err) => return Task::ready(Err(err)),
 296            };
 297
 298        cx.spawn(async move |this, cx| {
 299            let context_buffer = context_buffer_task.await;
 300
 301            this.update(cx, |this, cx| {
 302                this.insert_symbol(
 303                    make_context_symbol(
 304                        context_buffer,
 305                        project_path,
 306                        symbol_name,
 307                        symbol_range,
 308                        symbol_enclosing_range,
 309                    ),
 310                    cx,
 311                )
 312            })?;
 313            anyhow::Ok(true)
 314        })
 315    }
 316
 317    fn insert_symbol(&mut self, context_symbol: ContextSymbol, cx: &mut Context<Self>) {
 318        let id = self.next_context_id.post_inc();
 319        self.symbols.insert(context_symbol.id.clone(), id);
 320        self.symbols_by_path
 321            .entry(context_symbol.id.path.clone())
 322            .or_insert_with(Vec::new)
 323            .push(context_symbol.id.clone());
 324        self.symbol_buffers
 325            .insert(context_symbol.id.clone(), context_symbol.buffer.clone());
 326        self.context.push(AssistantContext::Symbol(SymbolContext {
 327            id,
 328            context_symbol,
 329        }));
 330        cx.notify();
 331    }
 332
 333    pub fn add_thread(
 334        &mut self,
 335        thread: Entity<Thread>,
 336        remove_if_exists: bool,
 337        cx: &mut Context<Self>,
 338    ) {
 339        if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {
 340            if remove_if_exists {
 341                self.remove_context(context_id, cx);
 342            }
 343        } else {
 344            self.insert_thread(thread, cx);
 345        }
 346    }
 347
 348    pub fn wait_for_summaries(&mut self, cx: &App) -> Task<()> {
 349        let tasks = std::mem::take(&mut self.thread_summary_tasks);
 350
 351        cx.spawn(async move |_cx| {
 352            join_all(tasks).await;
 353        })
 354    }
 355
 356    fn insert_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
 357        if let Some(summary_task) =
 358            thread.update(cx, |thread, cx| thread.generate_detailed_summary(cx))
 359        {
 360            let thread = thread.clone();
 361            let thread_store = self.thread_store.clone();
 362
 363            self.thread_summary_tasks.push(cx.spawn(async move |_, cx| {
 364                summary_task.await;
 365
 366                if let Some(thread_store) = thread_store {
 367                    // Save thread so its summary can be reused later
 368                    let save_task = thread_store
 369                        .update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx));
 370
 371                    if let Some(save_task) = save_task.ok() {
 372                        save_task.await.log_err();
 373                    }
 374                }
 375            }));
 376        }
 377
 378        let id = self.next_context_id.post_inc();
 379
 380        let text = thread.read(cx).latest_detailed_summary_or_text();
 381
 382        self.threads.insert(thread.read(cx).id().clone(), id);
 383        self.context
 384            .push(AssistantContext::Thread(ThreadContext { id, thread, text }));
 385        cx.notify();
 386    }
 387
 388    pub fn add_rules(
 389        &mut self,
 390        prompt_id: UserPromptId,
 391        title: impl Into<SharedString>,
 392        text: impl Into<SharedString>,
 393        remove_if_exists: bool,
 394        cx: &mut Context<ContextStore>,
 395    ) {
 396        if let Some(context_id) = self.includes_user_rules(&prompt_id) {
 397            if remove_if_exists {
 398                self.remove_context(context_id, cx);
 399            }
 400        } else {
 401            self.insert_user_rules(prompt_id, title, text, cx);
 402        }
 403    }
 404
 405    pub fn insert_user_rules(
 406        &mut self,
 407        prompt_id: UserPromptId,
 408        title: impl Into<SharedString>,
 409        text: impl Into<SharedString>,
 410        cx: &mut Context<ContextStore>,
 411    ) {
 412        let id = self.next_context_id.post_inc();
 413
 414        self.user_rules.insert(prompt_id, id);
 415        self.context.push(AssistantContext::Rules(RulesContext {
 416            id,
 417            prompt_id,
 418            title: title.into(),
 419            text: text.into(),
 420        }));
 421        cx.notify();
 422    }
 423
 424    pub fn add_fetched_url(
 425        &mut self,
 426        url: String,
 427        text: impl Into<SharedString>,
 428        cx: &mut Context<ContextStore>,
 429    ) {
 430        if self.includes_url(&url).is_none() {
 431            self.insert_fetched_url(url, text, cx);
 432        }
 433    }
 434
 435    fn insert_fetched_url(
 436        &mut self,
 437        url: String,
 438        text: impl Into<SharedString>,
 439        cx: &mut Context<ContextStore>,
 440    ) {
 441        let id = self.next_context_id.post_inc();
 442
 443        self.fetched_urls.insert(url.clone(), id);
 444        self.context
 445            .push(AssistantContext::FetchedUrl(FetchedUrlContext {
 446                id,
 447                url: url.into(),
 448                text: text.into(),
 449            }));
 450        cx.notify();
 451    }
 452
 453    pub fn add_image(&mut self, image: Arc<Image>, cx: &mut Context<ContextStore>) {
 454        let image_task = LanguageModelImage::from_image(image.clone(), cx).shared();
 455        let id = self.next_context_id.post_inc();
 456        self.context.push(AssistantContext::Image(ImageContext {
 457            id,
 458            original_image: image,
 459            image_task,
 460        }));
 461        cx.notify();
 462    }
 463
 464    pub fn wait_for_images(&self, cx: &App) -> Task<()> {
 465        let tasks = self
 466            .context
 467            .iter()
 468            .filter_map(|ctx| match ctx {
 469                AssistantContext::Image(ctx) => Some(ctx.image_task.clone()),
 470                _ => None,
 471            })
 472            .collect::<Vec<_>>();
 473
 474        cx.spawn(async move |_cx| {
 475            join_all(tasks).await;
 476        })
 477    }
 478
 479    pub fn add_selection(
 480        &mut self,
 481        buffer: Entity<Buffer>,
 482        range: Range<Anchor>,
 483        cx: &mut Context<ContextStore>,
 484    ) -> Task<Result<()>> {
 485        cx.spawn(async move |this, cx| {
 486            let (line_range, context_buffer_task) = this.update(cx, |_, cx| {
 487                load_context_buffer_range(buffer, range.clone(), cx)
 488            })??;
 489
 490            let context_buffer = context_buffer_task.await;
 491
 492            this.update(cx, |this, cx| {
 493                this.insert_selection(context_buffer, range, line_range, cx)
 494            })?;
 495
 496            anyhow::Ok(())
 497        })
 498    }
 499
 500    fn insert_selection(
 501        &mut self,
 502        context_buffer: ContextBuffer,
 503        range: Range<Anchor>,
 504        line_range: Range<Point>,
 505        cx: &mut Context<Self>,
 506    ) {
 507        let id = self.next_context_id.post_inc();
 508        self.context
 509            .push(AssistantContext::Selection(SelectionContext {
 510                id,
 511                range,
 512                line_range,
 513                context_buffer,
 514            }));
 515        cx.notify();
 516    }
 517
 518    pub fn accept_suggested_context(
 519        &mut self,
 520        suggested: &SuggestedContext,
 521        cx: &mut Context<ContextStore>,
 522    ) -> Task<Result<()>> {
 523        match suggested {
 524            SuggestedContext::File {
 525                buffer,
 526                icon_path: _,
 527                name: _,
 528            } => {
 529                if let Some(buffer) = buffer.upgrade() {
 530                    return self.add_file_from_buffer(buffer, cx);
 531                };
 532            }
 533            SuggestedContext::Thread { thread, name: _ } => {
 534                if let Some(thread) = thread.upgrade() {
 535                    self.insert_thread(thread, cx);
 536                };
 537            }
 538        }
 539        Task::ready(Ok(()))
 540    }
 541
 542    pub fn remove_context(&mut self, id: ContextId, cx: &mut Context<Self>) {
 543        let Some(ix) = self.context.iter().position(|context| context.id() == id) else {
 544            return;
 545        };
 546
 547        match self.context.remove(ix) {
 548            AssistantContext::File(_) => {
 549                self.files.retain(|_, context_id| *context_id != id);
 550            }
 551            AssistantContext::Directory(_) => {
 552                self.directories.retain(|_, context_id| *context_id != id);
 553            }
 554            AssistantContext::Symbol(symbol) => {
 555                if let Some(symbols_in_path) =
 556                    self.symbols_by_path.get_mut(&symbol.context_symbol.id.path)
 557                {
 558                    symbols_in_path.retain(|s| {
 559                        self.symbols
 560                            .get(s)
 561                            .map_or(false, |context_id| *context_id != id)
 562                    });
 563                }
 564                self.symbol_buffers.remove(&symbol.context_symbol.id);
 565                self.symbols.retain(|_, context_id| *context_id != id);
 566            }
 567            AssistantContext::Selection(_) => {}
 568            AssistantContext::FetchedUrl(_) => {
 569                self.fetched_urls.retain(|_, context_id| *context_id != id);
 570            }
 571            AssistantContext::Thread(_) => {
 572                self.threads.retain(|_, context_id| *context_id != id);
 573            }
 574            AssistantContext::Rules(RulesContext { prompt_id, .. }) => {
 575                self.user_rules.remove(&prompt_id);
 576            }
 577            AssistantContext::Image(_) => {}
 578        }
 579
 580        cx.notify();
 581    }
 582
 583    /// Returns whether the buffer is already included directly in the context, or if it will be
 584    /// included in the context via a directory. Directory inclusion is based on paths rather than
 585    /// buffer IDs as the directory will be re-scanned.
 586    pub fn will_include_buffer(
 587        &self,
 588        buffer_id: BufferId,
 589        project_path: &ProjectPath,
 590    ) -> Option<FileInclusion> {
 591        if let Some(context_id) = self.files.get(&buffer_id) {
 592            return Some(FileInclusion::Direct(*context_id));
 593        }
 594
 595        self.will_include_file_path_via_directory(project_path)
 596    }
 597
 598    /// Returns whether this file path is already included directly in the context, or if it will be
 599    /// included in the context via a directory.
 600    pub fn will_include_file_path(
 601        &self,
 602        project_path: &ProjectPath,
 603        cx: &App,
 604    ) -> Option<FileInclusion> {
 605        if !self.files.is_empty() {
 606            let found_file_context = self.context.iter().find(|context| match &context {
 607                AssistantContext::File(file_context) => {
 608                    let buffer = file_context.context_buffer.buffer.read(cx);
 609                    if let Some(context_path) = buffer.project_path(cx) {
 610                        &context_path == project_path
 611                    } else {
 612                        false
 613                    }
 614                }
 615                _ => false,
 616            });
 617            if let Some(context) = found_file_context {
 618                return Some(FileInclusion::Direct(context.id()));
 619            }
 620        }
 621
 622        self.will_include_file_path_via_directory(project_path)
 623    }
 624
 625    fn will_include_file_path_via_directory(
 626        &self,
 627        project_path: &ProjectPath,
 628    ) -> Option<FileInclusion> {
 629        if self.directories.is_empty() {
 630            return None;
 631        }
 632
 633        let mut path_buf = project_path.path.to_path_buf();
 634
 635        while path_buf.pop() {
 636            // TODO: This isn't very efficient. Consider using a better representation of the
 637            // directories map.
 638            let directory_project_path = ProjectPath {
 639                worktree_id: project_path.worktree_id,
 640                path: path_buf.clone().into(),
 641            };
 642            if let Some(_) = self.directories.get(&directory_project_path) {
 643                return Some(FileInclusion::InDirectory(directory_project_path));
 644            }
 645        }
 646
 647        None
 648    }
 649
 650    pub fn includes_directory(&self, project_path: &ProjectPath) -> Option<FileInclusion> {
 651        if let Some(context_id) = self.directories.get(project_path) {
 652            return Some(FileInclusion::Direct(*context_id));
 653        }
 654
 655        self.will_include_file_path_via_directory(project_path)
 656    }
 657
 658    pub fn included_symbol(&self, symbol_id: &ContextSymbolId) -> Option<ContextId> {
 659        self.symbols.get(symbol_id).copied()
 660    }
 661
 662    pub fn included_symbols_by_path(&self) -> &HashMap<ProjectPath, Vec<ContextSymbolId>> {
 663        &self.symbols_by_path
 664    }
 665
 666    pub fn buffer_for_symbol(&self, symbol_id: &ContextSymbolId) -> Option<Entity<Buffer>> {
 667        self.symbol_buffers.get(symbol_id).cloned()
 668    }
 669
 670    pub fn includes_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
 671        self.threads.get(thread_id).copied()
 672    }
 673
 674    pub fn includes_user_rules(&self, prompt_id: &UserPromptId) -> Option<ContextId> {
 675        self.user_rules.get(prompt_id).copied()
 676    }
 677
 678    pub fn includes_url(&self, url: &str) -> Option<ContextId> {
 679        self.fetched_urls.get(url).copied()
 680    }
 681
 682    /// Replaces the context that matches the ID of the new context, if any match.
 683    fn replace_context(&mut self, new_context: AssistantContext) {
 684        let id = new_context.id();
 685        for context in self.context.iter_mut() {
 686            if context.id() == id {
 687                *context = new_context;
 688                break;
 689            }
 690        }
 691    }
 692
 693    pub fn file_paths(&self, cx: &App) -> HashSet<ProjectPath> {
 694        self.context
 695            .iter()
 696            .filter_map(|context| match context {
 697                AssistantContext::File(file) => {
 698                    let buffer = file.context_buffer.buffer.read(cx);
 699                    buffer.project_path(cx)
 700                }
 701                AssistantContext::Directory(_)
 702                | AssistantContext::Symbol(_)
 703                | AssistantContext::Selection(_)
 704                | AssistantContext::FetchedUrl(_)
 705                | AssistantContext::Thread(_)
 706                | AssistantContext::Rules(_)
 707                | AssistantContext::Image(_) => None,
 708            })
 709            .collect()
 710    }
 711
 712    pub fn thread_ids(&self) -> HashSet<ThreadId> {
 713        self.threads.keys().cloned().collect()
 714    }
 715}
 716
 717pub enum FileInclusion {
 718    Direct(ContextId),
 719    InDirectory(ProjectPath),
 720}
 721
 722fn make_context_symbol(
 723    context_buffer: ContextBuffer,
 724    path: ProjectPath,
 725    name: SharedString,
 726    range: Range<Anchor>,
 727    enclosing_range: Range<Anchor>,
 728) -> ContextSymbol {
 729    ContextSymbol {
 730        id: ContextSymbolId { name, range, path },
 731        buffer_version: context_buffer.version,
 732        enclosing_range,
 733        buffer: context_buffer.buffer,
 734        text: context_buffer.text,
 735    }
 736}
 737
 738fn load_context_buffer_range(
 739    buffer: Entity<Buffer>,
 740    range: Range<Anchor>,
 741    cx: &App,
 742) -> Result<(Range<Point>, Task<ContextBuffer>)> {
 743    let buffer_ref = buffer.read(cx);
 744    let id = buffer_ref.remote_id();
 745
 746    let file = buffer_ref.file().context("context buffer missing path")?;
 747    let full_path = file.full_path(cx);
 748
 749    // Important to collect version at the same time as content so that staleness logic is correct.
 750    let version = buffer_ref.version();
 751    let content = buffer_ref.text_for_range(range.clone()).collect::<Rope>();
 752    let line_range = range.to_point(&buffer_ref.snapshot());
 753
 754    // Build the text on a background thread.
 755    let task = cx.background_spawn({
 756        let line_range = line_range.clone();
 757        async move {
 758            let text = to_fenced_codeblock(&full_path, content, Some(line_range));
 759            ContextBuffer {
 760                id,
 761                buffer,
 762                last_full_path: full_path.into(),
 763                version,
 764                text,
 765            }
 766        }
 767    });
 768
 769    Ok((line_range, task))
 770}
 771
 772fn load_context_buffer(buffer: Entity<Buffer>, cx: &App) -> Result<Task<ContextBuffer>> {
 773    let buffer_ref = buffer.read(cx);
 774    let id = buffer_ref.remote_id();
 775
 776    let file = buffer_ref.file().context("context buffer missing path")?;
 777    let full_path = file.full_path(cx);
 778
 779    // Important to collect version at the same time as content so that staleness logic is correct.
 780    let version = buffer_ref.version();
 781    let content = buffer_ref.as_rope().clone();
 782
 783    // Build the text on a background thread.
 784    Ok(cx.background_spawn(async move {
 785        let text = to_fenced_codeblock(&full_path, content, None);
 786        ContextBuffer {
 787            id,
 788            buffer,
 789            last_full_path: full_path.into(),
 790            version,
 791            text,
 792        }
 793    }))
 794}
 795
 796fn to_fenced_codeblock(
 797    path: &Path,
 798    content: Rope,
 799    line_range: Option<Range<Point>>,
 800) -> SharedString {
 801    let line_range_text = line_range.map(|range| {
 802        if range.start.row == range.end.row {
 803            format!(":{}", range.start.row + 1)
 804        } else {
 805            format!(":{}-{}", range.start.row + 1, range.end.row + 1)
 806        }
 807    });
 808
 809    let path_extension = path.extension().and_then(|ext| ext.to_str());
 810    let path_string = path.to_string_lossy();
 811    let capacity = 3
 812        + path_extension.map_or(0, |extension| extension.len() + 1)
 813        + path_string.len()
 814        + line_range_text.as_ref().map_or(0, |text| text.len())
 815        + 1
 816        + content.len()
 817        + 5;
 818    let mut buffer = String::with_capacity(capacity);
 819
 820    buffer.push_str("```");
 821
 822    if let Some(extension) = path_extension {
 823        buffer.push_str(extension);
 824        buffer.push(' ');
 825    }
 826    buffer.push_str(&path_string);
 827
 828    if let Some(line_range_text) = line_range_text {
 829        buffer.push_str(&line_range_text);
 830    }
 831
 832    buffer.push('\n');
 833    for chunk in content.chunks() {
 834        buffer.push_str(&chunk);
 835    }
 836
 837    if !buffer.ends_with('\n') {
 838        buffer.push('\n');
 839    }
 840
 841    buffer.push_str("```\n");
 842
 843    debug_assert!(
 844        buffer.len() == capacity - 1 || buffer.len() == capacity,
 845        "to_fenced_codeblock calculated capacity of {}, but length was {}",
 846        capacity,
 847        buffer.len(),
 848    );
 849
 850    buffer.into()
 851}
 852
 853fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
 854    let mut files = Vec::new();
 855
 856    for entry in worktree.child_entries(path) {
 857        if entry.is_dir() {
 858            files.extend(collect_files_in_path(worktree, &entry.path));
 859        } else if entry.is_file() {
 860            files.push(entry.path.clone());
 861        }
 862    }
 863
 864    files
 865}
 866
 867pub fn refresh_context_store_text(
 868    context_store: Entity<ContextStore>,
 869    changed_buffers: &HashSet<Entity<Buffer>>,
 870    cx: &App,
 871) -> impl Future<Output = Vec<ContextId>> + use<> {
 872    let mut tasks = Vec::new();
 873
 874    for context in &context_store.read(cx).context {
 875        let id = context.id();
 876
 877        let task = maybe!({
 878            match context {
 879                AssistantContext::File(file_context) => {
 880                    // TODO: Should refresh if the path has changed, as it's in the text.
 881                    if changed_buffers.is_empty()
 882                        || changed_buffers.contains(&file_context.context_buffer.buffer)
 883                    {
 884                        let context_store = context_store.clone();
 885                        return refresh_file_text(context_store, file_context, cx);
 886                    }
 887                }
 888                AssistantContext::Directory(directory_context) => {
 889                    let directory_path = directory_context.project_path(cx)?;
 890                    let should_refresh = directory_path.path != directory_context.last_path
 891                        || changed_buffers.is_empty()
 892                        || changed_buffers.iter().any(|buffer| {
 893                            let Some(buffer_path) = buffer.read(cx).project_path(cx) else {
 894                                return false;
 895                            };
 896                            buffer_path.starts_with(&directory_path)
 897                        });
 898
 899                    if should_refresh {
 900                        let context_store = context_store.clone();
 901                        return refresh_directory_text(
 902                            context_store,
 903                            directory_context,
 904                            directory_path,
 905                            cx,
 906                        );
 907                    }
 908                }
 909                AssistantContext::Symbol(symbol_context) => {
 910                    // TODO: Should refresh if the path has changed, as it's in the text.
 911                    if changed_buffers.is_empty()
 912                        || changed_buffers.contains(&symbol_context.context_symbol.buffer)
 913                    {
 914                        let context_store = context_store.clone();
 915                        return refresh_symbol_text(context_store, symbol_context, cx);
 916                    }
 917                }
 918                AssistantContext::Selection(selection_context) => {
 919                    // TODO: Should refresh if the path has changed, as it's in the text.
 920                    if changed_buffers.is_empty()
 921                        || changed_buffers.contains(&selection_context.context_buffer.buffer)
 922                    {
 923                        let context_store = context_store.clone();
 924                        return refresh_selection_text(context_store, selection_context, cx);
 925                    }
 926                }
 927                AssistantContext::Thread(thread_context) => {
 928                    if changed_buffers.is_empty() {
 929                        let context_store = context_store.clone();
 930                        return Some(refresh_thread_text(context_store, thread_context, cx));
 931                    }
 932                }
 933                // Intentionally omit refreshing fetched URLs as it doesn't seem all that useful,
 934                // and doing the caching properly could be tricky (unless it's already handled by
 935                // the HttpClient?).
 936                AssistantContext::FetchedUrl(_) => {}
 937                AssistantContext::Rules(user_rules_context) => {
 938                    let context_store = context_store.clone();
 939                    return Some(refresh_user_rules(context_store, user_rules_context, cx));
 940                }
 941                AssistantContext::Image(_) => {}
 942            }
 943
 944            None
 945        });
 946
 947        if let Some(task) = task {
 948            tasks.push(task.map(move |_| id));
 949        }
 950    }
 951
 952    future::join_all(tasks)
 953}
 954
 955fn refresh_file_text(
 956    context_store: Entity<ContextStore>,
 957    file_context: &FileContext,
 958    cx: &App,
 959) -> Option<Task<()>> {
 960    let id = file_context.id;
 961    let task = refresh_context_buffer(&file_context.context_buffer, cx);
 962    if let Some(task) = task {
 963        Some(cx.spawn(async move |cx| {
 964            let context_buffer = task.await;
 965            context_store
 966                .update(cx, |context_store, _| {
 967                    let new_file_context = FileContext { id, context_buffer };
 968                    context_store.replace_context(AssistantContext::File(new_file_context));
 969                })
 970                .ok();
 971        }))
 972    } else {
 973        None
 974    }
 975}
 976
 977fn refresh_directory_text(
 978    context_store: Entity<ContextStore>,
 979    directory_context: &DirectoryContext,
 980    directory_path: ProjectPath,
 981    cx: &App,
 982) -> Option<Task<()>> {
 983    let mut stale = false;
 984    let futures = directory_context
 985        .context_buffers
 986        .iter()
 987        .map(|context_buffer| {
 988            if let Some(refresh_task) = refresh_context_buffer(context_buffer, cx) {
 989                stale = true;
 990                future::Either::Left(refresh_task)
 991            } else {
 992                future::Either::Right(future::ready((*context_buffer).clone()))
 993            }
 994        })
 995        .collect::<Vec<_>>();
 996
 997    if !stale {
 998        return None;
 999    }
1000
1001    let context_buffers = future::join_all(futures);
1002
1003    let id = directory_context.id;
1004    let worktree = directory_context.worktree.clone();
1005    let entry_id = directory_context.entry_id;
1006    let last_path = directory_path.path;
1007    Some(cx.spawn(async move |cx| {
1008        let context_buffers = context_buffers.await;
1009        context_store
1010            .update(cx, |context_store, _| {
1011                let new_directory_context = DirectoryContext {
1012                    id,
1013                    worktree,
1014                    entry_id,
1015                    last_path,
1016                    context_buffers,
1017                };
1018                context_store.replace_context(AssistantContext::Directory(new_directory_context));
1019            })
1020            .ok();
1021    }))
1022}
1023
1024fn refresh_symbol_text(
1025    context_store: Entity<ContextStore>,
1026    symbol_context: &SymbolContext,
1027    cx: &App,
1028) -> Option<Task<()>> {
1029    let id = symbol_context.id;
1030    let task = refresh_context_symbol(&symbol_context.context_symbol, cx);
1031    if let Some(task) = task {
1032        Some(cx.spawn(async move |cx| {
1033            let context_symbol = task.await;
1034            context_store
1035                .update(cx, |context_store, _| {
1036                    let new_symbol_context = SymbolContext { id, context_symbol };
1037                    context_store.replace_context(AssistantContext::Symbol(new_symbol_context));
1038                })
1039                .ok();
1040        }))
1041    } else {
1042        None
1043    }
1044}
1045
1046fn refresh_selection_text(
1047    context_store: Entity<ContextStore>,
1048    selection_context: &SelectionContext,
1049    cx: &App,
1050) -> Option<Task<()>> {
1051    let id = selection_context.id;
1052    let range = selection_context.range.clone();
1053    let task = refresh_context_excerpt(&selection_context.context_buffer, range.clone(), cx);
1054    if let Some(task) = task {
1055        Some(cx.spawn(async move |cx| {
1056            let (line_range, context_buffer) = task.await;
1057            context_store
1058                .update(cx, |context_store, _| {
1059                    let new_selection_context = SelectionContext {
1060                        id,
1061                        range,
1062                        line_range,
1063                        context_buffer,
1064                    };
1065                    context_store
1066                        .replace_context(AssistantContext::Selection(new_selection_context));
1067                })
1068                .ok();
1069        }))
1070    } else {
1071        None
1072    }
1073}
1074
1075fn refresh_thread_text(
1076    context_store: Entity<ContextStore>,
1077    thread_context: &ThreadContext,
1078    cx: &App,
1079) -> Task<()> {
1080    let id = thread_context.id;
1081    let thread = thread_context.thread.clone();
1082    cx.spawn(async move |cx| {
1083        context_store
1084            .update(cx, |context_store, cx| {
1085                let text = thread.read(cx).latest_detailed_summary_or_text();
1086                context_store.replace_context(AssistantContext::Thread(ThreadContext {
1087                    id,
1088                    thread,
1089                    text,
1090                }));
1091            })
1092            .ok();
1093    })
1094}
1095
1096fn refresh_user_rules(
1097    context_store: Entity<ContextStore>,
1098    user_rules_context: &RulesContext,
1099    cx: &App,
1100) -> Task<()> {
1101    let id = user_rules_context.id;
1102    let prompt_id = user_rules_context.prompt_id;
1103    let Some(thread_store) = context_store.read(cx).thread_store.as_ref() else {
1104        return Task::ready(());
1105    };
1106    let Ok(load_task) = thread_store.read_with(cx, |thread_store, cx| {
1107        thread_store.load_rules(prompt_id, cx)
1108    }) else {
1109        return Task::ready(());
1110    };
1111    cx.spawn(async move |cx| {
1112        if let Ok((metadata, text)) = load_task.await {
1113            if let Some(title) = metadata.title.clone() {
1114                context_store
1115                    .update(cx, |context_store, _cx| {
1116                        context_store.replace_context(AssistantContext::Rules(RulesContext {
1117                            id,
1118                            prompt_id,
1119                            title,
1120                            text: text.into(),
1121                        }));
1122                    })
1123                    .ok();
1124                return;
1125            }
1126        }
1127        context_store
1128            .update(cx, |context_store, cx| {
1129                context_store.remove_context(id, cx);
1130            })
1131            .ok();
1132    })
1133}
1134
1135fn refresh_context_buffer(context_buffer: &ContextBuffer, cx: &App) -> Option<Task<ContextBuffer>> {
1136    let buffer = context_buffer.buffer.read(cx);
1137    if buffer.version.changed_since(&context_buffer.version) {
1138        load_context_buffer(context_buffer.buffer.clone(), cx).log_err()
1139    } else {
1140        None
1141    }
1142}
1143
1144fn refresh_context_excerpt(
1145    context_buffer: &ContextBuffer,
1146    range: Range<Anchor>,
1147    cx: &App,
1148) -> Option<impl Future<Output = (Range<Point>, ContextBuffer)> + use<>> {
1149    let buffer = context_buffer.buffer.read(cx);
1150    if buffer.version.changed_since(&context_buffer.version) {
1151        let (line_range, context_buffer_task) =
1152            load_context_buffer_range(context_buffer.buffer.clone(), range, cx).log_err()?;
1153        Some(context_buffer_task.map(move |context_buffer| (line_range, context_buffer)))
1154    } else {
1155        None
1156    }
1157}
1158
1159fn refresh_context_symbol(
1160    context_symbol: &ContextSymbol,
1161    cx: &App,
1162) -> Option<impl Future<Output = ContextSymbol> + use<>> {
1163    let buffer = context_symbol.buffer.read(cx);
1164    let project_path = buffer.project_path(cx)?;
1165    if buffer.version.changed_since(&context_symbol.buffer_version) {
1166        let (_line_range, context_buffer_task) = load_context_buffer_range(
1167            context_symbol.buffer.clone(),
1168            context_symbol.enclosing_range.clone(),
1169            cx,
1170        )
1171        .log_err()?;
1172        let name = context_symbol.id.name.clone();
1173        let range = context_symbol.id.range.clone();
1174        let enclosing_range = context_symbol.enclosing_range.clone();
1175        Some(context_buffer_task.map(move |context_buffer| {
1176            make_context_symbol(context_buffer, project_path, name, range, enclosing_range)
1177        }))
1178    } else {
1179        None
1180    }
1181}