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