@ondeinference / ed
Code Commits Issues Pulls Cloud Agent CI/CD
main
main release/1.3.1
rs 208 lines 6.63 KB
Raw
1 //! Approvals answered by the webview.
2
3 use std::collections::HashMap;
4 use std::time::Duration;
5
6 use async_trait::async_trait;
7 use ed_agent::{ApprovalDecision, ApprovalHandler, ApprovalRequest};
8 use log::{error, warn};
9 use serde::Serialize;
10 use tauri::{AppHandle, Emitter};
11 use tokio::sync::{oneshot, Mutex};
12
13 use crate::events::{
14 ApprovalRequestedPayload, ApprovalResolvedPayload, EVENT_CHAT_APPROVAL_REQUESTED,
15 EVENT_CHAT_APPROVAL_RESOLVED,
16 };
17 use crate::scope::TurnScope;
18
19 /// How long to wait for the user before denying.
20 pub const DEFAULT_APPROVAL_TIMEOUT: Duration = Duration::from_secs(120);
21
22 /// An [`ApprovalHandler`] that asks the webview.
23 ///
24 /// For each request it emits `chat_approval_requested`, waits for the webview
25 /// to call `chat_respond_approval` (or [`respond`](TauriApprovals::respond)),
26 /// and emits `chat_approval_resolved` when the wait ends, however it ends.
27 /// A timeout, a closed window or [`cancel_all`](TauriApprovals::cancel_all)
28 /// all mean *deny*.
29 type Emit = Box<dyn Fn(&str, serde_json::Value) -> Result<(), String> + Send + Sync>;
30
31 pub struct TauriApprovals {
32 emit: Emit,
33 scope: TurnScope,
34 timeout: Duration,
35 pending: Mutex<HashMap<String, oneshot::Sender<ApprovalDecision>>>,
36 }
37
38 impl TauriApprovals {
39 pub fn new(app: AppHandle, timeout: Duration) -> Self {
40 Self::scoped(app, TurnScope::default(), timeout)
41 }
42
43 pub(crate) fn scoped(app: AppHandle, scope: TurnScope, timeout: Duration) -> Self {
44 let emit: Emit =
45 Box::new(move |event, payload| app.emit(event, payload).map_err(|err| err.to_string()));
46 Self::with_emitter(emit, scope, timeout)
47 }
48
49 fn with_emitter(emit: Emit, scope: TurnScope, timeout: Duration) -> Self {
50 Self {
51 emit,
52 scope,
53 timeout,
54 pending: Mutex::new(HashMap::new()),
55 }
56 }
57
58 fn emit(&self, event: &str, payload: impl Serialize) -> Result<(), String> {
59 let value = serde_json::to_value(payload).map_err(|err| err.to_string())?;
60 (self.emit)(event, value)
61 }
62
63 /// Answer a pending request. Errors if it already timed out or was
64 /// answered.
65 pub async fn respond(
66 &self,
67 request_id: &str,
68 decision: ApprovalDecision,
69 ) -> Result<(), String> {
70 let gone = || "This approval request is no longer pending.".to_string();
71 let sender = self
72 .pending
73 .lock()
74 .await
75 .remove(request_id)
76 .ok_or_else(gone)?;
77 sender.send(decision).map_err(|_| gone())
78 }
79
80 /// Deny everything that is waiting. Use it when the conversation the
81 /// requests belong to goes away.
82 pub async fn cancel_all(&self) {
83 let pending: Vec<_> = self.pending.lock().await.drain().collect();
84 for (_, sender) in pending {
85 // Each waiting `approve` call emits its own resolved event once it
86 // sees the denial, so every request closes exactly once.
87 let _ = sender.send(ApprovalDecision::Deny);
88 }
89 }
90 }
91
92 #[async_trait]
93 impl ApprovalHandler for TauriApprovals {
94 async fn approve(&self, request: ApprovalRequest) -> ApprovalDecision {
95 let request_id = uuid::Uuid::new_v4().to_string();
96 let (sender, receiver) = oneshot::channel();
97 self.pending.lock().await.insert(request_id.clone(), sender);
98
99 let (session, _) = self.scope.get();
100 let payload = ApprovalRequestedPayload {
101 request_id: request_id.clone(),
102 session,
103 call: request.call,
104 risk: request.risk,
105 };
106 if let Err(err) = self.emit(EVENT_CHAT_APPROVAL_REQUESTED, payload) {
107 self.pending.lock().await.remove(&request_id);
108 error!("ed-tauri: could not request approval: {err}");
109 return ApprovalDecision::Deny;
110 }
111
112 let decision = match tokio::time::timeout(self.timeout, receiver).await {
113 Ok(Ok(decision)) => decision,
114 Ok(Err(_)) => ApprovalDecision::Deny,
115 Err(_) => {
116 warn!("ed-tauri: approval {request_id} timed out; denying");
117 ApprovalDecision::Deny
118 }
119 };
120 self.pending.lock().await.remove(&request_id);
121 if let Err(err) = self.emit(
122 EVENT_CHAT_APPROVAL_RESOLVED,
123 ApprovalResolvedPayload { request_id },
124 ) {
125 error!("ed-tauri: could not close approval: {err}");
126 }
127 decision
128 }
129 }
130
131 #[cfg(test)]
132 mod tests {
133 use std::sync::Arc;
134
135 use ed_agent::{ToolCall, ToolRisk};
136
137 use super::*;
138
139 fn request() -> ApprovalRequest {
140 ApprovalRequest {
141 call: ToolCall {
142 id: "1".into(),
143 name: "send_message".into(),
144 arguments: "{}".into(),
145 },
146 risk: ToolRisk::Mutating,
147 }
148 }
149
150 fn approvals(timeout: Duration) -> Arc<TauriApprovals> {
151 Arc::new(TauriApprovals::with_emitter(
152 Box::new(|_, _| Ok(())),
153 TurnScope::default(),
154 timeout,
155 ))
156 }
157
158 async fn pending_id(approvals: &TauriApprovals) -> String {
159 loop {
160 if let Some(id) = approvals.pending.lock().await.keys().next().cloned() {
161 return id;
162 }
163 tokio::task::yield_now().await;
164 }
165 }
166
167 #[tokio::test]
168 async fn the_webview_can_allow() {
169 let approvals = approvals(Duration::from_secs(5));
170 let waiting = tokio::spawn({
171 let approvals = approvals.clone();
172 async move { approvals.approve(request()).await }
173 });
174 let id = pending_id(&approvals).await;
175
176 approvals
177 .respond(&id, ApprovalDecision::AllowForSession)
178 .await
179 .unwrap();
180
181 assert_eq!(waiting.await.unwrap(), ApprovalDecision::AllowForSession);
182 assert!(approvals
183 .respond(&id, ApprovalDecision::Deny)
184 .await
185 .is_err());
186 }
187
188 #[tokio::test]
189 async fn cancel_all_denies_everything_pending() {
190 let approvals = approvals(Duration::from_secs(5));
191 let waiting = tokio::spawn({
192 let approvals = approvals.clone();
193 async move { approvals.approve(request()).await }
194 });
195 pending_id(&approvals).await;
196
197 approvals.cancel_all().await;
198
199 assert_eq!(waiting.await.unwrap(), ApprovalDecision::Deny);
200 }
201
202 #[tokio::test]
203 async fn an_unanswered_request_times_out_as_a_denial() {
204 let approvals = approvals(Duration::from_millis(20));
205 assert_eq!(approvals.approve(request()).await, ApprovalDecision::Deny);
206 assert!(approvals.pending.lock().await.is_empty());
207 }
208 }