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