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