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