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}