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