atman_runtime/
tool.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use tokio_util::sync::CancellationToken;
7
8use crate::error::RuntimeError;
9use crate::value::Value;
10
11// Not `Send`: proc-macro2 spans hold `Rc<()>`. Eval / exec await inline, never spawn.
12pub type BoxFut<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
13
14pub type ToolResult = Result<Value, RuntimeError>;
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum Tier {
18    Zero,
19    One,
20    Two,
21    Three,
22    Four,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
26pub enum ApprovalLevel {
27    Auto,
28    Approve,
29    Dangerous,
30}
31
32impl ApprovalLevel {
33    pub fn from_tier(tier: Tier) -> Self {
34        match tier {
35            Tier::Zero => ApprovalLevel::Auto,
36            Tier::One | Tier::Two => ApprovalLevel::Approve,
37            Tier::Three | Tier::Four => ApprovalLevel::Dangerous,
38        }
39    }
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum CancelBehavior {
44    AbortSafe,
45    Revertible,
46    Atomic,
47    Irreversible,
48}
49
50#[derive(Debug, Default, Clone)]
51pub struct ToolArgs {
52    pub positional: Vec<Value>,
53    pub named: Vec<(String, Value)>,
54}
55
56impl ToolArgs {
57    pub fn positional(&self, index: usize) -> Result<&Value, RuntimeError> {
58        self.positional
59            .get(index)
60            .ok_or_else(|| RuntimeError::MissingArg(format!("positional[{index}]")))
61    }
62
63    pub fn named(&self, name: &str) -> Option<&Value> {
64        self.named.iter().find(|(k, _)| k == name).map(|(_, v)| v)
65    }
66}
67
68#[derive(Clone, Default)]
69pub struct ToolCtx {
70    pub cancel: CancellationToken,
71    pub turn_id: Option<crate::event::TurnId>,
72    pub flow_run_id: Option<crate::event::FlowRunId>,
73    pub event_seq: Option<u64>,
74    pub prompt_resolver: Option<std::sync::Arc<dyn crate::rendezvous::PromptResolver>>,
75    pub registry: Option<std::sync::Arc<ToolRegistry>>,
76    pub sandbox: Option<std::sync::Arc<dyn crate::sandbox::Sandbox>>,
77    pub events: Option<crate::event::EventSink>,
78    pub stdout_broadcast: Option<tokio::sync::broadcast::Sender<String>>,
79    pub session_messages: Option<std::sync::Arc<Vec<crate::message::Message>>>,
80    pub session_messages_handle:
81        Option<std::sync::Arc<std::sync::Mutex<Vec<crate::message::Message>>>>,
82    pub compact_lock_handle: Option<std::sync::Arc<tokio::sync::Mutex<()>>>,
83    pub current_node_id: Option<String>,
84    pub stream_tx: Option<tokio::sync::broadcast::Sender<crate::stream::StreamFrame>>,
85    pub read_files:
86        Option<std::sync::Arc<std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>>>,
87    pub approval: Option<std::sync::Arc<crate::session::ApprovalRegistry>>,
88    pub forms: Option<std::sync::Arc<crate::session::FormRegistry>>,
89    pub providers: Option<std::sync::Arc<crate::provider::ProviderRegistry>>,
90    pub session_dir: Option<std::path::PathBuf>,
91    pub data_root: Option<std::path::PathBuf>,
92    pub project_index: Option<std::sync::Arc<crate::index::AnchorIndex>>,
93    pub fs_access: crate::fs_access::FsAccessPolicy,
94    pub lifecycle_fire_tx:
95        Option<tokio::sync::mpsc::UnboundedSender<atman_dsl::ast::LifecycleEvent>>,
96    pub bg_registry: Option<std::sync::Arc<crate::tools::bash_bg::BgRegistry>>,
97    pub term_registry: Option<std::sync::Arc<crate::tools::term::TermRegistry>>,
98    pub session_id: Option<String>,
99    pub trust: Option<crate::trust::TrustConfig>,
100    pub current_model: Option<String>,
101}
102
103impl ToolCtx {
104    pub fn new() -> Self {
105        Self::default()
106    }
107
108    pub fn with_anchors(
109        mut self,
110        turn_id: Option<crate::event::TurnId>,
111        flow_run_id: Option<crate::event::FlowRunId>,
112        event_seq: Option<u64>,
113    ) -> Self {
114        self.turn_id = turn_id;
115        self.flow_run_id = flow_run_id;
116        self.event_seq = event_seq;
117        self
118    }
119
120    pub fn with_registry(mut self, registry: std::sync::Arc<ToolRegistry>) -> Self {
121        self.registry = Some(registry);
122        self
123    }
124
125    pub fn with_sandbox(mut self, sandbox: std::sync::Arc<dyn crate::sandbox::Sandbox>) -> Self {
126        self.sandbox = Some(sandbox);
127        self
128    }
129
130    pub fn with_events(mut self, events: crate::event::EventSink) -> Self {
131        self.events = Some(events);
132        self
133    }
134
135    pub fn with_stdout_broadcast(mut self, tx: tokio::sync::broadcast::Sender<String>) -> Self {
136        self.stdout_broadcast = Some(tx);
137        self
138    }
139
140    pub fn with_session_messages(
141        mut self,
142        msgs: std::sync::Arc<Vec<crate::message::Message>>,
143    ) -> Self {
144        self.session_messages = Some(msgs);
145        self
146    }
147
148    pub fn with_session_messages_handle(
149        mut self,
150        handle: std::sync::Arc<std::sync::Mutex<Vec<crate::message::Message>>>,
151    ) -> Self {
152        self.session_messages_handle = Some(handle);
153        self
154    }
155
156    pub fn with_compact_lock_handle(
157        mut self,
158        handle: std::sync::Arc<tokio::sync::Mutex<()>>,
159    ) -> Self {
160        self.compact_lock_handle = Some(handle);
161        self
162    }
163
164    pub fn with_current_node(mut self, node_id: Option<String>) -> Self {
165        self.current_node_id = node_id;
166        self
167    }
168
169    pub fn with_read_files(
170        mut self,
171        set: std::sync::Arc<std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>>,
172    ) -> Self {
173        self.read_files = Some(set);
174        self
175    }
176
177    pub fn with_providers(
178        mut self,
179        providers: std::sync::Arc<crate::provider::ProviderRegistry>,
180    ) -> Self {
181        self.providers = Some(providers);
182        self
183    }
184
185    pub fn with_session_dir(mut self, dir: std::path::PathBuf) -> Self {
186        self.session_dir = Some(dir);
187        self
188    }
189
190    pub fn with_data_root(mut self, dir: std::path::PathBuf) -> Self {
191        self.data_root = Some(dir);
192        self
193    }
194
195    pub fn with_approval(
196        mut self,
197        approval: std::sync::Arc<crate::session::ApprovalRegistry>,
198    ) -> Self {
199        self.approval = Some(approval);
200        self
201    }
202
203    pub fn with_fs_access(mut self, policy: crate::fs_access::FsAccessPolicy) -> Self {
204        self.fs_access = policy;
205        self
206    }
207
208    pub fn with_lifecycle_fire_tx(
209        mut self,
210        tx: tokio::sync::mpsc::UnboundedSender<atman_dsl::ast::LifecycleEvent>,
211    ) -> Self {
212        self.lifecycle_fire_tx = Some(tx);
213        self
214    }
215
216    pub fn with_forms(mut self, forms: std::sync::Arc<crate::session::FormRegistry>) -> Self {
217        self.forms = Some(forms);
218        self
219    }
220
221    pub fn with_bg_registry(
222        mut self,
223        registry: std::sync::Arc<crate::tools::bash_bg::BgRegistry>,
224    ) -> Self {
225        self.bg_registry = Some(registry);
226        self
227    }
228
229    pub fn with_session_id(mut self, id: impl Into<String>) -> Self {
230        self.session_id = Some(id.into());
231        self
232    }
233
234    pub fn with_trust(mut self, trust: crate::trust::TrustConfig) -> Self {
235        self.trust = Some(trust);
236        self
237    }
238
239    pub fn with_current_model(mut self, model: impl Into<String>) -> Self {
240        self.current_model = Some(model.into());
241        self
242    }
243
244    pub fn with_term_registry(
245        mut self,
246        registry: std::sync::Arc<crate::tools::term::TermRegistry>,
247    ) -> Self {
248        self.term_registry = Some(registry);
249        self
250    }
251
252    pub fn note_read(&self, path: &std::path::Path) {
253        if let Some(set) = &self.read_files
254            && let Ok(mut lock) = set.lock()
255        {
256            lock.insert(path.to_path_buf());
257        }
258    }
259
260    pub fn has_read(&self, path: &std::path::Path) -> bool {
261        self.read_files
262            .as_ref()
263            .and_then(|set| set.lock().ok().map(|lock| lock.contains(path)))
264            .unwrap_or(false)
265    }
266
267    pub fn with_project_index(mut self, idx: std::sync::Arc<crate::index::AnchorIndex>) -> Self {
268        self.project_index = Some(idx);
269        self
270    }
271
272    pub fn with_stream_tx(
273        mut self,
274        tx: tokio::sync::broadcast::Sender<crate::stream::StreamFrame>,
275    ) -> Self {
276        self.stream_tx = Some(tx);
277        self
278    }
279}
280
281pub trait Tool: Send + Sync {
282    fn name(&self) -> &str;
283    fn tier(&self) -> Tier;
284    fn approval_level(&self, _args: &ToolArgs, _ctx: &ToolCtx) -> ApprovalLevel {
285        ApprovalLevel::from_tier(self.tier())
286    }
287    fn cancel_behavior(&self) -> CancelBehavior {
288        CancelBehavior::AbortSafe
289    }
290    fn description(&self) -> Option<&str> {
291        None
292    }
293    fn input_schema(&self) -> serde_json::Value {
294        serde_json::json!({"type": "object"})
295    }
296    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult>;
297    fn preview_call<'a>(
298        &'a self,
299        _args: &'a ToolArgs,
300        _ctx: &'a ToolCtx,
301    ) -> BoxFut<'a, Option<String>> {
302        Box::pin(async { None })
303    }
304}
305
306pub fn tool_spec(tool: &dyn Tool) -> ToolSpec {
307    ToolSpec {
308        name: tool.name().to_string(),
309        description: tool.description().map(str::to_string),
310        input_schema: tool.input_schema(),
311    }
312}
313
314#[derive(Debug, Clone, serde::Serialize)]
315pub struct ToolSpec {
316    pub name: String,
317    #[serde(skip_serializing_if = "Option::is_none")]
318    pub description: Option<String>,
319    pub input_schema: serde_json::Value,
320}
321
322#[derive(Default, Clone)]
323pub struct ToolRegistry {
324    tools: HashMap<String, Arc<dyn Tool>>,
325}
326
327impl ToolRegistry {
328    pub fn new() -> Self {
329        Self::default()
330    }
331
332    pub fn register(&mut self, tool: Arc<dyn Tool>) {
333        self.tools.insert(tool.name().to_string(), tool);
334    }
335
336    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
337        self.tools.get(name).cloned()
338    }
339
340    pub fn has(&self, name: &str) -> bool {
341        self.tools.contains_key(name)
342    }
343
344    pub fn names(&self) -> impl Iterator<Item = &str> {
345        self.tools.keys().map(String::as_str)
346    }
347
348    pub fn iter(&self) -> impl Iterator<Item = (&String, &Arc<dyn Tool>)> {
349        self.tools.iter()
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356
357    #[test]
358    fn approval_level_default_maps_from_tier() {
359        assert_eq!(ApprovalLevel::from_tier(Tier::Zero), ApprovalLevel::Auto);
360        assert_eq!(ApprovalLevel::from_tier(Tier::One), ApprovalLevel::Approve);
361        assert_eq!(ApprovalLevel::from_tier(Tier::Two), ApprovalLevel::Approve);
362        assert_eq!(
363            ApprovalLevel::from_tier(Tier::Three),
364            ApprovalLevel::Dangerous
365        );
366        assert_eq!(
367            ApprovalLevel::from_tier(Tier::Four),
368            ApprovalLevel::Dangerous
369        );
370    }
371
372    #[test]
373    fn approval_level_ordered_auto_lt_approve_lt_dangerous() {
374        assert!(ApprovalLevel::Auto < ApprovalLevel::Approve);
375        assert!(ApprovalLevel::Approve < ApprovalLevel::Dangerous);
376    }
377}
atman_runtime/tool.rs

atman_runtime/
tool.rs
