1use std::fmt::Write as _;
2use std::path::{Path, PathBuf};
3
4use collections::HashMap;
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 name = path.to_string_lossy().into_owned().into();
64
65 let mut text = String::new();
66 push_fenced_codeblock(path, buffer.text(), &mut text);
67
68 self.context.push(Context {
69 id,
70 name,
71 kind: ContextKind::File,
72 text: text.into(),
73 });
74 }
75
76 pub fn insert_directory(&mut self, path: &Path, text: impl Into<SharedString>) {
77 let id = self.next_context_id.post_inc();
78 self.directories.insert(path.to_path_buf(), id);
79
80 let name = path.to_string_lossy().into_owned().into();
81
82 self.context.push(Context {
83 id,
84 name,
85 kind: ContextKind::Directory,
86 text: text.into(),
87 });
88 }
89
90 pub fn insert_thread(&mut self, thread: &Thread) {
91 let context_id = self.next_context_id.post_inc();
92 self.threads.insert(thread.id().clone(), context_id);
93
94 self.context.push(Context {
95 id: context_id,
96 name: thread.summary().unwrap_or("New thread".into()),
97 kind: ContextKind::Thread,
98 text: thread.text().into(),
99 });
100 }
101
102 pub fn insert_fetched_url(&mut self, url: String, text: impl Into<SharedString>) {
103 let context_id = self.next_context_id.post_inc();
104 self.fetched_urls.insert(url.clone(), context_id);
105
106 self.context.push(Context {
107 id: context_id,
108 name: url.into(),
109 kind: ContextKind::FetchedUrl,
110 text: text.into(),
111 });
112 }
113
114 pub fn remove_context(&mut self, id: &ContextId) {
115 let Some(ix) = self.context.iter().position(|c| c.id == *id) else {
116 return;
117 };
118
119 match self.context.remove(ix).kind {
120 ContextKind::File => {
121 self.files.retain(|_, p_id| p_id != id);
122 }
123 ContextKind::Directory => {
124 self.directories.retain(|_, p_id| p_id != id);
125 }
126 ContextKind::FetchedUrl => {
127 self.fetched_urls.retain(|_, p_id| p_id != id);
128 }
129 ContextKind::Thread => {
130 self.threads.retain(|_, p_id| p_id != id);
131 }
132 }
133 }
134
135 pub fn included_file(&self, path: &Path) -> Option<IncludedFile> {
136 if let Some(id) = self.files.get(path) {
137 return Some(IncludedFile::Direct(*id));
138 }
139
140 if self.directories.is_empty() {
141 return None;
142 }
143
144 let mut buf = path.to_path_buf();
145
146 while buf.pop() {
147 if let Some(_) = self.directories.get(&buf) {
148 return Some(IncludedFile::InDirectory(buf));
149 }
150 }
151
152 None
153 }
154
155 pub fn included_directory(&self, path: &Path) -> Option<ContextId> {
156 self.directories.get(path).copied()
157 }
158
159 pub fn included_thread(&self, thread_id: &ThreadId) -> Option<ContextId> {
160 self.threads.get(thread_id).copied()
161 }
162
163 pub fn included_url(&self, url: &str) -> Option<ContextId> {
164 self.fetched_urls.get(url).copied()
165 }
166}
167
168pub enum IncludedFile {
169 Direct(ContextId),
170 InDirectory(PathBuf),
171}
172
173pub(crate) fn push_fenced_codeblock(path: &Path, content: String, buf: &mut String) {
174 buf.reserve(content.len() + 64);
175
176 write!(buf, "```").unwrap();
177
178 if let Some(extension) = path.extension().and_then(|ext| ext.to_str()) {
179 write!(buf, "{} ", extension).unwrap();
180 }
181
182 write!(buf, "{}", path.display()).unwrap();
183
184 buf.push('\n');
185 buf.push_str(&content);
186
187 if !buf.ends_with('\n') {
188 buf.push('\n');
189 }
190
191 buf.push_str("```\n");
192}