johnandersen777.bsky.social/asyncai

Switch repository

Loading…

//! Pool sandbox: each REPL scope is a checkout of a `monty` worker
//! subprocess (`monty-pool`). A crash, hang, or heap overrun kills only the
//! worker; the worker enforces `max_memory` with `monty-alloc`.

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use monty_pool::{
    on_print_sync, Checkout, Pool, PoolConfig, PoolError, ReplConfig, ResumeValue, TurnEvent,
};
use monty_types::{ExcType, MontyObject, NameLookupResult, ResourceLimits};
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};

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

const SCRIPT_NAME: &str = "run_code.py";
const REQUEST_GRACE: Duration = Duration::from_secs(10);

/// Builds the process-wide worker pool.
pub async fn connect(worker: PathBuf, limits: &CodeModeLimits) -> Result<Arc<Pool>, String> {
    let mut config = PoolConfig::subprocess(worker);
    config.min_processes = 0;
    config.request_timeout = Some(limits.max_duration + REQUEST_GRACE);
    Pool::new(config)
        .await
        .map(Arc::new)
        .map_err(|error| format!("monty pool failed to start: {error}"))
}

/// One REPL scope backed by a pool checkout, driven by a tokio task.
pub struct PoolSession {
    jobs: UnboundedSender<Job>,
}

impl PoolSession {
    pub fn spawn(
        pool: Arc<Pool>,
        catalog: Arc<Catalog>,
        limits: CodeModeLimits,
        preamble: Arc<String>,
    ) -> Self {
        let (jobs, mut receiver) = unbounded_channel::<Job>();
        let repl = ReplConfig {
            script_name: SCRIPT_NAME.to_string(),
            limits: Some(
                ResourceLimits::default()
                    .max_feed_duration(limits.max_duration)
                    .max_suspensions(limits.max_suspensions)
                    .max_memory(limits.max_memory),
            ),
            ..ReplConfig::default()
        };

        tokio::spawn(async move {
            let mut checkout: Option<Checkout> = None;
            while let Some(mut job) = receiver.recv().await {
                if job.restart {
                    checkout = None;
                }
                if checkout.is_none() {
                    match start(&pool, &repl, &catalog, &limits, &preamble).await {
                        Ok(started) => checkout = Some(started),
                        Err(message) => {
                            let _ = job.events.send(HostEvent::Done(Err(message)));
                            continue;
                        }
                    }
                }
                let session = checkout.as_mut().expect("checkout started");
                let (outcome, keep) = run_snippet(session, &catalog, &limits, &mut job).await;
                if !keep {
                    checkout = None;
                }
                let _ = job.events.send(HostEvent::Done(outcome));
            }
            if let Some(session) = checkout {
                let _ = session.finish().await;
            }
        });

        Self { jobs }
    }
}

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

async fn start(
    pool: &Pool,
    repl: &ReplConfig,
    catalog: &Catalog,
    limits: &CodeModeLimits,
    preamble: &str,
) -> Result<Checkout, String> {
    let mut checkout = pool
        .checkout(repl)
        .await
        .map_err(|error| format!("monty worker checkout failed: {error}"))?;
    if !preamble.trim().is_empty() {
        let (events, _events_rx) = unbounded_channel();
        let (_results_tx, results) = unbounded_channel();
        let mut job = Job {
            code: preamble.to_string(),
            restart: false,
            events,
            results,
        };
        if let (Err(error), _) = run_snippet(&mut checkout, catalog, limits, &mut job).await {
            return Err(format!("code mode preamble failed: {error}"));
        }
    }
    Ok(checkout)
}

/// Runs one snippet. The flag is false when the worker was lost and the
/// checkout must be replaced.
async fn run_snippet(
    checkout: &mut Checkout,
    catalog: &Catalog,
    limits: &CodeModeLimits,
    job: &mut Job,
) -> (Result<Completion, String>, bool) {
    let code = job.code.clone();
    let max_print = limits.max_print_bytes;
    let mut prints = String::new();
    let mut truncated = false;
    let mut sink = on_print_sync(|_, text: &str| {
        if prints.len() + text.len() > max_print {
            truncated = true;
        } else {
            prints.push_str(text);
        }
    });
    let mut link = HostLink::new(catalog, limits, job);

    let mut step = checkout
        .feed(
            code,
            Vec::<(String, MontyObject)>::new(),
            Vec::new(),
            true,
            &mut sink,
        )
        .await;

    loop {
        let event = match step {
            Ok(event) => event,
            Err(error) => {
                drop(sink);
                return pool_error(error);
            }
        };
        if !matches!(event, TurnEvent::Complete(_)) {
            if let Err(error) = link.suspend() {
                drop(sink);
                return (Err(error.to_string()), false);
            }
        }

        step = match event {
            TurnEvent::Complete(value) => {
                drop(sink);
                let value = if value.type_name() == "NoneType" {
                    None
                } else {
                    Some(monty_to_json(&value))
                };
                if truncated {
                    prints.push_str("\n...[print output truncated]");
                }
                return (Ok(Completion { value, prints }), true);
            }
            TurnEvent::NameLookup {
                name, object_id, ..
            } => {
                let result = match (object_id, link.lookup(&name)) {
                    (None, Some(function)) => NameLookupResult::Value(function),
                    _ => NameLookupResult::Undefined,
                };
                checkout.resume_name_lookup(result, &mut sink).await
            }
            TurnEvent::FunctionCall {
                function_name,
                args,
                call_id,
                object_id,
                ..
            } => {
                let answer = match link.call(&function_name, object_id.is_some(), &args, call_id) {
                    CallDecision::Pending => ResumeValue::Future,
                    CallDecision::Raise(error) => ResumeValue::Error(error),
                    CallDecision::NotFound => ResumeValue::NotFound,
                };
                checkout.resume(answer, &mut sink).await
            }
            TurnEvent::ResolveFutures {
                pending_call_ids, ..
            } => {
                let mut ready = link.take_ready(&pending_call_ids);
                while ready.is_empty() {
                    if let Err(message) = link.wait().await {
                        drop(sink);
                        return (Err(message), false);
                    }
                    ready = link.take_ready(&pending_call_ids);
                }
                let results = ready
                    .into_iter()
                    .map(|(call_id, value)| {
                        let answer = match value {
                            Ok(object) => ResumeValue::Return(object),
                            Err(error) => ResumeValue::Error(error),
                        };
                        (call_id, answer)
                    })
                    .collect();
                checkout.resume_futures(results, &mut sink).await
            }
            TurnEvent::OsCall { .. } => {
                checkout
                    .resume(
                        ResumeValue::Error(exception(
                            ExcType::PermissionError,
                            "os access disabled in asyncai code mode",
                        )),
                        &mut sink,
                    )
                    .await
            }
        };
    }
}

/// Maps a pool failure to an error result and whether the checkout survives.
fn pool_error(error: PoolError) -> (Result<Completion, String>, bool) {
    match error {
        PoolError::Runtime(exception) => {
            let memory = exception.exc_type() == ExcType::MemoryError;
            let mut message = exception.to_string();
            if memory {
                message
                    .push_str("\nThe sandbox worker hit its memory limit; REPL state was reset.");
            }
            (Err(message), !memory)
        }
        PoolError::Typing(message) => (Err(message), true),
        other => (
            Err(format!(
                "sandbox worker failed: {other}\nREPL state was reset; recreate any variables."
            )),
            false,
        ),
    }
}