johnandersen777.bsky.social/asyncai

Switch repository

Loading…

//! In-process sandbox: a dedicated OS thread owns the Monty REPL.
//!
//! `monty` values use `Rc` internally, so the REPL and its suspended snapshots
//! stay on one thread for their whole life. No heap cap: see ADR 0004 A5.

use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::Arc;

use monty::{MontyRepl, ReplProgress, ReplStartError};
use monty_types::{
    CompileOptions, ExcType, ExtFunctionResult, MontyException, MontyObject, NameLookupResult,
    PrintWriter, ResourceLimits, ResourceTracker,
};

use super::catalog::Catalog;
use super::convert::monty_to_json;
use super::protocol::{
    exception, CallDecision, Completion, HostEvent, HostLink, Job, SandboxSession,
};
use super::CodeModeLimits;

const ENGINE_STACK_BYTES: usize = 16 * 1024 * 1024;
const SCRIPT_NAME: &str = "run_code.py";

/// Handle to an engine thread. Dropping it stops the thread.
pub struct Engine {
    jobs: Sender<Job>,
}

impl Engine {
    pub fn spawn(
        catalog: Arc<Catalog>,
        limits: CodeModeLimits,
        preamble: Arc<String>,
    ) -> std::io::Result<Self> {
        let (jobs, receiver) = mpsc::channel::<Job>();
        std::thread::Builder::new()
            .name("asyncai-code-mode".to_string())
            .stack_size(ENGINE_STACK_BYTES)
            .spawn(move || engine_main(catalog, limits, preamble, receiver))?;
        Ok(Self { jobs })
    }
}

impl SandboxSession for Engine {
    fn submit(&self, job: Job) -> Result<(), String> {
        self.jobs
            .send(job)
            .map_err(|_| "code mode engine stopped".to_string())
    }
}

fn engine_main(
    catalog: Arc<Catalog>,
    limits: CodeModeLimits,
    preamble: Arc<String>,
    jobs: Receiver<Job>,
) {
    let mut repl: Option<MontyRepl> = None;
    while let Ok(mut job) = jobs.recv() {
        if job.restart {
            repl = None;
        }
        let session = match repl.take() {
            Some(session) => session,
            None => match new_repl(&catalog, &limits, &preamble) {
                Ok(session) => session,
                Err(failed) => {
                    let (session, error) = *failed;
                    repl = Some(session);
                    let _ = job.events.send(HostEvent::Done(Err(format!(
                        "code mode preamble failed: {error}"
                    ))));
                    continue;
                }
            },
        };
        let (session, outcome) = run_snippet(session, &catalog, &limits, &mut job);
        repl = Some(session);
        let _ = job.events.send(HostEvent::Done(outcome));
    }
}

/// A fresh REPL with the helpers and operator preamble already fed.
fn new_repl(
    catalog: &Catalog,
    limits: &CodeModeLimits,
    preamble: &str,
) -> Result<MontyRepl, Box<(MontyRepl, String)>> {
    let resource_limits = ResourceLimits::default()
        .max_feed_duration(limits.max_duration)
        .max_suspensions(limits.max_suspensions);
    let repl = MontyRepl::new(
        SCRIPT_NAME,
        ResourceTracker::new(resource_limits),
        CompileOptions::default(),
    );
    if preamble.trim().is_empty() {
        return Ok(repl);
    }

    let (events, _events_rx) = tokio::sync::mpsc::unbounded_channel();
    let (_results_tx, results) = tokio::sync::mpsc::unbounded_channel();
    let mut job = Job {
        code: preamble.to_string(),
        restart: false,
        events,
        results,
    };
    match run_snippet(repl, catalog, limits, &mut job) {
        (repl, Ok(_)) => Ok(repl),
        (repl, Err(error)) => Err(Box::new((repl, error))),
    }
}

fn run_snippet(
    repl: MontyRepl,
    catalog: &Catalog,
    limits: &CodeModeLimits,
    job: &mut Job,
) -> (MontyRepl, Result<Completion, String>) {
    let code = job.code.clone();
    let mut prints = String::new();
    let mut link = HostLink::new(catalog, limits, job);
    let max_print = Some(limits.max_print_bytes);

    let mut step = repl.feed_start(
        &code,
        Vec::<(String, MontyObject)>::new(),
        PrintWriter::CollectString(&mut prints, max_print),
    );

    loop {
        let progress = match step {
            Ok(progress) => progress,
            Err(error) => {
                let ReplStartError { repl, error } = *error;
                return (repl, Err(error.to_string()));
            }
        };

        if !matches!(progress, ReplProgress::Complete { .. }) {
            if let Err(error) = link.suspend() {
                step = abort(
                    progress,
                    error,
                    PrintWriter::CollectString(&mut prints, max_print),
                );
                continue;
            }
        }

        step = match progress {
            ReplProgress::Complete { repl, value } => {
                let value = if value.type_name() == "NoneType" {
                    None
                } else {
                    Some(monty_to_json(&value))
                };
                return (repl, Ok(Completion { value, prints }));
            }
            ReplProgress::NameLookup(lookup) => {
                let result = match link.lookup(&lookup.name) {
                    Some(function) => NameLookupResult::Value(function),
                    None => NameLookupResult::Undefined,
                };
                lookup.resume(result, PrintWriter::CollectString(&mut prints, max_print))
            }
            ReplProgress::FunctionCall(call) => {
                let decision = link.call(
                    &call.function_name,
                    call.object_id.is_some(),
                    &call.args,
                    call.call_id,
                );
                let print = PrintWriter::CollectString(&mut prints, max_print);
                match decision {
                    CallDecision::Pending => call.resume_pending(print),
                    CallDecision::Raise(error) => {
                        call.resume(ExtFunctionResult::Error(error), print)
                    }
                    CallDecision::NotFound => {
                        let name = call.function_name.clone();
                        call.resume(ExtFunctionResult::NotFound(name), print)
                    }
                }
            }
            ReplProgress::ResolveFutures(futures) => {
                let pending = futures.pending_call_ids().to_vec();
                let mut ready = link.take_ready(&pending);
                let mut failure = None;
                while ready.is_empty() {
                    if let Err(message) = link.wait_blocking() {
                        failure = Some(message);
                        break;
                    }
                    ready = link.take_ready(&pending);
                }
                let print = PrintWriter::CollectString(&mut prints, max_print);
                match failure {
                    Some(message) => {
                        futures.abort(exception(ExcType::RuntimeError, message), print)
                    }
                    None => futures.resume(
                        ready
                            .into_iter()
                            .map(|(call_id, value)| {
                                let result = match value {
                                    Ok(object) => ExtFunctionResult::Return(object),
                                    Err(error) => ExtFunctionResult::Error(error),
                                };
                                (call_id, result)
                            })
                            .collect(),
                        print,
                    ),
                }
            }
            ReplProgress::OsCall(call) => call.resume(
                ExtFunctionResult::Error(exception(
                    ExcType::PermissionError,
                    "os access disabled in asyncai code mode",
                )),
                PrintWriter::CollectString(&mut prints, max_print),
            ),
        };
    }
}

fn abort(
    progress: ReplProgress,
    error: MontyException,
    print: PrintWriter<'_>,
) -> Result<ReplProgress, Box<ReplStartError>> {
    match progress {
        ReplProgress::FunctionCall(call) => call.abort(error, print),
        ReplProgress::OsCall(call) => call.abort(error, print),
        ReplProgress::NameLookup(lookup) => lookup.abort(error, print),
        ReplProgress::ResolveFutures(futures) => futures.abort(error, print),
        ReplProgress::Complete { repl, .. } => Err(Box::new(ReplStartError { repl, error })),
    }
}