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