1use std::fmt::Write as _;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use anyhow::{anyhow, bail, Result};
6use collections::{BTreeMap, HashMap};
7use gpui::{AppContext, Model, ModelContext, SharedString, Task, WeakView};
8use language::Buffer;
9use project::{ProjectPath, Worktree};
10use text::BufferId;
11use workspace::Workspace;
12
13use crate::context::{
14 Context, ContextId, ContextKind, ContextSnapshot, DirectoryContext, FetchedUrlContext,
15 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(
65 &mut self,
66 project_path: ProjectPath,
67 cx: &mut ModelContext<Self>,
68 ) -> Task<Result<()>> {
69 let workspace = self.workspace.clone();
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 = open_buffer_task.await?;
83 let buffer_id = buffer.update(&mut cx, |buffer, _cx| buffer.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 this.update(&mut cx, |this, cx| {
101 this.insert_file(buffer, cx);
102 })?;
103
104 anyhow::Ok(())
105 })
106 }
107
108 pub fn insert_file(&mut self, buffer_model: Model<Buffer>, cx: &AppContext) {
109 let buffer = buffer_model.read(cx);
110 let Some(file) = buffer.file() else {
111 return;
112 };
113
114 let mut text = String::new();
115 push_fenced_codeblock(file.path(), buffer.text(), &mut text);
116
117 let id = self.next_context_id.post_inc();
118 self.files.insert(buffer.remote_id(), id);
119 self.context.push(Context::File(FileContext {
120 id,
121 buffer: buffer_model,
122 version: buffer.version.clone(),
123 text: text.into(),
124 }));
125 }
126
127 pub fn add_directory(
128 &mut self,
129 project_path: ProjectPath,
130 cx: &mut ModelContext<Self>,
131 ) -> Task<Result<()>> {
132 let workspace = self.workspace.clone();
133 let Some(project) = workspace
134 .upgrade()
135 .map(|workspace| workspace.read(cx).project().clone())
136 else {
137 return Task::ready(Err(anyhow!("failed to read project")));
138 };
139
140 let already_included = if let Some(context_id) = self.includes_directory(&project_path.path)
141 {
142 self.remove_context(context_id);
143 true
144 } else {
145 false
146 };
147 if already_included {
148 return Task::ready(Ok(()));
149 }
150
151 let worktree_id = project_path.worktree_id;
152 cx.spawn(|this, mut cx| async move {
153 let worktree = project.update(&mut cx, |project, cx| {
154 project
155 .worktree_for_id(worktree_id, cx)
156 .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))
157 })??;
158
159 let files = worktree.update(&mut cx, |worktree, _cx| {
160 collect_files_in_path(worktree, &project_path.path)
161 })?;
162
163 let open_buffer_tasks = project.update(&mut cx, |project, cx| {
164 files
165 .into_iter()
166 .map(|file_path| {
167 project.open_buffer(
168 ProjectPath {
169 worktree_id,
170 path: file_path.clone(),
171 },
172 cx,
173 )
174 })
175 .collect::<Vec<_>>()
176 })?;
177
178 let buffers = futures::future::join_all(open_buffer_tasks).await;
179
180 this.update(&mut cx, |this, cx| {
181 let mut text = String::new();
182 let mut directory_buffers = BTreeMap::new();
183 for buffer_model in buffers {
184 let buffer_model = buffer_model?;
185 let buffer = buffer_model.read(cx);
186 let path = buffer.file().map_or(&project_path.path, |file| file.path());
187 push_fenced_codeblock(&path, buffer.text(), &mut text);
188 directory_buffers
189 .insert(buffer.remote_id(), (buffer_model, buffer.version.clone()));
190 }
191
192 if directory_buffers.is_empty() {
193 bail!(
194 "could not read any text files from {}",
195 &project_path.path.display()
196 );
197 }
198
199 this.insert_directory(&project_path.path, directory_buffers, text);
200
201 anyhow::Ok(())
202 })??;
203
204 anyhow::Ok(())
205 })
206 }
207
208 pub fn insert_directory(
209 &mut self,
210 path: &Path,
211 buffers: BTreeMap<BufferId, (Model<Buffer>, clock::Global)>,
212 text: impl Into<SharedString>,
213 ) {
214 let id = self.next_context_id.post_inc();
215 self.directories.insert(path.to_path_buf(), id);
216
217 let full_path: SharedString = path.to_string_lossy().into_owned().into();
218
219 let name = match path.file_name() {
220 Some(name) => name.to_string_lossy().into_owned().into(),
221 None => full_path.clone(),
222 };
223
224 let parent = path
225 .parent()
226 .and_then(|p| p.file_name())
227 .map(|p| p.to_string_lossy().into_owned().into());
228
229 self.context.push(Context::Directory(DirectoryContext {
230 path: path.into(),
231 buffers,
232 snapshot: ContextSnapshot {
233 id,
234 name,
235 parent,
236 tooltip: Some(full_path),
237 kind: ContextKind::Directory,
238 text: text.into(),
239 },
240 }));
241 }
242
243 pub fn add_thread(&mut self, thread: Model<Thread>, cx: &mut ModelContext<Self>) {
244 if let Some(context_id) = self.includes_thread(&thread.read(cx).id()) {
245 self.remove_context(context_id);
246 } else {
247 self.insert_thread(thread, cx);
248 }
249 }
250
251 pub fn insert_thread(&mut self, thread: Model<Thread>, cx: &AppContext) {
252 let id = self.next_context_id.post_inc();
253 let thread_ref = thread.read(cx);
254 let text = thread_ref.text().into();
255
256 self.threads.insert(thread_ref.id().clone(), id);
257 self.context
258 .push(Context::Thread(ThreadContext { id, thread, text }));
259 }
260
261 pub fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
262 let id = self.next_context_id.post_inc();
263
264 self.fetched_urls.insert(url.clone(), id);
265 self.context.push(Context::FetchedUrl(FetchedUrlContext {
266 id,
267 url: url.into(),
268 text: text.into(),
269 }));
270 }
271
272 pub fn remove_context(&mut self, id: ContextId) {
273 let Some(ix) = self.context.iter().position(|context| context.id() == id) else {
274 return;
275 };
276
277 match self.context.remove(ix) {
278 Context::File(_) => {
279 self.files.retain(|_, context_id| *context_id != id);
280 }
281 Context::Directory(_) => {
282 self.directories.retain(|_, context_id| *context_id != id);
283 }
284 Context::FetchedUrl(_) => {
285 self.fetched_urls.retain(|_, context_id| *context_id != id);
286 }
287 Context::Thread(_) => {
288 self.threads.retain(|_, context_id| *context_id != id);
289 }
290 }
291 }
292
293 /// Returns whether the buffer is already included directly in the context, or if it will be
294 /// included in the context via a directory. Directory inclusion is based on paths rather than
295 /// buffer IDs as the directory will be re-scanned.
296 pub fn will_include_buffer(&self, buffer_id: BufferId, path: &Path) -> Option<FileInclusion> {
297 if let Some(context_id) = self.files.get(&buffer_id) {
298 return Some(FileInclusion::Direct(*context_id));
299 }
300
301 self.will_include_file_path_via_directory(path)
302 }
303
304 /// Returns whether this file path is already included directly in the context, or if it will be
305 /// included in the context via a directory.
306 pub fn will_include_file_path(&self, path: &Path, cx: &AppContext) -> Option<FileInclusion> {
307 if !self.files.is_empty() {
308 let found_file_context = self.context.iter().find(|context| match &context {
309 Context::File(file_context) => {
310 if let Some(file_path) = file_context.path(cx) {
311 *file_path == *path
312 } else {
313 false
314 }
315 }
316 _ => false,
317 });
318 if let Some(context) = found_file_context {
319 return Some(FileInclusion::Direct(context.id()));
320 }
321 }
322
323 self.will_include_file_path_via_directory(path)
324 }
325
326 fn will_include_file_path_via_directory(&self, path: &Path) -> Option<FileInclusion> {
327 if self.directories.is_empty() {
328 return None;
329 }
330
331 let mut buf = path.to_path_buf();
332
333 while buf.pop() {
334 if let Some(_) = self.directories.get(&buf) {
335 return Some(FileInclusion::InDirectory(buf));
336 }
337 }
338
339 None
340 }
341
342 pub fn includes_directory(&self, path: &Path) -> Option<ContextId> {
343 self.directories.get(path).copied()
344 }
345
346 pub fn includes_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
347 self.threads.get(thread_id).copied()
348 }
349
350 pub fn includes_url(&self, url: &str) -> Option<ContextId> {
351 self.fetched_urls.get(url).copied()
352 }
353}
354
355pub enum FileInclusion {
356 Direct(ContextId),
357 InDirectory(PathBuf),
358}
359
360pub(crate) fn push_fenced_codeblock(path: &Path, content: String, buffer: &mut String) {
361 buffer.reserve(content.len() + 64);
362
363 write!(buffer, "```").unwrap();
364
365 if let Some(extension) = path.extension().and_then(|ext| ext.to_str()) {
366 write!(buffer, "{} ", extension).unwrap();
367 }
368
369 write!(buffer, "{}", path.display()).unwrap();
370
371 buffer.push('\n');
372 buffer.push_str(&content);
373
374 if !buffer.ends_with('\n') {
375 buffer.push('\n');
376 }
377
378 buffer.push_str("```\n");
379}
380
381fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
382 let mut files = Vec::new();
383
384 for entry in worktree.child_entries(path) {
385 if entry.is_dir() {
386 files.extend(collect_files_in_path(worktree, &entry.path));
387 } else if entry.is_file() {
388 files.push(entry.path.clone());
389 }
390 }
391
392 files
393}