gemini_adk_rs/code_mode/
monty.rs

1//! The in-process sandbox: Monty on a blocking thread.
2
3use std::sync::mpsc;
4
5use async_trait::async_trait;
6use monty::{MontyRun, RunProgress};
7use monty_types::{
8    CompileOptions, ExcType, ExtFunctionResult, MontyException, MontyObject, NameLookupResult,
9    ObjectRef, PrintWriter, ResourceLimits, ResourceTracker,
10};
11use serde_json::{Map, Value, json};
12use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
13
14use super::CodeLimits;
15use super::sandbox::{CodeSandbox, Execution, Suspension};
16
17/// Runs programs in [Monty](https://github.com/pydantic/monty) on a blocking
18/// thread of this process.
19///
20/// Monty bounds a program's interpreter time, but in-process it cannot bound
21/// memory: a program that asks for a huge allocation can abort the process,
22/// or hold the thread past its time limit while the allocation completes.
23/// Use [`MontyPool`](super::MontyPool) where programs must not be able to
24/// take the process down.
25#[derive(Debug, Clone, Copy, Default)]
26pub struct InProcess;
27
28#[async_trait]
29impl CodeSandbox for InProcess {
30    async fn start(
31        &self,
32        program: &str,
33        functions: &[String],
34        limits: &CodeLimits,
35    ) -> (Box<dyn Execution>, Suspension) {
36        let (stops, stopped) = unbounded_channel();
37        let (answers, answered) = mpsc::channel();
38        let program = program.to_string();
39        let functions = functions.to_vec();
40        let limits = *limits;
41        tokio::task::spawn_blocking(move || run(program, functions, limits, &stops, &answered));
42        let mut execution = Thread { stopped, answers };
43        let first = execution.next().await;
44        (Box::new(execution), first)
45    }
46}
47
48/// The host's end of a program running on a blocking thread.
49struct Thread {
50    stopped: UnboundedReceiver<Suspension>,
51    answers: mpsc::Sender<Result<Value, String>>,
52}
53
54impl Thread {
55    async fn next(&mut self) -> Suspension {
56        self.stopped
57            .recv()
58            .await
59            .unwrap_or_else(|| Suspension::failed("the sandbox stopped"))
60    }
61}
62
63#[async_trait]
64impl Execution for Thread {
65    async fn resume(&mut self, answer: Result<Value, String>) -> Suspension {
66        if self.answers.send(answer).is_err() {
67            return Suspension::failed("the program has ended");
68        }
69        self.next().await
70    }
71}
72
73/// Run `program` to its end, reporting each call on `stops` and waiting on
74/// `answers` for its result. Returns early when the host drops its end.
75fn run(
76    program: String,
77    functions: Vec<String>,
78    limits: CodeLimits,
79    stops: &UnboundedSender<Suspension>,
80    answers: &mpsc::Receiver<Result<Value, String>>,
81) {
82    // Monty raises MemoryError past the print cap, which would end a program
83    // that only printed too much; collect with headroom and let the response
84    // cut what the model reads.
85    let cap = Some(limits.output_chars.saturating_mul(16).max(1 << 16));
86    let mut printed = String::new();
87    let finish = |result: Result<Value, String>, printed: String| {
88        let _ = stops.send(Suspension::Finished { result, printed });
89    };
90    let runner = match MontyRun::new(
91        program,
92        "main.py",
93        functions.clone(),
94        CompileOptions::default(),
95    ) {
96        Ok(runner) => runner,
97        Err(e) => return finish(Err(e.to_string()), printed),
98    };
99    let tracker = ResourceTracker::new(ResourceLimits {
100        max_feed_duration: Some(limits.cpu),
101        ..ResourceLimits::default()
102    });
103    let inputs = functions
104        .iter()
105        .map(|name| MontyObject::function(name.clone(), None))
106        .collect();
107    let mut progress = runner.start(
108        inputs,
109        tracker,
110        PrintWriter::CollectString(&mut printed, cap),
111    );
112    loop {
113        let step = match progress {
114            Err(e) => return finish(Err(e.to_string()), printed),
115            Ok(step) => step,
116        };
117        progress = match step {
118            RunProgress::Complete(value) => {
119                return finish(Ok(to_json(value.as_ref())), printed);
120            }
121            // Monty reports a call to any unknown name as an external call;
122            // only the functions given are the host's.
123            RunProgress::FunctionCall(call) if !functions.contains(&call.function_name) => {
124                let name = call.function_name.clone();
125                call.resume(
126                    ExtFunctionResult::NotFound(name),
127                    PrintWriter::CollectString(&mut printed, cap),
128                )
129            }
130            RunProgress::FunctionCall(call) => {
131                let positional = call.args.args().map(to_json).collect();
132                let keywords: Map<String, Value> = call
133                    .args
134                    .kwargs()
135                    .map(|(k, v)| (key(k), to_json(v)))
136                    .collect();
137                let stop = Suspension::Call {
138                    function: call.function_name.clone(),
139                    positional,
140                    keywords,
141                };
142                if stops.send(stop).is_err() {
143                    return;
144                }
145                // The host dropped its end: the program is cancelled.
146                let Ok(answer) = answers.recv() else {
147                    return;
148                };
149                let print = PrintWriter::CollectString(&mut printed, cap);
150                match answer {
151                    Ok(value) => call.resume(from_json(&value), print),
152                    Err(reason) => call.resume(
153                        MontyException::new(ExcType::RuntimeError, Some(reason)),
154                        print,
155                    ),
156                }
157            }
158            RunProgress::NameLookup(lookup) => lookup.resume(
159                NameLookupResult::Undefined,
160                PrintWriter::CollectString(&mut printed, cap),
161            ),
162            RunProgress::OsCall(os) => os.resume(
163                MontyException::new(
164                    ExcType::PermissionError,
165                    Some("this sandbox has no file system, network or clock access".into()),
166                ),
167                PrintWriter::CollectString(&mut printed, cap),
168            ),
169            RunProgress::ResolveFutures(_) => {
170                return finish(
171                    Err("async code is not supported; call the functions directly".into()),
172                    printed,
173                );
174            }
175        };
176    }
177}
178
179pub(super) fn key(k: ObjectRef<'_>) -> String {
180    k.as_str().map_or_else(|| k.py_repr(), str::to_string)
181}
182
183/// A Monty value as JSON. Tuples become lists; a value JSON cannot hold is
184/// given as its Python repr.
185pub(super) fn to_json(o: ObjectRef<'_>) -> Value {
186    match o.type_name() {
187        "NoneType" => Value::Null,
188        "bool" => o.as_bool().map_or(Value::Null, Value::Bool),
189        "int" => o
190            .as_int()
191            .map_or_else(|| Value::String(o.py_repr()), |i| json!(i)),
192        "float" => o
193            .as_float()
194            .and_then(serde_json::Number::from_f64)
195            .map_or(Value::Null, Value::Number),
196        "str" => o
197            .as_str()
198            .map_or(Value::Null, |s| Value::String(s.to_string())),
199        "list" | "tuple" => Value::Array(
200            o.items()
201                .unwrap_or_default()
202                .into_iter()
203                .map(to_json)
204                .collect(),
205        ),
206        "dict" => Value::Object(
207            o.pairs()
208                .unwrap_or_default()
209                .into_iter()
210                .map(|(k, v)| (key(k), to_json(v)))
211                .collect(),
212        ),
213        _ => Value::String(o.py_repr()),
214    }
215}
216
217/// JSON as a Monty value.
218pub(super) fn from_json(v: &Value) -> MontyObject {
219    match v {
220        Value::Null => MontyObject::none(),
221        Value::Bool(b) => MontyObject::bool(*b),
222        Value::Number(n) => n.as_i64().map_or_else(
223            || MontyObject::float(n.as_f64().unwrap_or(f64::NAN)),
224            MontyObject::int,
225        ),
226        Value::String(s) => MontyObject::string(s.clone()),
227        Value::Array(a) => MontyObject::list(a.iter().map(from_json)),
228        Value::Object(m) => MontyObject::dict(
229            m.iter()
230                .map(|(k, v)| (MontyObject::string(k.clone()), from_json(v))),
231        ),
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use std::time::Duration;
238
239    use super::*;
240
241    async fn finish(program: &str) -> Suspension {
242        let (_, stop) = InProcess.start(program, &[], &CodeLimits::default()).await;
243        stop
244    }
245
246    #[tokio::test]
247    async fn a_program_returns_its_last_expression_as_json() {
248        let stop = finish("x = [1, 2.5, 'a', None, True]\n{'x': x, 'n': len(x)}").await;
249        assert_eq!(
250            stop,
251            Suspension::Finished {
252                result: Ok(json!({"x": [1, 2.5, "a", null, true], "n": 5})),
253                printed: String::new(),
254            }
255        );
256    }
257
258    #[tokio::test]
259    async fn dates_and_printing_work() {
260        let stop = finish(
261            "import datetime\nd = datetime.date(2026, 10, 11)\nprint('day', d.weekday())\nd.isoformat()",
262        )
263        .await;
264        assert_eq!(
265            stop,
266            Suspension::Finished {
267                result: Ok(json!("2026-10-11")),
268                printed: "day 6\n".into(),
269            }
270        );
271    }
272
273    #[tokio::test]
274    async fn a_loop_ends_at_the_time_limit() {
275        let limits = CodeLimits {
276            cpu: Duration::from_millis(200),
277            ..CodeLimits::default()
278        };
279        let (_, stop) = InProcess.start("while True:\n    pass", &[], &limits).await;
280        let Suspension::Finished { result: Err(e), .. } = stop else {
281            panic!("expected a timeout, got {stop:?}");
282        };
283        assert!(e.contains("TimeoutError"), "{e}");
284    }
285
286    #[tokio::test]
287    async fn the_file_system_is_refused() {
288        let stop = finish("from pathlib import Path\nPath('/etc/passwd').read_text()").await;
289        let Suspension::Finished { result: Err(e), .. } = stop else {
290            panic!("expected a refusal, got {stop:?}");
291        };
292        assert!(e.contains("PermissionError"), "{e}");
293    }
294
295    #[tokio::test]
296    async fn a_call_suspends_until_the_host_answers() {
297        let functions = vec!["double".to_string()];
298        let (mut run, stop) = InProcess
299            .start(
300                "try:\n    double(1)\nexcept RuntimeError as e:\n    r = str(e)\n[double(n=21), r]",
301                &functions,
302                &CodeLimits::default(),
303            )
304            .await;
305        assert_eq!(
306            stop,
307            Suspension::Call {
308                function: "double".into(),
309                positional: vec![json!(1)],
310                keywords: Map::new(),
311            }
312        );
313        let stop = run.resume(Err("not yet".into())).await;
314        let Suspension::Call { keywords, .. } = stop else {
315            panic!("expected a second call, got {stop:?}");
316        };
317        assert_eq!(keywords.get("n"), Some(&json!(21)));
318        let stop = run.resume(Ok(json!({"value": 42}))).await;
319        assert_eq!(
320            stop,
321            Suspension::Finished {
322                result: Ok(json!([{"value": 42}, "not yet"])),
323                printed: String::new(),
324            }
325        );
326    }
327
328    #[tokio::test]
329    async fn an_undefined_name_raises_name_error() {
330        let stop = finish("lookup_everything()").await;
331        let Suspension::Finished { result: Err(e), .. } = stop else {
332            panic!("expected a NameError, got {stop:?}");
333        };
334        assert!(e.contains("NameError"), "{e}");
335    }
336}
gemini_adk_rs/code_mode/monty.rs

gemini_adk_rs/code_mode/
monty.rs
