1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use anyhow::{anyhow, bail, Result};
5use collections::{BTreeMap, HashMap};
6use gpui::{AppContext, AsyncAppContext, Model, ModelContext, SharedString, Task, WeakView};
7use language::Buffer;
8use project::{ProjectPath, Worktree};
9use rope::Rope;
10use text::BufferId;
11use workspace::Workspace;
12
13use crate::context::{
14 Context, ContextBuffer, ContextId, ContextKind, ContextSnapshot, DirectoryContext,
15 FetchedUrlContext, FileContext, ThreadContext,
16};
17use crate::thread::{Thread, ThreadId};
18
19pub struct ContextStore {
20 workspace: WeakView<Workspace>,
21 context: Vec<Context>,
22 // TODO: If an EntityId is used for all context types (like BufferId), can remove ContextId.
23 next_context_id: ContextId,
24 files: BTreeMap<BufferId, ContextId>,
25 directories: HashMap<PathBuf, ContextId>,
26 threads: HashMap<ThreadId, ContextId>,
27 fetched_urls: HashMap<String, ContextId>,
28}
29
30impl ContextStore {
31 pub fn new(workspace: WeakView<Workspace>) -> Self {
32 Self {
33 workspace,
34 context: Vec::new(),
35 next_context_id: ContextId(0),
36 files: BTreeMap::default(),
37 directories: HashMap::default(),
38 threads: HashMap::default(),
39 fetched_urls: HashMap::default(),
40 }
41 }
42
43 pub fn snapshot<'a>(
44 &'a self,
45 cx: &'a AppContext,
46 ) -> 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<Context> {
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 ModelContext<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(|this, mut cx| async move {
79 let open_buffer_task = project.update(&mut cx, |project, cx| {
80 project.open_buffer(project_path.clone(), cx)
81 })?;
82
83 let buffer_model = open_buffer_task.await?;
84 let buffer_id = this.update(&mut cx, |_, cx| buffer_model.read(cx).remote_id())?;
85
86 let already_included = this.update(&mut 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(&mut cx, |_, cx| {
102 let buffer = buffer_model.read(cx);
103 collect_buffer_info_and_text(
104 project_path.path.clone(),
105 buffer_model,
106 buffer,
107 &cx.to_async(),
108 )
109 })?;
110
111 let text = text_task.await;
112
113 this.update(&mut 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_model: Model<Buffer>,
124 cx: &mut ModelContext<Self>,
125 ) -> Task<Result<()>> {
126 cx.spawn(|this, mut cx| async move {
127 let (buffer_info, text_task) = this.update(&mut cx, |_, cx| {
128 let buffer = buffer_model.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_model,
135 buffer,
136 &cx.to_async(),
137 ))
138 })??;
139
140 let text = text_task.await;
141
142 this.update(&mut cx, |this, _cx| {
143 this.insert_file(make_context_buffer(buffer_info, text))
144 })?;
145
146 anyhow::Ok(())
147 })
148 }
149
150 pub 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.push(Context::File(FileContext {
154 id,
155 buffer: context_buffer,
156 }));
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_buffer_tasks = project.update(&mut cx, |project, cx| {
196 files
197 .iter()
198 .map(|file_path| {
199 project.open_buffer(
200 ProjectPath {
201 worktree_id,
202 path: file_path.clone(),
203 },
204 cx,
205 )
206 })
207 .collect::<Vec<_>>()
208 })?;
209
210 let buffers = futures::future::join_all(open_buffer_tasks).await;
211
212 let mut buffer_infos = Vec::new();
213 let mut text_tasks = Vec::new();
214 this.update(&mut cx, |_, cx| {
215 for (path, buffer_model) in files.into_iter().zip(buffers) {
216 let buffer_model = buffer_model?;
217 let buffer = buffer_model.read(cx);
218 let (buffer_info, text_task) =
219 collect_buffer_info_and_text(path, buffer_model, buffer, &cx.to_async());
220 buffer_infos.push(buffer_info);
221 text_tasks.push(text_task);
222 }
223 anyhow::Ok(())
224 })??;
225
226 let buffer_texts = futures::future::join_all(text_tasks).await;
227 let directory_buffers = buffer_infos
228 .into_iter()
229 .zip(buffer_texts.iter())
230 .map(|(info, text)| make_context_buffer(info, text.clone()))
231 .collect::<Vec<_>>();
232
233 if directory_buffers.is_empty() {
234 bail!("No text files found in {}", &project_path.path.display());
235 }
236
237 // TODO: include directory path in text?
238
239 this.update(&mut cx, |this, _| {
240 this.insert_directory(&project_path.path, directory_buffers, buffer_texts.into());
241 })?;
242
243 anyhow::Ok(())
244 })
245 }
246
247 pub fn insert_directory(
248 &mut self,
249 path: &Path,
250 buffers: Vec<ContextBuffer>,
251 text: Box<[SharedString]>,
252 ) {
253 let id = self.next_context_id.post_inc();
254 self.directories.insert(path.to_path_buf(), id);
255
256 let full_path: SharedString = path.to_string_lossy().into_owned().into();
257
258 let name = match path.file_name() {
259 Some(name) => name.to_string_lossy().into_owned().into(),
260 None => full_path.clone(),
261 };
262
263 let parent = path
264 .parent()
265 .and_then(|p| p.file_name())
266 .map(|p| p.to_string_lossy().into_owned().into());
267
268 self.context.push(Context::Directory(DirectoryContext {
269 path: path.into(),
270 buffers,
271 snapshot: ContextSnapshot {
272 id,
273 name,
274 parent,
275 tooltip: Some(full_path),
276 icon_path: None,
277 kind: ContextKind::Directory,
278 text,
279 },
280 }));
281 }
282
283 pub fn add_thread(&mut self, thread: Model<Thread>, cx: &mut ModelContext<Self>) {
284 if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {
285 self.remove_context(context_id);
286 } else {
287 self.insert_thread(thread, cx);
288 }
289 }
290
291 pub fn insert_thread(&mut self, thread: Model<Thread>, cx: &AppContext) {
292 let id = self.next_context_id.post_inc();
293 let thread_ref = thread.read(cx);
294 let text = thread_ref.text().into();
295
296 self.threads.insert(thread_ref.id().clone(), id);
297 self.context
298 .push(Context::Thread(ThreadContext { id, thread, text }));
299 }
300
301 pub fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
302 let id = self.next_context_id.post_inc();
303
304 self.fetched_urls.insert(url.clone(), id);
305 self.context.push(Context::FetchedUrl(FetchedUrlContext {
306 id,
307 url: url.into(),
308 text: text.into(),
309 }));
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 if let Some(file_path) = file_context.path(cx) {
351 *file_path == *path
352 } else {
353 false
354 }
355 }
356 _ => false,
357 });
358 if let Some(context) = found_file_context {
359 return Some(FileInclusion::Direct(context.id()));
360 }
361 }
362
363 self.will_include_file_path_via_directory(path)
364 }
365
366 fn will_include_file_path_via_directory(&self, path: &Path) -> Option<FileInclusion> {
367 if self.directories.is_empty() {
368 return None;
369 }
370
371 let mut buf = path.to_path_buf();
372
373 while buf.pop() {
374 if let Some(_) = self.directories.get(&buf) {
375 return Some(FileInclusion::InDirectory(buf));
376 }
377 }
378
379 None
380 }
381
382 pub fn includes_directory(&self, path: &Path) -> Option<ContextId> {
383 self.directories.get(path).copied()
384 }
385
386 pub fn includes_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
387 self.threads.get(thread_id).copied()
388 }
389
390 pub fn includes_url(&self, url: &str) -> Option<ContextId> {
391 self.fetched_urls.get(url).copied()
392 }
393}
394
395pub enum FileInclusion {
396 Direct(ContextId),
397 InDirectory(PathBuf),
398}
399
400// ContextBuffer without text.
401struct BufferInfo {
402 buffer_model: Model<Buffer>,
403 id: BufferId,
404 version: clock::Global,
405}
406
407fn make_context_buffer(info: BufferInfo, text: SharedString) -> ContextBuffer {
408 ContextBuffer {
409 id: info.id,
410 buffer: info.buffer_model,
411 version: info.version,
412 text,
413 }
414}
415
416fn collect_buffer_info_and_text(
417 path: Arc<Path>,
418 buffer_model: Model<Buffer>,
419 buffer: &Buffer,
420 cx: &AsyncAppContext,
421) -> (BufferInfo, Task<SharedString>) {
422 let buffer_info = BufferInfo {
423 id: buffer.remote_id(),
424 buffer_model,
425 version: buffer.version(),
426 };
427 // Important to collect version at the same time as content so that staleness logic is correct.
428 let content = buffer.as_rope().clone();
429 let text_task = cx
430 .background_executor()
431 .spawn(async move { to_fenced_codeblock(&path, content) });
432 (buffer_info, text_task)
433}
434
435fn to_fenced_codeblock(path: &Path, content: Rope) -> SharedString {
436 let path_extension = path.extension().and_then(|ext| ext.to_str());
437 let path_string = path.to_string_lossy();
438 let capacity = 3
439 + path_extension.map_or(0, |extension| extension.len() + 1)
440 + path_string.len()
441 + 1
442 + content.len()
443 + 5;
444 let mut buffer = String::with_capacity(capacity);
445
446 buffer.push_str("```");
447
448 if let Some(extension) = path_extension {
449 buffer.push_str(extension);
450 buffer.push(' ');
451 }
452 buffer.push_str(&path_string);
453
454 buffer.push('\n');
455 for chunk in content.chunks() {
456 buffer.push_str(&chunk);
457 }
458
459 if !buffer.ends_with('\n') {
460 buffer.push('\n');
461 }
462
463 buffer.push_str("```\n");
464
465 if buffer.len() > capacity {
466 log::error!(
467 "to_fenced_codeblock calculated capacity {} but length was {}",
468 capacity,
469 buffer.len()
470 );
471 }
472
473 buffer.into()
474}
475
476fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
477 let mut files = Vec::new();
478
479 for entry in worktree.child_entries(path) {
480 if entry.is_dir() {
481 files.extend(collect_files_in_path(worktree, &entry.path));
482 } else if entry.is_file() {
483 files.push(entry.path.clone());
484 }
485 }
486
487 files
488}