gemini_adk_rs/code_mode/
pool.rs

1//! The worker-process sandbox: Monty in a pool of subprocesses.
2
3use std::path::PathBuf;
4use std::sync::Arc;
5use std::time::Duration;
6
7use async_trait::async_trait;
8use monty_pool::{
9    Checkout, Pool, PoolConfig, PoolError, ReplConfig, ResumeValue, TurnEvent, on_print_sync,
10};
11use monty_types::{ExcType, MontyException, MontyObject, NameLookupResult, ResourceLimits};
12use serde_json::{Map, Value};
13
14use super::CodeLimits;
15use super::monty::{from_json, key, to_json};
16use super::sandbox::{CodeSandbox, Execution, Suspension};
17
18/// Runs programs in a pool of Monty worker processes.
19///
20/// Each program gets a worker of its own. A crash, an allocation past the
21/// memory limit or a hang kills only that worker, which the pool replaces;
22/// the session's process and every other call are untouched. Workers are the
23/// `monty` binary: `cargo install monty-runtime`, or `pip install
24/// pydantic-monty-runtime`.
25#[derive(Clone)]
26pub struct MontyPool {
27    pool: Arc<Pool>,
28}
29
30impl std::fmt::Debug for MontyPool {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("MontyPool").finish_non_exhaustive()
33    }
34}
35
36impl MontyPool {
37    /// Start a pool of up to `workers` processes of the `monty` binary at
38    /// `binary`, one kept warm. A turn that outlives 30 s kills its worker,
39    /// whatever the program's own limits say.
40    pub async fn spawn(binary: impl Into<PathBuf>, workers: usize) -> Result<Self, String> {
41        let mut config = PoolConfig::subprocess(binary);
42        config.min_processes = 1.min(workers);
43        config.max_processes = workers.max(1);
44        config.request_timeout = Some(Duration::from_secs(30));
45        config.checkout_timeout = Some(Duration::from_secs(5));
46        Self::with_config(config).await
47    }
48
49    /// A pool built from `config`, for a remote relay or other settings.
50    pub async fn with_config(config: PoolConfig) -> Result<Self, String> {
51        let pool = Pool::new(config).await.map_err(|e| e.to_string())?;
52        Ok(Self {
53            pool: Arc::new(pool),
54        })
55    }
56}
57
58#[async_trait]
59impl CodeSandbox for MontyPool {
60    async fn start(
61        &self,
62        program: &str,
63        functions: &[String],
64        limits: &CodeLimits,
65    ) -> (Box<dyn Execution>, Suspension) {
66        let repl = ReplConfig {
67            limits: Some(ResourceLimits {
68                max_feed_duration: Some(limits.cpu),
69                max_memory: Some(limits.memory),
70                ..ResourceLimits::default()
71            }),
72            ..ReplConfig::default()
73        };
74        let mut execution = Worker {
75            checkout: None,
76            functions: functions.to_vec(),
77            printed: String::new(),
78            cap: limits.output_chars,
79        };
80        let mut checkout = match self.pool.checkout(&repl).await {
81            Ok(checkout) => checkout,
82            Err(e) => {
83                return (
84                    Box::new(execution),
85                    Suspension::failed(format!("no sandbox worker: {e}")),
86                );
87            }
88        };
89        let inputs: Vec<(String, MontyObject)> = functions
90            .iter()
91            .map(|name| (name.clone(), MontyObject::function(name.clone(), None)))
92            .collect();
93        let event = {
94            let printed = &mut execution.printed;
95            let cap = execution.cap;
96            let mut on_print = on_print_sync(move |_, text| append(printed, text, cap));
97            checkout
98                .feed(program, inputs, Vec::new(), true, &mut on_print)
99                .await
100        };
101        execution.checkout = Some(checkout);
102        let first = execution.settle(event).await;
103        (Box::new(execution), first)
104    }
105}
106
107/// One program's worker, held from its start to its end.
108struct Worker {
109    /// `None` once the program has ended and the worker is returned or
110    /// discarded.
111    checkout: Option<Checkout>,
112    functions: Vec<String>,
113    printed: String,
114    cap: usize,
115}
116
117/// How to answer an event the host does not see.
118enum Next {
119    Resume(ResumeValue),
120    Undefined,
121}
122
123fn append(printed: &mut String, text: &str, cap: usize) {
124    if printed.chars().count() < cap {
125        printed.push_str(text);
126    }
127}
128
129impl Worker {
130    /// Answer every event the host need not see (an unknown name, an OS
131    /// call) and return where the program stops for the host.
132    async fn settle(&mut self, mut event: Result<TurnEvent, PoolError>) -> Suspension {
133        loop {
134            let next = match event {
135                Ok(TurnEvent::FunctionCall {
136                    function_name,
137                    args,
138                    ..
139                }) => {
140                    if self.functions.contains(&function_name) {
141                        let positional = args.args().map(to_json).collect();
142                        let keywords: Map<String, Value> =
143                            args.kwargs().map(|(k, v)| (key(k), to_json(v))).collect();
144                        return Suspension::Call {
145                            function: function_name,
146                            positional,
147                            keywords,
148                        };
149                    }
150                    Next::Resume(ResumeValue::Error(MontyException::new(
151                        ExcType::NameError,
152                        Some(format!("name '{function_name}' is not defined")),
153                    )))
154                }
155                Ok(TurnEvent::OsCall { .. }) => {
156                    Next::Resume(ResumeValue::Error(MontyException::new(
157                        ExcType::PermissionError,
158                        Some("this sandbox has no file system, network or clock access".into()),
159                    )))
160                }
161                Ok(TurnEvent::NameLookup { .. }) => Next::Undefined,
162                Ok(TurnEvent::ResolveFutures { .. }) => {
163                    // Mid-execution state is not returned to the pool.
164                    self.checkout = None;
165                    return self.finished(Err(
166                        "async code is not supported; call the functions directly".into(),
167                    ));
168                }
169                Ok(TurnEvent::Complete(value)) => {
170                    let result = to_json(value.as_ref());
171                    self.release().await;
172                    return self.finished(Ok(result));
173                }
174                Err(PoolError::Runtime(exception)) => {
175                    // After a time or memory limit the worker's heap is not
176                    // to be trusted: discard it rather than return it.
177                    if matches!(
178                        exception.exc_type(),
179                        ExcType::TimeoutError | ExcType::MemoryError
180                    ) {
181                        self.checkout = None;
182                    } else {
183                        self.release().await;
184                    }
185                    return self.finished(Err(exception.to_string()));
186                }
187                Err(error) => {
188                    self.checkout = None;
189                    return self.finished(Err(format!("the sandbox worker failed: {error}")));
190                }
191            };
192            let Some(checkout) = self.checkout.as_mut() else {
193                return Suspension::failed("the program has ended");
194            };
195            let printed = &mut self.printed;
196            let cap = self.cap;
197            let mut on_print = on_print_sync(move |_, text| append(printed, text, cap));
198            event = match next {
199                Next::Resume(value) => checkout.resume(value, &mut on_print).await,
200                Next::Undefined => {
201                    checkout
202                        .resume_name_lookup(NameLookupResult::Undefined, &mut on_print)
203                        .await
204                }
205            };
206        }
207    }
208
209    /// Return the worker to the pool.
210    async fn release(&mut self) {
211        if let Some(checkout) = self.checkout.take() {
212            let _ = checkout.finish().await;
213        }
214    }
215
216    fn finished(&mut self, result: Result<Value, String>) -> Suspension {
217        Suspension::Finished {
218            result,
219            printed: std::mem::take(&mut self.printed),
220        }
221    }
222}
223
224#[async_trait]
225impl Execution for Worker {
226    async fn resume(&mut self, answer: Result<Value, String>) -> Suspension {
227        let Some(checkout) = self.checkout.as_mut() else {
228            return Suspension::failed("the program has ended");
229        };
230        let value = match answer {
231            Ok(value) => ResumeValue::Return(from_json(&value)),
232            Err(reason) => {
233                ResumeValue::Error(MontyException::new(ExcType::RuntimeError, Some(reason)))
234            }
235        };
236        let event = {
237            let printed = &mut self.printed;
238            let cap = self.cap;
239            let mut on_print = on_print_sync(move |_, text| append(printed, text, cap));
240            checkout.resume(value, &mut on_print).await
241        };
242        self.settle(event).await
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    /// The worker binary, from `MONTY_BIN`. Without it these tests pass
251    /// without running.
252    fn binary() -> Option<PathBuf> {
253        std::env::var_os("MONTY_BIN").map(PathBuf::from)
254    }
255
256    #[tokio::test]
257    async fn a_program_runs_in_a_worker_and_calls_back() {
258        let Some(bin) = binary() else { return };
259        let pool = MontyPool::spawn(bin, 2).await.expect("pool");
260        let functions = vec!["double".to_string()];
261        let (mut run, stop) = pool
262            .start(
263                "print('hi')\n[double(21), sum(range(4))]",
264                &functions,
265                &CodeLimits::default(),
266            )
267            .await;
268        assert!(matches!(stop, Suspension::Call { .. }), "{stop:?}");
269        let stop = run.resume(Ok(serde_json::json!(42))).await;
270        assert_eq!(
271            stop,
272            Suspension::Finished {
273                result: Ok(serde_json::json!([42, 6])),
274                printed: "hi\n".into(),
275            }
276        );
277    }
278
279    #[tokio::test]
280    async fn a_huge_allocation_kills_only_the_worker() {
281        let Some(bin) = binary() else { return };
282        let pool = MontyPool::spawn(bin, 1).await.expect("pool");
283        let (_, stop) = pool
284            .start("len('a' * 10**10)", &[], &CodeLimits::default())
285            .await;
286        let Suspension::Finished { result: Err(e), .. } = stop else {
287            panic!("expected a memory error, got {stop:?}");
288        };
289        assert!(e.contains("MemoryError") || e.contains("memory"), "{e}");
290        // The pool replaced the worker: the next program runs.
291        let (_, stop) = pool.start("1 + 1", &[], &CodeLimits::default()).await;
292        assert_eq!(
293            stop,
294            Suspension::Finished {
295                result: Ok(serde_json::json!(2)),
296                printed: String::new(),
297            }
298        );
299    }
300}
gemini_adk_rs/code_mode/pool.rs

gemini_adk_rs/code_mode/
pool.rs
