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

pawan/
subagent.rs
