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