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/// Sendable boxed future used by provider and tool traits.
11pub type BoxFut<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
12
13pub type ToolResult = Result<Value, RuntimeError>;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
16#[serde(rename_all = "snake_case")]
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, Clone, Copy, PartialEq, Eq)]
51pub enum InvocationPlane {
52    Ordinary,
53    PermissionControl,
54}
55
56#[derive(Debug, Default, Clone)]
57pub struct ToolArgs {
58    pub positional: Vec<Value>,
59    pub named: Vec<(String, Value)>,
60}
61
62impl ToolArgs {
63    pub fn positional(&self, index: usize) -> Result<&Value, RuntimeError> {
64        self.positional
65            .get(index)
66            .ok_or_else(|| RuntimeError::MissingArg(format!("positional[{index}]")))
67    }
68
69    pub fn named(&self, name: &str) -> Option<&Value> {
70        self.named.iter().find(|(k, _)| k == name).map(|(_, v)| v)
71    }
72}
73
74#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
75pub enum HistorySegment {
76    #[default]
77    Root,
78    Spawned,
79}
80
81#[derive(Clone, Copy, Debug, PartialEq, Eq)]
82pub enum PathOrigin {
83    Omitted,
84    Relative,
85    ExplicitInside,
86    ExplicitExternal,
87    Unbound,
88}
89
90#[derive(Clone, Debug, PartialEq, Eq)]
91pub struct ResolvedPath {
92    pub path: std::path::PathBuf,
93    pub origin: PathOrigin,
94}
95
96#[derive(Clone, Default)]
97pub struct ToolCtx {
98    pub cancel: CancellationToken,
99    pub turn_id: Option<crate::event::TurnId>,
100    pub flow_run_id: Option<crate::event::FlowRunId>,
101    pub history_segment: HistorySegment,
102    pub event_seq: Option<u64>,
103    pub prompt_resolver: Option<std::sync::Arc<dyn crate::rendezvous::PromptResolver>>,
104    pub registry: Option<std::sync::Arc<ToolRegistry>>,
105    pub sandbox: Option<std::sync::Arc<dyn crate::sandbox::Sandbox>>,
106    pub events: Option<crate::event::EventSink>,
107    pub stdout_broadcast: Option<tokio::sync::broadcast::Sender<String>>,
108    pub session_messages: Option<std::sync::Arc<Vec<crate::message::Message>>>,
109    pub session_messages_handle:
110        Option<std::sync::Arc<std::sync::Mutex<Vec<crate::message::Message>>>>,
111    pub session_runtime: Option<std::sync::Arc<crate::session::Session>>,
112    pub(crate) deferred_input_session: Option<std::sync::Arc<crate::session::Session>>,
113    pub compact_lock_handle: Option<std::sync::Arc<tokio::sync::Mutex<()>>>,
114    pub(crate) context_epoch_handle: Option<std::sync::Arc<std::sync::atomic::AtomicU64>>,
115    pub(crate) context_prefix_tracker:
116        Option<std::sync::Arc<std::sync::Mutex<crate::context_plan::ContextPrefixTracker>>>,
117    pub current_node_id: Option<String>,
118    pub stream_tx: Option<tokio::sync::broadcast::Sender<crate::stream::StreamFrame>>,
119    pub read_files:
120        Option<std::sync::Arc<std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>>>,
121    pub approval: Option<std::sync::Arc<crate::session::ApprovalRegistry>>,
122    pub permission_broker: Option<std::sync::Arc<crate::permission::PermissionBroker>>,
123    pub(crate) invocation_authorization: Option<crate::permission::InvocationAuthorization>,
124    pub forms: Option<std::sync::Arc<crate::session::FormRegistry>>,
125    pub providers: Option<std::sync::Arc<crate::provider::ProviderRegistry>>,
126    pub session_dir: Option<std::path::PathBuf>,
127    pub output_store: Option<std::sync::Arc<crate::tools::tool_output::OutputStore>>,
128    pub data_root: Option<std::path::PathBuf>,
129    pub project_index: Option<std::sync::Arc<crate::index::AnchorIndex>>,
130    pub fs_access: crate::fs_access::FsAccessPolicy,
131    pub workspace: Option<crate::git_workspace::WorkspaceBinding>,
132    pub flow_workspace_service: Option<std::sync::Arc<crate::flow_workspace::FlowWorkspaceService>>,
133    pub lifecycle_fire_tx:
134        Option<tokio::sync::mpsc::UnboundedSender<atman_dsl::ast::LifecycleEvent>>,
135    pub bg_registry: Option<std::sync::Arc<crate::tools::bash_bg::BgRegistry>>,
136    pub term_registry: Option<std::sync::Arc<crate::tools::term::TermRegistry>>,
137    pub watch_hub: Option<std::sync::Arc<crate::watch::WatchHub>>,
138    pub flow_registry: Option<std::sync::Arc<crate::tools::agent_ctrl::FlowRegistry>>,
139    pub flow_identity: Option<std::sync::Arc<crate::flow_authority::FlowIdentity>>,
140    pub task_registry: Option<crate::task_registry::TaskRegistry>,
141    pub session_id: Option<String>,
142    pub trust: Option<crate::trust::TrustConfig>,
143    pub safety: Option<crate::safety::SafetyConfig>,
144    pub current_model: Option<String>,
145    pub call_intent: Option<crate::message::ToolCallIntent>,
146    pub tool_use_id: Option<String>,
147    pub(crate) model_tool_exposures: Option<ToolExposureRegistry>,
148    pub(crate) invocation_env: crate::invocation_env::InvocationEnv,
149    pub watch_rules: Option<crate::streaming::WatchRules>,
150    /// Called when memory.recent_turns is invoked, with the count of returned messages.
151    pub on_memory_recent: Option<std::sync::Arc<dyn Fn(u16) + Send + Sync>>,
152    pub history_store: Option<std::sync::Arc<dyn crate::history_store::HistoryStore>>,
153    pub agent_entry: Option<std::sync::Arc<crate::tools::agent_ctrl::FlowEntry>>,
154    pub tool_output_budget: crate::tools::tool_output::ToolOutputBudget,
155}
156
157#[derive(Clone, Default)]
158pub(crate) struct ToolExposureRegistry {
159    pending: std::sync::Arc<std::sync::Mutex<std::collections::HashMap<ToolExposureKey, usize>>>,
160}
161
162#[derive(Clone, Debug, PartialEq, Eq, Hash)]
163struct ToolExposureKey {
164    flow_run_id: Option<crate::event::FlowRunId>,
165    tool_use_id: String,
166    tool_name: String,
167}
168
169impl ToolExposureRegistry {
170    pub(crate) fn register_response<I, S>(
171        &self,
172        flow_run_id: Option<&crate::event::FlowRunId>,
173        message: &crate::message::Message,
174        exposed_names: I,
175    ) where
176        I: IntoIterator<Item = S>,
177        S: AsRef<str>,
178    {
179        let exposed_names: std::collections::HashSet<String> = exposed_names
180            .into_iter()
181            .map(|name| name.as_ref().to_string())
182            .collect();
183        let mut id_counts = std::collections::HashMap::new();
184        for part in &message.parts {
185            if let crate::message::MessagePart::ToolUse { id, .. } = part {
186                *id_counts.entry(id.as_str()).or_insert(0usize) += 1;
187            }
188        }
189
190        let mut pending = self.pending.lock().unwrap();
191        for part in &message.parts {
192            let crate::message::MessagePart::ToolUse { id, name, .. } = part else {
193                continue;
194            };
195            if id_counts.get(id.as_str()) != Some(&1) || !exposed_names.contains(name) {
196                continue;
197            }
198            *pending
199                .entry(ToolExposureKey {
200                    flow_run_id: flow_run_id.cloned(),
201                    tool_use_id: id.clone(),
202                    tool_name: name.clone(),
203                })
204                .or_insert(0) += 1;
205        }
206    }
207
208    pub(crate) fn claim(
209        &self,
210        flow_run_id: Option<&crate::event::FlowRunId>,
211        tool_use_id: &str,
212        tool_name: &str,
213    ) -> bool {
214        let key = ToolExposureKey {
215            flow_run_id: flow_run_id.cloned(),
216            tool_use_id: tool_use_id.to_string(),
217            tool_name: tool_name.to_string(),
218        };
219        let mut pending = self.pending.lock().unwrap();
220        let Some(count) = pending.get_mut(&key) else {
221            return false;
222        };
223        *count -= 1;
224        if *count == 0 {
225            pending.remove(&key);
226        }
227        true
228    }
229}
230
231impl ToolCtx {
232    pub fn new() -> Self {
233        Self::default()
234    }
235
236    pub fn with_anchors(
237        mut self,
238        turn_id: Option<crate::event::TurnId>,
239        flow_run_id: Option<crate::event::FlowRunId>,
240        event_seq: Option<u64>,
241    ) -> Self {
242        self.turn_id = turn_id;
243        self.flow_run_id = flow_run_id;
244        self.event_seq = event_seq;
245        self
246    }
247
248    pub fn with_history_segment(mut self, segment: HistorySegment) -> Self {
249        self.history_segment = segment;
250        self
251    }
252
253    pub fn with_call_intent(mut self, call_intent: Option<crate::message::ToolCallIntent>) -> Self {
254        self.call_intent = call_intent;
255        self
256    }
257
258    pub fn with_tool_use_id(mut self, tool_use_id: impl Into<String>) -> Self {
259        self.tool_use_id = Some(tool_use_id.into());
260        self
261    }
262
263    pub(crate) fn with_invocation_env(
264        mut self,
265        invocation_env: crate::invocation_env::InvocationEnv,
266    ) -> Self {
267        self.invocation_env = invocation_env;
268        self
269    }
270
271    pub fn message_flow_run_id(&self) -> Option<crate::event::FlowRunId> {
272        match self.history_segment {
273            HistorySegment::Root => None,
274            HistorySegment::Spawned => self.flow_run_id.clone(),
275        }
276    }
277
278    pub fn with_registry(mut self, registry: std::sync::Arc<ToolRegistry>) -> Self {
279        self.registry = Some(registry);
280        self
281    }
282
283    pub fn with_sandbox(mut self, sandbox: std::sync::Arc<dyn crate::sandbox::Sandbox>) -> Self {
284        self.sandbox = Some(sandbox);
285        self
286    }
287
288    pub fn with_events(mut self, events: crate::event::EventSink) -> Self {
289        self.events = Some(events);
290        self
291    }
292
293    pub fn with_stdout_broadcast(mut self, tx: tokio::sync::broadcast::Sender<String>) -> Self {
294        self.stdout_broadcast = Some(tx);
295        self
296    }
297
298    pub fn with_session_messages(
299        mut self,
300        msgs: std::sync::Arc<Vec<crate::message::Message>>,
301    ) -> Self {
302        self.session_messages = Some(msgs);
303        self
304    }
305
306    pub fn with_session_messages_handle(
307        mut self,
308        handle: std::sync::Arc<std::sync::Mutex<Vec<crate::message::Message>>>,
309    ) -> Self {
310        self.session_messages_handle = Some(handle);
311        self
312    }
313
314    pub fn with_session_runtime(
315        mut self,
316        session: std::sync::Arc<crate::session::Session>,
317    ) -> Self {
318        self.deferred_input_session = Some(session.clone());
319        self.session_runtime = Some(session);
320        self
321    }
322
323    pub fn with_compact_lock_handle(
324        mut self,
325        handle: std::sync::Arc<tokio::sync::Mutex<()>>,
326    ) -> Self {
327        self.compact_lock_handle = Some(handle);
328        self
329    }
330
331    pub(crate) fn context_epoch_seed(&self) -> Option<String> {
332        self.context_epoch_handle.as_ref().map(|epoch| {
333            format!(
334                "generation:{}",
335                epoch.load(std::sync::atomic::Ordering::Relaxed)
336            )
337        })
338    }
339
340    pub(crate) fn advance_context_epoch(&self) {
341        if let Some(epoch) = self.context_epoch_handle.as_ref() {
342            epoch.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
343        }
344    }
345
346    pub fn with_current_node(mut self, node_id: Option<String>) -> Self {
347        self.current_node_id = node_id;
348        self
349    }
350
351    pub fn with_read_files(
352        mut self,
353        set: std::sync::Arc<std::sync::Mutex<std::collections::HashSet<std::path::PathBuf>>>,
354    ) -> Self {
355        self.read_files = Some(set);
356        self
357    }
358
359    pub fn with_providers(
360        mut self,
361        providers: std::sync::Arc<crate::provider::ProviderRegistry>,
362    ) -> Self {
363        self.providers = Some(providers);
364        self
365    }
366
367    pub fn with_session_dir(mut self, dir: std::path::PathBuf) -> Self {
368        self.output_store = Some(std::sync::Arc::new(
369            crate::tools::tool_output::OutputStore::at(dir.clone()),
370        ));
371        self.session_dir = Some(dir);
372        self
373    }
374
375    pub fn with_output_store(
376        mut self,
377        store: std::sync::Arc<crate::tools::tool_output::OutputStore>,
378    ) -> Self {
379        self.output_store = Some(store);
380        self
381    }
382
383    pub fn with_data_root(mut self, dir: std::path::PathBuf) -> Self {
384        self.data_root = Some(dir);
385        self
386    }
387
388    pub fn with_approval(
389        mut self,
390        approval: std::sync::Arc<crate::session::ApprovalRegistry>,
391    ) -> Self {
392        self.approval = Some(approval);
393        self
394    }
395
396    pub fn with_permission_broker(
397        mut self,
398        broker: std::sync::Arc<crate::permission::PermissionBroker>,
399    ) -> Self {
400        self.permission_broker = Some(broker);
401        self
402    }
403
404    // A per-call clone prevents concurrent dispatch entries from sharing permits.
405    pub(crate) fn authorized_for(
406        &self,
407        authorization: crate::permission::InvocationAuthorization,
408    ) -> Self {
409        let mut ctx = self.clone();
410        ctx.invocation_authorization = Some(authorization);
411        ctx
412    }
413
414    pub(crate) fn invocation_authorization(
415        &self,
416    ) -> Option<&crate::permission::InvocationAuthorization> {
417        self.invocation_authorization.as_ref()
418    }
419
420    pub(crate) fn invocation_authorization_for(
421        &self,
422        tool_name: &str,
423    ) -> Result<&crate::permission::InvocationAuthorization, crate::error::RuntimeError> {
424        let authorization = self.invocation_authorization.as_ref().ok_or_else(|| {
425            crate::error::RuntimeError::ToolFailed(format!(
426                "{tool_name}: missing invocation authorization"
427            ))
428        })?;
429        if authorization.tool_name() != tool_name {
430            return Err(crate::error::RuntimeError::ToolFailed(format!(
431                "{tool_name}: invocation authorization belongs to {}",
432                authorization.tool_name()
433            )));
434        }
435        Ok(authorization)
436    }
437
438    pub fn with_fs_access(mut self, policy: crate::fs_access::FsAccessPolicy) -> Self {
439        self.fs_access = policy;
440        self
441    }
442
443    pub fn with_workspace(mut self, binding: crate::git_workspace::WorkspaceBinding) -> Self {
444        self.fs_access.workspace = Some(binding.path.clone());
445        self.workspace = Some(binding);
446        self
447    }
448
449    pub fn with_flow_workspace_service(
450        mut self,
451        service: std::sync::Arc<crate::flow_workspace::FlowWorkspaceService>,
452    ) -> Self {
453        self.flow_workspace_service = Some(service);
454        self
455    }
456
457    pub fn resolve_cwd(
458        &self,
459        explicit: Option<&std::path::Path>,
460    ) -> Result<std::path::PathBuf, RuntimeError> {
461        Ok(self.resolve_cwd_with_origin(explicit)?.path)
462    }
463
464    pub fn resolve_cwd_with_origin(
465        &self,
466        explicit: Option<&std::path::Path>,
467    ) -> Result<ResolvedPath, RuntimeError> {
468        match explicit {
469            Some(path) => self.resolve_path_with_origin(path),
470            None => {
471                let mut resolved = self.resolve_path_with_origin(std::path::Path::new("."))?;
472                resolved.origin = if self.workspace.is_some() {
473                    PathOrigin::Omitted
474                } else {
475                    PathOrigin::Unbound
476                };
477                Ok(resolved)
478            }
479        }
480    }
481
482    pub fn resolve_path(&self, path: &std::path::Path) -> Result<std::path::PathBuf, RuntimeError> {
483        Ok(self.resolve_path_with_origin(path)?.path)
484    }
485
486    pub fn resolve_path_with_origin(
487        &self,
488        path: &std::path::Path,
489    ) -> Result<ResolvedPath, RuntimeError> {
490        let Some(binding) = &self.workspace else {
491            let base = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
492            let candidate = if path.is_absolute() {
493                path.to_path_buf()
494            } else {
495                base.join(path)
496            };
497            return Ok(ResolvedPath {
498                path: crate::fs_access::canonicalize_stable(&candidate),
499                origin: PathOrigin::Unbound,
500            });
501        };
502        let root = crate::fs_access::canonicalize_stable(&binding.path);
503        let candidate = if path.is_absolute() {
504            path.to_path_buf()
505        } else {
506            binding.path.join(path)
507        };
508        let resolved = crate::fs_access::canonicalize_stable(&candidate);
509        if path.is_absolute() {
510            return Ok(ResolvedPath {
511                origin: if resolved.starts_with(&root) {
512                    PathOrigin::ExplicitInside
513                } else {
514                    PathOrigin::ExplicitExternal
515                },
516                path: resolved,
517            });
518        }
519        if !resolved.starts_with(&root) {
520            return Err(RuntimeError::ToolFailed(format!(
521                "managed workspace path {} escapes workspace root {}",
522                path.display(),
523                binding.path.display()
524            )));
525        }
526        Ok(ResolvedPath {
527            path: resolved,
528            origin: PathOrigin::Relative,
529        })
530    }
531
532    pub fn with_lifecycle_fire_tx(
533        mut self,
534        tx: tokio::sync::mpsc::UnboundedSender<atman_dsl::ast::LifecycleEvent>,
535    ) -> Self {
536        self.lifecycle_fire_tx = Some(tx);
537        self
538    }
539
540    pub fn with_forms(mut self, forms: std::sync::Arc<crate::session::FormRegistry>) -> Self {
541        self.forms = Some(forms);
542        self
543    }
544
545    pub fn with_bg_registry(
546        mut self,
547        registry: std::sync::Arc<crate::tools::bash_bg::BgRegistry>,
548    ) -> Self {
549        self.bg_registry = Some(registry);
550        self
551    }
552
553    pub fn with_session_id(mut self, id: impl Into<String>) -> Self {
554        self.session_id = Some(id.into());
555        self
556    }
557
558    pub fn with_trust(mut self, trust: crate::trust::TrustConfig) -> Self {
559        self.trust = Some(trust);
560        self
561    }
562
563    /// Freezes the active trust policy for one tool invocation. The broker uses
564    /// this snapshot to mint an invocation authorization carrying the selected
565    /// execution boundary.
566    pub fn for_tool_invocation(mut self, _tier: Tier) -> Self {
567        if let Some(session) = self.session_runtime.as_ref() {
568            self.trust = Some(session.trust_config());
569        }
570        self
571    }
572
573    pub fn with_safety(mut self, safety: crate::safety::SafetyConfig) -> Self {
574        self.safety = Some(safety);
575        self
576    }
577
578    pub fn with_current_model(mut self, model: impl Into<String>) -> Self {
579        self.current_model = Some(model.into());
580        self
581    }
582
583    pub fn with_watch_rules(mut self, rules: crate::streaming::WatchRules) -> Self {
584        self.watch_rules = Some(rules);
585        self
586    }
587
588    pub fn with_term_registry(
589        mut self,
590        registry: std::sync::Arc<crate::tools::term::TermRegistry>,
591    ) -> Self {
592        self.term_registry = Some(registry);
593        self
594    }
595
596    pub fn with_watch_hub(mut self, hub: std::sync::Arc<crate::watch::WatchHub>) -> Self {
597        self.watch_hub = Some(hub);
598        self
599    }
600
601    pub fn with_flow_registry(
602        mut self,
603        registry: std::sync::Arc<crate::tools::agent_ctrl::FlowRegistry>,
604    ) -> Self {
605        self.flow_registry = Some(registry);
606        self
607    }
608
609    pub fn with_task_registry(mut self, registry: crate::task_registry::TaskRegistry) -> Self {
610        self.task_registry = Some(registry);
611        self
612    }
613
614    pub fn note_read(&self, path: &std::path::Path) {
615        if let Some(set) = &self.read_files
616            && let Ok(mut lock) = set.lock()
617        {
618            lock.insert(path.to_path_buf());
619        }
620    }
621
622    pub fn has_read(&self, path: &std::path::Path) -> bool {
623        self.read_files
624            .as_ref()
625            .and_then(|set| set.lock().ok().map(|lock| lock.contains(path)))
626            .unwrap_or(false)
627    }
628
629    pub fn with_project_index(mut self, idx: std::sync::Arc<crate::index::AnchorIndex>) -> Self {
630        self.project_index = Some(idx);
631        self
632    }
633
634    pub fn with_history_store(
635        mut self,
636        store: std::sync::Arc<dyn crate::history_store::HistoryStore>,
637    ) -> Self {
638        self.history_store = Some(store);
639        self
640    }
641
642    pub fn with_agent_entry(
643        mut self,
644        entry: std::sync::Arc<crate::tools::agent_ctrl::FlowEntry>,
645    ) -> Self {
646        self.agent_entry = Some(entry);
647        self
648    }
649
650    pub fn with_stream_tx(
651        mut self,
652        tx: tokio::sync::broadcast::Sender<crate::stream::StreamFrame>,
653    ) -> Self {
654        self.stream_tx = Some(tx);
655        self
656    }
657}
658
659pub trait Tool: Send + Sync {
660    fn name(&self) -> &str;
661    fn tier(&self) -> Tier;
662    fn requires_call_intent(&self) -> bool {
663        true
664    }
665    fn invocation_plane(&self) -> InvocationPlane {
666        InvocationPlane::Ordinary
667    }
668    fn approval_level(&self, _args: &ToolArgs, _ctx: &ToolCtx) -> ApprovalLevel {
669        ApprovalLevel::from_tier(self.tier())
670    }
671    fn cancel_behavior(&self) -> CancelBehavior {
672        CancelBehavior::AbortSafe
673    }
674    fn description(&self) -> Option<&str> {
675        None
676    }
677    fn input_schema(&self) -> serde_json::Value {
678        serde_json::json!({"type": "object"})
679    }
680    // Defaulting to none avoids treating arbitrary command or URL arguments as paths.
681    fn invocation_provenance(
682        &self,
683        _args: &ToolArgs,
684        _ctx: &ToolCtx,
685    ) -> Result<crate::permission::ResourceProvenance, RuntimeError> {
686        Ok(crate::permission::ResourceProvenance::none())
687    }
688    fn model_followups(&self, _result: &Value, _ctx: &ToolCtx) -> Vec<crate::message::Message> {
689        Vec::new()
690    }
691    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult>;
692    fn preview_call<'a>(
693        &'a self,
694        _args: &'a ToolArgs,
695        _ctx: &'a ToolCtx,
696    ) -> BoxFut<'a, Option<String>> {
697        Box::pin(async { None })
698    }
699}
700
701pub fn tool_spec(tool: &dyn Tool) -> ToolSpec {
702    let mut input_schema = tool.input_schema();
703    if tool.requires_call_intent() {
704        decorate_tool_input_schema(&mut input_schema);
705    }
706    canonicalize_json_object_keys(&mut input_schema);
707    ToolSpec {
708        name: tool.name().to_string(),
709        description: tool.description().map(str::to_string),
710        input_schema,
711    }
712}
713
714fn canonicalize_json_object_keys(value: &mut serde_json::Value) {
715    match value {
716        serde_json::Value::Object(object) => {
717            let mut entries: Vec<_> = std::mem::take(object).into_iter().collect();
718            for (_, value) in &mut entries {
719                canonicalize_json_object_keys(value);
720            }
721            entries.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
722            object.extend(entries);
723        }
724        serde_json::Value::Array(items) => {
725            for item in items {
726                canonicalize_json_object_keys(item);
727            }
728        }
729        _ => {}
730    }
731}
732
733fn tool_call_intent_schema() -> serde_json::Value {
734    serde_json::json!({
735        "type": "string",
736        "minLength": 1,
737        "maxLength": crate::message::TOOL_CALL_INTENT_MAX_CHARS,
738        "pattern": "\\S"
739    })
740}
741
742fn decorate_tool_input_schema(schema: &mut serde_json::Value) {
743    let Some(root) = schema.as_object_mut() else {
744        return;
745    };
746    if root.get("type").and_then(serde_json::Value::as_str) != Some("object") {
747        return;
748    }
749    if root
750        .get("properties")
751        .and_then(serde_json::Value::as_object)
752        .is_some_and(|properties| properties.contains_key(crate::message::TOOL_CALL_INTENT_FIELD))
753    {
754        return;
755    }
756    {
757        let properties = root
758            .entry("properties")
759            .or_insert_with(|| serde_json::Value::Object(Default::default()));
760        let Some(properties) = properties.as_object_mut() else {
761            return;
762        };
763        properties.insert(
764            crate::message::TOOL_CALL_INTENT_FIELD.into(),
765            tool_call_intent_schema(),
766        );
767    }
768    let required = root
769        .entry("required")
770        .or_insert_with(|| serde_json::Value::Array(Vec::new()));
771    let Some(required) = required.as_array_mut() else {
772        return;
773    };
774    if !required
775        .iter()
776        .any(|field| field.as_str() == Some(crate::message::TOOL_CALL_INTENT_FIELD))
777    {
778        required.push(crate::message::TOOL_CALL_INTENT_FIELD.into());
779    }
780}
781
782fn tool_spec_call_intent_support(tool_name: &str, tools: &[ToolSpec]) -> Option<bool> {
783    tools
784        .iter()
785        .find(|tool| tool.name == tool_name)
786        .map(|tool| {
787            tool.input_schema
788                .get("properties")
789                .and_then(serde_json::Value::as_object)
790                .and_then(|properties| properties.get(crate::message::TOOL_CALL_INTENT_FIELD))
791                == Some(&tool_call_intent_schema())
792        })
793}
794
795pub fn tool_spec_supports_call_intent(tool_name: &str, tools: &[ToolSpec]) -> bool {
796    tool_spec_call_intent_support(tool_name, tools) == Some(true)
797}
798
799pub fn tool_spec_blocks_call_intent(tool_name: &str, tools: &[ToolSpec]) -> bool {
800    tool_spec_call_intent_support(tool_name, tools) == Some(false)
801}
802
803pub fn tool_schema_uses_call_intent_field(schema: &serde_json::Value) -> bool {
804    schema
805        .get("properties")
806        .and_then(serde_json::Value::as_object)
807        .is_some_and(|properties| properties.contains_key(crate::message::TOOL_CALL_INTENT_FIELD))
808}
809
810#[derive(Debug, Clone, serde::Serialize)]
811pub struct ToolSpec {
812    pub name: String,
813    #[serde(skip_serializing_if = "Option::is_none")]
814    pub description: Option<String>,
815    pub input_schema: serde_json::Value,
816}
817
818#[derive(Default, Clone)]
819pub struct ToolRegistry {
820    tools: std::sync::Arc<std::sync::RwLock<HashMap<String, std::sync::Arc<dyn Tool>>>>,
821}
822
823impl ToolRegistry {
824    pub fn new() -> Self {
825        Self::default()
826    }
827
828    pub fn register(&self, tool: std::sync::Arc<dyn Tool>) {
829        assert!(
830            !crate::eval::is_evaluator_intrinsic(tool.name()),
831            "tool name `{}` is reserved for an evaluator intrinsic",
832            tool.name()
833        );
834        self.tools
835            .write()
836            .unwrap()
837            .insert(tool.name().to_string(), tool);
838    }
839
840    pub fn get(&self, name: &str) -> Option<std::sync::Arc<dyn Tool>> {
841        self.tools.read().unwrap().get(name).cloned()
842    }
843
844    pub fn has(&self, name: &str) -> bool {
845        self.tools.read().unwrap().contains_key(name)
846    }
847
848    pub fn names(&self) -> Vec<String> {
849        self.tools.read().unwrap().keys().cloned().collect()
850    }
851
852    pub fn iter(&self) -> Vec<(String, std::sync::Arc<dyn Tool>)> {
853        self.tools
854            .read()
855            .unwrap()
856            .iter()
857            .map(|(k, v)| (k.clone(), v.clone()))
858            .collect()
859    }
860
861    /// Replace one qualified tool namespace while holding a single write lock.
862    pub fn replace_namespace(&self, prefix: &str, tools: Vec<std::sync::Arc<dyn Tool>>) {
863        assert!(!prefix.is_empty(), "tool namespace prefix cannot be empty");
864        assert!(
865            tools.iter().all(|tool| tool.name().starts_with(prefix)),
866            "replacement tools must belong to namespace `{prefix}`"
867        );
868        let mut registry = self.tools.write().unwrap();
869        registry.retain(|name, _| !name.starts_with(prefix));
870        for tool in tools {
871            registry.insert(tool.name().to_string(), tool);
872        }
873    }
874
875    /// Remove namespaced tools that are no longer backed by an enabled source.
876    pub fn retain_namespaces(&self, root_prefix: &str, retained_prefixes: &[String]) {
877        assert!(
878            retained_prefixes
879                .iter()
880                .all(|prefix| prefix.starts_with(root_prefix)),
881            "retained namespaces must belong to root `{root_prefix}`"
882        );
883        self.tools.write().unwrap().retain(|name, _| {
884            !name.starts_with(root_prefix)
885                || retained_prefixes
886                    .iter()
887                    .any(|prefix| name.starts_with(prefix))
888        });
889    }
890
891    /// Remove all tools whose name starts with `prefix` (e.g. `"mcp."`).
892    pub fn unregister_prefix(&self, prefix: &str) {
893        self.tools
894            .write()
895            .unwrap()
896            .retain(|k, _| !k.starts_with(prefix));
897    }
898}
899
900#[cfg(test)]
901mod tests {
902    use super::*;
903
904    struct NamedTool(&'static str);
905
906    impl Tool for NamedTool {
907        fn name(&self) -> &str {
908            self.0
909        }
910
911        fn tier(&self) -> Tier {
912            Tier::Zero
913        }
914
915        fn call<'a>(&'a self, _args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
916            Box::pin(async { Ok(crate::value::Value::Unit) })
917        }
918    }
919
920    #[test]
921    fn namespace_replacement_removes_stale_tools_without_touching_peers() {
922        let registry = ToolRegistry::new();
923        registry.register(std::sync::Arc::new(NamedTool("mcp.alpha.old")));
924        registry.register(std::sync::Arc::new(NamedTool("mcp.beta.keep")));
925
926        registry.replace_namespace(
927            "mcp.alpha.",
928            vec![std::sync::Arc::new(NamedTool("mcp.alpha.new"))],
929        );
930
931        assert!(!registry.has("mcp.alpha.old"));
932        assert!(registry.has("mcp.alpha.new"));
933        assert!(registry.has("mcp.beta.keep"));
934    }
935
936    #[test]
937    fn namespace_retention_only_removes_disabled_sources() {
938        let registry = ToolRegistry::new();
939        registry.register(std::sync::Arc::new(NamedTool("mcp.alpha.keep")));
940        registry.register(std::sync::Arc::new(NamedTool("mcp.beta.remove")));
941        registry.register(std::sync::Arc::new(NamedTool("fs.read")));
942
943        registry.retain_namespaces("mcp.", &["mcp.alpha.".to_string()]);
944
945        assert!(registry.has("mcp.alpha.keep"));
946        assert!(!registry.has("mcp.beta.remove"));
947        assert!(registry.has("fs.read"));
948    }
949
950    #[test]
951    fn model_tool_exposure_is_scoped_exact_and_single_use() {
952        let exposures = ToolExposureRegistry::default();
953        let run = crate::event::FlowRunId::now();
954        let other_run = crate::event::FlowRunId::now();
955        let response = crate::message::Message {
956            role: crate::message::MessageRole::Assistant,
957            parts: vec![crate::message::MessagePart::ToolUse {
958                id: "call-1".into(),
959                name: "allowed.probe".into(),
960                input: serde_json::json!({}),
961                intent: None,
962            }],
963            turn_id: crate::event::TurnId::now(),
964            origin: crate::message::MessageOrigin::User,
965        };
966        exposures.register_response(Some(&run), &response, ["allowed.probe"]);
967
968        assert!(!exposures.claim(Some(&other_run), "call-1", "allowed.probe"));
969        assert!(!exposures.claim(Some(&run), "call-1", "changed.probe"));
970        assert!(exposures.claim(Some(&run), "call-1", "allowed.probe"));
971        assert!(!exposures.claim(Some(&run), "call-1", "allowed.probe"));
972    }
973
974    #[test]
975    fn model_tool_exposure_rejects_unexposed_and_duplicate_response_ids() {
976        let exposures = ToolExposureRegistry::default();
977        let run = crate::event::FlowRunId::now();
978        let response = crate::message::Message {
979            role: crate::message::MessageRole::Assistant,
980            parts: vec![
981                crate::message::MessagePart::ToolUse {
982                    id: "duplicate".into(),
983                    name: "allowed.probe".into(),
984                    input: serde_json::json!({}),
985                    intent: None,
986                },
987                crate::message::MessagePart::ToolUse {
988                    id: "duplicate".into(),
989                    name: "allowed.probe".into(),
990                    input: serde_json::json!({}),
991                    intent: None,
992                },
993                crate::message::MessagePart::ToolUse {
994                    id: "hidden".into(),
995                    name: "hidden.probe".into(),
996                    input: serde_json::json!({}),
997                    intent: None,
998                },
999            ],
1000            turn_id: crate::event::TurnId::now(),
1001            origin: crate::message::MessageOrigin::User,
1002        };
1003        exposures.register_response(Some(&run), &response, ["allowed.probe"]);
1004
1005        assert!(!exposures.claim(Some(&run), "duplicate", "allowed.probe"));
1006        assert!(!exposures.claim(Some(&run), "hidden", "hidden.probe"));
1007    }
1008
1009    struct ReservedEnvTool;
1010
1011    struct ObjectTool;
1012
1013    impl Tool for ObjectTool {
1014        fn name(&self) -> &str {
1015            "probe"
1016        }
1017
1018        fn tier(&self) -> Tier {
1019            Tier::Zero
1020        }
1021
1022        fn input_schema(&self) -> serde_json::Value {
1023            serde_json::json!({
1024                "type": "object",
1025                "properties": {"value": {"type": "integer"}},
1026                "required": ["value"]
1027            })
1028        }
1029
1030        fn call<'a>(&'a self, _args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
1031            Box::pin(async { Ok(Value::Unit) })
1032        }
1033    }
1034
1035    struct CollidingTool;
1036
1037    impl Tool for CollidingTool {
1038        fn name(&self) -> &str {
1039            "collision"
1040        }
1041
1042        fn tier(&self) -> Tier {
1043            Tier::Zero
1044        }
1045
1046        fn input_schema(&self) -> serde_json::Value {
1047            serde_json::json!({
1048                "type": "object",
1049                "properties": {"_atman_intent": {"type": "integer"}}
1050            })
1051        }
1052
1053        fn call<'a>(&'a self, _args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
1054            Box::pin(async { Ok(Value::Unit) })
1055        }
1056    }
1057
1058    impl Tool for ReservedEnvTool {
1059        fn name(&self) -> &str {
1060            "env"
1061        }
1062
1063        fn tier(&self) -> Tier {
1064            Tier::Zero
1065        }
1066
1067        fn call<'a>(&'a self, _args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
1068            Box::pin(async { Ok(Value::Unit) })
1069        }
1070    }
1071
1072    #[test]
1073    #[should_panic(expected = "tool name `env` is reserved for an evaluator intrinsic")]
1074    fn evaluator_intrinsic_names_cannot_be_registered_as_tools() {
1075        ToolRegistry::new().register(std::sync::Arc::new(ReservedEnvTool));
1076    }
1077
1078    #[test]
1079    fn tool_spec_requires_call_intent_after_business_fields() {
1080        let spec = tool_spec(&ObjectTool);
1081        assert_eq!(
1082            spec.input_schema["required"],
1083            serde_json::json!(["value", "_atman_intent"])
1084        );
1085        assert_eq!(
1086            spec.input_schema["properties"][crate::message::TOOL_CALL_INTENT_FIELD],
1087            tool_call_intent_schema()
1088        );
1089        assert!(tool_spec_supports_call_intent("probe", &[spec]));
1090        assert_eq!(
1091            serde_json::to_vec(&tool_spec(&ObjectTool)).unwrap(),
1092            serde_json::to_vec(&tool_spec(&ObjectTool)).unwrap()
1093        );
1094    }
1095
1096    #[test]
1097    fn tool_spec_preserves_business_field_collision() {
1098        let spec = tool_spec(&CollidingTool);
1099        assert_eq!(
1100            spec.input_schema["properties"][crate::message::TOOL_CALL_INTENT_FIELD],
1101            serde_json::json!({"type": "integer"})
1102        );
1103        assert!(!tool_spec_supports_call_intent("collision", &[spec]));
1104    }
1105
1106    #[test]
1107    fn canonical_json_orders_nested_object_keys() {
1108        let mut left: serde_json::Value =
1109            serde_json::from_str(r#"{"zeta":{"y":1,"x":2},"alpha":0}"#).unwrap();
1110        let mut right: serde_json::Value =
1111            serde_json::from_str(r#"{"alpha":0,"zeta":{"x":2,"y":1}}"#).unwrap();
1112        canonicalize_json_object_keys(&mut left);
1113        canonicalize_json_object_keys(&mut right);
1114        assert_eq!(
1115            serde_json::to_vec(&left).unwrap(),
1116            serde_json::to_vec(&right).unwrap()
1117        );
1118    }
1119
1120    #[test]
1121    fn tool_call_input_codec_round_trips_metadata_and_preserves_collisions() {
1122        let spec = tool_spec(&ObjectTool);
1123        let wire = serde_json::json!({
1124            "value": 7,
1125            "_atman_intent": "  inspect\nstate  "
1126        });
1127        let (clean, intent) =
1128            crate::message::decode_tool_call_input(wire, "probe", std::slice::from_ref(&spec));
1129        assert_eq!(clean, serde_json::json!({"value": 7}));
1130        assert_eq!(
1131            intent.as_ref().map(|value| value.as_str()),
1132            Some("inspect state")
1133        );
1134        assert_eq!(
1135            crate::message::encode_tool_call_input(&clean, intent.as_ref(), "probe", &[spec]),
1136            serde_json::json!({"value": 7, "_atman_intent": "inspect state"})
1137        );
1138        assert_eq!(
1139            crate::message::encode_tool_call_input(&clean, intent.as_ref(), "probe", &[]),
1140            serde_json::json!({"value": 7, "_atman_intent": "inspect state"})
1141        );
1142
1143        let spec = tool_spec(&ObjectTool);
1144        let (clean, intent) = crate::message::decode_tool_call_input(
1145            serde_json::json!({"value": 7, "_atman_intent": 42}),
1146            "probe",
1147            &[spec],
1148        );
1149        assert_eq!(clean, serde_json::json!({"value": 7}));
1150        assert!(intent.is_none());
1151
1152        let collision = tool_spec(&CollidingTool);
1153        let business_input = serde_json::json!({"_atman_intent": 42});
1154        let (clean, intent) = crate::message::decode_tool_call_input(
1155            business_input.clone(),
1156            "collision",
1157            &[collision],
1158        );
1159        assert_eq!(clean, business_input);
1160        assert!(intent.is_none());
1161    }
1162
1163    #[test]
1164    fn approval_level_default_maps_from_tier() {
1165        assert_eq!(ApprovalLevel::from_tier(Tier::Zero), ApprovalLevel::Auto);
1166        assert_eq!(ApprovalLevel::from_tier(Tier::One), ApprovalLevel::Approve);
1167        assert_eq!(ApprovalLevel::from_tier(Tier::Two), ApprovalLevel::Approve);
1168        assert_eq!(
1169            ApprovalLevel::from_tier(Tier::Three),
1170            ApprovalLevel::Dangerous
1171        );
1172        assert_eq!(
1173            ApprovalLevel::from_tier(Tier::Four),
1174            ApprovalLevel::Dangerous
1175        );
1176    }
1177
1178    #[test]
1179    fn approval_level_ordered_auto_lt_approve_lt_dangerous() {
1180        assert!(ApprovalLevel::Auto < ApprovalLevel::Approve);
1181        assert!(ApprovalLevel::Approve < ApprovalLevel::Dangerous);
1182    }
1183
1184    #[test]
1185    fn invocation_snapshot_tracks_session_trust() {
1186        let dir = tempfile::tempdir().unwrap();
1187        let initial = crate::trust::TrustConfig::default();
1188        let session = std::sync::Arc::new(
1189            crate::session::Session::open_with_trust(dir.path(), initial.clone()).unwrap(),
1190        );
1191        let sandbox: std::sync::Arc<dyn crate::sandbox::Sandbox> =
1192            std::sync::Arc::new(crate::sandbox::SandboxExec::new(dir.path()));
1193        let flows = std::sync::Arc::new(crate::tools::agent_ctrl::FlowRegistry::new());
1194        let identity = flows
1195            .register_root(
1196                session.id().to_string(),
1197                crate::event::FlowRunId::now(),
1198                crate::flow_authority::EffectiveAuthority::root(&initial, true, None),
1199            )
1200            .unwrap();
1201        let mut base = ToolCtx::new()
1202            .with_trust(crate::trust::TrustConfig {
1203                mode: crate::trust::TrustMode::Reckless,
1204                ..crate::trust::TrustConfig::default()
1205            })
1206            .with_session_runtime(std::sync::Arc::clone(&session))
1207            .with_sandbox(sandbox);
1208        base.flow_identity = Some(identity);
1209
1210        let controlled = base.clone().for_tool_invocation(Tier::Four);
1211        assert_eq!(controlled.trust, Some(initial));
1212        assert!(controlled.sandbox.is_some());
1213
1214        let reckless = crate::trust::TrustConfig {
1215            mode: crate::trust::TrustMode::Reckless,
1216            ..crate::trust::TrustConfig::default()
1217        };
1218        session.update_trust(reckless.clone(), |_| Ok(())).unwrap();
1219        let unrestricted = base.for_tool_invocation(Tier::Four);
1220
1221        assert_eq!(unrestricted.trust, Some(reckless));
1222        assert!(unrestricted.sandbox.is_some());
1223        assert!(controlled.sandbox.is_some());
1224    }
1225
1226    #[test]
1227    fn invocation_context_retains_sandbox_across_control_tools() {
1228        let dir = tempfile::tempdir().unwrap();
1229        let sandbox: std::sync::Arc<dyn crate::sandbox::Sandbox> =
1230            std::sync::Arc::new(crate::sandbox::SandboxExec::new(dir.path()));
1231        let base = ToolCtx::new().with_sandbox(sandbox);
1232
1233        let control_ctx = base.for_tool_invocation(Tier::Two);
1234        assert!(control_ctx.sandbox.is_some());
1235        assert!(
1236            control_ctx
1237                .for_tool_invocation(Tier::Four)
1238                .sandbox
1239                .is_some()
1240        );
1241    }
1242
1243    fn binding(path: std::path::PathBuf) -> crate::git_workspace::WorkspaceBinding {
1244        crate::git_workspace::WorkspaceBinding {
1245            workspace_id: "workspace".into(),
1246            repository_root: path.clone(),
1247            path,
1248            branch: None,
1249        }
1250    }
1251
1252    #[test]
1253    fn workspace_resolver_rebinds_policy_without_changing_process_cwd() {
1254        let workspace = tempfile::tempdir().unwrap();
1255        let process_cwd = std::env::current_dir().unwrap();
1256        let ctx = ToolCtx::new()
1257            .with_fs_access(crate::fs_access::FsAccessPolicy {
1258                mode: crate::fs_access::FsAccessMode::ReadOnly,
1259                workspace: Some(process_cwd.clone()),
1260            })
1261            .with_workspace(binding(workspace.path().to_path_buf()));
1262
1263        let canonical_workspace = crate::fs_access::canonicalize_stable(workspace.path());
1264        let omitted = ctx.resolve_cwd_with_origin(None).unwrap();
1265        assert_eq!(omitted.path, canonical_workspace);
1266        assert_eq!(omitted.origin, PathOrigin::Omitted);
1267
1268        let relative = ctx
1269            .resolve_path_with_origin(std::path::Path::new("nested/file"))
1270            .unwrap();
1271        assert_eq!(relative.path, canonical_workspace.join("nested/file"));
1272        assert_eq!(relative.origin, PathOrigin::Relative);
1273
1274        let inside = ctx
1275            .resolve_path_with_origin(&workspace.path().join("inside"))
1276            .unwrap();
1277        assert_eq!(inside.origin, PathOrigin::ExplicitInside);
1278
1279        let external = ctx.resolve_path_with_origin(&process_cwd).unwrap();
1280        assert_eq!(external.origin, PathOrigin::ExplicitExternal);
1281        assert_eq!(ctx.fs_access.mode, crate::fs_access::FsAccessMode::ReadOnly);
1282        assert_eq!(ctx.fs_access.workspace.as_deref(), Some(workspace.path()));
1283        assert_eq!(std::env::current_dir().unwrap(), process_cwd);
1284    }
1285
1286    #[test]
1287    fn workspace_resolver_rejects_parent_escape() {
1288        let workspace = tempfile::tempdir().unwrap();
1289        let ctx = ToolCtx::new().with_workspace(binding(workspace.path().to_path_buf()));
1290        let error = ctx
1291            .resolve_path(std::path::Path::new("../outside"))
1292            .unwrap_err();
1293        assert!(error.to_string().contains("escapes workspace root"));
1294    }
1295
1296    #[cfg(unix)]
1297    #[test]
1298    fn workspace_resolver_rejects_symlink_escape() {
1299        let workspace = tempfile::tempdir().unwrap();
1300        let outside = tempfile::tempdir().unwrap();
1301        std::os::unix::fs::symlink(outside.path(), workspace.path().join("link")).unwrap();
1302        let ctx = ToolCtx::new().with_workspace(binding(workspace.path().to_path_buf()));
1303
1304        let error = ctx
1305            .resolve_path(std::path::Path::new("link/file"))
1306            .unwrap_err();
1307        assert!(error.to_string().contains("escapes workspace root"));
1308    }
1309
1310    #[test]
1311    fn ordinary_context_keeps_process_cwd_semantics() {
1312        let process_cwd = std::env::current_dir().unwrap();
1313        let ctx = ToolCtx::new();
1314        assert_eq!(ctx.resolve_cwd(None).unwrap(), process_cwd);
1315        assert_eq!(
1316            ctx.resolve_path(std::path::Path::new("child")).unwrap(),
1317            process_cwd.join("child")
1318        );
1319    }
1320}
atman_runtime/tool.rs

atman_runtime/
tool.rs
