1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::{Arc, Condvar, Mutex};
5use std::time::{Duration, Instant};
6
7use serde_json::Value;
8
9use crate::tools::ApprovalDecision;
10use crate::types::{Metadata, ToolCall};
11
12pub type ApprovalFuture<T> = Pin<Box<dyn Future<Output = Result<T, ApprovalError>> + Send>>;
13
14pub trait ApprovalProvider: Send + Sync {
15 fn should_request(&self, request: &ApprovalRequest) -> bool;
16 fn decide(&self, request: &ApprovalRequest) -> ApprovalFuture<Option<ApprovalDecision>>;
17}
18
19#[derive(Debug, Clone, PartialEq)]
20pub struct ApprovalRequest {
21 pub request_id: String,
22 pub run_id: String,
23 pub trace_id: String,
24 pub agent_name: String,
25 pub cycle_index: u32,
26 pub tool_call_id: String,
27 pub tool_name: String,
28 pub arguments: Value,
29 pub preview: String,
30 pub metadata: Metadata,
31}
32
33impl ApprovalRequest {
34 pub fn for_tool_call(
35 run_id: impl Into<String>,
36 trace_id: impl Into<String>,
37 agent_name: impl Into<String>,
38 cycle_index: u32,
39 call: &ToolCall,
40 ) -> Self {
41 let run_id = run_id.into();
42 let trace_id = trace_id.into();
43 let agent_name = agent_name.into();
44 let arguments = Value::Object(call.arguments.clone().into_iter().collect());
45 Self {
46 request_id: format!("approval:{run_id}:{}:{}", call.name, call.id),
47 run_id,
48 trace_id,
49 agent_name,
50 cycle_index,
51 tool_call_id: call.id.clone(),
52 tool_name: call.name.clone(),
53 preview: format!("{} {}", call.name, arguments),
54 arguments,
55 metadata: Metadata::new(),
56 }
57 }
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub struct ApprovalError {
62 message: String,
63}
64
65impl ApprovalError {
66 pub fn new(message: impl Into<String>) -> Self {
67 Self {
68 message: message.into(),
69 }
70 }
71
72 pub fn message(&self) -> &str {
73 &self.message
74 }
75}
76
77impl std::fmt::Display for ApprovalError {
78 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79 formatter.write_str(&self.message)
80 }
81}
82
83impl std::error::Error for ApprovalError {}
84
85#[derive(Clone, Default)]
86pub struct ApprovalBroker {
87 inner: Arc<ApprovalBrokerInner>,
88}
89
90#[derive(Default)]
91struct ApprovalBrokerInner {
92 pending: Mutex<HashMap<String, PendingApproval>>,
93 changed: Condvar,
94}
95
96struct PendingApproval {
97 request: ApprovalRequest,
98 decision: Option<ApprovalDecision>,
99}
100
101impl ApprovalBroker {
102 pub fn register(&self, request: ApprovalRequest) -> Result<(), ApprovalError> {
103 let mut pending = self
104 .inner
105 .pending
106 .lock()
107 .map_err(|_| ApprovalError::new("approval broker lock poisoned"))?;
108 pending.insert(
109 request.request_id.clone(),
110 PendingApproval {
111 request,
112 decision: None,
113 },
114 );
115 self.inner.changed.notify_all();
116 Ok(())
117 }
118
119 pub fn resolve(
120 &self,
121 request_id: impl AsRef<str>,
122 decision: ApprovalDecision,
123 ) -> Result<(), ApprovalError> {
124 let mut pending = self
125 .inner
126 .pending
127 .lock()
128 .map_err(|_| ApprovalError::new("approval broker lock poisoned"))?;
129 let request_id = request_id.as_ref();
130 let Some(entry) = pending.get_mut(request_id) else {
131 return Err(ApprovalError::new(format!(
132 "unknown approval request: {request_id}"
133 )));
134 };
135 entry.decision = Some(decision);
136 self.inner.changed.notify_all();
137 Ok(())
138 }
139
140 pub fn pending_request(&self, request_id: impl AsRef<str>) -> Option<ApprovalRequest> {
141 self.inner.pending.lock().ok().and_then(|pending| {
142 pending
143 .get(request_id.as_ref())
144 .map(|entry| entry.request.clone())
145 })
146 }
147
148 pub(crate) fn wait_blocking(
149 &self,
150 request_id: &str,
151 timeout: Option<Duration>,
152 ) -> Result<ApprovalDecision, ApprovalError> {
153 let started = Instant::now();
154 let mut pending = self
155 .inner
156 .pending
157 .lock()
158 .map_err(|_| ApprovalError::new("approval broker lock poisoned"))?;
159 loop {
160 if let Some(decision) = pending
161 .get(request_id)
162 .and_then(|entry| entry.decision.clone())
163 {
164 pending.remove(request_id);
165 return Ok(decision);
166 }
167
168 if let Some(timeout) = timeout {
169 let elapsed = started.elapsed();
170 if elapsed >= timeout {
171 pending.remove(request_id);
172 return Ok(ApprovalDecision::timeout("approval timed out"));
173 }
174 let remaining = timeout.saturating_sub(elapsed);
175 let (next_pending, wait_result) = self
176 .inner
177 .changed
178 .wait_timeout(pending, remaining)
179 .map_err(|_| ApprovalError::new("approval broker lock poisoned"))?;
180 pending = next_pending;
181 if wait_result.timed_out() {
182 pending.remove(request_id);
183 return Ok(ApprovalDecision::timeout("approval timed out"));
184 }
185 } else {
186 pending = self
187 .inner
188 .changed
189 .wait(pending)
190 .map_err(|_| ApprovalError::new("approval broker lock poisoned"))?;
191 }
192 }
193 }
194}
195
196pub(crate) fn block_on_approval_future<T: Send + 'static>(
197 future: ApprovalFuture<T>,
198) -> Result<T, ApprovalError> {
199 if let Ok(handle) = tokio::runtime::Handle::try_current() {
200 if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
201 tokio::task::block_in_place(|| handle.block_on(future))
202 } else {
203 std::thread::spawn(move || {
204 tokio::runtime::Builder::new_current_thread()
205 .enable_all()
206 .build()
207 .map_err(|error| ApprovalError::new(error.to_string()))?
208 .block_on(future)
209 })
210 .join()
211 .map_err(|_| ApprovalError::new("approval future thread panicked"))?
212 }
213 } else {
214 tokio::runtime::Builder::new_current_thread()
215 .enable_all()
216 .build()
217 .map_err(|error| ApprovalError::new(error.to_string()))?
218 .block_on(future)
219 }
220}