atman_runtime/
tool.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4
5use tokio_util::sync::CancellationToken;
6
7use crate::error::RuntimeError;
8use crate::value::Value;
9
10// `dyn Future` is now `Send` — AST Span was replaced with a custom `Copy + Send + Sync`
11// struct, removing the `proc_macro2::Span` (which held `Rc<()>`).  This means
12// providers, tools, and classifiers can be spawned with `tokio::spawn` instead
13// of requiring `spawn_local` + `LocalSet`.
14pub type BoxFut<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
15
16pub type ToolResult = Result<Value, RuntimeError>;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum Tier {
20    Zero,
21    One,
22    Two,
23    Three,
24    Four,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
28pub enum ApprovalLevel {
29    Auto,
30    Approve,
31    Dangerous,
32}
33
34impl ApprovalLevel {
35    pub fn from_tier(tier: Tier) -> Self {
36        match tier {
37            Tier::Zero => ApprovalLevel::Auto,
38            Tier::One | Tier::Two => ApprovalLevel::Approve,
39            Tier::Three | Tier::Four => ApprovalLevel::Dangerous,
40        }
41    }
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum CancelBehavior {
46    AbortSafe,
47    Revertible,
48    Atomic,
49    Irreversible,
50}
51
52#[derive(Debug, Default, Clone)]
53pub struct ToolArgs {
54    pub positional: Vec<Value>,
55    pub named: Vec<(String, Value)>,
56}
57
58impl ToolArgs {
59    pub fn positional(&self, index: usize) -> Result<&Value, RuntimeError> {
60        self.positional
61            .get(index)
62            .ok_or_else(|| RuntimeError::MissingArg(format!("positional[{index}]")))
63    }
64
65    pub fn named(&self, name: &str) -> Option<&Value> {
66        self.named.iter().find(|(k, _)| k == name).map(|(_, v)| v)
67    }
68}
69
70#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
71pub enum HistorySegment {
72    #[default]
73    Root,
74    Spawned,
75}
76
77#[derive(Clone, Default)]
78pub struct ToolCtx {
79    pub cancel: CancellationToken,
80    pub turn_id: Option<crate::event::TurnId>,
81    pub flow_run_id: Option<crate::event::FlowRunId>,
82    pub history_segment: HistorySegment,
83    pub event_seq: Option<u64>,
84    pub prompt_resolver: Option<std::sync::Arc<dyn crate::rendezvous::PromptResolver>>,
85    pub registry: Option<std::sync::Arc<ToolRegistry>>,
86    pub sandbox: Option<std::sync::Arc<dyn crate::sandbox::Sandbox>>,
87    pub events: Option<crate::event::EventSink>,
88    pub stdout_broadcast: Option<tokio::sync::broadcast::Sender<String>>,
89    pub session_messages: Option<std::sync::Arc<Vec<crate::message::Message>>>,
90    pub session_messages_handle:
91        Option<std::sync::Arc<std::sync::Mutex<Vec<crate::message::Message>>>>,
92    pub session_runtime: Option<std::sync::Arc<crate::session::Session>>,
93    pub compact_lock_handle: Option<std::sync::Arc<tokio::sync::Mutex<()>>>,
94    pub current_node_id: Option<String>,
95    pub stream_tx: Option<tokio::sync::broadcast::Sender<crate::stream::StreamFrame>>,
96    pub read_files:
97        Option<std::sync::Arc<std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>>>,
98    pub approval: Option<std::sync::Arc<crate::session::ApprovalRegistry>>,
99    pub forms: Option<std::sync::Arc<crate::session::FormRegistry>>,
100    pub providers: Option<std::sync::Arc<crate::provider::ProviderRegistry>>,
101    pub session_dir: Option<std::path::PathBuf>,
102    pub output_store: Option<std::sync::Arc<crate::tools::tool_output::OutputStore>>,
103    pub data_root: Option<std::path::PathBuf>,
104    pub project_index: Option<std::sync::Arc<crate::index::AnchorIndex>>,
105    pub fs_access: crate::fs_access::FsAccessPolicy,
106    pub lifecycle_fire_tx:
107        Option<tokio::sync::mpsc::UnboundedSender<atman_dsl::ast::LifecycleEvent>>,
108    pub bg_registry: Option<std::sync::Arc<crate::tools::bash_bg::BgRegistry>>,
109    pub term_registry: Option<std::sync::Arc<crate::tools::term::TermRegistry>>,
110    pub watch_hub: Option<std::sync::Arc<crate::watch::WatchHub>>,
111    pub flow_registry: Option<std::sync::Arc<crate::tools::agent_ctrl::FlowRegistry>>,
112    pub task_registry: Option<crate::task_registry::TaskRegistry>,
113    pub session_id: Option<String>,
114    pub trust: Option<crate::trust::TrustConfig>,
115    pub safety: Option<crate::safety::SafetyConfig>,
116    pub current_model: Option<String>,
117    pub watch_rules: Option<crate::streaming::WatchRules>,
118    /// Called when memory.recent_turns is invoked, with the count of returned messages.
119    pub on_memory_recent: Option<std::sync::Arc<dyn Fn(u16) + Send + Sync>>,
120    pub history_store: Option<std::sync::Arc<dyn crate::history_store::HistoryStore>>,
121    pub agent_entry: Option<std::sync::Arc<crate::tools::agent_ctrl::FlowEntry>>,
122    pub tool_output_budget: crate::tools::tool_output::ToolOutputBudget,
123}
124
125impl ToolCtx {
126    pub fn new() -> Self {
127        Self::default()
128    }
129
130    pub fn with_anchors(
131        mut self,
132        turn_id: Option<crate::event::TurnId>,
133        flow_run_id: Option<crate::event::FlowRunId>,
134        event_seq: Option<u64>,
135    ) -> Self {
136        self.turn_id = turn_id;
137        self.flow_run_id = flow_run_id;
138        self.event_seq = event_seq;
139        self
140    }
141
142    pub fn with_history_segment(mut self, segment: HistorySegment) -> Self {
143        self.history_segment = segment;
144        self
145    }
146
147    pub fn message_flow_run_id(&self) -> Option<crate::event::FlowRunId> {
148        match self.history_segment {
149            HistorySegment::Root => None,
150            HistorySegment::Spawned => self.flow_run_id.clone(),
151        }
152    }
153
154    pub fn with_registry(mut self, registry: std::sync::Arc<ToolRegistry>) -> Self {
155        self.registry = Some(registry);
156        self
157    }
158
159    pub fn with_sandbox(mut self, sandbox: std::sync::Arc<dyn crate::sandbox::Sandbox>) -> Self {
160        self.sandbox = Some(sandbox);
161        self
162    }
163
164    pub fn with_events(mut self, events: crate::event::EventSink) -> Self {
165        self.events = Some(events);
166        self
167    }
168
169    pub fn with_stdout_broadcast(mut self, tx: tokio::sync::broadcast::Sender<String>) -> Self {
170        self.stdout_broadcast = Some(tx);
171        self
172    }
173
174    pub fn with_session_messages(
175        mut self,
176        msgs: std::sync::Arc<Vec<crate::message::Message>>,
177    ) -> Self {
178        self.session_messages = Some(msgs);
179        self
180    }
181
182    pub fn with_session_messages_handle(
183        mut self,
184        handle: std::sync::Arc<std::sync::Mutex<Vec<crate::message::Message>>>,
185    ) -> Self {
186        self.session_messages_handle = Some(handle);
187        self
188    }
189
190    pub fn with_session_runtime(
191        mut self,
192        session: std::sync::Arc<crate::session::Session>,
193    ) -> Self {
194        self.session_runtime = Some(session);
195        self
196    }
197
198    pub fn with_compact_lock_handle(
199        mut self,
200        handle: std::sync::Arc<tokio::sync::Mutex<()>>,
201    ) -> Self {
202        self.compact_lock_handle = Some(handle);
203        self
204    }
205
206    pub fn with_current_node(mut self, node_id: Option<String>) -> Self {
207        self.current_node_id = node_id;
208        self
209    }
210
211    pub fn with_read_files(
212        mut self,
213        set: std::sync::Arc<std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>>,
214    ) -> Self {
215        self.read_files = Some(set);
216        self
217    }
218
219    pub fn with_providers(
220        mut self,
221        providers: std::sync::Arc<crate::provider::ProviderRegistry>,
222    ) -> Self {
223        self.providers = Some(providers);
224        self
225    }
226
227    pub fn with_session_dir(mut self, dir: std::path::PathBuf) -> Self {
228        self.output_store = Some(std::sync::Arc::new(
229            crate::tools::tool_output::OutputStore::at(dir.clone()),
230        ));
231        self.session_dir = Some(dir);
232        self
233    }
234
235    pub fn with_output_store(
236        mut self,
237        store: std::sync::Arc<crate::tools::tool_output::OutputStore>,
238    ) -> Self {
239        self.output_store = Some(store);
240        self
241    }
242
243    pub fn with_data_root(mut self, dir: std::path::PathBuf) -> Self {
244        self.data_root = Some(dir);
245        self
246    }
247
248    pub fn with_approval(
249        mut self,
250        approval: std::sync::Arc<crate::session::ApprovalRegistry>,
251    ) -> Self {
252        self.approval = Some(approval);
253        self
254    }
255
256    pub fn with_fs_access(mut self, policy: crate::fs_access::FsAccessPolicy) -> Self {
257        self.fs_access = policy;
258        self
259    }
260
261    pub fn with_lifecycle_fire_tx(
262        mut self,
263        tx: tokio::sync::mpsc::UnboundedSender<atman_dsl::ast::LifecycleEvent>,
264    ) -> Self {
265        self.lifecycle_fire_tx = Some(tx);
266        self
267    }
268
269    pub fn with_forms(mut self, forms: std::sync::Arc<crate::session::FormRegistry>) -> Self {
270        self.forms = Some(forms);
271        self
272    }
273
274    pub fn with_bg_registry(
275        mut self,
276        registry: std::sync::Arc<crate::tools::bash_bg::BgRegistry>,
277    ) -> Self {
278        self.bg_registry = Some(registry);
279        self
280    }
281
282    pub fn with_session_id(mut self, id: impl Into<String>) -> Self {
283        self.session_id = Some(id.into());
284        self
285    }
286
287    pub fn with_trust(mut self, trust: crate::trust::TrustConfig) -> Self {
288        self.trust = Some(trust);
289        self
290    }
291
292    pub fn with_safety(mut self, safety: crate::safety::SafetyConfig) -> Self {
293        self.safety = Some(safety);
294        self
295    }
296
297    pub fn with_current_model(mut self, model: impl Into<String>) -> Self {
298        self.current_model = Some(model.into());
299        self
300    }
301
302    pub fn with_watch_rules(mut self, rules: crate::streaming::WatchRules) -> Self {
303        self.watch_rules = Some(rules);
304        self
305    }
306
307    pub fn with_term_registry(
308        mut self,
309        registry: std::sync::Arc<crate::tools::term::TermRegistry>,
310    ) -> Self {
311        self.term_registry = Some(registry);
312        self
313    }
314
315    pub fn with_watch_hub(mut self, hub: std::sync::Arc<crate::watch::WatchHub>) -> Self {
316        self.watch_hub = Some(hub);
317        self
318    }
319
320    pub fn with_flow_registry(
321        mut self,
322        registry: std::sync::Arc<crate::tools::agent_ctrl::FlowRegistry>,
323    ) -> Self {
324        self.flow_registry = Some(registry);
325        self
326    }
327
328    pub fn with_task_registry(mut self, registry: crate::task_registry::TaskRegistry) -> Self {
329        self.task_registry = Some(registry);
330        self
331    }
332
333    pub fn note_read(&self, path: &std::path::Path) {
334        if let Some(set) = &self.read_files
335            && let Ok(mut lock) = set.lock()
336        {
337            lock.insert(path.to_path_buf());
338        }
339    }
340
341    pub fn has_read(&self, path: &std::path::Path) -> bool {
342        self.read_files
343            .as_ref()
344            .and_then(|set| set.lock().ok().map(|lock| lock.contains(path)))
345            .unwrap_or(false)
346    }
347
348    pub fn with_project_index(mut self, idx: std::sync::Arc<crate::index::AnchorIndex>) -> Self {
349        self.project_index = Some(idx);
350        self
351    }
352
353    pub fn with_history_store(
354        mut self,
355        store: std::sync::Arc<dyn crate::history_store::HistoryStore>,
356    ) -> Self {
357        self.history_store = Some(store);
358        self
359    }
360
361    pub fn with_agent_entry(
362        mut self,
363        entry: std::sync::Arc<crate::tools::agent_ctrl::FlowEntry>,
364    ) -> Self {
365        self.agent_entry = Some(entry);
366        self
367    }
368
369    pub fn with_stream_tx(
370        mut self,
371        tx: tokio::sync::broadcast::Sender<crate::stream::StreamFrame>,
372    ) -> Self {
373        self.stream_tx = Some(tx);
374        self
375    }
376}
377
378pub trait Tool: Send + Sync {
379    fn name(&self) -> &str;
380    fn tier(&self) -> Tier;
381    fn approval_level(&self, _args: &ToolArgs, _ctx: &ToolCtx) -> ApprovalLevel {
382        ApprovalLevel::from_tier(self.tier())
383    }
384    fn cancel_behavior(&self) -> CancelBehavior {
385        CancelBehavior::AbortSafe
386    }
387    fn description(&self) -> Option<&str> {
388        None
389    }
390    fn input_schema(&self) -> serde_json::Value {
391        serde_json::json!({"type": "object"})
392    }
393    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult>;
394    fn preview_call<'a>(
395        &'a self,
396        _args: &'a ToolArgs,
397        _ctx: &'a ToolCtx,
398    ) -> BoxFut<'a, Option<String>> {
399        Box::pin(async { None })
400    }
401}
402
403pub fn tool_spec(tool: &dyn Tool) -> ToolSpec {
404    ToolSpec {
405        name: tool.name().to_string(),
406        description: tool.description().map(str::to_string),
407        input_schema: tool.input_schema(),
408    }
409}
410
411#[derive(Debug, Clone, serde::Serialize)]
412pub struct ToolSpec {
413    pub name: String,
414    #[serde(skip_serializing_if = "Option::is_none")]
415    pub description: Option<String>,
416    pub input_schema: serde_json::Value,
417}
418
419#[derive(Default, Clone)]
420pub struct ToolRegistry {
421    tools: std::sync::Arc<std::sync::RwLock<HashMap<String, std::sync::Arc<dyn Tool>>>>,
422}
423
424impl ToolRegistry {
425    pub fn new() -> Self {
426        Self::default()
427    }
428
429    pub fn register(&self, tool: std::sync::Arc<dyn Tool>) {
430        self.tools
431            .write()
432            .unwrap()
433            .insert(tool.name().to_string(), tool);
434    }
435
436    pub fn get(&self, name: &str) -> Option<std::sync::Arc<dyn Tool>> {
437        self.tools.read().unwrap().get(name).cloned()
438    }
439
440    pub fn has(&self, name: &str) -> bool {
441        self.tools.read().unwrap().contains_key(name)
442    }
443
444    pub fn names(&self) -> Vec<String> {
445        self.tools.read().unwrap().keys().cloned().collect()
446    }
447
448    pub fn iter(&self) -> Vec<(String, std::sync::Arc<dyn Tool>)> {
449        self.tools
450            .read()
451            .unwrap()
452            .iter()
453            .map(|(k, v)| (k.clone(), v.clone()))
454            .collect()
455    }
456
457    /// Remove all tools whose name starts with `prefix` (e.g. `"mcp."`).
458    pub fn unregister_prefix(&self, prefix: &str) {
459        self.tools
460            .write()
461            .unwrap()
462            .retain(|k, _| !k.starts_with(prefix));
463    }
464}
465
466#[cfg(test)]
467mod tests {
468    use super::*;
469
470    #[test]
471    fn approval_level_default_maps_from_tier() {
472        assert_eq!(ApprovalLevel::from_tier(Tier::Zero), ApprovalLevel::Auto);
473        assert_eq!(ApprovalLevel::from_tier(Tier::One), ApprovalLevel::Approve);
474        assert_eq!(ApprovalLevel::from_tier(Tier::Two), ApprovalLevel::Approve);
475        assert_eq!(
476            ApprovalLevel::from_tier(Tier::Three),
477            ApprovalLevel::Dangerous
478        );
479        assert_eq!(
480            ApprovalLevel::from_tier(Tier::Four),
481            ApprovalLevel::Dangerous
482        );
483    }
484
485    #[test]
486    fn approval_level_ordered_auto_lt_approve_lt_dangerous() {
487        assert!(ApprovalLevel::Auto < ApprovalLevel::Approve);
488        assert!(ApprovalLevel::Approve < ApprovalLevel::Dangerous);
489    }
490}
atman_runtime/tool.rs

atman_runtime/
tool.rs
