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::{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 insert_thread(&mut self, thread: &Thread) {
233 let context_id = self.next_context_id.post_inc();
234 self.threads.insert(thread.id().clone(), context_id);
235
236 self.context.push(Context {
237 id: context_id,
238 name: thread.summary().unwrap_or("New thread".into()),
239 parent: None,
240 tooltip: None,
241 kind: ContextKind::Thread,
242 text: thread.text().into(),
243 });
244 }
245
246 pub fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
247 let context_id = self.next_context_id.post_inc();
248 self.fetched_urls.insert(url.clone(), context_id);
249
250 self.context.push(Context {
251 id: context_id,
252 name: url.into(),
253 parent: None,
254 tooltip: None,
255 kind: ContextKind::FetchedUrl,
256 text: text.into(),
257 });
258 }
259
260 pub fn remove_context(&mut self, id: &ContextId) {
261 let Some(ix) = self.context.iter().position(|context| context.id == *id) else {
262 return;
263 };
264
265 match self.context.remove(ix).kind {
266 ContextKind::File => {
267 self.files.retain(|_, context_id| context_id != id);
268 }
269 ContextKind::Directory => {
270 self.directories.retain(|_, context_id| context_id != id);
271 }
272 ContextKind::FetchedUrl => {
273 self.fetched_urls.retain(|_, context_id| context_id != id);
274 }
275 ContextKind::Thread => {
276 self.threads.retain(|_, context_id| context_id != id);
277 }
278 }
279 }
280
281 pub fn included_file(&self, path: &Path) -> Option<IncludedFile> {
282 if let Some(id) = self.files.get(path) {
283 return Some(IncludedFile::Direct(*id));
284 }
285
286 if self.directories.is_empty() {
287 return None;
288 }
289
290 let mut buf = path.to_path_buf();
291
292 while buf.pop() {
293 if let Some(_) = self.directories.get(&buf) {
294 return Some(IncludedFile::InDirectory(buf));
295 }
296 }
297
298 None
299 }
300
301 pub fn included_directory(&self, path: &Path) -> Option<ContextId> {
302 self.directories.get(path).copied()
303 }
304
305 pub fn included_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
306 self.threads.get(thread_id).copied()
307 }
308
309 pub fn included_url(&self, url: &str) -> Option<ContextId> {
310 self.fetched_urls.get(url).copied()
311 }
312
313 pub fn duplicated_names(&self) -> HashSet<SharedString> {
314 let mut seen = HashSet::default();
315 let mut dupes = HashSet::default();
316
317 for context in self.context().iter() {
318 if !seen.insert(&context.name) {
319 dupes.insert(context.name.clone());
320 }
321 }
322
323 dupes
324 }
325}
326
327pub enum IncludedFile {
328 Direct(ContextId),
329 InDirectory(PathBuf),
330}
331
332pub(crate) fn push_fenced_codeblock(path: &Path, content: String, buffer: &mut String) {
333 buffer.reserve(content.len() + 64);
334
335 write!(buffer, "```").unwrap();
336
337 if let Some(extension) = path.extension().and_then(|ext| ext.to_str()) {
338 write!(buffer, "{} ", extension).unwrap();
339 }
340
341 write!(buffer, "{}", path.display()).unwrap();
342
343 buffer.push('\n');
344 buffer.push_str(&content);
345
346 if !buffer.ends_with('\n') {
347 buffer.push('\n');
348 }
349
350 buffer.push_str("```\n");
351}
352
353fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
354 let mut files = Vec::new();
355
356 for entry in worktree.child_entries(path) {
357 if entry.is_dir() {
358 files.extend(collect_files_in_path(worktree, &entry.path));
359 } else if entry.is_file() {
360 files.push(entry.path.clone());
361 }
362 }
363
364 files
365}