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