agent_base/tool/
policy.rs

1use async_trait::async_trait;
2use serde_json::Value;
3
4use super::{Content, ToolContext};
5use crate::types::{AgentResult, ApprovalRequest, RiskLevel};
6
7/// Policy-based control over tool execution.
8///
9/// Implement this trait to customise how tools are approved, monitored, and
10/// validated during an agent run. The pipeline calls each hook at a specific
11/// point in the tool lifecycle:
12///
13/// ```text
14/// evaluate_approval  →  before_call  →  (tool executes)  →  after_call
15/// ```
16///
17/// # Example: auto-approve read-only tools
18///
19/// ```ignore
20/// struct ReadOnlyPolicy;
21///
22/// #[async_trait]
23/// impl ToolPolicy for ReadOnlyPolicy {
24///     async fn evaluate_approval(&self, tool_name: &str, _args: &Value) -> Option<ApprovalRequest> {
25///         if tool_name == "read_file" || tool_name == "search" {
26///             None  // auto-approve — no prompt for user
27///         } else {
28///             Some(ApprovalRequest { message: format!("Allow {}?", tool_name) })
29///         }
30///     }
31/// }
32/// ```
33///
34/// All hooks have default no-op implementations, so you only need to override
35/// the ones you care about.
36#[async_trait]
37pub trait ToolPolicy: Send + Sync {
38    /// Called **before** every tool call.
39    ///
40    /// Return `None` to auto-approve (skip the approval handler entirely).
41    /// Return `Some(ApprovalRequest)` to defer to the configured
42    /// [`ApprovalHandler`](crate::ApprovalHandler).
43    ///
44    /// This is the primary hook for implementing permission guards — for
45    /// example, auto-approving read-only operations while prompting for
46    /// destructive ones.
47    async fn evaluate_approval(&self, tool_name: &str, args: &Value) -> Option<ApprovalRequest>;
48
49    /// Called immediately **before** a tool executes, after approval has been
50    /// granted (or auto-approved).
51    ///
52    /// Use this for:
53    /// - Input validation (reject malformed arguments before the tool runs).
54    /// - Auditing / logging the raw call.
55    /// - Rate-limiting or quota enforcement.
56    ///
57    /// Return an `Err` to **cancel** the tool call before execution. The error
58    /// message is surfaced to the LLM so it can correct its approach.
59    fn before_call(&self, tool_name: &str, args: &Value, ctx: &ToolContext) -> AgentResult<()> {
60        let _ = (tool_name, args, ctx);
61        Ok(())
62    }
63
64    /// Called immediately **after** a tool executes, before the result is
65    /// returned to the LLM.
66    ///
67    /// Use this for:
68    /// - Output scrubbing / redaction (strip secrets from tool results).
69    /// - Truncation or formatting of large outputs.
70    /// - Recording metrics or audit trails.
71    ///
72    /// The `result` is the raw [`Vec<Content>`](Content) produced by the tool.
73    /// You can inspect it but not modify it through this hook — if you need to
74    /// transform the output, use a middleware instead.
75    ///
76    /// Return an `Err` to **reject** the result. The error message is surfaced
77    /// to the LLM.
78    fn after_call(
79        &self,
80        tool_name: &str,
81        args: &Value,
82        result: &[Content],
83        ctx: &ToolContext,
84    ) -> AgentResult<()> {
85        let _ = (tool_name, args, result, ctx);
86        Ok(())
87    }
88}
89
90/// A [`ToolPolicy`] that requires approval for every tool call.
91///
92/// Combined with a deny-all approval handler, this denies every tool for the
93/// agent that carries it. Used by multi-agent runtimes as a fallback when a
94/// child should have "no permission" but its parent has no [`ToolPolicy`] of
95/// its own to inherit (so there is no notion of which tools are "dangerous").
96#[derive(Debug, Clone, Default)]
97pub struct DenyAllToolPolicy;
98
99#[async_trait]
100impl ToolPolicy for DenyAllToolPolicy {
101    async fn evaluate_approval(&self, tool_name: &str, _args: &Value) -> Option<ApprovalRequest> {
102        Some(ApprovalRequest {
103            title: format!("Permission required: {tool_name}"),
104            message: format!("This agent has no permission to call `{tool_name}`."),
105            action_key: None,
106            risk_level: RiskLevel::Destructive,
107            raw: None,
108        })
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use serde_json::json;
116
117    /// Minimal policy that auto-approves everything; exercises the default
118    /// no-op `before_call` / `after_call` hooks.
119    struct NoopPolicy;
120
121    #[async_trait]
122    impl ToolPolicy for NoopPolicy {
123        async fn evaluate_approval(
124            &self,
125            _tool_name: &str,
126            _args: &Value,
127        ) -> Option<ApprovalRequest> {
128            None
129        }
130    }
131
132    #[test]
133    fn default_hooks_are_noop() {
134        let p = NoopPolicy;
135        let ctx = ToolContext::for_test();
136        let args = json!({ "x": 1 });
137        assert!(p.before_call("echo", &args, &ctx).is_ok());
138        assert!(p.after_call("echo", &args, &[], &ctx).is_ok());
139    }
140
141    #[tokio::test]
142    async fn deny_all_policy_requires_approval_for_every_tool() {
143        let p = DenyAllToolPolicy;
144        let req = p
145            .evaluate_approval("any_tool", &json!({}))
146            .await
147            .expect("every tool should require approval");
148        assert_eq!(req.risk_level, RiskLevel::Destructive);
149        assert!(req.title.contains("any_tool"));
150    }
151}
agent_base/tool/policy.rs

agent_base/tool/
policy.rs
