johnandersen777.bsky.social/asyncai

Switch repository

Loading…

//! The host protocol every sandbox speaks: jobs in, `HostEvent`s out, call
//! results back in. Sandbox-independent call handling lives here.

use std::collections::HashMap;

use monty_types::{CallArgs, ExcType, MontyException, MontyObject};
use serde_json::Value;
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};

use super::catalog::Catalog;
use super::convert::{object_ref_to_json, tool_text_to_monty};
use super::dispatch::{budget_exceeded, check_call};
use super::CodeModeLimits;
use crate::tools::ToolCallOutput;

/// Outcome of one nested tool call, as delivered back to the sandbox.
pub type CallOutcome = Result<ToolCallOutput, String>;

/// Event sent from a sandbox to the async host while a snippet runs.
#[derive(Debug)]
pub enum HostEvent {
    /// The snippet called a catalog function; run it and send the result.
    Call {
        call_id: u32,
        py_name: String,
        mcp_name: String,
        arguments: Value,
        mutating: bool,
    },
    /// The snippet finished.
    Done(Result<Completion, String>),
}

/// A finished snippet: its final value (`None` for a `None` result) and prints.
#[derive(Debug, Clone, PartialEq)]
pub struct Completion {
    pub value: Option<Value>,
    pub prints: String,
}

/// One `run_code` request.
pub struct Job {
    pub code: String,
    pub restart: bool,
    pub events: UnboundedSender<HostEvent>,
    pub results: UnboundedReceiver<(u32, CallOutcome)>,
}

/// A started sandbox scope (one REPL). Dropping it ends the scope.
pub trait SandboxSession: Send + Sync {
    fn submit(&self, job: Job) -> Result<(), String>;
}

/// What the sandbox should do with a function call.
pub enum CallDecision {
    /// Handed to the host; resume with a pending future.
    Pending,
    /// Raise this exception at the call site.
    Raise(MontyException),
    /// Not a catalog function: raise `NameError`.
    NotFound,
}

pub fn exception(kind: ExcType, message: impl Into<String>) -> MontyException {
    MontyException::new(kind, Some(message.into()))
}

/// Per-snippet call handling shared by all sandboxes.
pub struct HostLink<'a> {
    catalog: &'a Catalog,
    limits: &'a CodeModeLimits,
    job: &'a mut Job,
    nested_calls: usize,
    suspensions: usize,
    buffered: HashMap<u32, CallOutcome>,
}

impl<'a> HostLink<'a> {
    pub fn new(catalog: &'a Catalog, limits: &'a CodeModeLimits, job: &'a mut Job) -> Self {
        Self {
            catalog,
            limits,
            job,
            nested_calls: 0,
            suspensions: 0,
            buffered: HashMap::new(),
        }
    }

    /// Counts a suspension; `Err` once the per-snippet budget is spent.
    pub fn suspend(&mut self) -> Result<(), MontyException> {
        self.suspensions += 1;
        if self.suspensions > self.limits.max_suspensions {
            return Err(exception(
                ExcType::RuntimeError,
                format!(
                    "suspension limit of {} exceeded",
                    self.limits.max_suspensions
                ),
            ));
        }
        Ok(())
    }

    /// The value for an undefined name: a catalog function, or `None`.
    pub fn lookup(&self, name: &str) -> Option<MontyObject> {
        self.catalog
            .get(name)
            .map(|entry| MontyObject::function(entry.py_name.clone(), entry.doc.clone()))
    }

    pub fn call(
        &mut self,
        name: &str,
        is_method: bool,
        args: &CallArgs,
        call_id: u32,
    ) -> CallDecision {
        if is_method {
            return CallDecision::Raise(exception(
                ExcType::TypeError,
                format!("{name}: host objects are not supported"),
            ));
        }
        let Some(entry) = self.catalog.get(name) else {
            return CallDecision::NotFound;
        };
        if self.nested_calls >= self.limits.max_tool_calls {
            let rejection = budget_exceeded(self.limits.max_tool_calls);
            return CallDecision::Raise(exception(rejection.kind, rejection.message));
        }
        let kwargs = args
            .kwargs()
            .map(|(key, value)| {
                let key = match key.as_str() {
                    Some(text) => text.to_string(),
                    None => key.py_repr(),
                };
                (key, object_ref_to_json(value))
            })
            .collect();
        let arguments = match check_call(entry, args.args().len(), kwargs) {
            Ok(arguments) => arguments,
            Err(rejection) => {
                return CallDecision::Raise(exception(rejection.kind, rejection.message))
            }
        };

        self.nested_calls += 1;
        let sent = self.job.events.send(HostEvent::Call {
            call_id,
            py_name: entry.py_name.clone(),
            mcp_name: entry.mcp_name.clone(),
            arguments: Value::Object(arguments),
            mutating: entry.mutating,
        });
        if sent.is_err() {
            return CallDecision::Raise(exception(ExcType::RuntimeError, "code mode host stopped"));
        }
        CallDecision::Pending
    }

    /// Results already delivered for any of `pending`.
    pub fn take_ready(
        &mut self,
        pending: &[u32],
    ) -> Vec<(u32, Result<MontyObject, MontyException>)> {
        while let Ok((call_id, outcome)) = self.job.results.try_recv() {
            self.buffered.insert(call_id, outcome);
        }
        pending
            .iter()
            .filter_map(|call_id| {
                self.buffered
                    .remove(call_id)
                    .map(|outcome| (*call_id, outcome_to_value(outcome)))
            })
            .collect()
    }

    /// Blocks the current (non-async) thread for one more result.
    pub fn wait_blocking(&mut self) -> Result<(), String> {
        match self.job.results.blocking_recv() {
            Some((call_id, outcome)) => {
                self.buffered.insert(call_id, outcome);
                Ok(())
            }
            None => Err("code mode host stopped".to_string()),
        }
    }

    /// Awaits one more result.
    pub async fn wait(&mut self) -> Result<(), String> {
        match self.job.results.recv().await {
            Some((call_id, outcome)) => {
                self.buffered.insert(call_id, outcome);
                Ok(())
            }
            None => Err("code mode host stopped".to_string()),
        }
    }
}

/// A nested call's outcome as a Python value or exception.
pub fn outcome_to_value(outcome: CallOutcome) -> Result<MontyObject, MontyException> {
    match outcome {
        Ok(output) if output.is_error => Err(exception(ExcType::RuntimeError, output.content)),
        Ok(output) => Ok(tool_text_to_monty(&output.content)),
        Err(message) => Err(exception(
            ExcType::RuntimeError,
            format!("tool error: {message}"),
        )),
    }
}