atman_runtime/
task_registry.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use std::time::Instant;
4
5use serde::{Deserialize, Serialize};
6use tokio::sync::broadcast;
7use tokio_util::sync::CancellationToken;
8use uuid::Uuid;
9
10#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
11#[serde(transparent)]
12pub struct TaskId(pub Uuid);
13
14impl TaskId {
15    pub fn now() -> Self {
16        Self(Uuid::now_v7())
17    }
18}
19
20impl std::fmt::Display for TaskId {
21    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22        self.0.fmt(f)
23    }
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
27#[serde(rename_all = "snake_case")]
28pub enum TaskKind {
29    Bash,
30    Terminal,
31    Flow,
32}
33
34impl TaskKind {
35    pub fn label(self) -> &'static str {
36        match self {
37            TaskKind::Bash => "Bash",
38            TaskKind::Terminal => "Terminal",
39            TaskKind::Flow => "Flow",
40        }
41    }
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
45#[serde(rename_all = "snake_case")]
46pub enum TaskStatus {
47    Running,
48    Killing,
49    Ok,
50    Err,
51    Killed,
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
55#[serde(rename_all = "snake_case")]
56pub enum TaskTermination {
57    Killed,
58    Suicide,
59}
60
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub enum KillOutcome {
63    Killed { termination: TaskTermination },
64    NotFound,
65    NotRunning,
66    SelfKillRejected,
67}
68
69impl TaskStatus {
70    pub fn display_label(self) -> &'static str {
71        match self {
72            TaskStatus::Running => "running",
73            TaskStatus::Killing => "stopping",
74            TaskStatus::Ok => "completed",
75            TaskStatus::Err => "failed",
76            TaskStatus::Killed => "stopped",
77        }
78    }
79
80    pub fn is_terminal(self) -> bool {
81        matches!(self, TaskStatus::Ok | TaskStatus::Err | TaskStatus::Killed)
82    }
83
84    pub fn is_running(self) -> bool {
85        matches!(self, TaskStatus::Running | TaskStatus::Killing)
86    }
87}
88
89pub fn normalize_task_label(value: &str) -> Option<String> {
90    crate::message::ToolCallIntent::new(value).map(|label| label.as_str().to_owned())
91}
92
93#[derive(Debug, Clone)]
94pub struct TaskSnapshot {
95    pub id: TaskId,
96    pub kind: TaskKind,
97    pub label: String,
98    pub command: Option<String>,
99    pub status: TaskStatus,
100    pub started_at: Instant,
101    pub ended_at: Option<Instant>,
102    pub source_handle: String,
103    pub session_id: String,
104    pub workspace_id: Option<String>,
105    pub flow_run_id: Option<crate::event::FlowRunId>,
106    pub termination: Option<TaskTermination>,
107}
108
109#[derive(Debug, Clone, PartialEq, Eq)]
110pub struct TaskDisplay {
111    pub label: String,
112    pub command: Option<String>,
113}
114
115impl From<String> for TaskDisplay {
116    fn from(label: String) -> Self {
117        Self {
118            label,
119            command: None,
120        }
121    }
122}
123
124impl From<&str> for TaskDisplay {
125    fn from(label: &str) -> Self {
126        label.to_owned().into()
127    }
128}
129
130impl TaskSnapshot {
131    pub fn elapsed_ms(&self) -> u64 {
132        self.ended_at
133            .unwrap_or_else(Instant::now)
134            .duration_since(self.started_at)
135            .as_millis() as u64
136    }
137
138    pub fn is_running(&self) -> bool {
139        self.status.is_running()
140    }
141}
142
143#[derive(Debug, Clone)]
144pub enum TaskEvent {
145    Registered(TaskSnapshot),
146    StatusChanged {
147        id: TaskId,
148        kind: TaskKind,
149        old: TaskStatus,
150        new: TaskStatus,
151        termination: Option<TaskTermination>,
152    },
153    Reaped {
154        id: TaskId,
155    },
156}
157
158#[derive(Debug, Clone, Default)]
159pub struct TaskFilter {
160    pub kind: Option<TaskKind>,
161    pub status: Option<TaskStatus>,
162    pub session_id: Option<String>,
163}
164
165impl TaskFilter {
166    pub fn all() -> Self {
167        Self::default()
168    }
169
170    pub fn running() -> Self {
171        Self {
172            status: Some(TaskStatus::Running),
173            ..Default::default()
174        }
175    }
176
177    pub fn matches(&self, snap: &TaskSnapshot) -> bool {
178        if let Some(k) = self.kind
179            && snap.kind != k
180        {
181            return false;
182        }
183        if let Some(s) = self.status
184            && snap.status != s
185        {
186            return false;
187        }
188        if let Some(ref sid) = self.session_id
189            && snap.session_id != *sid
190        {
191            return false;
192        }
193        true
194    }
195}
196
197struct TaskEntry {
198    snapshot: TaskSnapshot,
199    cancel: CancellationToken,
200    kill_hook: Option<std::sync::Arc<dyn Fn() + Send + Sync>>,
201}
202
203/// Central management layer for all observable tasks.
204///
205/// Sub-registries (BgRegistry, TermRegistry, Executor, agent_ctrl) register
206/// tasks here on spawn and call `finish` when the task ends. Typed operations
207/// (bash.output, term.input, term.capture) stay on the sub-registries and are
208/// looked up by `source_handle`.
209#[derive(Clone)]
210pub struct TaskRegistry {
211    inner: Arc<std::sync::Mutex<HashMap<TaskId, TaskEntry>>>,
212    event_tx: broadcast::Sender<TaskEvent>,
213}
214
215impl Default for TaskRegistry {
216    fn default() -> Self {
217        let (event_tx, _) = broadcast::channel(256);
218        Self {
219            inner: Arc::new(std::sync::Mutex::new(HashMap::new())),
220            event_tx,
221        }
222    }
223}
224
225fn running_snapshot(
226    kind: TaskKind,
227    display: TaskDisplay,
228    source_handle: String,
229    session_id: String,
230    workspace_id: Option<String>,
231    flow_run_id: Option<crate::event::FlowRunId>,
232) -> TaskSnapshot {
233    TaskSnapshot {
234        id: TaskId::now(),
235        kind,
236        label: display.label,
237        command: display.command,
238        status: TaskStatus::Running,
239        started_at: Instant::now(),
240        ended_at: None,
241        source_handle,
242        session_id,
243        workspace_id,
244        flow_run_id,
245        termination: None,
246    }
247}
248
249impl TaskRegistry {
250    pub fn new() -> Self {
251        Self::default()
252    }
253
254    pub fn register(
255        &self,
256        kind: TaskKind,
257        display: TaskDisplay,
258        source_handle: String,
259        session_id: String,
260        cancel: CancellationToken,
261    ) -> TaskId {
262        self.register_snapshot(
263            running_snapshot(kind, display, source_handle, session_id, None, None),
264            cancel,
265            None,
266        )
267    }
268
269    pub fn register_flow(
270        &self,
271        label: String,
272        source_handle: String,
273        session_id: String,
274        cancel: CancellationToken,
275        workspace_id: Option<String>,
276    ) -> TaskId {
277        self.register_snapshot(
278            running_snapshot(
279                TaskKind::Flow,
280                label.into(),
281                source_handle,
282                session_id,
283                workspace_id,
284                None,
285            ),
286            cancel,
287            None,
288        )
289    }
290
291    pub fn register_flow_with_run_id(
292        &self,
293        label: String,
294        source_handle: String,
295        session_id: String,
296        cancel: CancellationToken,
297        workspace_id: Option<String>,
298        flow_run_id: crate::event::FlowRunId,
299    ) -> TaskId {
300        self.register_snapshot(
301            running_snapshot(
302                TaskKind::Flow,
303                label.into(),
304                source_handle,
305                session_id,
306                workspace_id,
307                Some(flow_run_id),
308            ),
309            cancel,
310            None,
311        )
312    }
313
314    pub fn register_with_kill_hook(
315        &self,
316        kind: TaskKind,
317        display: TaskDisplay,
318        source_handle: String,
319        session_id: String,
320        cancel: CancellationToken,
321        kill_hook: Option<std::sync::Arc<dyn Fn() + Send + Sync>>,
322    ) -> TaskId {
323        self.register_snapshot(
324            running_snapshot(kind, display, source_handle, session_id, None, None),
325            cancel,
326            kill_hook,
327        )
328    }
329
330    fn register_snapshot(
331        &self,
332        snapshot: TaskSnapshot,
333        cancel: CancellationToken,
334        kill_hook: Option<std::sync::Arc<dyn Fn() + Send + Sync>>,
335    ) -> TaskId {
336        let id = snapshot.id.clone();
337        let entry = TaskEntry {
338            snapshot: snapshot.clone(),
339            cancel,
340            kill_hook,
341        };
342        self.inner.lock().unwrap().insert(id.clone(), entry);
343        let _ = self.event_tx.send(TaskEvent::Registered(snapshot));
344        id
345    }
346
347    pub fn lookup(&self, id: &TaskId) -> Option<TaskSnapshot> {
348        self.inner
349            .lock()
350            .unwrap()
351            .get(id)
352            .map(|e| e.snapshot.clone())
353    }
354
355    /// Find a task by its source handle (e.g. "bg_1", "term_2", run_id).
356    pub fn lookup_by_handle(&self, handle: &str) -> Option<TaskSnapshot> {
357        self.inner
358            .lock()
359            .unwrap()
360            .values()
361            .find(|e| e.snapshot.source_handle == handle)
362            .map(|e| e.snapshot.clone())
363    }
364
365    pub fn list(&self, filter: &TaskFilter) -> Vec<TaskSnapshot> {
366        let inner = self.inner.lock().unwrap();
367        let mut out: Vec<TaskSnapshot> = inner
368            .values()
369            .map(|e| e.snapshot.clone())
370            .filter(|s| filter.matches(s))
371            .collect();
372        out.sort_by_key(|s| s.started_at);
373        out
374    }
375
376    /// Kill a task at the request of an external operator, such as the TUI.
377    ///
378    /// Operator actions have no Flow identity, so they can never be mistaken
379    /// for a Flow killing itself.
380    pub fn kill_from_operator(&self, id: &TaskId) -> KillOutcome {
381        self.kill_from(id, None, false)
382    }
383
384    pub fn kill_from(
385        &self,
386        id: &TaskId,
387        caller_flow_run_id: Option<&crate::event::FlowRunId>,
388        suicide: bool,
389    ) -> KillOutcome {
390        let mut inner = self.inner.lock().unwrap();
391        let Some(entry) = inner.get_mut(id) else {
392            return KillOutcome::NotFound;
393        };
394        if entry.snapshot.status.is_terminal() {
395            return KillOutcome::NotRunning;
396        }
397        let self_targeting = caller_flow_run_id.is_some()
398            && entry.snapshot.flow_run_id.as_ref() == caller_flow_run_id;
399        if self_targeting && !suicide {
400            return KillOutcome::SelfKillRejected;
401        }
402        let termination = if self_targeting {
403            TaskTermination::Suicide
404        } else {
405            TaskTermination::Killed
406        };
407        let old_status = entry.snapshot.status;
408        let kind = entry.snapshot.kind;
409        entry.snapshot.status = TaskStatus::Killing;
410        entry.snapshot.termination = Some(termination);
411        let cancel = entry.cancel.clone();
412        let hook = entry.kill_hook.clone();
413        drop(inner);
414        cancel.cancel();
415        if let Some(hook) = hook {
416            hook();
417        }
418        let _ = self.event_tx.send(TaskEvent::StatusChanged {
419            id: id.clone(),
420            kind,
421            old: old_status,
422            new: TaskStatus::Killing,
423            termination: Some(termination),
424        });
425        KillOutcome::Killed { termination }
426    }
427
428    /// Transition a task to a terminal status. Called by the owning
429    /// sub-registry when the task finishes.
430    pub fn finish(&self, id: &TaskId, status: TaskStatus) {
431        let mut inner = self.inner.lock().unwrap();
432        let Some(entry) = inner.get_mut(id) else {
433            return;
434        };
435        if entry.snapshot.status.is_terminal() {
436            return;
437        }
438        let old = entry.snapshot.status;
439        entry.snapshot.status = status;
440        entry.snapshot.ended_at = Some(Instant::now());
441        let kind = entry.snapshot.kind;
442        let termination = entry.snapshot.termination;
443        drop(inner);
444        let _ = self.event_tx.send(TaskEvent::StatusChanged {
445            id: id.clone(),
446            kind,
447            old,
448            new: status,
449            termination,
450        });
451    }
452
453    pub fn reap(&self, id: &TaskId) {
454        let mut inner = self.inner.lock().unwrap();
455        let should_remove = inner
456            .get(id)
457            .map(|e| e.snapshot.status.is_terminal())
458            .unwrap_or(false);
459        if should_remove {
460            inner.remove(id);
461            drop(inner);
462            let _ = self.event_tx.send(TaskEvent::Reaped { id: id.clone() });
463        }
464    }
465
466    pub fn subscribe(&self) -> broadcast::Receiver<TaskEvent> {
467        self.event_tx.subscribe()
468    }
469
470    pub fn running_count(&self) -> usize {
471        self.inner
472            .lock()
473            .unwrap()
474            .values()
475            .filter(|e| e.snapshot.status.is_running())
476            .count()
477    }
478}
479
480#[cfg(test)]
481mod tests {
482    use super::*;
483
484    fn cancel() -> CancellationToken {
485        CancellationToken::new()
486    }
487
488    #[test]
489    fn register_and_lookup() {
490        let reg = TaskRegistry::new();
491        let id = reg.register(
492            TaskKind::Bash,
493            "cargo build".into(),
494            "bg_1".into(),
495            "sess".into(),
496            cancel(),
497        );
498        let snap = reg.lookup(&id).expect("found");
499        assert_eq!(snap.kind, TaskKind::Bash);
500        assert_eq!(snap.status, TaskStatus::Running);
501        assert!(snap.ended_at.is_none());
502        assert!(snap.command.is_none());
503    }
504
505    #[test]
506    fn command_is_independent_from_the_user_facing_label() {
507        let reg = TaskRegistry::new();
508        let id = reg.register(
509            TaskKind::Bash,
510            TaskDisplay {
511                label: "运行项目测试".into(),
512                command: Some("cargo test --workspace".into()),
513            },
514            "bg_1".into(),
515            "sess".into(),
516            cancel(),
517        );
518        let snap = reg.lookup(&id).expect("found");
519        assert_eq!(snap.label, "运行项目测试");
520        assert_eq!(snap.command.as_deref(), Some("cargo test --workspace"));
521    }
522
523    #[test]
524    fn lookup_by_handle() {
525        let reg = TaskRegistry::new();
526        let _id = reg.register(
527            TaskKind::Terminal,
528            "vim".into(),
529            "term_1".into(),
530            "sess".into(),
531            cancel(),
532        );
533        let snap = reg.lookup_by_handle("term_1").expect("found");
534        assert_eq!(snap.kind, TaskKind::Terminal);
535        assert!(reg.lookup_by_handle("nope").is_none());
536    }
537
538    #[test]
539    fn list_filters_by_kind_and_status() {
540        let reg = TaskRegistry::new();
541        let b1 = reg.register(
542            TaskKind::Bash,
543            "a".into(),
544            "bg_1".into(),
545            "s".into(),
546            cancel(),
547        );
548        let _t1 = reg.register(
549            TaskKind::Terminal,
550            "vim".into(),
551            "term_1".into(),
552            "s".into(),
553            cancel(),
554        );
555        let _b2 = reg.register(
556            TaskKind::Bash,
557            "ls".into(),
558            "bg_2".into(),
559            "s".into(),
560            cancel(),
561        );
562
563        let bash_only = reg.list(&TaskFilter {
564            kind: Some(TaskKind::Bash),
565            ..Default::default()
566        });
567        assert_eq!(bash_only.len(), 2);
568
569        reg.finish(&b1, TaskStatus::Ok);
570        let running = reg.list(&TaskFilter::running());
571        assert_eq!(running.len(), 2);
572    }
573
574    #[test]
575    fn kill_cancels_token() {
576        let reg = TaskRegistry::new();
577        let tok = cancel();
578        let id = reg.register(
579            TaskKind::Bash,
580            "x".into(),
581            "bg".into(),
582            "s".into(),
583            tok.clone(),
584        );
585        assert_eq!(
586            reg.kill_from_operator(&id),
587            KillOutcome::Killed {
588                termination: TaskTermination::Killed
589            }
590        );
591        assert!(tok.is_cancelled());
592    }
593
594    #[test]
595    fn self_kill_requires_suicide_confirmation() {
596        let reg = TaskRegistry::new();
597        let run_id = crate::event::FlowRunId::now();
598        let token = cancel();
599        let id = reg.register_flow_with_run_id(
600            "flow".into(),
601            "agent".into(),
602            "s".into(),
603            token.clone(),
604            None,
605            run_id.clone(),
606        );
607
608        assert_eq!(
609            reg.kill_from(&id, Some(&run_id), false),
610            KillOutcome::SelfKillRejected
611        );
612        assert!(!token.is_cancelled());
613        assert_eq!(reg.lookup(&id).unwrap().status, TaskStatus::Running);
614    }
615
616    #[test]
617    fn confirmed_self_kill_records_suicide_cause() {
618        let reg = TaskRegistry::new();
619        let run_id = crate::event::FlowRunId::now();
620        let token = cancel();
621        let id = reg.register_flow_with_run_id(
622            "flow".into(),
623            "agent".into(),
624            "s".into(),
625            token.clone(),
626            None,
627            run_id.clone(),
628        );
629
630        assert_eq!(
631            reg.kill_from(&id, Some(&run_id), true),
632            KillOutcome::Killed {
633                termination: TaskTermination::Suicide
634            }
635        );
636        assert!(token.is_cancelled());
637        assert_eq!(
638            reg.lookup(&id).unwrap().termination,
639            Some(TaskTermination::Suicide)
640        );
641    }
642
643    #[test]
644    fn operator_kill_is_not_classified_as_suicide() {
645        let reg = TaskRegistry::new();
646        let target_run_id = crate::event::FlowRunId::now();
647        let id = reg.register_flow_with_run_id(
648            "flow".into(),
649            "agent".into(),
650            "s".into(),
651            cancel(),
652            None,
653            target_run_id,
654        );
655
656        assert_eq!(
657            reg.kill_from_operator(&id),
658            KillOutcome::Killed {
659                termination: TaskTermination::Killed
660            }
661        );
662        assert_eq!(
663            reg.lookup(&id).unwrap().termination,
664            Some(TaskTermination::Killed)
665        );
666    }
667
668    #[test]
669    fn another_flow_can_kill_without_suicide_confirmation() {
670        let reg = TaskRegistry::new();
671        let target_run_id = crate::event::FlowRunId::now();
672        let caller_run_id = crate::event::FlowRunId::now();
673        let id = reg.register_flow_with_run_id(
674            "flow".into(),
675            "agent".into(),
676            "s".into(),
677            cancel(),
678            None,
679            target_run_id,
680        );
681
682        assert_eq!(
683            reg.kill_from(&id, Some(&caller_run_id), false),
684            KillOutcome::Killed {
685                termination: TaskTermination::Killed
686            }
687        );
688    }
689
690    #[test]
691    fn kill_returns_false_for_terminal() {
692        let reg = TaskRegistry::new();
693        let id = reg.register(
694            TaskKind::Bash,
695            "x".into(),
696            "bg".into(),
697            "s".into(),
698            cancel(),
699        );
700        reg.finish(&id, TaskStatus::Ok);
701        assert_eq!(reg.kill_from_operator(&id), KillOutcome::NotRunning);
702    }
703
704    #[test]
705    fn finish_is_idempotent() {
706        let reg = TaskRegistry::new();
707        let id = reg.register(
708            TaskKind::Bash,
709            "x".into(),
710            "bg".into(),
711            "s".into(),
712            cancel(),
713        );
714        reg.finish(&id, TaskStatus::Ok);
715        reg.finish(&id, TaskStatus::Err);
716        let snap = reg.lookup(&id).unwrap();
717        assert_eq!(snap.status, TaskStatus::Ok);
718    }
719
720    #[test]
721    fn reap_removes_terminal_only() {
722        let reg = TaskRegistry::new();
723        let id = reg.register(
724            TaskKind::Bash,
725            "x".into(),
726            "bg".into(),
727            "s".into(),
728            cancel(),
729        );
730        reg.reap(&id);
731        assert!(reg.lookup(&id).is_some());
732        reg.finish(&id, TaskStatus::Ok);
733        reg.reap(&id);
734        assert!(reg.lookup(&id).is_none());
735    }
736
737    #[test]
738    fn subscribe_receives_registered_event() {
739        let reg = TaskRegistry::new();
740        let mut rx = reg.subscribe();
741        let _id = reg.register(
742            TaskKind::Bash,
743            "x".into(),
744            "bg".into(),
745            "s".into(),
746            cancel(),
747        );
748        let ev = rx.try_recv().expect("got event");
749        match ev {
750            TaskEvent::Registered(s) => assert_eq!(s.kind, TaskKind::Bash),
751            _ => panic!("wrong event"),
752        }
753    }
754
755    #[test]
756    fn subscribe_receives_status_changed() {
757        let reg = TaskRegistry::new();
758        let mut rx = reg.subscribe();
759        let id = reg.register(
760            TaskKind::Bash,
761            "x".into(),
762            "bg".into(),
763            "s".into(),
764            cancel(),
765        );
766        let _ = rx.try_recv();
767        reg.finish(&id, TaskStatus::Ok);
768        let ev = rx.try_recv().expect("got status event");
769        match ev {
770            TaskEvent::StatusChanged { new, .. } => assert_eq!(new, TaskStatus::Ok),
771            _ => panic!("wrong event"),
772        }
773    }
774
775    #[test]
776    fn filter_matches_combines() {
777        let snap = TaskSnapshot {
778            id: TaskId::now(),
779            kind: TaskKind::Terminal,
780            label: "vim".into(),
781            command: None,
782            status: TaskStatus::Running,
783            started_at: Instant::now(),
784            ended_at: None,
785            source_handle: "term_1".into(),
786            session_id: "sess_a".into(),
787            workspace_id: None,
788            flow_run_id: None,
789            termination: None,
790        };
791        let f = TaskFilter {
792            kind: Some(TaskKind::Terminal),
793            status: Some(TaskStatus::Running),
794            session_id: Some("sess_a".into()),
795        };
796        assert!(f.matches(&snap));
797
798        let f2 = TaskFilter {
799            kind: Some(TaskKind::Bash),
800            ..Default::default()
801        };
802        assert!(!f2.matches(&snap));
803    }
804
805    #[test]
806    fn task_status_has_stable_display_labels() {
807        assert_eq!(TaskStatus::Running.display_label(), "running");
808        assert_eq!(TaskStatus::Killing.display_label(), "stopping");
809        assert_eq!(TaskStatus::Ok.display_label(), "completed");
810        assert_eq!(TaskStatus::Err.display_label(), "failed");
811        assert_eq!(TaskStatus::Killed.display_label(), "stopped");
812    }
813
814    #[test]
815    fn task_label_normalization_matches_tool_intent_bounds() {
816        assert_eq!(
817            normalize_task_label("  检查   当前状态  ").as_deref(),
818            Some("检查 当前状态")
819        );
820        assert!(normalize_task_label(" \n\t ").is_none());
821        assert_eq!(
822            normalize_task_label(&"x".repeat(121))
823                .unwrap()
824                .chars()
825                .count(),
826            120
827        );
828    }
829}
atman_runtime/task_registry.rs

atman_runtime/
task_registry.rs
