Skip to main content

defect_agent/session/
permissions.rs

1//! Write-back channel for permission requests.
2//!
3//! `Session::resolve_permission` sends the client response to the ACP reverse request
4//! `session/request_permission` back to the main loop, which waits using
5//! [`PermissionGate::wait`].
6//!
7//! Permission management — see session and turn-loop designs.
8
9use agent_client_protocol_schema::ToolCallId;
10use dashmap::DashMap;
11use tokio::sync::oneshot;
12use tokio_util::sync::CancellationToken;
13
14use crate::event::PermissionResolution;
15
16/// A registry of pending permission requests.
17///
18/// Each in-flight turn holds a shared `Arc<PermissionGate>`:
19/// - The main loop registers a waiter and awaits via [`Self::wait`]
20/// - The ACP bridge layer calls [`Self::resolve`] after receiving the client response
21#[derive(Default)]
22pub struct PermissionGate {
23    waiters: DashMap<ToolCallId, oneshot::Sender<PermissionResolution>>,
24}
25
26impl PermissionGate {
27    pub fn new() -> Self {
28        Self::default()
29    }
30
31    /// Register a waiter and await until [`Self::resolve`] is called or `cancel` fires.
32    ///
33    /// When `cancel` fires, returns [`PermissionResolution::Cancelled`] — the main loop
34    /// handles this as "User cancelled".
35    ///
36    /// If a waiter already exists for the same `id`, the old sender is dropped (the old
37    /// wait receives [`PermissionResolution::Cancelled`], avoiding a hang). This path
38    /// should theoretically never be hit — the main loop only calls `wait` once per
39    /// tool_use.
40    pub async fn wait(&self, id: ToolCallId, cancel: CancellationToken) -> PermissionResolution {
41        let (tx, rx) = oneshot::channel();
42        if let Some(prev) = self.waiters.insert(id.clone(), tx) {
43            // This should not happen: `wait` called twice for the same `id`. Wake the old
44            // waiter with `Cancelled` to prevent it from hanging forever.
45            tracing::warn!(
46                tool_call_id = %id,
47                "PermissionGate::wait called twice for same id; cancelling previous waiter"
48            );
49            let _ = prev.send(PermissionResolution::Cancelled);
50        }
51
52        tokio::select! {
53            biased;
54            () = cancel.cancelled() => {
55                // Remove our registration if it is still present; resolve may race with
56                // cancel.
57                self.waiters.remove(&id);
58                PermissionResolution::Cancelled
59            }
60            recv = rx => match recv {
61                Ok(outcome) => outcome,
62                // Sender was replaced or gate was dropped; use cancellation semantics.
63                Err(_) => PermissionResolution::Cancelled,
64            }
65        }
66    }
67
68    /// Deliver the outcome to the waiter. If `id` has no waiter (already removed by
69    /// cancel, or the main loop hasn't called wait yet), silently no-op — the ACP bridge
70    /// layer is unaware of main-loop timing, and duplicate or late resolves must not
71    /// corrupt the turn.
72    pub fn resolve(&self, id: &ToolCallId, outcome: PermissionResolution) {
73        if let Some((_, tx)) = self.waiters.remove(id) {
74            // Ignore if the receiver has been dropped — the main loop may have already
75            // returned via the cancel path.
76            let _ = tx.send(outcome);
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests;