rab/agent/
footer_data_provider.rs1use std::collections::BTreeMap;
2use std::fs;
3use std::path::{Path, PathBuf};
4
5pub struct FooterDataProvider {
19 cwd: PathBuf,
20 git_branch: Option<String>,
21 extension_statuses: BTreeMap<String, String>,
22 available_provider_count: usize,
23}
24
25impl FooterDataProvider {
26 pub fn new(cwd: PathBuf) -> Self {
27 let mut provider = Self {
28 cwd,
29 git_branch: None,
30 extension_statuses: BTreeMap::new(),
31 available_provider_count: 1,
32 };
33 provider.refresh_git_branch();
34 provider
35 }
36
37 pub fn get_git_branch(&self) -> Option<&str> {
40 self.git_branch.as_deref()
41 }
42
43 pub fn refresh_git_branch(&mut self) {
45 self.git_branch = resolve_git_branch(&self.cwd);
46 }
47
48 pub fn set_cwd(&mut self, cwd: PathBuf) {
49 self.cwd = cwd;
50 self.refresh_git_branch();
51 }
52
53 pub fn get_extension_statuses(&self) -> &BTreeMap<String, String> {
56 &self.extension_statuses
57 }
58
59 pub fn set_extension_status(&mut self, key: &str, text: Option<&str>) {
60 if let Some(text) = text {
61 self.extension_statuses
62 .insert(key.to_string(), text.to_string());
63 } else {
64 self.extension_statuses.remove(key);
65 }
66 }
67
68 pub fn clear_extension_statuses(&mut self) {
69 self.extension_statuses.clear();
70 }
71
72 pub fn get_available_provider_count(&self) -> usize {
75 self.available_provider_count
76 }
77
78 pub fn set_available_provider_count(&mut self, count: usize) {
79 self.available_provider_count = count;
80 }
81
82 #[cfg(test)]
84 pub fn set_test_git_branch(&mut self, branch: Option<&str>) {
85 self.git_branch = branch.map(|s| s.to_string());
86 }
87}
88
89#[cfg(test)]
92mod tests {
93 use super::*;
94
95 #[test]
96 fn test_new_provider_refreshes_git_branch() {
97 let provider = FooterDataProvider::new(PathBuf::from("/tmp"));
98 assert!(provider.get_git_branch().is_none());
100 }
101
102 #[test]
103 fn test_set_test_git_branch() {
104 let mut provider = FooterDataProvider::new(PathBuf::from("/tmp"));
105 provider.set_test_git_branch(Some("main"));
106 assert_eq!(provider.get_git_branch(), Some("main"));
107 }
108
109 #[test]
110 fn test_set_test_git_branch_none() {
111 let mut provider = FooterDataProvider::new(PathBuf::from("/tmp"));
112 provider.set_test_git_branch(Some("feature"));
113 provider.set_test_git_branch(None);
114 assert!(provider.get_git_branch().is_none());
115 }
116
117 #[test]
118 fn test_extension_statuses() {
119 let mut provider = FooterDataProvider::new(PathBuf::from("/tmp"));
120 assert!(provider.get_extension_statuses().is_empty());
121
122 provider.set_extension_status("bash", Some("ready"));
123 assert_eq!(
124 provider.get_extension_statuses().get("bash"),
125 Some(&"ready".to_string())
126 );
127
128 provider.set_extension_status("bash", None);
129 assert!(provider.get_extension_statuses().is_empty());
130 }
131
132 #[test]
133 fn test_extension_statuses_sorted() {
134 let mut provider = FooterDataProvider::new(PathBuf::from("/tmp"));
135 provider.set_extension_status("zzz", Some("last"));
136 provider.set_extension_status("aaa", Some("first"));
137 provider.set_extension_status("mmm", Some("middle"));
138
139 let keys: Vec<&String> = provider.get_extension_statuses().keys().collect();
140 assert_eq!(keys, vec!["aaa", "mmm", "zzz"]);
141 }
142
143 #[test]
144 fn test_clear_extension_statuses() {
145 let mut provider = FooterDataProvider::new(PathBuf::from("/tmp"));
146 provider.set_extension_status("bash", Some("ready"));
147 provider.clear_extension_statuses();
148 assert!(provider.get_extension_statuses().is_empty());
149 }
150
151 #[test]
152 fn test_provider_count() {
153 let mut provider = FooterDataProvider::new(PathBuf::from("/tmp"));
154 assert_eq!(provider.get_available_provider_count(), 1);
155 provider.set_available_provider_count(3);
156 assert_eq!(provider.get_available_provider_count(), 3);
157 }
158
159 #[test]
160 fn test_set_cwd_refreshes_git_branch() {
161 let mut provider = FooterDataProvider::new(PathBuf::from("/tmp"));
162 provider.set_test_git_branch(Some("old-branch"));
163 provider.set_cwd(PathBuf::from("/nonexistent"));
165 assert!(provider.get_git_branch().is_none());
166 }
167
168 #[test]
171 fn test_find_git_paths_no_git() {
172 let tmp = std::env::temp_dir().join(format!("rab-test-{}", uuid::Uuid::new_v4()));
173 std::fs::create_dir_all(&tmp).unwrap();
174 let result = find_git_paths(&tmp);
175 assert!(result.is_none());
176 let _ = std::fs::remove_dir_all(&tmp);
177 }
178
179 #[test]
180 fn test_find_git_paths_regular_repo() {
181 let tmp = std::env::temp_dir().join(format!("rab-test-{}", uuid::Uuid::new_v4()));
182 std::fs::create_dir_all(&tmp.join(".git")).unwrap();
183 std::fs::write(&tmp.join(".git").join("HEAD"), "ref: refs/heads/main\n").unwrap();
184
185 let result = find_git_paths(&tmp);
186 assert!(result.is_some());
187 let paths = result.unwrap();
188 assert_eq!(paths.head_path, tmp.join(".git").join("HEAD"));
189
190 let _ = std::fs::remove_dir_all(&tmp);
191 }
192
193 #[test]
194 fn test_find_git_paths_walk_up() {
195 let tmp = std::env::temp_dir().join(format!("rab-test-{}", uuid::Uuid::new_v4()));
196 std::fs::create_dir_all(&tmp.join("sub").join("deep")).unwrap();
197 std::fs::create_dir_all(&tmp.join(".git")).unwrap();
198 std::fs::write(&tmp.join(".git").join("HEAD"), "ref: refs/heads/main\n").unwrap();
199
200 let result = find_git_paths(&tmp.join("sub").join("deep"));
202 assert!(result.is_some());
203
204 let _ = std::fs::remove_dir_all(&tmp);
205 }
206
207 #[test]
208 fn test_resolve_git_branch_from_head() {
209 let tmp = std::env::temp_dir().join(format!("rab-test-{}", uuid::Uuid::new_v4()));
210 std::fs::create_dir_all(&tmp.join(".git")).unwrap();
211 std::fs::write(
212 &tmp.join(".git").join("HEAD"),
213 "ref: refs/heads/feature-branch\n",
214 )
215 .unwrap();
216
217 let result = resolve_git_branch(&tmp);
218 assert_eq!(result.as_deref(), Some("feature-branch"));
219
220 let _ = std::fs::remove_dir_all(&tmp);
221 }
222
223 #[test]
224 fn test_resolve_git_branch_detached() {
225 let tmp = std::env::temp_dir().join(format!("rab-test-{}", uuid::Uuid::new_v4()));
226 std::fs::create_dir_all(&tmp.join(".git")).unwrap();
227 std::fs::write(&tmp.join(".git").join("HEAD"), "abc123def456\n").unwrap();
228
229 let result = resolve_git_branch(&tmp);
230 assert_eq!(result.as_deref(), Some("detached"));
231
232 let _ = std::fs::remove_dir_all(&tmp);
233 }
234
235 #[test]
236 fn test_resolve_git_branch_no_git() {
237 let tmp = std::env::temp_dir().join(format!("rab-test-{}", uuid::Uuid::new_v4()));
238 std::fs::create_dir_all(&tmp).unwrap();
239
240 let result = resolve_git_branch(&tmp);
241 assert!(result.is_none());
242
243 let _ = std::fs::remove_dir_all(&tmp);
244 }
245}
246
247struct GitPaths {
248 _repo_dir: PathBuf,
249 head_path: PathBuf,
250}
251
252fn find_git_paths(cwd: &Path) -> Option<GitPaths> {
254 let mut dir = Some(cwd.to_path_buf());
255 while let Some(ref d) = dir {
256 let git_path = d.join(".git");
257 if git_path.exists() {
258 if git_path.is_file() {
259 let content = fs::read_to_string(&git_path).ok()?;
261 let content = content.trim();
262 if let Some(git_dir_str) = content.strip_prefix("gitdir: ") {
263 let git_dir = d.join(git_dir_str);
264 let head_path = git_dir.join("HEAD");
265 if head_path.exists() {
266 return Some(GitPaths {
267 _repo_dir: d.clone(),
268 head_path,
269 });
270 }
271 }
272 } else if git_path.is_dir() {
273 let head_path = git_path.join("HEAD");
275 if head_path.exists() {
276 return Some(GitPaths {
277 _repo_dir: d.clone(),
278 head_path,
279 });
280 }
281 }
282 }
283 dir = d.parent().map(|p| p.to_path_buf());
284 }
285 None
286}
287
288fn resolve_git_branch(cwd: &Path) -> Option<String> {
290 let paths = find_git_paths(cwd)?;
291 let content = fs::read_to_string(&paths.head_path).ok()?;
292 let content = content.trim();
293
294 if let Some(branch) = content.strip_prefix("ref: refs/heads/") {
295 if branch == ".invalid" {
296 resolve_branch_with_git(&paths._repo_dir)
298 } else {
299 Some(branch.to_string())
300 }
301 } else {
302 Some("detached".to_string())
304 }
305}
306
307fn resolve_branch_with_git(repo_dir: &Path) -> Option<String> {
309 let output = std::process::Command::new("git")
310 .args([
311 "--no-optional-locks",
312 "symbolic-ref",
313 "--quiet",
314 "--short",
315 "HEAD",
316 ])
317 .current_dir(repo_dir)
318 .output()
319 .ok()?;
320 if output.status.success() {
321 let branch = String::from_utf8_lossy(&output.stdout).trim().to_string();
322 if !branch.is_empty() {
323 return Some(branch);
324 }
325 }
326 Some("detached".to_string())
327}