Skip to main content

atman_runtime/tools/
task_ops.rs

1use crate::error::RuntimeError;
2use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
3use crate::value::Value;
4
5pub struct TaskList;
6
7fn extract_optional_string(args: &ToolArgs, name: &str) -> Option<String> {
8    match args.named(name)? {
9        Value::Str(s) => Some(s.clone()),
10        _ => None,
11    }
12}
13
14fn extract_optional_bool(args: &ToolArgs, name: &str) -> Option<bool> {
15    match args.named(name)? {
16        Value::Bool(b) => Some(*b),
17        _ => None,
18    }
19}
20
21impl Tool for TaskList {
22    fn name(&self) -> &str {
23        "task.list"
24    }
25
26    fn tier(&self) -> Tier {
27        Tier::One
28    }
29
30    fn description(&self) -> Option<&str> {
31        Some(
32            "List all background tasks (bash, terminal, flow). Returns array of {id, kind, label, status, elapsed_ms, source_handle}. Filter by kind/status. Default: running only; pass all=true to include completed.",
33        )
34    }
35
36    fn input_schema(&self) -> serde_json::Value {
37        serde_json::json!({
38            "type": "object",
39            "properties": {
40                "kind": {"type": "string", "description": "Filter by kind: bash|terminal|flow"},
41                "status": {"type": "string", "description": "Filter by status: running|ok|err|killed"},
42                "all": {"type": "boolean", "description": "Include completed tasks (default false)"}
43            }
44        })
45    }
46
47    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
48        Box::pin(async move {
49            let kind_filter = extract_optional_string(&args, "kind");
50            let status_filter = extract_optional_string(&args, "status");
51            let all = extract_optional_bool(&args, "all").unwrap_or(false);
52
53            let registry = ctx.task_registry.clone().ok_or_else(|| {
54                RuntimeError::ToolFailed("task.list: registry not available".into())
55            })?;
56
57            let mut filter = crate::task_registry::TaskFilter::all();
58            if !all {
59                filter.status = Some(crate::task_registry::TaskStatus::Running);
60            }
61            if let Some(k) = &kind_filter {
62                filter.kind = parse_kind(k);
63            }
64            if let Some(s) = &status_filter {
65                filter.status = parse_status(s);
66            }
67
68            let snapshots = registry.list(&filter);
69            let arr: Vec<Value> = snapshots
70                .iter()
71                .map(|s| {
72                    Value::from_json(serde_json::json!({
73                        "id": s.id.0.to_string(),
74                        "kind": kind_to_str(s.kind),
75                        "label": s.label,
76                        "status": status_to_str(s.status),
77                        "elapsed_ms": s.elapsed_ms(),
78                        "source_handle": s.source_handle,
79                        "termination": s.termination.map(|termination| match termination {
80                            crate::task_registry::TaskTermination::Killed => "killed",
81                            crate::task_registry::TaskTermination::Suicide => "suicide",
82                        }),
83                    }))
84                })
85                .collect();
86            Ok(Value::List(arr))
87        })
88    }
89}
90
91pub struct TaskKill;
92
93impl Tool for TaskKill {
94    fn name(&self) -> &str {
95        "task.kill"
96    }
97
98    fn tier(&self) -> Tier {
99        Tier::Three
100    }
101
102    fn description(&self) -> Option<&str> {
103        Some(
104            "Kill a background task by id. A Flow killing its own task must pass suicide=true. Returns {killed, termination}.",
105        )
106    }
107
108    fn input_schema(&self) -> serde_json::Value {
109        serde_json::json!({
110            "type": "object",
111            "properties": {
112                "id": {"type": "string", "description": "Task id to kill"},
113                "suicide": {"type": "boolean", "description": "Required when a Flow targets its own task id (default false)"}
114            },
115            "required": ["id"]
116        })
117    }
118
119    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
120        Box::pin(async move {
121            let id_str = extract_optional_string(&args, "id").ok_or_else(|| {
122                RuntimeError::ToolFailed("task.kill: missing 'id' argument".into())
123            })?;
124            let id =
125                crate::task_registry::TaskId(uuid::Uuid::parse_str(&id_str).map_err(|e| {
126                    RuntimeError::ToolFailed(format!("task.kill: invalid id: {e}"))
127                })?);
128            let registry = ctx.task_registry.clone().ok_or_else(|| {
129                RuntimeError::ToolFailed("task.kill: registry not available".into())
130            })?;
131            let suicide = extract_optional_bool(&args, "suicide").unwrap_or(false);
132            let outcome = registry.kill_from(&id, ctx.flow_run_id.as_ref(), suicide);
133            match outcome {
134                crate::task_registry::KillOutcome::Killed { termination } => {
135                    Ok(Value::from_json(serde_json::json!({
136                        "killed": true,
137                        "termination": match termination {
138                            crate::task_registry::TaskTermination::Killed => "killed",
139                            crate::task_registry::TaskTermination::Suicide => "suicide",
140                        }
141                    })))
142                }
143                crate::task_registry::KillOutcome::SelfKillRejected => {
144                    Err(RuntimeError::ToolFailed(
145                        "task.kill: this Flow is killing itself; pass suicide=true to confirm self-termination"
146                            .into(),
147                    ))
148                }
149                crate::task_registry::KillOutcome::NotFound
150                | crate::task_registry::KillOutcome::NotRunning => Ok(Value::from_json(
151                    serde_json::json!({"killed": false, "termination": null}),
152                )),
153            }
154        })
155    }
156}
157
158fn parse_kind(s: &str) -> Option<crate::task_registry::TaskKind> {
159    use crate::task_registry::TaskKind;
160    match s {
161        "bash" => Some(TaskKind::Bash),
162        "terminal" => Some(TaskKind::Terminal),
163        "flow" => Some(TaskKind::Flow),
164        _ => None,
165    }
166}
167
168fn parse_status(s: &str) -> Option<crate::task_registry::TaskStatus> {
169    use crate::task_registry::TaskStatus;
170    match s {
171        "running" => Some(TaskStatus::Running),
172        "ok" => Some(TaskStatus::Ok),
173        "err" => Some(TaskStatus::Err),
174        "killed" => Some(TaskStatus::Killed),
175        _ => None,
176    }
177}
178
179fn kind_to_str(k: crate::task_registry::TaskKind) -> &'static str {
180    use crate::task_registry::TaskKind;
181    match k {
182        TaskKind::Bash => "bash",
183        TaskKind::Terminal => "terminal",
184        TaskKind::Flow => "flow",
185    }
186}
187
188fn status_to_str(s: crate::task_registry::TaskStatus) -> &'static str {
189    use crate::task_registry::TaskStatus;
190    match s {
191        TaskStatus::Running => "running",
192        TaskStatus::Killing => "killing",
193        TaskStatus::Ok => "ok",
194        TaskStatus::Err => "err",
195        TaskStatus::Killed => "killed",
196    }
197}