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