pawan/
subagent.rs

1//! Live subagent run tracking for tools and the TUI queue strip.
2
3use std::collections::HashMap;
4use std::sync::{Mutex, OnceLock};
5use std::time::Instant;
6
7/// Lifecycle state exposed to the TUI.
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum SubagentState {
10    Running,
11    Done,
12    Failed,
13}
14
15/// One tracked subagent invocation.
16#[derive(Debug, Clone)]
17pub struct SubagentRun {
18    pub id: String,
19    pub label: String,
20    pub source: String,
21    pub agent_type: Option<String>,
22    pub state: SubagentState,
23    pub current_tool: Option<String>,
24    started_at: Instant,
25    pub duration_ms: Option<u64>,
26    pub error: Option<String>,
27}
28
29/// Handle returned when a subagent run starts; updates the global tracker.
30#[derive(Debug, Clone)]
31pub struct SubagentHandle {
32    id: String,
33}
34
35fn tracker() -> &'static Mutex<HashMap<String, SubagentRun>> {
36    static TRACKER: OnceLock<Mutex<HashMap<String, SubagentRun>>> = OnceLock::new();
37    TRACKER.get_or_init(|| Mutex::new(HashMap::new()))
38}
39
40impl SubagentHandle {
41    pub fn start(label: impl Into<String>, source: &str, agent_type: Option<String>) -> Self {
42        let id = uuid::Uuid::new_v4().to_string();
43        let short_id = id[..8.min(id.len())].to_string();
44        let run = SubagentRun {
45            id: short_id.clone(),
46            label: label.into(),
47            source: source.to_string(),
48            agent_type,
49            state: SubagentState::Running,
50            current_tool: None,
51            started_at: Instant::now(),
52            duration_ms: None,
53            error: None,
54        };
55        tracker().lock().unwrap().insert(short_id.clone(), run);
56        Self { id: short_id }
57    }
58
59    pub fn id(&self) -> &str {
60        &self.id
61    }
62
63    pub fn set_tool(&self, tool_name: &str) {
64        if let Some(run) = tracker().lock().unwrap().get_mut(&self.id) {
65            run.current_tool = Some(tool_name.to_string());
66        }
67    }
68
69    pub fn clear_tool(&self) {
70        if let Some(run) = tracker().lock().unwrap().get_mut(&self.id) {
71            run.current_tool = None;
72        }
73    }
74
75    pub fn complete_ok(&self) {
76        if let Some(run) = tracker().lock().unwrap().get_mut(&self.id) {
77            run.state = SubagentState::Done;
78            run.duration_ms = Some(run.started_at.elapsed().as_millis() as u64);
79            run.current_tool = None;
80        }
81    }
82
83    pub fn complete_err(&self, error: impl Into<String>) {
84        if let Some(run) = tracker().lock().unwrap().get_mut(&self.id) {
85            run.state = SubagentState::Failed;
86            run.error = Some(error.into());
87            run.duration_ms = Some(run.started_at.elapsed().as_millis() as u64);
88            run.current_tool = None;
89        }
90    }
91
92    pub fn dismiss(self) {
93        tracker().lock().unwrap().remove(&self.id);
94    }
95}
96
97/// Snapshot of runs still visible in the queue strip.
98pub fn snapshot_queue(max_age_ms: u64) -> Vec<SubagentRun> {
99    let now = Instant::now();
100    let mut guard = tracker().lock().unwrap();
101    guard.retain(|_, run| {
102        if run.state == SubagentState::Running {
103            return true;
104        }
105        let age = run
106            .duration_ms
107            .unwrap_or_else(|| now.duration_since(run.started_at).as_millis() as u64);
108        age < max_age_ms
109    });
110    let mut runs: Vec<_> = guard.values().cloned().collect();
111    runs.sort_by_key(|a| a.started_at);
112    runs
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn handle_lifecycle_updates_state() {
121        let h = SubagentHandle::start("explore auth", "task", Some("explore".into()));
122        h.set_tool("grep_search");
123        h.clear_tool();
124        h.complete_ok();
125        let snap = snapshot_queue(60_000);
126        assert!(snap
127            .iter()
128            .any(|r| r.id == h.id() && r.state == SubagentState::Done));
129        h.dismiss();
130        assert!(snapshot_queue(60_000).is_empty());
131    }
132
133    #[test]
134    fn complete_err_sets_failed_state() {
135        let h = SubagentHandle::start("failing run", "task", None);
136        h.complete_err("connection reset");
137        let snap = snapshot_queue(60_000);
138        let run = snap
139            .iter()
140            .find(|r| r.id == h.id())
141            .expect("failed run should remain in queue");
142        assert_eq!(run.state, SubagentState::Failed);
143        assert_eq!(run.error.as_deref(), Some("connection reset"));
144        assert!(run.duration_ms.is_some());
145        h.dismiss();
146    }
147
148    #[test]
149    fn snapshot_queue_evicts_old_done_runs() {
150        let h = SubagentHandle::start("slow done", "task", None);
151        std::thread::sleep(std::time::Duration::from_millis(100));
152        h.complete_ok();
153        let snap = snapshot_queue(50);
154        assert!(
155            !snap.iter().any(|r| r.id == h.id()),
156            "done run whose duration exceeds max_age_ms should be evicted"
157        );
158        h.dismiss();
159    }
160
161    #[test]
162    fn set_tool_and_clear_tool_update_running_run() {
163        let h = SubagentHandle::start("tooling", "task", None);
164        h.set_tool("read");
165        let with_tool = snapshot_queue(60_000)
166            .into_iter()
167            .find(|r| r.id == h.id())
168            .expect("running run should be visible");
169        assert_eq!(with_tool.state, SubagentState::Running);
170        assert_eq!(with_tool.current_tool.as_deref(), Some("read"));
171
172        h.clear_tool();
173        let cleared = snapshot_queue(60_000)
174            .into_iter()
175            .find(|r| r.id == h.id())
176            .expect("running run should still be visible");
177        assert_eq!(cleared.current_tool, None);
178        h.dismiss();
179    }
180
181    #[test]
182    fn multiple_concurrent_handles_tracked() {
183        let h1 = SubagentHandle::start("first", "task", Some("explore".into()));
184        let h2 = SubagentHandle::start("second", "task", Some("quick_task".into()));
185        let snap = snapshot_queue(60_000);
186        assert!(snap.iter().any(|r| r.id == h1.id() && r.label == "first"));
187        assert!(snap.iter().any(|r| r.id == h2.id() && r.label == "second"));
188        assert_ne!(h1.id(), h2.id());
189        h1.dismiss();
190        h2.dismiss();
191    }
192
193    #[test]
194    fn complete_ok_sets_done_state_and_duration() {
195        let h = SubagentHandle::start("ok run", "task", None);
196        h.complete_ok();
197        let run = snapshot_queue(60_000)
198            .into_iter()
199            .find(|r| r.id == h.id())
200            .expect("done run");
201        assert_eq!(run.state, SubagentState::Done);
202        assert!(run.duration_ms.is_some());
203        h.dismiss();
204    }
205
206    #[test]
207    fn dismiss_removes_run_from_snapshot() {
208        let h = SubagentHandle::start("ephemeral", "task", None);
209        assert!(snapshot_queue(60_000).iter().any(|r| r.id == h.id()));
210        let id = h.id().to_string();
211        h.dismiss();
212        assert!(!snapshot_queue(60_000).iter().any(|r| r.id == id));
213    }
214
215    #[test]
216    fn snapshot_queue_retains_running_past_ttl() {
217        let h = SubagentHandle::start("still running", "task", None);
218        std::thread::sleep(std::time::Duration::from_millis(100));
219        let snap = snapshot_queue(1);
220        assert!(snap
221            .iter()
222            .any(|r| r.id == h.id() && r.state == SubagentState::Running));
223        h.dismiss();
224    }
225
226    #[test]
227    fn complete_err_clears_current_tool() {
228        let h = SubagentHandle::start("tool then fail", "task", None);
229        h.set_tool("bash");
230        h.complete_err("boom");
231        let run = snapshot_queue(60_000)
232            .into_iter()
233            .find(|r| r.id == h.id())
234            .expect("failed run");
235        assert_eq!(run.current_tool, None);
236        h.dismiss();
237    }
238
239    #[test]
240    fn start_records_label_source_and_agent_type() {
241        let h = SubagentHandle::start("my label", "quick_task", Some("explore".into()));
242        let run = snapshot_queue(60_000)
243            .into_iter()
244            .find(|r| r.id == h.id())
245            .expect("running run");
246        assert_eq!(run.label, "my label");
247        assert_eq!(run.source, "quick_task");
248        assert_eq!(run.agent_type.as_deref(), Some("explore"));
249        h.dismiss();
250    }
251
252    #[test]
253    fn concurrent_handles_receive_distinct_ids() {
254        let handles: Vec<_> = (0..5)
255            .map(|i| SubagentHandle::start(format!("run-{i}"), "task", None))
256            .collect();
257        let ids: std::collections::HashSet<_> = handles.iter().map(SubagentHandle::id).collect();
258        assert_eq!(ids.len(), handles.len());
259        for h in handles {
260            h.dismiss();
261        }
262    }
263}
pawan/subagent.rs

pawan/
subagent.rs
