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, Default)]
71pub struct ToolCtx {
72    pub cancel: CancellationToken,
73    pub turn_id: Option<crate::event::TurnId>,
74    pub flow_run_id: Option<crate::event::FlowRunId>,
75    pub event_seq: Option<u64>,
76    pub prompt_resolver: Option<std::sync::Arc<dyn crate::rendezvous::PromptResolver>>,
77    pub registry: Option<std::sync::Arc<ToolRegistry>>,
78    pub sandbox: Option<std::sync::Arc<dyn crate::sandbox::Sandbox>>,
79    pub events: Option<crate::event::EventSink>,
80    pub stdout_broadcast: Option<tokio::sync::broadcast::Sender<String>>,
81    pub session_messages: Option<std::sync::Arc<Vec<crate::message::Message>>>,
82    pub session_messages_handle:
83        Option<std::sync::Arc<std::sync::Mutex<Vec<crate::message::Message>>>>,
84    pub session_runtime: Option<std::sync::Arc<crate::session::Session>>,
85    pub compact_lock_handle: Option<std::sync::Arc<tokio::sync::Mutex<()>>>,
86    pub current_node_id: Option<String>,
87    pub stream_tx: Option<tokio::sync::broadcast::Sender<crate::stream::StreamFrame>>,
88    pub read_files:
89        Option<std::sync::Arc<std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>>>,
90    pub approval: Option<std::sync::Arc<crate::session::ApprovalRegistry>>,
91    pub forms: Option<std::sync::Arc<crate::session::FormRegistry>>,
92    pub providers: Option<std::sync::Arc<crate::provider::ProviderRegistry>>,
93    pub session_dir: Option<std::path::PathBuf>,
94    pub data_root: Option<std::path::PathBuf>,
95    pub project_index: Option<std::sync::Arc<crate::index::AnchorIndex>>,
96    pub fs_access: crate::fs_access::FsAccessPolicy,
97    pub lifecycle_fire_tx:
98        Option<tokio::sync::mpsc::UnboundedSender<atman_dsl::ast::LifecycleEvent>>,
99    pub bg_registry: Option<std::sync::Arc<crate::tools::bash_bg::BgRegistry>>,
100    pub term_registry: Option<std::sync::Arc<crate::tools::term::TermRegistry>>,
101    pub session_id: Option<String>,
102    pub trust: Option<crate::trust::TrustConfig>,
103    pub current_model: Option<String>,
104    /// Called when memory.recent_turns is invoked, with the count of returned messages.
105    pub on_memory_recent: Option<std::sync::Arc<dyn Fn(u16) + Send + Sync>>,
106}
107
108impl ToolCtx {
109    pub fn new() -> Self {
110        Self::default()
111    }
112
113    pub fn with_anchors(
114        mut self,
115        turn_id: Option<crate::event::TurnId>,
116        flow_run_id: Option<crate::event::FlowRunId>,
117        event_seq: Option<u64>,
118    ) -> Self {
119        self.turn_id = turn_id;
120        self.flow_run_id = flow_run_id;
121        self.event_seq = event_seq;
122        self
123    }
124
125    pub fn with_registry(mut self, registry: std::sync::Arc<ToolRegistry>) -> Self {
126        self.registry = Some(registry);
127        self
128    }
129
130    pub fn with_sandbox(mut self, sandbox: std::sync::Arc<dyn crate::sandbox::Sandbox>) -> Self {
131        self.sandbox = Some(sandbox);
132        self
133    }
134
135    pub fn with_events(mut self, events: crate::event::EventSink) -> Self {
136        self.events = Some(events);
137        self
138    }
139
140    pub fn with_stdout_broadcast(mut self, tx: tokio::sync::broadcast::Sender<String>) -> Self {
141        self.stdout_broadcast = Some(tx);
142        self
143    }
144
145    pub fn with_session_messages(
146        mut self,
147        msgs: std::sync::Arc<Vec<crate::message::Message>>,
148    ) -> Self {
149        self.session_messages = Some(msgs);
150        self
151    }
152
153    pub fn with_session_messages_handle(
154        mut self,
155        handle: std::sync::Arc<std::sync::Mutex<Vec<crate::message::Message>>>,
156    ) -> Self {
157        self.session_messages_handle = Some(handle);
158        self
159    }
160
161    pub fn with_session_runtime(
162        mut self,
163        session: std::sync::Arc<crate::session::Session>,
164    ) -> Self {
165        self.session_runtime = Some(session);
166        self
167    }
168
169    pub fn with_compact_lock_handle(
170        mut self,
171        handle: std::sync::Arc<tokio::sync::Mutex<()>>,
172    ) -> Self {
173        self.compact_lock_handle = Some(handle);
174        self
175    }
176
177    pub fn with_current_node(mut self, node_id: Option<String>) -> Self {
178        self.current_node_id = node_id;
179        self
180    }
181
182    pub fn with_read_files(
183        mut self,
184        set: std::sync::Arc<std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>>,
185    ) -> Self {
186        self.read_files = Some(set);
187        self
188    }
189
190    pub fn with_providers(
191        mut self,
192        providers: std::sync::Arc<crate::provider::ProviderRegistry>,
193    ) -> Self {
194        self.providers = Some(providers);
195        self
196    }
197
198    pub fn with_session_dir(mut self, dir: std::path::PathBuf) -> Self {
199        self.session_dir = Some(dir);
200        self
201    }
202
203    pub fn with_data_root(mut self, dir: std::path::PathBuf) -> Self {
204        self.data_root = Some(dir);
205        self
206    }
207
208    pub fn with_approval(
209        mut self,
210        approval: std::sync::Arc<crate::session::ApprovalRegistry>,
211    ) -> Self {
212        self.approval = Some(approval);
213        self
214    }
215
216    pub fn with_fs_access(mut self, policy: crate::fs_access::FsAccessPolicy) -> Self {
217        self.fs_access = policy;
218        self
219    }
220
221    pub fn with_lifecycle_fire_tx(
222        mut self,
223        tx: tokio::sync::mpsc::UnboundedSender<atman_dsl::ast::LifecycleEvent>,
224    ) -> Self {
225        self.lifecycle_fire_tx = Some(tx);
226        self
227    }
228
229    pub fn with_forms(mut self, forms: std::sync::Arc<crate::session::FormRegistry>) -> Self {
230        self.forms = Some(forms);
231        self
232    }
233
234    pub fn with_bg_registry(
235        mut self,
236        registry: std::sync::Arc<crate::tools::bash_bg::BgRegistry>,
237    ) -> Self {
238        self.bg_registry = Some(registry);
239        self
240    }
241
242    pub fn with_session_id(mut self, id: impl Into<String>) -> Self {
243        self.session_id = Some(id.into());
244        self
245    }
246
247    pub fn with_trust(mut self, trust: crate::trust::TrustConfig) -> Self {
248        self.trust = Some(trust);
249        self
250    }
251
252    pub fn with_current_model(mut self, model: impl Into<String>) -> Self {
253        self.current_model = Some(model.into());
254        self
255    }
256
257    pub fn with_term_registry(
258        mut self,
259        registry: std::sync::Arc<crate::tools::term::TermRegistry>,
260    ) -> Self {
261        self.term_registry = Some(registry);
262        self
263    }
264
265    pub fn note_read(&self, path: &std::path::Path) {
266        if let Some(set) = &self.read_files
267            && let Ok(mut lock) = set.lock()
268        {
269            lock.insert(path.to_path_buf());
270        }
271    }
272
273    pub fn has_read(&self, path: &std::path::Path) -> bool {
274        self.read_files
275            .as_ref()
276            .and_then(|set| set.lock().ok().map(|lock| lock.contains(path)))
277            .unwrap_or(false)
278    }
279
280    pub fn with_project_index(mut self, idx: std::sync::Arc<crate::index::AnchorIndex>) -> Self {
281        self.project_index = Some(idx);
282        self
283    }
284
285    pub fn with_stream_tx(
286        mut self,
287        tx: tokio::sync::broadcast::Sender<crate::stream::StreamFrame>,
288    ) -> Self {
289        self.stream_tx = Some(tx);
290        self
291    }
292}
293
294pub trait Tool: Send + Sync {
295    fn name(&self) -> &str;
296    fn tier(&self) -> Tier;
297    fn approval_level(&self, _args: &ToolArgs, _ctx: &ToolCtx) -> ApprovalLevel {
298        ApprovalLevel::from_tier(self.tier())
299    }
300    fn cancel_behavior(&self) -> CancelBehavior {
301        CancelBehavior::AbortSafe
302    }
303    fn description(&self) -> Option<&str> {
304        None
305    }
306    fn input_schema(&self) -> serde_json::Value {
307        serde_json::json!({"type": "object"})
308    }
309    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult>;
310    fn preview_call<'a>(
311        &'a self,
312        _args: &'a ToolArgs,
313        _ctx: &'a ToolCtx,
314    ) -> BoxFut<'a, Option<String>> {
315        Box::pin(async { None })
316    }
317}
318
319pub fn tool_spec(tool: &dyn Tool) -> ToolSpec {
320    ToolSpec {
321        name: tool.name().to_string(),
322        description: tool.description().map(str::to_string),
323        input_schema: tool.input_schema(),
324    }
325}
326
327#[derive(Debug, Clone, serde::Serialize)]
328pub struct ToolSpec {
329    pub name: String,
330    #[serde(skip_serializing_if = "Option::is_none")]
331    pub description: Option<String>,
332    pub input_schema: serde_json::Value,
333}
334
335#[derive(Default, Clone)]
336pub struct ToolRegistry {
337    tools: std::sync::Arc<std::sync::RwLock<HashMap<String, std::sync::Arc<dyn Tool>>>>,
338}
339
340impl ToolRegistry {
341    pub fn new() -> Self {
342        Self::default()
343    }
344
345    pub fn register(&self, tool: std::sync::Arc<dyn Tool>) {
346        self.tools
347            .write()
348            .unwrap()
349            .insert(tool.name().to_string(), tool);
350    }
351
352    pub fn get(&self, name: &str) -> Option<std::sync::Arc<dyn Tool>> {
353        self.tools.read().unwrap().get(name).cloned()
354    }
355
356    pub fn has(&self, name: &str) -> bool {
357        self.tools.read().unwrap().contains_key(name)
358    }
359
360    pub fn names(&self) -> Vec<String> {
361        self.tools.read().unwrap().keys().cloned().collect()
362    }
363
364    pub fn iter(&self) -> Vec<(String, std::sync::Arc<dyn Tool>)> {
365        self.tools
366            .read()
367            .unwrap()
368            .iter()
369            .map(|(k, v)| (k.clone(), v.clone()))
370            .collect()
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377
378    #[test]
379    fn approval_level_default_maps_from_tier() {
380        assert_eq!(ApprovalLevel::from_tier(Tier::Zero), ApprovalLevel::Auto);
381        assert_eq!(ApprovalLevel::from_tier(Tier::One), ApprovalLevel::Approve);
382        assert_eq!(ApprovalLevel::from_tier(Tier::Two), ApprovalLevel::Approve);
383        assert_eq!(
384            ApprovalLevel::from_tier(Tier::Three),
385            ApprovalLevel::Dangerous
386        );
387        assert_eq!(
388            ApprovalLevel::from_tier(Tier::Four),
389            ApprovalLevel::Dangerous
390        );
391    }
392
393    #[test]
394    fn approval_level_ordered_auto_lt_approve_lt_dangerous() {
395        assert!(ApprovalLevel::Auto < ApprovalLevel::Approve);
396        assert!(ApprovalLevel::Approve < ApprovalLevel::Dangerous);
397    }
398}
atman_runtime/tool.rs

atman_runtime/
tool.rs
