Skip to main content

codei_tools/
approval_gate.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use tokio::sync::{oneshot, Mutex};
5
6use crate::{ApprovalHandler, ApprovalRequest, ApprovalResponse};
7
8/// Pending approval surfaced to the UI layer.
9pub struct SharedApprovalGate {
10    inner: Arc<Mutex<GateState>>,
11}
12
13#[derive(Default)]
14struct GateState {
15    pending: Option<PendingApproval>,
16}
17
18struct PendingApproval {
19    request: ApprovalRequest,
20    respond: oneshot::Sender<bool>,
21}
22
23impl Default for SharedApprovalGate {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl SharedApprovalGate {
30    pub fn new() -> Self {
31        Self {
32            inner: Arc::new(Mutex::new(GateState::default())),
33        }
34    }
35
36    pub fn handler(self: &Arc<Self>) -> GateApprovalHandler {
37        GateApprovalHandler {
38            gate: Arc::clone(self),
39        }
40    }
41
42    pub async fn take_pending(&self) -> Option<ApprovalRequest> {
43        let guard = self.inner.lock().await;
44        guard.pending.as_ref().map(|p| p.request.clone())
45    }
46
47    pub async fn respond(&self, approved: bool) -> bool {
48        let respond = {
49            let mut guard = self.inner.lock().await;
50            guard.pending.take().map(|p| p.respond)
51        };
52        if let Some(tx) = respond {
53            tx.send(approved).is_ok()
54        } else {
55            false
56        }
57    }
58}
59
60pub struct GateApprovalHandler {
61    gate: Arc<SharedApprovalGate>,
62}
63
64#[async_trait]
65impl ApprovalHandler for GateApprovalHandler {
66    async fn approve(&self, request: ApprovalRequest) -> ApprovalResponse {
67        match request.tool_name.as_str() {
68            "write" | "edit" | "shell" => {}
69            _ => return ApprovalResponse { approved: true },
70        }
71
72        let (tx, rx) = oneshot::channel();
73        {
74            let mut guard = self.gate.inner.lock().await;
75            guard.pending = Some(PendingApproval {
76                request: request.clone(),
77                respond: tx,
78            });
79        }
80
81        let approved = rx.await.unwrap_or(false);
82        ApprovalResponse { approved }
83    }
84}