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

atman_runtime/
tool.rs
