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