agentstategraph_policy/
external.rs1use std::collections::HashMap;
26use std::io::{Read, Write};
27use std::process::{Child, Output};
28use std::sync::Arc;
29use std::thread::{self, JoinHandle};
30use std::time::{Duration, Instant};
31
32use crate::selector::Situation;
33use crate::types::{Decision, EvaluatorSource};
34
35#[derive(Debug, thiserror::Error)]
37pub enum ExternalError {
38 #[error("evaluator kind '{0}' not registered")]
42 NotRegistered(String),
43
44 #[error("source resolution failed: {0}")]
48 SourceResolution(String),
49
50 #[error("evaluator failed: {0}")]
53 Execution(String),
54
55 #[error(transparent)]
57 Other(#[from] Box<dyn std::error::Error + Send + Sync>),
58}
59
60impl ExternalError {
61 pub fn code(&self) -> &'static str {
66 match self {
67 ExternalError::NotRegistered(_) => "evaluator_not_registered",
68 ExternalError::SourceResolution(_) => "evaluator_source_unresolved",
69 ExternalError::Execution(_) => "evaluator_failed",
70 ExternalError::Other(e) if e.is::<TimedOut>() => "evaluator_timeout",
71 ExternalError::Other(_) => "evaluator_error",
72 }
73 }
74}
75
76pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
79
80#[derive(Debug, thiserror::Error)]
83#[error("evaluator timed out after {0:?}")]
84pub struct TimedOut(pub Duration);
85
86pub fn wait_with_timeout(
92 mut child: Child,
93 input: Option<Vec<u8>>,
94 timeout: Duration,
95) -> Result<Output, ExternalError> {
96 let started = Instant::now();
97 let stdin = child.stdin.take();
100 let writer = input.and_then(|input| {
101 let mut stdin = stdin?;
102 Some(thread::spawn(move || stdin.write_all(&input)))
103 });
104 let stdout = child.stdout.take().map(read_on_thread);
105 let stderr = child.stderr.take().map(read_on_thread);
106
107 let mut pause = Duration::from_millis(1);
108 let status = loop {
109 let exited = child
110 .try_wait()
111 .map_err(|e| ExternalError::Execution(format!("wait: {e}")))?;
112 if let Some(status) = exited {
113 break status;
114 }
115 let elapsed = started.elapsed();
116 if elapsed >= timeout {
117 let _ = child.kill();
118 let _ = child.wait();
119 return Err(ExternalError::Other(Box::new(TimedOut(timeout))));
120 }
121 thread::sleep(pause.min(timeout - elapsed));
122 pause = (pause * 2).min(Duration::from_millis(50));
123 };
124
125 if let Some(Ok(Err(e))) = writer.map(JoinHandle::join) {
126 if e.kind() != std::io::ErrorKind::BrokenPipe {
129 return Err(ExternalError::Execution(format!("write stdin: {e}")));
130 }
131 }
132 let collect = |pipe: Option<JoinHandle<Vec<u8>>>| {
133 pipe.and_then(|reader| reader.join().ok())
134 .unwrap_or_default()
135 };
136 Ok(Output {
137 status,
138 stdout: collect(stdout),
139 stderr: collect(stderr),
140 })
141}
142
143fn read_on_thread(mut pipe: impl Read + Send + 'static) -> JoinHandle<Vec<u8>> {
144 thread::spawn(move || {
145 let mut out = Vec::new();
146 let _ = pipe.read_to_end(&mut out);
147 out
148 })
149}
150
151pub trait ExternalEvaluator: Send + Sync {
157 fn kind(&self) -> &'static str;
161
162 fn evaluate(
167 &self,
168 source: &EvaluatorSource,
169 situation: &Situation,
170 action: &str,
171 agent_id: &str,
172 ) -> Result<Decision, ExternalError>;
173}
174
175#[derive(Default)]
184pub struct ExternalEvaluatorRegistry {
185 impls: HashMap<&'static str, Arc<dyn ExternalEvaluator>>,
186}
187
188impl ExternalEvaluatorRegistry {
189 pub fn new() -> Self {
191 Self {
192 impls: HashMap::new(),
193 }
194 }
195
196 pub fn register(&mut self, eval: Arc<dyn ExternalEvaluator>) {
199 self.impls.insert(eval.kind(), eval);
200 }
201
202 pub fn get(&self, kind: &str) -> Option<&Arc<dyn ExternalEvaluator>> {
204 self.impls.get(kind)
205 }
206
207 pub fn kinds(&self) -> Vec<&'static str> {
209 self.impls.keys().copied().collect()
210 }
211
212 pub fn len(&self) -> usize {
214 self.impls.len()
215 }
216
217 pub fn is_empty(&self) -> bool {
219 self.impls.is_empty()
220 }
221}
222
223impl std::fmt::Debug for ExternalEvaluatorRegistry {
224 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
225 f.debug_struct("ExternalEvaluatorRegistry")
226 .field("kinds", &self.kinds())
227 .finish()
228 }
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234
235 struct MockRunner;
236 impl ExternalEvaluator for MockRunner {
237 fn kind(&self) -> &'static str {
238 "mock"
239 }
240 fn evaluate(
241 &self,
242 _source: &EvaluatorSource,
243 _situation: &Situation,
244 _action: &str,
245 _agent_id: &str,
246 ) -> Result<Decision, ExternalError> {
247 Ok(Decision::NoPolicyMatch)
248 }
249 }
250
251 #[cfg(unix)]
252 #[test]
253 fn a_process_past_its_timeout_is_killed() {
254 use std::process::{Command, Stdio};
255 let child = Command::new("sleep")
256 .arg("5")
257 .stdout(Stdio::piped())
258 .stderr(Stdio::piped())
259 .spawn()
260 .unwrap();
261 let started = Instant::now();
262 let error = wait_with_timeout(child, None, Duration::from_millis(100)).unwrap_err();
263 assert!(started.elapsed() < Duration::from_secs(4));
264 assert_eq!(error.code(), "evaluator_timeout");
265 }
266
267 #[cfg(unix)]
268 #[test]
269 fn a_process_gets_its_input_and_returns_its_output() {
270 use std::process::{Command, Stdio};
271 let child = Command::new("cat")
272 .stdin(Stdio::piped())
273 .stdout(Stdio::piped())
274 .stderr(Stdio::piped())
275 .spawn()
276 .unwrap();
277 let input = vec![b'x'; 1 << 20];
279 let output = wait_with_timeout(child, Some(input.clone()), DEFAULT_TIMEOUT).unwrap();
280 assert!(output.status.success());
281 assert_eq!(output.stdout, input);
282 }
283
284 #[test]
285 fn registry_register_and_get() {
286 let mut reg = ExternalEvaluatorRegistry::new();
287 assert!(reg.is_empty());
288 reg.register(Arc::new(MockRunner));
289 assert_eq!(reg.len(), 1);
290 assert!(reg.get("mock").is_some());
291 assert!(reg.get("missing").is_none());
292 assert_eq!(reg.kinds(), vec!["mock"]);
293 }
294}