Skip to main content

atomcode_core/turn/
permission.rs

1use crate::tool::{ApprovalRequirement, PermissionDecision, ToolCall};
2use async_trait::async_trait;
3use tokio::sync::mpsc;
4
5/// Permission decision interface. TurnRunner calls this when a tool requires approval.
6/// Different implementations support interactive (main agent) and automatic (subagent) modes.
7#[async_trait]
8pub trait PermissionDecider: Send + Sync {
9    async fn decide(&self, call: &ToolCall, approval: &ApprovalRequirement) -> PermissionDecision;
10
11    /// Quick synchronous check: will this call be auto-approved without
12    /// user interaction?  Used by TurnRunner to skip the
13    /// `ApprovalRequested` event (and its associated TUI prompt row)
14    /// when the PermissionStore already has a session grant or override
15    /// that will cause `decide()` to return `Allow` immediately.
16    ///
17    /// Returning `false` does **not** mean the call will be denied —
18    /// only that it *might* need interactive approval.  Returning
19    /// `true` guarantees `decide()` will return `Allow` without
20    /// prompting.
21    fn will_auto_approve(&self, call: &ToolCall, approval: &ApprovalRequirement) -> bool;
22}
23
24/// Auto-permission modes for subagents
25#[derive(Debug, Clone)]
26pub enum AutoPermissionMode {
27    /// Allow all tools
28    BypassAll,
29    /// Allow edit tools (write_file, edit_file, search_replace), deny others
30    AcceptEdits,
31    /// Deny all tools that require approval
32    DenyAll,
33}
34
35const EDIT_TOOLS: &[&str] = &["create_file", "edit_file", "search_replace"];
36
37/// Automatic permission decider (used by SubagentLoop)
38pub struct AutoPermissionDecider {
39    mode: AutoPermissionMode,
40}
41
42impl AutoPermissionDecider {
43    pub fn new(mode: AutoPermissionMode) -> Self {
44        Self { mode }
45    }
46}
47
48#[async_trait]
49impl PermissionDecider for AutoPermissionDecider {
50    async fn decide(&self, call: &ToolCall, _approval: &ApprovalRequirement) -> PermissionDecision {
51        match self.mode {
52            AutoPermissionMode::BypassAll => PermissionDecision::Allow,
53            AutoPermissionMode::AcceptEdits => {
54                if EDIT_TOOLS.contains(&call.name.as_str()) {
55                    PermissionDecision::Allow
56                } else {
57                    PermissionDecision::Deny
58                }
59            }
60            AutoPermissionMode::DenyAll => PermissionDecision::Deny,
61        }
62    }
63
64    fn will_auto_approve(&self, call: &ToolCall, _approval: &ApprovalRequirement) -> bool {
65        // AutoPermissionDecider never prompts the user — it either
66        // allows or denies based on its mode.  Return true when the
67        // decision will be Allow (no interactive prompt involved).
68        match self.mode {
69            AutoPermissionMode::BypassAll => true,
70            AutoPermissionMode::AcceptEdits => EDIT_TOOLS.contains(&call.name.as_str()),
71            AutoPermissionMode::DenyAll => false,
72        }
73    }
74}
75
76/// Approval request sent to AgentLoop's command loop
77#[derive(Debug)]
78pub struct ApprovalRequest {
79    pub call: ToolCall,
80    pub reason: String,
81}
82
83/// Interactive permission decider (used by AgentLoop).
84/// Checks PermissionStore first (session grants, overrides),
85/// then falls back to sending approval request via channel.
86pub struct InteractivePermissionDecider {
87    request_tx: mpsc::UnboundedSender<ApprovalRequest>,
88    response_rx: tokio::sync::Mutex<mpsc::UnboundedReceiver<PermissionDecision>>,
89    /// Shared permission store — checked before sending interactive requests.
90    /// AgentLoop writes to this (grant_session on ApproveToolAlways),
91    /// TurnRunner reads from it (check before prompting user).
92    permission_store: std::sync::Arc<std::sync::RwLock<crate::tool::PermissionStore>>,
93}
94
95impl InteractivePermissionDecider {
96    pub fn new(
97        request_tx: mpsc::UnboundedSender<ApprovalRequest>,
98        response_rx: mpsc::UnboundedReceiver<PermissionDecision>,
99        permission_store: std::sync::Arc<std::sync::RwLock<crate::tool::PermissionStore>>,
100    ) -> Self {
101        Self {
102            request_tx,
103            response_rx: tokio::sync::Mutex::new(response_rx),
104            permission_store,
105        }
106    }
107}
108
109#[async_trait]
110impl PermissionDecider for InteractivePermissionDecider {
111    async fn decide(&self, call: &ToolCall, approval: &ApprovalRequirement) -> PermissionDecision {
112        // Check PermissionStore first — session grants and overrides
113        // take effect without prompting the user again. The full
114        // `ApprovalRequirement` (including `RequireApprovalAlways`) is
115        // passed through so the store can honor variants that must
116        // always prompt regardless of prior session grants.
117        if let Ok(store) = self.permission_store.read() {
118            match store.check(&call.name, approval) {
119                PermissionDecision::Allow => return PermissionDecision::Allow,
120                PermissionDecision::Deny => return PermissionDecision::Deny,
121                PermissionDecision::Ask(_) => {} // fall through to interactive
122            }
123        }
124
125        let reason = match approval {
126            ApprovalRequirement::RequireApproval(r)
127            | ApprovalRequirement::RequireApprovalAlways(r) => r.clone(),
128            ApprovalRequirement::AutoApprove => return PermissionDecision::Allow,
129        };
130
131        // Not in store — send interactive approval request
132        let request = ApprovalRequest {
133            call: call.clone(),
134            reason,
135        };
136        if self.request_tx.send(request).is_err() {
137            return PermissionDecision::Deny;
138        }
139        let mut rx = self.response_rx.lock().await;
140        rx.recv().await.unwrap_or(PermissionDecision::Deny)
141    }
142
143    fn will_auto_approve(&self, call: &ToolCall, approval: &ApprovalRequirement) -> bool {
144        // Mirror the PermissionStore check that `decide()` performs
145        // before falling through to the interactive channel. The full
146        // variant is forwarded so `RequireApprovalAlways` continues to
147        // prompt even when a session grant exists for the tool.
148        if matches!(approval, ApprovalRequirement::AutoApprove) {
149            return true;
150        }
151        if let Ok(store) = self.permission_store.read() {
152            matches!(store.check(&call.name, approval), PermissionDecision::Allow)
153        } else {
154            false
155        }
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    fn make_call(name: &str) -> ToolCall {
164        ToolCall {
165            id: "test".into(),
166            name: name.into(),
167            arguments: "{}".into(),
168        }
169    }
170
171    #[tokio::test]
172    async fn test_auto_bypass_allows_all() {
173        let d = AutoPermissionDecider::new(AutoPermissionMode::BypassAll);
174        assert!(matches!(
175            d.decide(&make_call("bash"), &ApprovalRequirement::RequireApproval("dangerous".into())).await,
176            PermissionDecision::Allow
177        ));
178    }
179
180    #[tokio::test]
181    async fn test_auto_deny_denies_all() {
182        let d = AutoPermissionDecider::new(AutoPermissionMode::DenyAll);
183        assert!(matches!(
184            d.decide(&make_call("bash"), &ApprovalRequirement::RequireApproval("dangerous".into())).await,
185            PermissionDecision::Deny
186        ));
187    }
188
189    #[tokio::test]
190    async fn test_auto_accept_edits_allows_write() {
191        let d = AutoPermissionDecider::new(AutoPermissionMode::AcceptEdits);
192        assert!(matches!(
193            d.decide(&make_call("create_file"), &ApprovalRequirement::RequireApproval("write".into())).await,
194            PermissionDecision::Allow
195        ));
196        assert!(matches!(
197            d.decide(&make_call("edit_file"), &ApprovalRequirement::RequireApproval("edit".into())).await,
198            PermissionDecision::Allow
199        ));
200        assert!(matches!(
201            d.decide(&make_call("search_replace"), &ApprovalRequirement::RequireApproval("sr".into())).await,
202            PermissionDecision::Allow
203        ));
204    }
205
206    #[tokio::test]
207    async fn test_auto_accept_edits_denies_bash() {
208        let d = AutoPermissionDecider::new(AutoPermissionMode::AcceptEdits);
209        assert!(matches!(
210            d.decide(&make_call("bash"), &ApprovalRequirement::RequireApproval("dangerous".into())).await,
211            PermissionDecision::Deny
212        ));
213    }
214
215    #[tokio::test]
216    async fn test_interactive_allow() {
217        let (req_tx, mut req_rx) = mpsc::unbounded_channel();
218        let (resp_tx, resp_rx) = mpsc::unbounded_channel();
219        let store =
220            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
221        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
222
223        let call = make_call("bash");
224        let approval = ApprovalRequirement::RequireApproval("dangerous".into());
225        let fut = d.decide(&call, &approval);
226
227        tokio::spawn(async move {
228            let _req = req_rx.recv().await.unwrap();
229            resp_tx.send(PermissionDecision::Allow).unwrap();
230        });
231
232        assert!(matches!(fut.await, PermissionDecision::Allow));
233    }
234
235    #[tokio::test]
236    async fn test_interactive_deny() {
237        let (req_tx, mut req_rx) = mpsc::unbounded_channel();
238        let (resp_tx, resp_rx) = mpsc::unbounded_channel();
239        let store =
240            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
241        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
242
243        let call = make_call("bash");
244        let approval = ApprovalRequirement::RequireApproval("dangerous".into());
245        let fut = d.decide(&call, &approval);
246
247        tokio::spawn(async move {
248            let _req = req_rx.recv().await.unwrap();
249            resp_tx.send(PermissionDecision::Deny).unwrap();
250        });
251
252        assert!(matches!(fut.await, PermissionDecision::Deny));
253    }
254
255    #[tokio::test]
256    async fn test_interactive_channel_closed_returns_deny() {
257        let (req_tx, req_rx) = mpsc::unbounded_channel();
258        let (_resp_tx, resp_rx) = mpsc::unbounded_channel();
259        let store =
260            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
261        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
262
263        drop(req_rx); // close request channel
264        let call = make_call("bash");
265        assert!(matches!(
266            d.decide(&call, &ApprovalRequirement::RequireApproval("dangerous".into())).await,
267            PermissionDecision::Deny
268        ));
269    }
270
271    #[tokio::test]
272    async fn test_interactive_session_grant_skips_channel() {
273        let (req_tx, _req_rx) = mpsc::unbounded_channel();
274        let (_resp_tx, resp_rx) = mpsc::unbounded_channel();
275        let store =
276            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
277
278        // Grant session permission for "bash" BEFORE creating the decider
279        store.write().unwrap().grant_session("bash");
280
281        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
282        let call = make_call("bash");
283
284        // Should return Allow immediately from PermissionStore,
285        // WITHOUT sending a request on the channel (channel is not even read).
286        let decision = d.decide(&call, &ApprovalRequirement::RequireApproval("dangerous".into())).await;
287        assert!(matches!(decision, PermissionDecision::Allow));
288    }
289
290    #[tokio::test]
291    async fn test_interactive_require_approval_always_with_grant_still_prompts() {
292        // RequireApprovalAlways must always prompt the user — even if [A] was
293        // pressed earlier for this tool. The session grant must NOT bypass it.
294        let (req_tx, mut req_rx) = mpsc::unbounded_channel();
295        let (resp_tx, resp_rx) = mpsc::unbounded_channel();
296        let store =
297            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
298        store.write().unwrap().grant_session("bash");
299        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
300        let call = make_call("bash");
301        let approval = ApprovalRequirement::RequireApprovalAlways("sensitive".into());
302        let fut = d.decide(&call, &approval);
303
304        // Expect a request on the channel (NOT auto-approved by the store).
305        // Deny it and confirm decide() returns Deny.
306        tokio::spawn(async move {
307            let _req = req_rx.recv().await.expect("channel must receive request");
308            resp_tx.send(PermissionDecision::Deny).unwrap();
309        });
310
311        assert!(matches!(fut.await, PermissionDecision::Deny));
312    }
313
314    // ── will_auto_approve tests ──
315
316    #[test]
317    fn test_will_auto_approve_auto_bypass() {
318        let d = AutoPermissionDecider::new(AutoPermissionMode::BypassAll);
319        let call = make_call("bash");
320        assert!(d.will_auto_approve(&call, &ApprovalRequirement::RequireApproval("dangerous".into())));
321    }
322
323    #[test]
324    fn test_will_auto_approve_auto_deny() {
325        let d = AutoPermissionDecider::new(AutoPermissionMode::DenyAll);
326        let call = make_call("bash");
327        assert!(!d.will_auto_approve(&call, &ApprovalRequirement::RequireApproval("dangerous".into())));
328    }
329
330    #[test]
331    fn test_will_auto_approve_auto_accept_edits() {
332        let d = AutoPermissionDecider::new(AutoPermissionMode::AcceptEdits);
333        let edit_call = make_call("edit_file");
334        let bash_call = make_call("bash");
335        assert!(d.will_auto_approve(&edit_call, &ApprovalRequirement::RequireApproval("write".into())));
336        assert!(!d.will_auto_approve(&bash_call, &ApprovalRequirement::RequireApproval("dangerous".into())));
337    }
338
339    #[test]
340    fn test_will_auto_approve_interactive_no_grant() {
341        let (req_tx, _req_rx) = mpsc::unbounded_channel();
342        let (_resp_tx, resp_rx) = mpsc::unbounded_channel();
343        let store =
344            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
345        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
346        let call = make_call("bash");
347        // No session grant → will NOT auto-approve
348        assert!(!d.will_auto_approve(&call, &ApprovalRequirement::RequireApproval("dangerous".into())));
349    }
350
351    #[test]
352    fn test_will_auto_approve_interactive_with_session_grant() {
353        let (req_tx, _req_rx) = mpsc::unbounded_channel();
354        let (_resp_tx, resp_rx) = mpsc::unbounded_channel();
355        let store =
356            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
357        store.write().unwrap().grant_session("bash");
358        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
359        let call = make_call("bash");
360        // Session grant exists → WILL auto-approve
361        assert!(d.will_auto_approve(&call, &ApprovalRequirement::RequireApproval("dangerous".into())));
362    }
363
364    #[test]
365    fn test_will_auto_approve_interactive_require_approval_always_with_grant() {
366        // RequireApprovalAlways means "always prompt, even if [A] was pressed
367        // earlier for this tool". A session grant must NOT bypass it.
368        let (req_tx, _req_rx) = mpsc::unbounded_channel();
369        let (_resp_tx, resp_rx) = mpsc::unbounded_channel();
370        let store =
371            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
372        store.write().unwrap().grant_session("bash");
373        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
374        let call = make_call("bash");
375        assert!(!d.will_auto_approve(&call, &ApprovalRequirement::RequireApprovalAlways("sensitive".into())));
376    }
377
378    #[test]
379    fn test_will_auto_approve_interactive_require_approval_always_no_grant() {
380        // RequireApprovalAlways WITHOUT a session grant: still needs prompt.
381        let (req_tx, _req_rx) = mpsc::unbounded_channel();
382        let (_resp_tx, resp_rx) = mpsc::unbounded_channel();
383        let store =
384            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
385        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
386        let call = make_call("bash");
387        assert!(!d.will_auto_approve(&call, &ApprovalRequirement::RequireApprovalAlways("sensitive".into())));
388    }
389
390    #[test]
391    fn test_will_auto_approve_interactive_different_tool_not_auto() {
392        // Session grant for "bash" should NOT auto-approve "mcp__zouwu__query"
393        let (req_tx, _req_rx) = mpsc::unbounded_channel();
394        let (_resp_tx, resp_rx) = mpsc::unbounded_channel();
395        let store =
396            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
397        store.write().unwrap().grant_session("bash");
398        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
399        let call = make_call("mcp__zouwu__query");
400        assert!(!d.will_auto_approve(&call, &ApprovalRequirement::RequireApproval("mcp tool".into())));
401    }
402
403    #[test]
404    fn test_will_auto_approve_interactive_same_tool_auto() {
405        // The key scenario from the bug: user presses [A] for an MCP tool,
406        // subsequent calls to the same tool should auto-approve.
407        let (req_tx, _req_rx) = mpsc::unbounded_channel();
408        let (_resp_tx, resp_rx) = mpsc::unbounded_channel();
409        let store =
410            std::sync::Arc::new(std::sync::RwLock::new(crate::tool::PermissionStore::new()));
411        // Simulate: user pressed [A] for this MCP tool in a previous call
412        store.write().unwrap().grant_session("mcp__zouwu-mcp-server__query_requirements");
413        let d = InteractivePermissionDecider::new(req_tx, resp_rx, store);
414        let call = make_call("mcp__zouwu-mcp-server__query_requirements");
415        assert!(d.will_auto_approve(&call, &ApprovalRequirement::RequireApproval("mcp tool".into())));
416    }
417}