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

atman_runtime/
tool.rs
