Skip to main content

agent_base/engine/
approval.rs

1use async_trait::async_trait;
2
3use crate::types::{AgentResult, ApprovalDecision, ApprovalRequest};
4
5/// Trait for handling tool approval requests.
6///
7/// # Cancellation Contract
8///
9/// Implementors MUST handle cancellation via the provided `cancel_token`.
10/// The token is cancelled when the user interrupts the agent (e.g., Ctrl+C).
11///
12/// - **Async handlers** (channels, timers): race against `cancel_token.cancelled()` using `tokio::select!`
13/// - **Blocking handlers** (stdin, file IO): use `tokio::task::spawn_blocking` + `tokio::select!` to avoid blocking the tokio runtime
14///
15/// # Timeout
16///
17/// The caller (`process_approval`) wraps this method with a 300-second timeout.
18/// If the handler does not return within 300s, the approval is denied automatically.
19///
20/// # Example (blocking stdin)
21///
22/// ```ignore
23/// async fn approve(&self, request: ApprovalRequest, cancel_token: CancellationToken) -> AgentResult<ApprovalDecision> {
24///     tokio::select! {
25///         _ = cancel_token.cancelled() => Err(AgentError::Cancelled),
26///         result = tokio::task::spawn_blocking(|| read_stdin()) => result?,
27///     }
28/// }
29/// ```
30#[async_trait]
31pub trait ApprovalHandler: Send + Sync {
32    async fn approve(
33        &self,
34        request: ApprovalRequest,
35        cancel_token: tokio_util::sync::CancellationToken,
36    ) -> AgentResult<ApprovalDecision>;
37}
38
39#[derive(Clone, Debug, Default)]
40pub struct DenyAllApprovalHandler;
41
42#[async_trait]
43impl ApprovalHandler for DenyAllApprovalHandler {
44    async fn approve(
45        &self,
46        _request: ApprovalRequest,
47        _cancel_token: tokio_util::sync::CancellationToken,
48    ) -> AgentResult<ApprovalDecision> {
49        Ok(ApprovalDecision::Deny)
50    }
51}
52
53#[derive(Clone, Debug, Default)]
54pub struct AllowAllApprovalHandler;
55
56#[async_trait]
57impl ApprovalHandler for AllowAllApprovalHandler {
58    async fn approve(
59        &self,
60        _request: ApprovalRequest,
61        _cancel_token: tokio_util::sync::CancellationToken,
62    ) -> AgentResult<ApprovalDecision> {
63        Ok(ApprovalDecision::AllowAlways)
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use crate::types::RiskLevel;
71    use tokio_util::sync::CancellationToken;
72
73    fn request() -> ApprovalRequest {
74        ApprovalRequest {
75            title: "Delete file".to_string(),
76            message: "Really delete?".to_string(),
77            action_key: Some("delete".to_string()),
78            risk_level: RiskLevel::Destructive,
79            raw: None,
80        }
81    }
82
83    #[tokio::test]
84    async fn deny_all_returns_deny() {
85        let decision = DenyAllApprovalHandler
86            .approve(request(), CancellationToken::new())
87            .await
88            .unwrap();
89        assert_eq!(decision, ApprovalDecision::Deny);
90    }
91
92    #[tokio::test]
93    async fn allow_all_returns_allow_always() {
94        let decision = AllowAllApprovalHandler
95            .approve(request(), CancellationToken::new())
96            .await
97            .unwrap();
98        assert_eq!(decision, ApprovalDecision::AllowAlways);
99    }
100}