codei_tools/
approval_gate.rs1use std::sync::Arc;
2
3use async_trait::async_trait;
4use tokio::sync::{oneshot, Mutex};
5
6use crate::{ApprovalHandler, ApprovalRequest, ApprovalResponse};
7
8pub 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}