context_store.rs

  1use std::ops::Range;
  2use std::path::{Path, PathBuf};
  3use std::sync::Arc;
  4
  5use anyhow::{Context as _, Result, anyhow};
  6use collections::{BTreeMap, HashMap, HashSet};
  7use futures::{self, Future, FutureExt, future};
  8use gpui::{App, AppContext as _, AsyncApp, Context, Entity, SharedString, Task, WeakEntity};
  9use language::{Buffer, File};
 10use project::{ProjectItem, ProjectPath, Worktree};
 11use rope::Rope;
 12use text::{Anchor, BufferId, OffsetRangeExt};
 13use util::{ResultExt, maybe};
 14use workspace::Workspace;
 15
 16use crate::context::{
 17    AssistantContext, ContextBuffer, ContextId, ContextSymbol, ContextSymbolId, DirectoryContext,
 18    FetchedUrlContext, FileContext, SymbolContext, ThreadContext,
 19};
 20use crate::context_strip::SuggestedContext;
 21use crate::thread::{Thread, ThreadId};
 22
 23pub struct ContextStore {
 24    workspace: WeakEntity<Workspace>,
 25    context: Vec<AssistantContext>,
 26    // TODO: If an EntityId is used for all context types (like BufferId), can remove ContextId.
 27    next_context_id: ContextId,
 28    files: BTreeMap<BufferId, ContextId>,
 29    directories: HashMap<PathBuf, ContextId>,
 30    symbols: HashMap<ContextSymbolId, ContextId>,
 31    symbol_buffers: HashMap<ContextSymbolId, Entity<Buffer>>,
 32    symbols_by_path: HashMap<ProjectPath, Vec<ContextSymbolId>>,
 33    threads: HashMap<ThreadId, ContextId>,
 34    fetched_urls: HashMap<String, ContextId>,
 35}
 36
 37impl ContextStore {
 38    pub fn new(workspace: WeakEntity<Workspace>) -> Self {
 39        Self {
 40            workspace,
 41            context: Vec::new(),
 42            next_context_id: ContextId(0),
 43            files: BTreeMap::default(),
 44            directories: HashMap::default(),
 45            symbols: HashMap::default(),
 46            symbol_buffers: HashMap::default(),
 47            symbols_by_path: HashMap::default(),
 48            threads: HashMap::default(),
 49            fetched_urls: HashMap::default(),
 50        }
 51    }
 52
 53    pub fn context(&self) -> &Vec<AssistantContext> {
 54        &self.context
 55    }
 56
 57    pub fn context_for_id(&self, id: ContextId) -> Option<&AssistantContext> {
 58        self.context().iter().find(|context| context.id() == id)
 59    }
 60
 61    pub fn clear(&mut self) {
 62        self.context.clear();
 63        self.files.clear();
 64        self.directories.clear();
 65        self.threads.clear();
 66        self.fetched_urls.clear();
 67    }
 68
 69    pub fn add_file_from_path(
 70        &mut self,
 71        project_path: ProjectPath,
 72        remove_if_exists: bool,
 73        cx: &mut Context<Self>,
 74    ) -> Task<Result<()>> {
 75        let workspace = self.workspace.clone();
 76
 77        let Some(project) = workspace
 78            .upgrade()
 79            .map(|workspace| workspace.read(cx).project().clone())
 80        else {
 81            return Task::ready(Err(anyhow!("failed to read project")));
 82        };
 83
 84        cx.spawn(async move |this, cx| {
 85            let open_buffer_task = project.update(cx, |project, cx| {
 86                project.open_buffer(project_path.clone(), cx)
 87            })?;
 88
 89            let buffer_entity = open_buffer_task.await?;
 90            let buffer_id = this.update(cx, |_, cx| buffer_entity.read(cx).remote_id())?;
 91
 92            let already_included = this.update(cx, |this, _cx| {
 93                match this.will_include_buffer(buffer_id, &project_path.path) {
 94                    Some(FileInclusion::Direct(context_id)) => {
 95                        if remove_if_exists {
 96                            this.remove_context(context_id);
 97                        }
 98                        true
 99                    }
100                    Some(FileInclusion::InDirectory(_)) => true,
101                    None => false,
102                }
103            })?;
104
105            if already_included {
106                return anyhow::Ok(());
107            }
108
109            let (buffer_info, text_task) = this.update(cx, |_, cx| {
110                let buffer = buffer_entity.read(cx);
111                collect_buffer_info_and_text(
112                    project_path.path.clone(),
113                    buffer_entity,
114                    buffer,
115                    None,
116                    cx.to_async(),
117                )
118            })??;
119
120            let text = text_task.await;
121
122            this.update(cx, |this, _cx| {
123                this.insert_file(make_context_buffer(buffer_info, text));
124            })?;
125
126            anyhow::Ok(())
127        })
128    }
129
130    pub fn add_file_from_buffer(
131        &mut self,
132        buffer_entity: Entity<Buffer>,
133        cx: &mut Context<Self>,
134    ) -> Task<Result<()>> {
135        cx.spawn(async move |this, cx| {
136            let (buffer_info, text_task) = this.update(cx, |_, cx| {
137                let buffer = buffer_entity.read(cx);
138                let Some(file) = buffer.file() else {
139                    return Err(anyhow!("Buffer has no path."));
140                };
141                collect_buffer_info_and_text(
142                    file.path().clone(),
143                    buffer_entity,
144                    buffer,
145                    None,
146                    cx.to_async(),
147                )
148            })??;
149
150            let text = text_task.await;
151
152            this.update(cx, |this, _cx| {
153                this.insert_file(make_context_buffer(buffer_info, text))
154            })?;
155
156            anyhow::Ok(())
157        })
158    }
159
160    fn insert_file(&mut self, context_buffer: ContextBuffer) {
161        let id = self.next_context_id.post_inc();
162        self.files.insert(context_buffer.id, id);
163        self.context.push(AssistantContext::File(FileContext {
164            id,
165            context_buffer: context_buffer,
166        }));
167    }
168
169    pub fn add_directory(
170        &mut self,
171        project_path: ProjectPath,
172        remove_if_exists: bool,
173        cx: &mut Context<Self>,
174    ) -> Task<Result<()>> {
175        let workspace = self.workspace.clone();
176        let Some(project) = workspace
177            .upgrade()
178            .map(|workspace| workspace.read(cx).project().clone())
179        else {
180            return Task::ready(Err(anyhow!("failed to read project")));
181        };
182
183        let already_included = match self.includes_directory(&project_path.path) {
184            Some(FileInclusion::Direct(context_id)) => {
185                if remove_if_exists {
186                    self.remove_context(context_id);
187                }
188                true
189            }
190            Some(FileInclusion::InDirectory(_)) => true,
191            None => false,
192        };
193        if already_included {
194            return Task::ready(Ok(()));
195        }
196
197        let worktree_id = project_path.worktree_id;
198        cx.spawn(async move |this, cx| {
199            let worktree = project.update(cx, |project, cx| {
200                project
201                    .worktree_for_id(worktree_id, cx)
202                    .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))
203            })??;
204
205            let files = worktree.update(cx, |worktree, _cx| {
206                collect_files_in_path(worktree, &project_path.path)
207            })?;
208
209            let open_buffers_task = project.update(cx, |project, cx| {
210                let tasks = files.iter().map(|file_path| {
211                    project.open_buffer(
212                        ProjectPath {
213                            worktree_id,
214                            path: file_path.clone(),
215                        },
216                        cx,
217                    )
218                });
219                future::join_all(tasks)
220            })?;
221
222            let buffers = open_buffers_task.await;
223
224            let mut buffer_infos = Vec::new();
225            let mut text_tasks = Vec::new();
226            this.update(cx, |_, cx| {
227                for (path, buffer_entity) in files.into_iter().zip(buffers) {
228                    // Skip all binary files and other non-UTF8 files
229                    if let Ok(buffer_entity) = buffer_entity {
230                        let buffer = buffer_entity.read(cx);
231                        if let Some((buffer_info, text_task)) = collect_buffer_info_and_text(
232                            path,
233                            buffer_entity,
234                            buffer,
235                            None,
236                            cx.to_async(),
237                        )
238                        .log_err()
239                        {
240                            buffer_infos.push(buffer_info);
241                            text_tasks.push(text_task);
242                        }
243                    }
244                }
245                anyhow::Ok(())
246            })??;
247
248            let buffer_texts = future::join_all(text_tasks).await;
249            let context_buffers = buffer_infos
250                .into_iter()
251                .zip(buffer_texts)
252                .map(|(info, text)| make_context_buffer(info, text))
253                .collect::<Vec<_>>();
254
255            if context_buffers.is_empty() {
256                return Err(anyhow!(
257                    "No text files found in {}",
258                    &project_path.path.display()
259                ));
260            }
261
262            this.update(cx, |this, _| {
263                this.insert_directory(project_path, context_buffers);
264            })?;
265
266            anyhow::Ok(())
267        })
268    }
269
270    fn insert_directory(&mut self, project_path: ProjectPath, context_buffers: Vec<ContextBuffer>) {
271        let id = self.next_context_id.post_inc();
272        self.directories.insert(project_path.path.to_path_buf(), id);
273
274        self.context
275            .push(AssistantContext::Directory(DirectoryContext {
276                id,
277                project_path,
278                context_buffers,
279            }));
280    }
281
282    pub fn add_symbol(
283        &mut self,
284        buffer: Entity<Buffer>,
285        symbol_name: SharedString,
286        symbol_range: Range<Anchor>,
287        symbol_enclosing_range: Range<Anchor>,
288        remove_if_exists: bool,
289        cx: &mut Context<Self>,
290    ) -> Task<Result<bool>> {
291        let buffer_ref = buffer.read(cx);
292        let Some(file) = buffer_ref.file() else {
293            return Task::ready(Err(anyhow!("Buffer has no path.")));
294        };
295
296        let Some(project_path) = buffer_ref.project_path(cx) else {
297            return Task::ready(Err(anyhow!("Buffer has no project path.")));
298        };
299
300        if let Some(symbols_for_path) = self.symbols_by_path.get(&project_path) {
301            let mut matching_symbol_id = None;
302            for symbol in symbols_for_path {
303                if &symbol.name == &symbol_name {
304                    let snapshot = buffer_ref.snapshot();
305                    if symbol.range.to_offset(&snapshot) == symbol_range.to_offset(&snapshot) {
306                        matching_symbol_id = self.symbols.get(symbol).cloned();
307                        break;
308                    }
309                }
310            }
311
312            if let Some(id) = matching_symbol_id {
313                if remove_if_exists {
314                    self.remove_context(id);
315                }
316                return Task::ready(Ok(false));
317            }
318        }
319
320        let (buffer_info, collect_content_task) = match collect_buffer_info_and_text(
321            file.path().clone(),
322            buffer,
323            buffer_ref,
324            Some(symbol_enclosing_range.clone()),
325            cx.to_async(),
326        ) {
327            Ok((buffer_info, collect_context_task)) => (buffer_info, collect_context_task),
328            Err(err) => return Task::ready(Err(err)),
329        };
330
331        cx.spawn(async move |this, cx| {
332            let content = collect_content_task.await;
333
334            this.update(cx, |this, _cx| {
335                this.insert_symbol(make_context_symbol(
336                    buffer_info,
337                    project_path,
338                    symbol_name,
339                    symbol_range,
340                    symbol_enclosing_range,
341                    content,
342                ))
343            })?;
344            anyhow::Ok(true)
345        })
346    }
347
348    fn insert_symbol(&mut self, context_symbol: ContextSymbol) {
349        let id = self.next_context_id.post_inc();
350        self.symbols.insert(context_symbol.id.clone(), id);
351        self.symbols_by_path
352            .entry(context_symbol.id.path.clone())
353            .or_insert_with(Vec::new)
354            .push(context_symbol.id.clone());
355        self.symbol_buffers
356            .insert(context_symbol.id.clone(), context_symbol.buffer.clone());
357        self.context.push(AssistantContext::Symbol(SymbolContext {
358            id,
359            context_symbol,
360        }));
361    }
362
363    pub fn add_thread(
364        &mut self,
365        thread: Entity<Thread>,
366        remove_if_exists: bool,
367        cx: &mut Context<Self>,
368    ) {
369        if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {
370            if remove_if_exists {
371                self.remove_context(context_id);
372            }
373        } else {
374            self.insert_thread(thread, cx);
375        }
376    }
377
378    fn insert_thread(&mut self, thread: Entity<Thread>, cx: &App) {
379        let id = self.next_context_id.post_inc();
380        let text = thread.read(cx).text().into();
381
382        self.threads.insert(thread.read(cx).id().clone(), id);
383        self.context
384            .push(AssistantContext::Thread(ThreadContext { id, thread, text }));
385    }
386
387    pub fn add_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
388        if self.includes_url(&url).is_none() {
389            self.insert_fetched_url(url, text);
390        }
391    }
392
393    fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
394        let id = self.next_context_id.post_inc();
395
396        self.fetched_urls.insert(url.clone(), id);
397        self.context
398            .push(AssistantContext::FetchedUrl(FetchedUrlContext {
399                id,
400                url: url.into(),
401                text: text.into(),
402            }));
403    }
404
405    pub fn accept_suggested_context(
406        &mut self,
407        suggested: &SuggestedContext,
408        cx: &mut Context<ContextStore>,
409    ) -> Task<Result<()>> {
410        match suggested {
411            SuggestedContext::File {
412                buffer,
413                icon_path: _,
414                name: _,
415            } => {
416                if let Some(buffer) = buffer.upgrade() {
417                    return self.add_file_from_buffer(buffer, cx);
418                };
419            }
420            SuggestedContext::Thread { thread, name: _ } => {
421                if let Some(thread) = thread.upgrade() {
422                    self.insert_thread(thread, cx);
423                };
424            }
425        }
426        Task::ready(Ok(()))
427    }
428
429    pub fn remove_context(&mut self, id: ContextId) {
430        let Some(ix) = self.context.iter().position(|context| context.id() == id) else {
431            return;
432        };
433
434        match self.context.remove(ix) {
435            AssistantContext::File(_) => {
436                self.files.retain(|_, context_id| *context_id != id);
437            }
438            AssistantContext::Directory(_) => {
439                self.directories.retain(|_, context_id| *context_id != id);
440            }
441            AssistantContext::Symbol(symbol) => {
442                if let Some(symbols_in_path) =
443                    self.symbols_by_path.get_mut(&symbol.context_symbol.id.path)
444                {
445                    symbols_in_path.retain(|s| {
446                        self.symbols
447                            .get(s)
448                            .map_or(false, |context_id| *context_id != id)
449                    });
450                }
451                self.symbol_buffers.remove(&symbol.context_symbol.id);
452                self.symbols.retain(|_, context_id| *context_id != id);
453            }
454            AssistantContext::FetchedUrl(_) => {
455                self.fetched_urls.retain(|_, context_id| *context_id != id);
456            }
457            AssistantContext::Thread(_) => {
458                self.threads.retain(|_, context_id| *context_id != id);
459            }
460        }
461    }
462
463    /// Returns whether the buffer is already included directly in the context, or if it will be
464    /// included in the context via a directory. Directory inclusion is based on paths rather than
465    /// buffer IDs as the directory will be re-scanned.
466    pub fn will_include_buffer(&self, buffer_id: BufferId, path: &Path) -> Option<FileInclusion> {
467        if let Some(context_id) = self.files.get(&buffer_id) {
468            return Some(FileInclusion::Direct(*context_id));
469        }
470
471        self.will_include_file_path_via_directory(path)
472    }
473
474    /// Returns whether this file path is already included directly in the context, or if it will be
475    /// included in the context via a directory.
476    pub fn will_include_file_path(&self, path: &Path, cx: &App) -> Option<FileInclusion> {
477        if !self.files.is_empty() {
478            let found_file_context = self.context.iter().find(|context| match &context {
479                AssistantContext::File(file_context) => {
480                    let buffer = file_context.context_buffer.buffer.read(cx);
481                    if let Some(file_path) = buffer_path_log_err(buffer, cx) {
482                        *file_path == *path
483                    } else {
484                        false
485                    }
486                }
487                _ => false,
488            });
489            if let Some(context) = found_file_context {
490                return Some(FileInclusion::Direct(context.id()));
491            }
492        }
493
494        self.will_include_file_path_via_directory(path)
495    }
496
497    fn will_include_file_path_via_directory(&self, path: &Path) -> Option<FileInclusion> {
498        if self.directories.is_empty() {
499            return None;
500        }
501
502        let mut buf = path.to_path_buf();
503
504        while buf.pop() {
505            if let Some(_) = self.directories.get(&buf) {
506                return Some(FileInclusion::InDirectory(buf));
507            }
508        }
509
510        None
511    }
512
513    pub fn includes_directory(&self, path: &Path) -> Option<FileInclusion> {
514        if let Some(context_id) = self.directories.get(path) {
515            return Some(FileInclusion::Direct(*context_id));
516        }
517
518        self.will_include_file_path_via_directory(path)
519    }
520
521    pub fn included_symbol(&self, symbol_id: &ContextSymbolId) -> Option<ContextId> {
522        self.symbols.get(symbol_id).copied()
523    }
524
525    pub fn included_symbols_by_path(&self) -> &HashMap<ProjectPath, Vec<ContextSymbolId>> {
526        &self.symbols_by_path
527    }
528
529    pub fn buffer_for_symbol(&self, symbol_id: &ContextSymbolId) -> Option<Entity<Buffer>> {
530        self.symbol_buffers.get(symbol_id).cloned()
531    }
532
533    pub fn includes_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
534        self.threads.get(thread_id).copied()
535    }
536
537    pub fn includes_url(&self, url: &str) -> Option<ContextId> {
538        self.fetched_urls.get(url).copied()
539    }
540
541    /// Replaces the context that matches the ID of the new context, if any match.
542    fn replace_context(&mut self, new_context: AssistantContext) {
543        let id = new_context.id();
544        for context in self.context.iter_mut() {
545            if context.id() == id {
546                *context = new_context;
547                break;
548            }
549        }
550    }
551
552    pub fn file_paths(&self, cx: &App) -> HashSet<PathBuf> {
553        self.context
554            .iter()
555            .filter_map(|context| match context {
556                AssistantContext::File(file) => {
557                    let buffer = file.context_buffer.buffer.read(cx);
558                    buffer_path_log_err(buffer, cx).map(|p| p.to_path_buf())
559                }
560                AssistantContext::Directory(_)
561                | AssistantContext::Symbol(_)
562                | AssistantContext::FetchedUrl(_)
563                | AssistantContext::Thread(_) => None,
564            })
565            .collect()
566    }
567
568    pub fn thread_ids(&self) -> HashSet<ThreadId> {
569        self.threads.keys().cloned().collect()
570    }
571}
572
573pub enum FileInclusion {
574    Direct(ContextId),
575    InDirectory(PathBuf),
576}
577
578// ContextBuffer without text.
579struct BufferInfo {
580    buffer_entity: Entity<Buffer>,
581    file: Arc<dyn File>,
582    id: BufferId,
583    version: clock::Global,
584}
585
586fn make_context_buffer(info: BufferInfo, text: SharedString) -> ContextBuffer {
587    ContextBuffer {
588        id: info.id,
589        buffer: info.buffer_entity,
590        file: info.file,
591        version: info.version,
592        text,
593    }
594}
595
596fn make_context_symbol(
597    info: BufferInfo,
598    path: ProjectPath,
599    name: SharedString,
600    range: Range<Anchor>,
601    enclosing_range: Range<Anchor>,
602    text: SharedString,
603) -> ContextSymbol {
604    ContextSymbol {
605        id: ContextSymbolId { name, range, path },
606        buffer_version: info.version,
607        enclosing_range,
608        buffer: info.buffer_entity,
609        text,
610    }
611}
612
613fn collect_buffer_info_and_text(
614    path: Arc<Path>,
615    buffer_entity: Entity<Buffer>,
616    buffer: &Buffer,
617    range: Option<Range<Anchor>>,
618    cx: AsyncApp,
619) -> Result<(BufferInfo, Task<SharedString>)> {
620    let buffer_info = BufferInfo {
621        id: buffer.remote_id(),
622        buffer_entity,
623        file: buffer
624            .file()
625            .context("buffer context must have a file")?
626            .clone(),
627        version: buffer.version(),
628    };
629    // Important to collect version at the same time as content so that staleness logic is correct.
630    let content = if let Some(range) = range {
631        buffer.text_for_range(range).collect::<Rope>()
632    } else {
633        buffer.as_rope().clone()
634    };
635    let text_task = cx.background_spawn(async move { to_fenced_codeblock(&path, content) });
636    Ok((buffer_info, text_task))
637}
638
639pub fn buffer_path_log_err(buffer: &Buffer, cx: &App) -> Option<Arc<Path>> {
640    if let Some(file) = buffer.file() {
641        let mut path = file.path().clone();
642        if path.as_os_str().is_empty() {
643            path = file.full_path(cx).into();
644        }
645        Some(path)
646    } else {
647        log::error!("Buffer that had a path unexpectedly no longer has a path.");
648        None
649    }
650}
651
652fn to_fenced_codeblock(path: &Path, content: Rope) -> SharedString {
653    let path_extension = path.extension().and_then(|ext| ext.to_str());
654    let path_string = path.to_string_lossy();
655    let capacity = 3
656        + path_extension.map_or(0, |extension| extension.len() + 1)
657        + path_string.len()
658        + 1
659        + content.len()
660        + 5;
661    let mut buffer = String::with_capacity(capacity);
662
663    buffer.push_str("```");
664
665    if let Some(extension) = path_extension {
666        buffer.push_str(extension);
667        buffer.push(' ');
668    }
669    buffer.push_str(&path_string);
670
671    buffer.push('\n');
672    for chunk in content.chunks() {
673        buffer.push_str(&chunk);
674    }
675
676    if !buffer.ends_with('\n') {
677        buffer.push('\n');
678    }
679
680    buffer.push_str("```\n");
681
682    debug_assert!(
683        buffer.len() == capacity - 1 || buffer.len() == capacity,
684        "to_fenced_codeblock calculated capacity of {}, but length was {}",
685        capacity,
686        buffer.len(),
687    );
688
689    buffer.into()
690}
691
692fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
693    let mut files = Vec::new();
694
695    for entry in worktree.child_entries(path) {
696        if entry.is_dir() {
697            files.extend(collect_files_in_path(worktree, &entry.path));
698        } else if entry.is_file() {
699            files.push(entry.path.clone());
700        }
701    }
702
703    files
704}
705
706pub fn refresh_context_store_text(
707    context_store: Entity<ContextStore>,
708    changed_buffers: &HashSet<Entity<Buffer>>,
709    cx: &App,
710) -> impl Future<Output = Vec<ContextId>> + use<> {
711    let mut tasks = Vec::new();
712
713    for context in &context_store.read(cx).context {
714        let id = context.id();
715
716        let task = maybe!({
717            match context {
718                AssistantContext::File(file_context) => {
719                    if changed_buffers.is_empty()
720                        || changed_buffers.contains(&file_context.context_buffer.buffer)
721                    {
722                        let context_store = context_store.clone();
723                        return refresh_file_text(context_store, file_context, cx);
724                    }
725                }
726                AssistantContext::Directory(directory_context) => {
727                    let should_refresh = changed_buffers.is_empty()
728                        || changed_buffers.iter().any(|buffer| {
729                            let buffer = buffer.read(cx);
730
731                            buffer_path_log_err(&buffer, cx).map_or(false, |path| {
732                                path.starts_with(&directory_context.project_path.path)
733                            })
734                        });
735
736                    if should_refresh {
737                        let context_store = context_store.clone();
738                        return refresh_directory_text(context_store, directory_context, cx);
739                    }
740                }
741                AssistantContext::Symbol(symbol_context) => {
742                    if changed_buffers.is_empty()
743                        || changed_buffers.contains(&symbol_context.context_symbol.buffer)
744                    {
745                        let context_store = context_store.clone();
746                        return refresh_symbol_text(context_store, symbol_context, cx);
747                    }
748                }
749                AssistantContext::Thread(thread_context) => {
750                    if changed_buffers.is_empty() {
751                        let context_store = context_store.clone();
752                        return Some(refresh_thread_text(context_store, thread_context, cx));
753                    }
754                }
755                // Intentionally omit refreshing fetched URLs as it doesn't seem all that useful,
756                // and doing the caching properly could be tricky (unless it's already handled by
757                // the HttpClient?).
758                AssistantContext::FetchedUrl(_) => {}
759            }
760
761            None
762        });
763
764        if let Some(task) = task {
765            tasks.push(task.map(move |_| id));
766        }
767    }
768
769    future::join_all(tasks)
770}
771
772fn refresh_file_text(
773    context_store: Entity<ContextStore>,
774    file_context: &FileContext,
775    cx: &App,
776) -> Option<Task<()>> {
777    let id = file_context.id;
778    let task = refresh_context_buffer(&file_context.context_buffer, cx);
779    if let Some(task) = task {
780        Some(cx.spawn(async move |cx| {
781            let context_buffer = task.await;
782            context_store
783                .update(cx, |context_store, _| {
784                    let new_file_context = FileContext { id, context_buffer };
785                    context_store.replace_context(AssistantContext::File(new_file_context));
786                })
787                .ok();
788        }))
789    } else {
790        None
791    }
792}
793
794fn refresh_directory_text(
795    context_store: Entity<ContextStore>,
796    directory_context: &DirectoryContext,
797    cx: &App,
798) -> Option<Task<()>> {
799    let mut stale = false;
800    let futures = directory_context
801        .context_buffers
802        .iter()
803        .map(|context_buffer| {
804            if let Some(refresh_task) = refresh_context_buffer(context_buffer, cx) {
805                stale = true;
806                future::Either::Left(refresh_task)
807            } else {
808                future::Either::Right(future::ready((*context_buffer).clone()))
809            }
810        })
811        .collect::<Vec<_>>();
812
813    if !stale {
814        return None;
815    }
816
817    let context_buffers = future::join_all(futures);
818
819    let id = directory_context.id;
820    let project_path = directory_context.project_path.clone();
821    Some(cx.spawn(async move |cx| {
822        let context_buffers = context_buffers.await;
823        context_store
824            .update(cx, |context_store, _| {
825                let new_directory_context = DirectoryContext {
826                    id,
827                    project_path,
828                    context_buffers,
829                };
830                context_store.replace_context(AssistantContext::Directory(new_directory_context));
831            })
832            .ok();
833    }))
834}
835
836fn refresh_symbol_text(
837    context_store: Entity<ContextStore>,
838    symbol_context: &SymbolContext,
839    cx: &App,
840) -> Option<Task<()>> {
841    let id = symbol_context.id;
842    let task = refresh_context_symbol(&symbol_context.context_symbol, cx);
843    if let Some(task) = task {
844        Some(cx.spawn(async move |cx| {
845            let context_symbol = task.await;
846            context_store
847                .update(cx, |context_store, _| {
848                    let new_symbol_context = SymbolContext { id, context_symbol };
849                    context_store.replace_context(AssistantContext::Symbol(new_symbol_context));
850                })
851                .ok();
852        }))
853    } else {
854        None
855    }
856}
857
858fn refresh_thread_text(
859    context_store: Entity<ContextStore>,
860    thread_context: &ThreadContext,
861    cx: &App,
862) -> Task<()> {
863    let id = thread_context.id;
864    let thread = thread_context.thread.clone();
865    cx.spawn(async move |cx| {
866        context_store
867            .update(cx, |context_store, cx| {
868                let text = thread.read(cx).text().into();
869                context_store.replace_context(AssistantContext::Thread(ThreadContext {
870                    id,
871                    thread,
872                    text,
873                }));
874            })
875            .ok();
876    })
877}
878
879fn refresh_context_buffer(
880    context_buffer: &ContextBuffer,
881    cx: &App,
882) -> Option<impl Future<Output = ContextBuffer> + use<>> {
883    let buffer = context_buffer.buffer.read(cx);
884    let path = buffer_path_log_err(buffer, cx)?;
885    if buffer.version.changed_since(&context_buffer.version) {
886        let (buffer_info, text_task) = collect_buffer_info_and_text(
887            path,
888            context_buffer.buffer.clone(),
889            buffer,
890            None,
891            cx.to_async(),
892        )
893        .log_err()?;
894        Some(text_task.map(move |text| make_context_buffer(buffer_info, text)))
895    } else {
896        None
897    }
898}
899
900fn refresh_context_symbol(
901    context_symbol: &ContextSymbol,
902    cx: &App,
903) -> Option<impl Future<Output = ContextSymbol> + use<>> {
904    let buffer = context_symbol.buffer.read(cx);
905    let path = buffer_path_log_err(buffer, cx)?;
906    let project_path = buffer.project_path(cx)?;
907    if buffer.version.changed_since(&context_symbol.buffer_version) {
908        let (buffer_info, text_task) = collect_buffer_info_and_text(
909            path,
910            context_symbol.buffer.clone(),
911            buffer,
912            Some(context_symbol.enclosing_range.clone()),
913            cx.to_async(),
914        )
915        .log_err()?;
916        let name = context_symbol.id.name.clone();
917        let range = context_symbol.id.range.clone();
918        let enclosing_range = context_symbol.enclosing_range.clone();
919        Some(text_task.map(move |text| {
920            make_context_symbol(
921                buffer_info,
922                project_path,
923                name,
924                range,
925                enclosing_range,
926                text,
927            )
928        }))
929    } else {
930        None
931    }
932}