1use std::fmt::Write as _;
2use std::path::{Path, PathBuf};
3
4use anyhow::{anyhow, Result};
5use collections::{HashMap, HashSet};
6use gpui::{ModelContext, SharedString, Task, WeakView};
7use language::Buffer;
8use project::ProjectPath;
9use workspace::Workspace;
10
11use crate::thread::Thread;
12use crate::{
13 context::{Context, ContextId, ContextKind},
14 thread::ThreadId,
15};
16
17pub struct ContextStore {
18 workspace: WeakView<Workspace>,
19 context: Vec<Context>,
20 next_context_id: ContextId,
21 files: HashMap<PathBuf, ContextId>,
22 directories: HashMap<PathBuf, ContextId>,
23 threads: HashMap<ThreadId, ContextId>,
24 fetched_urls: HashMap<String, ContextId>,
25}
26
27impl ContextStore {
28 pub fn new(workspace: WeakView<Workspace>) -> Self {
29 Self {
30 workspace,
31 context: Vec::new(),
32 next_context_id: ContextId(0),
33 files: HashMap::default(),
34 directories: HashMap::default(),
35 threads: HashMap::default(),
36 fetched_urls: HashMap::default(),
37 }
38 }
39
40 pub fn context(&self) -> &Vec<Context> {
41 &self.context
42 }
43
44 pub fn clear(&mut self) {
45 self.context.clear();
46 self.files.clear();
47 self.directories.clear();
48 self.threads.clear();
49 self.fetched_urls.clear();
50 }
51
52 pub fn add_file(
53 &mut self,
54 project_path: ProjectPath,
55 cx: &mut ModelContext<Self>,
56 ) -> Task<Result<()>> {
57 let workspace = self.workspace.clone();
58 let Some(project) = workspace
59 .upgrade()
60 .map(|workspace| workspace.read(cx).project().clone())
61 else {
62 return Task::ready(Err(anyhow!("failed to read project")));
63 };
64
65 let already_included = match self.included_file(&project_path.path) {
66 Some(IncludedFile::Direct(context_id)) => {
67 self.remove_context(&context_id);
68 true
69 }
70 Some(IncludedFile::InDirectory(_)) => true,
71 None => false,
72 };
73 if already_included {
74 return Task::ready(Ok(()));
75 }
76
77 cx.spawn(|this, mut cx| async move {
78 let open_buffer_task =
79 project.update(&mut cx, |project, cx| project.open_buffer(project_path, cx))?;
80
81 let buffer = open_buffer_task.await?;
82 this.update(&mut cx, |this, cx| {
83 this.insert_file(buffer.read(cx));
84 })?;
85
86 anyhow::Ok(())
87 })
88 }
89
90 pub fn insert_file(&mut self, buffer: &Buffer) {
91 let Some(file) = buffer.file() else {
92 return;
93 };
94
95 let path = file.path();
96
97 let id = self.next_context_id.post_inc();
98 self.files.insert(path.to_path_buf(), id);
99
100 let full_path: SharedString = path.to_string_lossy().into_owned().into();
101
102 let name = match path.file_name() {
103 Some(name) => name.to_string_lossy().into_owned().into(),
104 None => full_path.clone(),
105 };
106
107 let parent = path
108 .parent()
109 .and_then(|p| p.file_name())
110 .map(|p| p.to_string_lossy().into_owned().into());
111
112 let mut text = String::new();
113 push_fenced_codeblock(path, buffer.text(), &mut text);
114
115 self.context.push(Context {
116 id,
117 name,
118 parent,
119 tooltip: Some(full_path),
120 kind: ContextKind::File,
121 text: text.into(),
122 });
123 }
124
125 pub fn insert_directory(&mut self, path: &Path, text: impl Into<SharedString>) {
126 let id = self.next_context_id.post_inc();
127 self.directories.insert(path.to_path_buf(), id);
128
129 let full_path: SharedString = path.to_string_lossy().into_owned().into();
130
131 let name = match path.file_name() {
132 Some(name) => name.to_string_lossy().into_owned().into(),
133 None => full_path.clone(),
134 };
135
136 let parent = path
137 .parent()
138 .and_then(|p| p.file_name())
139 .map(|p| p.to_string_lossy().into_owned().into());
140
141 self.context.push(Context {
142 id,
143 name,
144 parent,
145 tooltip: Some(full_path),
146 kind: ContextKind::Directory,
147 text: text.into(),
148 });
149 }
150
151 pub fn insert_thread(&mut self, thread: &Thread) {
152 let context_id = self.next_context_id.post_inc();
153 self.threads.insert(thread.id().clone(), context_id);
154
155 self.context.push(Context {
156 id: context_id,
157 name: thread.summary().unwrap_or("New thread".into()),
158 parent: None,
159 tooltip: None,
160 kind: ContextKind::Thread,
161 text: thread.text().into(),
162 });
163 }
164
165 pub fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
166 let context_id = self.next_context_id.post_inc();
167 self.fetched_urls.insert(url.clone(), context_id);
168
169 self.context.push(Context {
170 id: context_id,
171 name: url.into(),
172 parent: None,
173 tooltip: None,
174 kind: ContextKind::FetchedUrl,
175 text: text.into(),
176 });
177 }
178
179 pub fn remove_context(&mut self, id: &ContextId) {
180 let Some(ix) = self.context.iter().position(|context| context.id == *id) else {
181 return;
182 };
183
184 match self.context.remove(ix).kind {
185 ContextKind::File => {
186 self.files.retain(|_, context_id| context_id != id);
187 }
188 ContextKind::Directory => {
189 self.directories.retain(|_, context_id| context_id != id);
190 }
191 ContextKind::FetchedUrl => {
192 self.fetched_urls.retain(|_, context_id| context_id != id);
193 }
194 ContextKind::Thread => {
195 self.threads.retain(|_, context_id| context_id != id);
196 }
197 }
198 }
199
200 pub fn included_file(&self, path: &Path) -> Option<IncludedFile> {
201 if let Some(id) = self.files.get(path) {
202 return Some(IncludedFile::Direct(*id));
203 }
204
205 if self.directories.is_empty() {
206 return None;
207 }
208
209 let mut buf = path.to_path_buf();
210
211 while buf.pop() {
212 if let Some(_) = self.directories.get(&buf) {
213 return Some(IncludedFile::InDirectory(buf));
214 }
215 }
216
217 None
218 }
219
220 pub fn included_directory(&self, path: &Path) -> Option<ContextId> {
221 self.directories.get(path).copied()
222 }
223
224 pub fn included_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
225 self.threads.get(thread_id).copied()
226 }
227
228 pub fn included_url(&self, url: &str) -> Option<ContextId> {
229 self.fetched_urls.get(url).copied()
230 }
231
232 pub fn duplicated_names(&self) -> HashSet<SharedString> {
233 let mut seen = HashSet::default();
234 let mut dupes = HashSet::default();
235
236 for context in self.context().iter() {
237 if !seen.insert(&context.name) {
238 dupes.insert(context.name.clone());
239 }
240 }
241
242 dupes
243 }
244}
245
246pub enum IncludedFile {
247 Direct(ContextId),
248 InDirectory(PathBuf),
249}
250
251pub(crate) fn push_fenced_codeblock(path: &Path, content: String, buffer: &mut String) {
252 buffer.reserve(content.len() + 64);
253
254 write!(buffer, "```").unwrap();
255
256 if let Some(extension) = path.extension().and_then(|ext| ext.to_str()) {
257 write!(buffer, "{} ", extension).unwrap();
258 }
259
260 write!(buffer, "{}", path.display()).unwrap();
261
262 buffer.push('\n');
263 buffer.push_str(&content);
264
265 if !buffer.ends_with('\n') {
266 buffer.push('\n');
267 }
268
269 buffer.push_str("```\n");
270}