1use std::collections::{BTreeMap, BTreeSet};
2use std::fs::OpenOptions;
3use std::io::ErrorKind;
4use std::path::{Path, PathBuf};
5use std::thread;
6use std::time::{Duration, Instant, SystemTime};
7
8use anyhow::{Context, Result, bail};
9use chrono::Utc;
10use serde::{Deserialize, Serialize};
11
12use super::config::RoleType;
13use super::hierarchy::resolve_hierarchy;
14use super::{team_config_dir, team_config_path};
15
16const BENCH_FILE_NAME: &str = "bench.yaml";
17const BENCH_LOCK_STALE_SECS: u64 = 30;
18const BENCH_LOCK_TIMEOUT_MS: u64 = 2_000;
19const BENCH_LOCK_RETRY_MS: u64 = 10;
20
21#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
22pub struct BenchState {
23 #[serde(default)]
24 pub benched: BTreeMap<String, BenchEntry>,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28pub struct BenchEntry {
29 pub timestamp: String,
30 #[serde(default, skip_serializing_if = "Option::is_none")]
31 pub reason: Option<String>,
32}
33
34pub fn bench_file_path(project_root: &Path) -> PathBuf {
35 team_config_dir(project_root).join(BENCH_FILE_NAME)
36}
37
38pub fn load_bench_state(project_root: &Path) -> Result<BenchState> {
39 load_bench_state_from_path(&bench_file_path(project_root))
40}
41
42pub fn benched_engineer_names(project_root: &Path) -> Result<BTreeSet<String>> {
43 Ok(load_bench_state(project_root)?
44 .benched
45 .keys()
46 .cloned()
47 .collect())
48}
49
50pub fn bench_engineer(
51 project_root: &Path,
52 engineer: &str,
53 reason: Option<&str>,
54) -> Result<BenchEntry> {
55 validate_engineer(project_root, engineer)?;
56 with_bench_lock(project_root, || {
57 let path = bench_file_path(project_root);
58 let mut state = load_bench_state_from_path(&path)?;
59 let entry = BenchEntry {
60 timestamp: Utc::now().to_rfc3339(),
61 reason: normalize_reason(reason),
62 };
63 state.benched.insert(engineer.to_string(), entry.clone());
64 write_bench_state(&path, &state)?;
65 Ok(entry)
66 })
67}
68
69pub fn unbench_engineer(project_root: &Path, engineer: &str) -> Result<bool> {
70 validate_engineer(project_root, engineer)?;
71 with_bench_lock(project_root, || {
72 let path = bench_file_path(project_root);
73 let mut state = load_bench_state_from_path(&path)?;
74 let removed = state.benched.remove(engineer).is_some();
75 write_bench_state(&path, &state)?;
76 Ok(removed)
77 })
78}
79
80pub fn format_benched_engineers_section(state: &BenchState) -> Option<String> {
81 if state.benched.is_empty() {
82 return None;
83 }
84
85 let mut lines = vec![
86 "Benched Engineers".to_string(),
87 format!("{:<20} {:<26} {}", "ENGINEER", "SINCE", "REASON"),
88 ];
89 for (engineer, entry) in &state.benched {
90 lines.push(format!(
91 "{:<20} {:<26} {}",
92 engineer,
93 entry.timestamp,
94 entry.reason.as_deref().unwrap_or("-"),
95 ));
96 }
97
98 Some(lines.join("\n"))
99}
100
101fn normalize_reason(reason: Option<&str>) -> Option<String> {
102 reason
103 .map(str::trim)
104 .filter(|reason| !reason.is_empty())
105 .map(str::to_string)
106}
107
108fn validate_engineer(project_root: &Path, engineer: &str) -> Result<()> {
109 let config = super::config::TeamConfig::load(&team_config_path(project_root))?;
110 let members = resolve_hierarchy(&config)?;
111 if members
112 .iter()
113 .any(|member| member.name == engineer && member.role_type == RoleType::Engineer)
114 {
115 Ok(())
116 } else {
117 bail!("unknown engineer '{engineer}'");
118 }
119}
120
121fn load_bench_state_from_path(path: &Path) -> Result<BenchState> {
122 if !path.exists() {
123 return Ok(BenchState::default());
124 }
125 let content = std::fs::read_to_string(path)
126 .with_context(|| format!("failed to read {}", path.display()))?;
127 if content.trim().is_empty() {
128 return Ok(BenchState::default());
129 }
130 serde_yaml::from_str(&content).with_context(|| format!("failed to parse {}", path.display()))
131}
132
133fn write_bench_state(path: &Path, state: &BenchState) -> Result<()> {
134 if let Some(parent) = path.parent() {
135 std::fs::create_dir_all(parent)
136 .with_context(|| format!("failed to create {}", parent.display()))?;
137 }
138 let yaml = serde_yaml::to_string(state).context("failed to serialize bench state")?;
139 let temp_path = path.with_extension(format!(
140 "yaml.tmp-{}-{}",
141 std::process::id(),
142 Utc::now().timestamp_nanos_opt().unwrap_or_default()
143 ));
144 std::fs::write(&temp_path, yaml)
145 .with_context(|| format!("failed to write {}", temp_path.display()))?;
146 std::fs::rename(&temp_path, path)
147 .with_context(|| format!("failed to replace {}", path.display()))?;
148 Ok(())
149}
150
151fn with_bench_lock<T>(project_root: &Path, operation: impl FnOnce() -> Result<T>) -> Result<T> {
152 let _guard = BenchLockGuard::acquire(project_root)?;
153 operation()
154}
155
156struct BenchLockGuard {
157 path: PathBuf,
158}
159
160impl BenchLockGuard {
161 fn acquire(project_root: &Path) -> Result<Self> {
162 let path = bench_file_path(project_root).with_extension("yaml.lock");
163 if let Some(parent) = path.parent() {
164 std::fs::create_dir_all(parent)
165 .with_context(|| format!("failed to create {}", parent.display()))?;
166 }
167 let started = Instant::now();
168 loop {
169 match OpenOptions::new().create_new(true).write(true).open(&path) {
170 Ok(_) => return Ok(Self { path }),
171 Err(error) if error.kind() == ErrorKind::AlreadyExists => {
172 if lock_is_stale(&path) {
173 let _ = std::fs::remove_file(&path);
174 continue;
175 }
176 if started.elapsed() >= Duration::from_millis(BENCH_LOCK_TIMEOUT_MS) {
177 bail!("timed out waiting for bench state lock");
178 }
179 thread::sleep(Duration::from_millis(BENCH_LOCK_RETRY_MS));
180 }
181 Err(error) => {
182 return Err(error)
183 .with_context(|| format!("failed to acquire {}", path.display()));
184 }
185 }
186 }
187 }
188}
189
190impl Drop for BenchLockGuard {
191 fn drop(&mut self) {
192 let _ = std::fs::remove_file(&self.path);
193 }
194}
195
196fn lock_is_stale(path: &Path) -> bool {
197 let Ok(metadata) = std::fs::metadata(path) else {
198 return false;
199 };
200 let Ok(modified) = metadata.modified() else {
201 return false;
202 };
203 let Ok(age) = SystemTime::now().duration_since(modified) else {
204 return false;
205 };
206 age.as_secs() >= BENCH_LOCK_STALE_SECS
207}
208
209#[cfg(test)]
210mod tests {
211 use std::sync::{Arc, Barrier};
212
213 use super::*;
214
215 fn write_team_config(root: &Path) {
216 let team_dir = root.join(".batty").join("team_config");
217 std::fs::create_dir_all(&team_dir).unwrap();
218 std::fs::write(
219 team_dir.join("team.yaml"),
220 r#"
221name: test
222agent: codex
223roles:
224 - name: architect
225 role_type: architect
226 - name: manager
227 role_type: manager
228 - name: engineer
229 role_type: engineer
230 instances: 2
231"#,
232 )
233 .unwrap();
234 }
235
236 #[test]
237 fn bench_engineer_persists_reason_and_timestamp() {
238 let tmp = tempfile::tempdir().unwrap();
239 write_team_config(tmp.path());
240
241 let entry = bench_engineer(tmp.path(), "eng-1-1", Some("session end")).unwrap();
242 let state = load_bench_state(tmp.path()).unwrap();
243
244 assert_eq!(state.benched.get("eng-1-1"), Some(&entry));
245 assert_eq!(entry.reason.as_deref(), Some("session end"));
246 assert!(!entry.timestamp.is_empty());
247 }
248
249 #[test]
250 fn unbench_engineer_removes_entry() {
251 let tmp = tempfile::tempdir().unwrap();
252 write_team_config(tmp.path());
253
254 bench_engineer(tmp.path(), "eng-1-1", Some("pause")).unwrap();
255 assert!(unbench_engineer(tmp.path(), "eng-1-1").unwrap());
256 assert!(
257 !load_bench_state(tmp.path())
258 .unwrap()
259 .benched
260 .contains_key("eng-1-1")
261 );
262 }
263
264 #[test]
265 fn bench_engineer_rejects_unknown_engineer() {
266 let tmp = tempfile::tempdir().unwrap();
267 write_team_config(tmp.path());
268
269 let error = bench_engineer(tmp.path(), "eng-9", Some("pause")).unwrap_err();
270 assert!(error.to_string().contains("unknown engineer 'eng-9'"));
271 }
272
273 #[test]
274 fn format_section_includes_timestamp_and_reason() {
275 let mut state = BenchState::default();
276 state.benched.insert(
277 "eng-1-1".to_string(),
278 BenchEntry {
279 timestamp: "2026-04-10T10:00:00Z".to_string(),
280 reason: Some("session end".to_string()),
281 },
282 );
283
284 let formatted = format_benched_engineers_section(&state).unwrap();
285 assert!(formatted.contains("Benched Engineers"));
286 assert!(formatted.contains("eng-1-1"));
287 assert!(formatted.contains("2026-04-10T10:00:00Z"));
288 assert!(formatted.contains("session end"));
289 }
290
291 #[test]
292 fn concurrent_bench_and_unbench_preserve_both_updates() {
293 let tmp = tempfile::tempdir().unwrap();
294 write_team_config(tmp.path());
295 bench_engineer(tmp.path(), "eng-1-1", Some("existing")).unwrap();
296
297 let barrier = Arc::new(Barrier::new(3));
298 let root_a = tmp.path().to_path_buf();
299 let barrier_a = Arc::clone(&barrier);
300 let add = std::thread::spawn(move || {
301 barrier_a.wait();
302 for _ in 0..50 {
303 bench_engineer(&root_a, "eng-1-2", Some("new")).unwrap();
304 }
305 });
306 let root_b = tmp.path().to_path_buf();
307 let barrier_b = Arc::clone(&barrier);
308 let remove = std::thread::spawn(move || {
309 barrier_b.wait();
310 for _ in 0..50 {
311 unbench_engineer(&root_b, "eng-1-1").unwrap();
312 }
313 });
314
315 barrier.wait();
316 add.join().unwrap();
317 remove.join().unwrap();
318 let state = load_bench_state(tmp.path()).unwrap();
319 assert!(!state.benched.contains_key("eng-1-1"));
320 assert_eq!(
321 state
322 .benched
323 .get("eng-1-2")
324 .and_then(|entry| entry.reason.as_deref()),
325 Some("new")
326 );
327 }
328}