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