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