1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use anyhow::{anyhow, bail, Result};
5use collections::{BTreeMap, HashMap, HashSet};
6use futures::{self, future, Future, FutureExt};
7use gpui::{App, AppContext as _, AsyncApp, Context, Entity, SharedString, Task, WeakEntity};
8use language::Buffer;
9use project::{ProjectPath, Worktree};
10use rope::Rope;
11use text::BufferId;
12use util::maybe;
13use workspace::Workspace;
14
15use crate::context::{
16 AssistantContext, ContextBuffer, ContextId, ContextSnapshot, DirectoryContext,
17 FetchedUrlContext, FileContext, ThreadContext,
18};
19use crate::context_strip::SuggestedContext;
20use crate::thread::{Thread, ThreadId};
21
22pub struct ContextStore {
23 workspace: WeakEntity<Workspace>,
24 context: Vec<AssistantContext>,
25 // TODO: If an EntityId is used for all context types (like BufferId), can remove ContextId.
26 next_context_id: ContextId,
27 files: BTreeMap<BufferId, ContextId>,
28 directories: HashMap<PathBuf, ContextId>,
29 threads: HashMap<ThreadId, ContextId>,
30 fetched_urls: HashMap<String, ContextId>,
31}
32
33impl ContextStore {
34 pub fn new(workspace: WeakEntity<Workspace>) -> Self {
35 Self {
36 workspace,
37 context: Vec::new(),
38 next_context_id: ContextId(0),
39 files: BTreeMap::default(),
40 directories: HashMap::default(),
41 threads: HashMap::default(),
42 fetched_urls: HashMap::default(),
43 }
44 }
45
46 pub fn snapshot<'a>(&'a self, cx: &'a App) -> impl Iterator<Item = ContextSnapshot> + 'a {
47 self.context()
48 .iter()
49 .flat_map(|context| context.snapshot(cx))
50 }
51
52 pub fn context(&self) -> &Vec<AssistantContext> {
53 &self.context
54 }
55
56 pub fn clear(&mut self) {
57 self.context.clear();
58 self.files.clear();
59 self.directories.clear();
60 self.threads.clear();
61 self.fetched_urls.clear();
62 }
63
64 pub fn add_file_from_path(
65 &mut self,
66 project_path: ProjectPath,
67 remove_if_exists: bool,
68 cx: &mut Context<Self>,
69 ) -> Task<Result<()>> {
70 let workspace = self.workspace.clone();
71
72 let Some(project) = workspace
73 .upgrade()
74 .map(|workspace| workspace.read(cx).project().clone())
75 else {
76 return Task::ready(Err(anyhow!("failed to read project")));
77 };
78
79 cx.spawn(async move |this, cx| {
80 let open_buffer_task = project.update(cx, |project, cx| {
81 project.open_buffer(project_path.clone(), cx)
82 })?;
83
84 let buffer_entity = open_buffer_task.await?;
85 let buffer_id = this.update(cx, |_, cx| buffer_entity.read(cx).remote_id())?;
86
87 let already_included = this.update(cx, |this, _cx| {
88 match this.will_include_buffer(buffer_id, &project_path.path) {
89 Some(FileInclusion::Direct(context_id)) => {
90 if remove_if_exists {
91 this.remove_context(context_id);
92 }
93 true
94 }
95 Some(FileInclusion::InDirectory(_)) => true,
96 None => false,
97 }
98 })?;
99
100 if already_included {
101 return anyhow::Ok(());
102 }
103
104 let (buffer_info, text_task) = this.update(cx, |_, cx| {
105 let buffer = buffer_entity.read(cx);
106 collect_buffer_info_and_text(
107 project_path.path.clone(),
108 buffer_entity,
109 buffer,
110 cx.to_async(),
111 )
112 })?;
113
114 let text = text_task.await;
115
116 this.update(cx, |this, _cx| {
117 this.insert_file(make_context_buffer(buffer_info, text));
118 })?;
119
120 anyhow::Ok(())
121 })
122 }
123
124 pub fn add_file_from_buffer(
125 &mut self,
126 buffer_entity: Entity<Buffer>,
127 cx: &mut Context<Self>,
128 ) -> Task<Result<()>> {
129 cx.spawn(async move |this, cx| {
130 let (buffer_info, text_task) = this.update(cx, |_, cx| {
131 let buffer = buffer_entity.read(cx);
132 let Some(file) = buffer.file() else {
133 return Err(anyhow!("Buffer has no path."));
134 };
135 Ok(collect_buffer_info_and_text(
136 file.path().clone(),
137 buffer_entity,
138 buffer,
139 cx.to_async(),
140 ))
141 })??;
142
143 let text = text_task.await;
144
145 this.update(cx, |this, _cx| {
146 this.insert_file(make_context_buffer(buffer_info, text))
147 })?;
148
149 anyhow::Ok(())
150 })
151 }
152
153 fn insert_file(&mut self, context_buffer: ContextBuffer) {
154 let id = self.next_context_id.post_inc();
155 self.files.insert(context_buffer.id, id);
156 self.context
157 .push(AssistantContext::File(FileContext { id, context_buffer }));
158 }
159
160 pub fn add_directory(
161 &mut self,
162 project_path: ProjectPath,
163 remove_if_exists: bool,
164 cx: &mut Context<Self>,
165 ) -> Task<Result<()>> {
166 let workspace = self.workspace.clone();
167 let Some(project) = workspace
168 .upgrade()
169 .map(|workspace| workspace.read(cx).project().clone())
170 else {
171 return Task::ready(Err(anyhow!("failed to read project")));
172 };
173
174 let already_included = if let Some(context_id) = self.includes_directory(&project_path.path)
175 {
176 if remove_if_exists {
177 self.remove_context(context_id);
178 }
179 true
180 } else {
181 false
182 };
183 if already_included {
184 return Task::ready(Ok(()));
185 }
186
187 let worktree_id = project_path.worktree_id;
188 cx.spawn(async move |this, cx| {
189 let worktree = project.update(cx, |project, cx| {
190 project
191 .worktree_for_id(worktree_id, cx)
192 .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))
193 })??;
194
195 let files = worktree.update(cx, |worktree, _cx| {
196 collect_files_in_path(worktree, &project_path.path)
197 })?;
198
199 let open_buffers_task = project.update(cx, |project, cx| {
200 let tasks = files.iter().map(|file_path| {
201 project.open_buffer(
202 ProjectPath {
203 worktree_id,
204 path: file_path.clone(),
205 },
206 cx,
207 )
208 });
209 future::join_all(tasks)
210 })?;
211
212 let buffers = open_buffers_task.await;
213
214 let mut buffer_infos = Vec::new();
215 let mut text_tasks = Vec::new();
216 this.update(cx, |_, cx| {
217 for (path, buffer_entity) in files.into_iter().zip(buffers) {
218 // Skip all binary files and other non-UTF8 files
219 if let Ok(buffer_entity) = buffer_entity {
220 let buffer = buffer_entity.read(cx);
221 let (buffer_info, text_task) = collect_buffer_info_and_text(
222 path,
223 buffer_entity,
224 buffer,
225 cx.to_async(),
226 );
227 buffer_infos.push(buffer_info);
228 text_tasks.push(text_task);
229 }
230 }
231 anyhow::Ok(())
232 })??;
233
234 let buffer_texts = future::join_all(text_tasks).await;
235 let context_buffers = buffer_infos
236 .into_iter()
237 .zip(buffer_texts)
238 .map(|(info, text)| make_context_buffer(info, text))
239 .collect::<Vec<_>>();
240
241 if context_buffers.is_empty() {
242 bail!("No text files found in {}", &project_path.path.display());
243 }
244
245 this.update(cx, |this, _| {
246 this.insert_directory(&project_path.path, context_buffers);
247 })?;
248
249 anyhow::Ok(())
250 })
251 }
252
253 fn insert_directory(&mut self, path: &Path, context_buffers: Vec<ContextBuffer>) {
254 let id = self.next_context_id.post_inc();
255 self.directories.insert(path.to_path_buf(), id);
256
257 self.context
258 .push(AssistantContext::Directory(DirectoryContext::new(
259 id,
260 path,
261 context_buffers,
262 )));
263 }
264
265 pub fn add_thread(
266 &mut self,
267 thread: Entity<Thread>,
268 remove_if_exists: bool,
269 cx: &mut Context<Self>,
270 ) {
271 if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {
272 if remove_if_exists {
273 self.remove_context(context_id);
274 }
275 } else {
276 self.insert_thread(thread, cx);
277 }
278 }
279
280 fn insert_thread(&mut self, thread: Entity<Thread>, cx: &App) {
281 let id = self.next_context_id.post_inc();
282 let text = thread.read(cx).text().into();
283
284 self.threads.insert(thread.read(cx).id().clone(), id);
285 self.context
286 .push(AssistantContext::Thread(ThreadContext { id, thread, text }));
287 }
288
289 pub fn add_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
290 if self.includes_url(&url).is_none() {
291 self.insert_fetched_url(url, text);
292 }
293 }
294
295 fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
296 let id = self.next_context_id.post_inc();
297
298 self.fetched_urls.insert(url.clone(), id);
299 self.context
300 .push(AssistantContext::FetchedUrl(FetchedUrlContext {
301 id,
302 url: url.into(),
303 text: text.into(),
304 }));
305 }
306
307 pub fn accept_suggested_context(
308 &mut self,
309 suggested: &SuggestedContext,
310 cx: &mut Context<ContextStore>,
311 ) -> Task<Result<()>> {
312 match suggested {
313 SuggestedContext::File {
314 buffer,
315 icon_path: _,
316 name: _,
317 } => {
318 if let Some(buffer) = buffer.upgrade() {
319 return self.add_file_from_buffer(buffer, cx);
320 };
321 }
322 SuggestedContext::Thread { thread, name: _ } => {
323 if let Some(thread) = thread.upgrade() {
324 self.insert_thread(thread, cx);
325 };
326 }
327 }
328 Task::ready(Ok(()))
329 }
330
331 pub fn remove_context(&mut self, id: ContextId) {
332 let Some(ix) = self.context.iter().position(|context| context.id() == id) else {
333 return;
334 };
335
336 match self.context.remove(ix) {
337 AssistantContext::File(_) => {
338 self.files.retain(|_, context_id| *context_id != id);
339 }
340 AssistantContext::Directory(_) => {
341 self.directories.retain(|_, context_id| *context_id != id);
342 }
343 AssistantContext::FetchedUrl(_) => {
344 self.fetched_urls.retain(|_, context_id| *context_id != id);
345 }
346 AssistantContext::Thread(_) => {
347 self.threads.retain(|_, context_id| *context_id != id);
348 }
349 }
350 }
351
352 /// Returns whether the buffer is already included directly in the context, or if it will be
353 /// included in the context via a directory. Directory inclusion is based on paths rather than
354 /// buffer IDs as the directory will be re-scanned.
355 pub fn will_include_buffer(&self, buffer_id: BufferId, path: &Path) -> Option<FileInclusion> {
356 if let Some(context_id) = self.files.get(&buffer_id) {
357 return Some(FileInclusion::Direct(*context_id));
358 }
359
360 self.will_include_file_path_via_directory(path)
361 }
362
363 /// Returns whether this file path is already included directly in the context, or if it will be
364 /// included in the context via a directory.
365 pub fn will_include_file_path(&self, path: &Path, cx: &App) -> Option<FileInclusion> {
366 if !self.files.is_empty() {
367 let found_file_context = self.context.iter().find(|context| match &context {
368 AssistantContext::File(file_context) => {
369 let buffer = file_context.context_buffer.buffer.read(cx);
370 if let Some(file_path) = buffer_path_log_err(buffer) {
371 *file_path == *path
372 } else {
373 false
374 }
375 }
376 _ => false,
377 });
378 if let Some(context) = found_file_context {
379 return Some(FileInclusion::Direct(context.id()));
380 }
381 }
382
383 self.will_include_file_path_via_directory(path)
384 }
385
386 fn will_include_file_path_via_directory(&self, path: &Path) -> Option<FileInclusion> {
387 if self.directories.is_empty() {
388 return None;
389 }
390
391 let mut buf = path.to_path_buf();
392
393 while buf.pop() {
394 if let Some(_) = self.directories.get(&buf) {
395 return Some(FileInclusion::InDirectory(buf));
396 }
397 }
398
399 None
400 }
401
402 pub fn includes_directory(&self, path: &Path) -> Option<ContextId> {
403 self.directories.get(path).copied()
404 }
405
406 pub fn includes_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
407 self.threads.get(thread_id).copied()
408 }
409
410 pub fn includes_url(&self, url: &str) -> Option<ContextId> {
411 self.fetched_urls.get(url).copied()
412 }
413
414 /// Replaces the context that matches the ID of the new context, if any match.
415 fn replace_context(&mut self, new_context: AssistantContext) {
416 let id = new_context.id();
417 for context in self.context.iter_mut() {
418 if context.id() == id {
419 *context = new_context;
420 break;
421 }
422 }
423 }
424
425 pub fn file_paths(&self, cx: &App) -> HashSet<PathBuf> {
426 self.context
427 .iter()
428 .filter_map(|context| match context {
429 AssistantContext::File(file) => {
430 let buffer = file.context_buffer.buffer.read(cx);
431 buffer_path_log_err(buffer).map(|p| p.to_path_buf())
432 }
433 AssistantContext::Directory(_)
434 | AssistantContext::FetchedUrl(_)
435 | AssistantContext::Thread(_) => None,
436 })
437 .collect()
438 }
439
440 pub fn thread_ids(&self) -> HashSet<ThreadId> {
441 self.threads.keys().cloned().collect()
442 }
443}
444
445pub enum FileInclusion {
446 Direct(ContextId),
447 InDirectory(PathBuf),
448}
449
450// ContextBuffer without text.
451struct BufferInfo {
452 buffer_entity: Entity<Buffer>,
453 id: BufferId,
454 version: clock::Global,
455}
456
457fn make_context_buffer(info: BufferInfo, text: SharedString) -> ContextBuffer {
458 ContextBuffer {
459 id: info.id,
460 buffer: info.buffer_entity,
461 version: info.version,
462 text,
463 }
464}
465
466fn collect_buffer_info_and_text(
467 path: Arc<Path>,
468 buffer_entity: Entity<Buffer>,
469 buffer: &Buffer,
470 cx: AsyncApp,
471) -> (BufferInfo, Task<SharedString>) {
472 let buffer_info = BufferInfo {
473 id: buffer.remote_id(),
474 buffer_entity,
475 version: buffer.version(),
476 };
477 // Important to collect version at the same time as content so that staleness logic is correct.
478 let content = buffer.as_rope().clone();
479 let text_task = cx.background_spawn(async move { to_fenced_codeblock(&path, content) });
480 (buffer_info, text_task)
481}
482
483pub fn buffer_path_log_err(buffer: &Buffer) -> Option<Arc<Path>> {
484 if let Some(file) = buffer.file() {
485 Some(file.path().clone())
486 } else {
487 log::error!("Buffer that had a path unexpectedly no longer has a path.");
488 None
489 }
490}
491
492fn to_fenced_codeblock(path: &Path, content: Rope) -> SharedString {
493 let path_extension = path.extension().and_then(|ext| ext.to_str());
494 let path_string = path.to_string_lossy();
495 let capacity = 3
496 + path_extension.map_or(0, |extension| extension.len() + 1)
497 + path_string.len()
498 + 1
499 + content.len()
500 + 5;
501 let mut buffer = String::with_capacity(capacity);
502
503 buffer.push_str("```");
504
505 if let Some(extension) = path_extension {
506 buffer.push_str(extension);
507 buffer.push(' ');
508 }
509 buffer.push_str(&path_string);
510
511 buffer.push('\n');
512 for chunk in content.chunks() {
513 buffer.push_str(&chunk);
514 }
515
516 if !buffer.ends_with('\n') {
517 buffer.push('\n');
518 }
519
520 buffer.push_str("```\n");
521
522 debug_assert!(
523 buffer.len() == capacity - 1 || buffer.len() == capacity,
524 "to_fenced_codeblock calculated capacity of {}, but length was {}",
525 capacity,
526 buffer.len(),
527 );
528
529 buffer.into()
530}
531
532fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
533 let mut files = Vec::new();
534
535 for entry in worktree.child_entries(path) {
536 if entry.is_dir() {
537 files.extend(collect_files_in_path(worktree, &entry.path));
538 } else if entry.is_file() {
539 files.push(entry.path.clone());
540 }
541 }
542
543 files
544}
545
546pub fn refresh_context_store_text(
547 context_store: Entity<ContextStore>,
548 changed_buffers: &HashSet<Entity<Buffer>>,
549 cx: &App,
550) -> impl Future<Output = Vec<ContextId>> {
551 let mut tasks = Vec::new();
552
553 for context in &context_store.read(cx).context {
554 let id = context.id();
555
556 let task = maybe!({
557 match context {
558 AssistantContext::File(file_context) => {
559 if changed_buffers.is_empty()
560 || changed_buffers.contains(&file_context.context_buffer.buffer)
561 {
562 let context_store = context_store.clone();
563 return refresh_file_text(context_store, file_context, cx);
564 }
565 }
566 AssistantContext::Directory(directory_context) => {
567 let should_refresh = changed_buffers.is_empty()
568 || changed_buffers.iter().any(|buffer| {
569 let buffer = buffer.read(cx);
570
571 buffer_path_log_err(&buffer)
572 .map_or(false, |path| path.starts_with(&directory_context.path))
573 });
574
575 if should_refresh {
576 let context_store = context_store.clone();
577 return refresh_directory_text(context_store, directory_context, cx);
578 }
579 }
580 AssistantContext::Thread(thread_context) => {
581 if changed_buffers.is_empty() {
582 let context_store = context_store.clone();
583 return Some(refresh_thread_text(context_store, thread_context, cx));
584 }
585 }
586 // Intentionally omit refreshing fetched URLs as it doesn't seem all that useful,
587 // and doing the caching properly could be tricky (unless it's already handled by
588 // the HttpClient?).
589 AssistantContext::FetchedUrl(_) => {}
590 }
591
592 None
593 });
594
595 if let Some(task) = task {
596 tasks.push(task.map(move |_| id));
597 }
598 }
599
600 future::join_all(tasks)
601}
602
603fn refresh_file_text(
604 context_store: Entity<ContextStore>,
605 file_context: &FileContext,
606 cx: &App,
607) -> Option<Task<()>> {
608 let id = file_context.id;
609 let task = refresh_context_buffer(&file_context.context_buffer, cx);
610 if let Some(task) = task {
611 Some(cx.spawn(async move |cx| {
612 let context_buffer = task.await;
613 context_store
614 .update(cx, |context_store, _| {
615 let new_file_context = FileContext { id, context_buffer };
616 context_store.replace_context(AssistantContext::File(new_file_context));
617 })
618 .ok();
619 }))
620 } else {
621 None
622 }
623}
624
625fn refresh_directory_text(
626 context_store: Entity<ContextStore>,
627 directory_context: &DirectoryContext,
628 cx: &App,
629) -> Option<Task<()>> {
630 let mut stale = false;
631 let futures = directory_context
632 .context_buffers
633 .iter()
634 .map(|context_buffer| {
635 if let Some(refresh_task) = refresh_context_buffer(context_buffer, cx) {
636 stale = true;
637 future::Either::Left(refresh_task)
638 } else {
639 future::Either::Right(future::ready((*context_buffer).clone()))
640 }
641 })
642 .collect::<Vec<_>>();
643
644 if !stale {
645 return None;
646 }
647
648 let context_buffers = future::join_all(futures);
649
650 let id = directory_context.snapshot.id;
651 let path = directory_context.path.clone();
652 Some(cx.spawn(async move |cx| {
653 let context_buffers = context_buffers.await;
654 context_store
655 .update(cx, |context_store, _| {
656 let new_directory_context = DirectoryContext::new(id, &path, context_buffers);
657 context_store.replace_context(AssistantContext::Directory(new_directory_context));
658 })
659 .ok();
660 }))
661}
662
663fn refresh_thread_text(
664 context_store: Entity<ContextStore>,
665 thread_context: &ThreadContext,
666 cx: &App,
667) -> Task<()> {
668 let id = thread_context.id;
669 let thread = thread_context.thread.clone();
670 cx.spawn(async move |cx| {
671 context_store
672 .update(cx, |context_store, cx| {
673 let text = thread.read(cx).text().into();
674 context_store.replace_context(AssistantContext::Thread(ThreadContext {
675 id,
676 thread,
677 text,
678 }));
679 })
680 .ok();
681 })
682}
683
684fn refresh_context_buffer(
685 context_buffer: &ContextBuffer,
686 cx: &App,
687) -> Option<impl Future<Output = ContextBuffer>> {
688 let buffer = context_buffer.buffer.read(cx);
689 let path = buffer_path_log_err(buffer)?;
690 if buffer.version.changed_since(&context_buffer.version) {
691 let (buffer_info, text_task) = collect_buffer_info_and_text(
692 path,
693 context_buffer.buffer.clone(),
694 buffer,
695 cx.to_async(),
696 );
697 Some(text_task.map(move |text| make_context_buffer(buffer_info, text)))
698 } else {
699 None
700 }
701}