1use std::fmt::Write as _;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use anyhow::{anyhow, bail, Result};
6use collections::{HashMap, HashSet};
7use gpui::{Model, ModelContext, SharedString, Task, WeakView};
8use language::Buffer;
9use project::{ProjectPath, Worktree};
10use workspace::Workspace;
11
12use crate::thread::Thread;
13use crate::{
14 context::{Context, ContextId, ContextKind},
15 thread::ThreadId,
16};
17
18pub struct ContextStore {
19 workspace: WeakView<Workspace>,
20 context: Vec<Context>,
21 next_context_id: ContextId,
22 files: HashMap<PathBuf, ContextId>,
23 directories: HashMap<PathBuf, ContextId>,
24 threads: HashMap<ThreadId, ContextId>,
25 fetched_urls: HashMap<String, ContextId>,
26}
27
28impl ContextStore {
29 pub fn new(workspace: WeakView<Workspace>) -> Self {
30 Self {
31 workspace,
32 context: Vec::new(),
33 next_context_id: ContextId(0),
34 files: HashMap::default(),
35 directories: HashMap::default(),
36 threads: HashMap::default(),
37 fetched_urls: HashMap::default(),
38 }
39 }
40
41 pub fn context(&self) -> &Vec<Context> {
42 &self.context
43 }
44
45 pub fn clear(&mut self) {
46 self.context.clear();
47 self.files.clear();
48 self.directories.clear();
49 self.threads.clear();
50 self.fetched_urls.clear();
51 }
52
53 pub fn add_file(
54 &mut self,
55 project_path: ProjectPath,
56 cx: &mut ModelContext<Self>,
57 ) -> Task<Result<()>> {
58 let workspace = self.workspace.clone();
59 let Some(project) = workspace
60 .upgrade()
61 .map(|workspace| workspace.read(cx).project().clone())
62 else {
63 return Task::ready(Err(anyhow!("failed to read project")));
64 };
65
66 let already_included = match self.included_file(&project_path.path) {
67 Some(IncludedFile::Direct(context_id)) => {
68 self.remove_context(&context_id);
69 true
70 }
71 Some(IncludedFile::InDirectory(_)) => true,
72 None => false,
73 };
74 if already_included {
75 return Task::ready(Ok(()));
76 }
77
78 cx.spawn(|this, mut cx| async move {
79 let open_buffer_task =
80 project.update(&mut cx, |project, cx| project.open_buffer(project_path, cx))?;
81
82 let buffer = open_buffer_task.await?;
83 this.update(&mut cx, |this, cx| {
84 this.insert_file(buffer.read(cx));
85 })?;
86
87 anyhow::Ok(())
88 })
89 }
90
91 pub fn insert_file(&mut self, buffer: &Buffer) {
92 let Some(file) = buffer.file() else {
93 return;
94 };
95
96 let path = file.path();
97
98 let id = self.next_context_id.post_inc();
99 self.files.insert(path.to_path_buf(), id);
100
101 let full_path: SharedString = path.to_string_lossy().into_owned().into();
102
103 let name = match path.file_name() {
104 Some(name) => name.to_string_lossy().into_owned().into(),
105 None => full_path.clone(),
106 };
107
108 let parent = path
109 .parent()
110 .and_then(|p| p.file_name())
111 .map(|p| p.to_string_lossy().into_owned().into());
112
113 let mut text = String::new();
114 push_fenced_codeblock(path, buffer.text(), &mut text);
115
116 self.context.push(Context {
117 id,
118 name,
119 parent,
120 tooltip: Some(full_path),
121 kind: ContextKind::File,
122 text: text.into(),
123 });
124 }
125
126 pub fn add_directory(
127 &mut self,
128 project_path: ProjectPath,
129 cx: &mut ModelContext<Self>,
130 ) -> Task<Result<()>> {
131 let workspace = self.workspace.clone();
132 let Some(project) = workspace
133 .upgrade()
134 .map(|workspace| workspace.read(cx).project().clone())
135 else {
136 return Task::ready(Err(anyhow!("failed to read project")));
137 };
138
139 let already_included = if let Some(context_id) = self.included_directory(&project_path.path)
140 {
141 self.remove_context(&context_id);
142 true
143 } else {
144 false
145 };
146 if already_included {
147 return Task::ready(Ok(()));
148 }
149
150 let worktree_id = project_path.worktree_id;
151 cx.spawn(|this, mut cx| async move {
152 let worktree = project.update(&mut cx, |project, cx| {
153 project
154 .worktree_for_id(worktree_id, cx)
155 .ok_or_else(|| anyhow!("no worktree found for {worktree_id:?}"))
156 })??;
157
158 let files = worktree.update(&mut cx, |worktree, _cx| {
159 collect_files_in_path(worktree, &project_path.path)
160 })?;
161
162 let open_buffer_tasks = project.update(&mut cx, |project, cx| {
163 files
164 .into_iter()
165 .map(|file_path| {
166 project.open_buffer(
167 ProjectPath {
168 worktree_id,
169 path: file_path.clone(),
170 },
171 cx,
172 )
173 })
174 .collect::<Vec<_>>()
175 })?;
176
177 let buffers = futures::future::join_all(open_buffer_tasks).await;
178
179 this.update(&mut cx, |this, cx| {
180 let mut text = String::new();
181 let mut added_files = 0;
182
183 for buffer in buffers.into_iter().flatten() {
184 let buffer = buffer.read(cx);
185 let path = buffer.file().map_or(&project_path.path, |file| file.path());
186 push_fenced_codeblock(&path, buffer.text(), &mut text);
187 added_files += 1;
188 }
189
190 if added_files == 0 {
191 bail!(
192 "could not read any text files from {}",
193 &project_path.path.display()
194 );
195 }
196
197 this.insert_directory(&project_path.path, text);
198
199 anyhow::Ok(())
200 })??;
201
202 anyhow::Ok(())
203 })
204 }
205
206 pub fn insert_directory(&mut self, path: &Path, text: impl Into<SharedString>) {
207 let id = self.next_context_id.post_inc();
208 self.directories.insert(path.to_path_buf(), id);
209
210 let full_path: SharedString = path.to_string_lossy().into_owned().into();
211
212 let name = match path.file_name() {
213 Some(name) => name.to_string_lossy().into_owned().into(),
214 None => full_path.clone(),
215 };
216
217 let parent = path
218 .parent()
219 .and_then(|p| p.file_name())
220 .map(|p| p.to_string_lossy().into_owned().into());
221
222 self.context.push(Context {
223 id,
224 name,
225 parent,
226 tooltip: Some(full_path),
227 kind: ContextKind::Directory,
228 text: text.into(),
229 });
230 }
231
232 pub fn add_thread(&mut self, thread: Model<Thread>, cx: &mut ModelContext<Self>) {
233 if let Some(context_id) = self.included_thread(&thread.read(cx).id()) {
234 self.remove_context(&context_id);
235 } else {
236 self.insert_thread(thread.read(cx));
237 }
238 }
239
240 pub fn insert_thread(&mut self, thread: &Thread) {
241 let context_id = self.next_context_id.post_inc();
242 self.threads.insert(thread.id().clone(), context_id);
243
244 self.context.push(Context {
245 id: context_id,
246 name: thread.summary().unwrap_or("New thread".into()),
247 parent: None,
248 tooltip: None,
249 kind: ContextKind::Thread,
250 text: thread.text().into(),
251 });
252 }
253
254 pub fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
255 let context_id = self.next_context_id.post_inc();
256 self.fetched_urls.insert(url.clone(), context_id);
257
258 self.context.push(Context {
259 id: context_id,
260 name: url.into(),
261 parent: None,
262 tooltip: None,
263 kind: ContextKind::FetchedUrl,
264 text: text.into(),
265 });
266 }
267
268 pub fn remove_context(&mut self, id: &ContextId) {
269 let Some(ix) = self.context.iter().position(|context| context.id == *id) else {
270 return;
271 };
272
273 match self.context.remove(ix).kind {
274 ContextKind::File => {
275 self.files.retain(|_, context_id| context_id != id);
276 }
277 ContextKind::Directory => {
278 self.directories.retain(|_, context_id| context_id != id);
279 }
280 ContextKind::FetchedUrl => {
281 self.fetched_urls.retain(|_, context_id| context_id != id);
282 }
283 ContextKind::Thread => {
284 self.threads.retain(|_, context_id| context_id != id);
285 }
286 }
287 }
288
289 pub fn included_file(&self, path: &Path) -> Option<IncludedFile> {
290 if let Some(id) = self.files.get(path) {
291 return Some(IncludedFile::Direct(*id));
292 }
293
294 if self.directories.is_empty() {
295 return None;
296 }
297
298 let mut buf = path.to_path_buf();
299
300 while buf.pop() {
301 if let Some(_) = self.directories.get(&buf) {
302 return Some(IncludedFile::InDirectory(buf));
303 }
304 }
305
306 None
307 }
308
309 pub fn included_directory(&self, path: &Path) -> Option<ContextId> {
310 self.directories.get(path).copied()
311 }
312
313 pub fn included_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
314 self.threads.get(thread_id).copied()
315 }
316
317 pub fn included_url(&self, url: &str) -> Option<ContextId> {
318 self.fetched_urls.get(url).copied()
319 }
320
321 pub fn duplicated_names(&self) -> HashSet<SharedString> {
322 let mut seen = HashSet::default();
323 let mut dupes = HashSet::default();
324
325 for context in self.context().iter() {
326 if !seen.insert(&context.name) {
327 dupes.insert(context.name.clone());
328 }
329 }
330
331 dupes
332 }
333}
334
335pub enum IncludedFile {
336 Direct(ContextId),
337 InDirectory(PathBuf),
338}
339
340pub(crate) fn push_fenced_codeblock(path: &Path, content: String, buffer: &mut String) {
341 buffer.reserve(content.len() + 64);
342
343 write!(buffer, "```").unwrap();
344
345 if let Some(extension) = path.extension().and_then(|ext| ext.to_str()) {
346 write!(buffer, "{} ", extension).unwrap();
347 }
348
349 write!(buffer, "{}", path.display()).unwrap();
350
351 buffer.push('\n');
352 buffer.push_str(&content);
353
354 if !buffer.ends_with('\n') {
355 buffer.push('\n');
356 }
357
358 buffer.push_str("```\n");
359}
360
361fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
362 let mut files = Vec::new();
363
364 for entry in worktree.child_entries(path) {
365 if entry.is_dir() {
366 files.extend(collect_files_in_path(worktree, &entry.path));
367 } else if entry.is_file() {
368 files.push(entry.path.clone());
369 }
370 }
371
372 files
373}