~ knotless Search/
↑↓ move↵ open⇧↵ all results

johnandersen777.bsky.social/asyncai

0
0
main
156 lines · 4.5 KB Raw History
//! Validation shared by every sandbox: a nested call either becomes an MCP
//! request or a Python exception, before anything leaves the sandbox.

use monty_types::ExcType;
use serde_json::{Map, Value};

use super::catalog::CatalogEntry;

/// A Python exception to raise at the call site.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Rejection {
    pub kind: ExcType,
    pub message: String,
}

impl Rejection {
    fn type_error(message: String) -> Self {
        Self {
            kind: ExcType::TypeError,
            message,
        }
    }
}

/// Checks one call against the tool's declared parameters and returns the
/// MCP arguments, keyed by MCP name. Optional arguments passed as `None` are
/// dropped.
pub fn check_call(
    entry: &CatalogEntry,
    positional: usize,
    kwargs: Vec<(String, Value)>,
) -> Result<Map<String, Value>, Rejection> {
    let name = &entry.py_name;
    if positional > 0 {
        return Err(Rejection::type_error(format!(
            "{name}() takes keyword arguments only, e.g. {name}(key=value)"
        )));
    }

    let mut arguments = Map::new();
    for (key, value) in kwargs {
        let Some(param) = entry.param(&key) else {
            let accepted = entry
                .params
                .iter()
                .map(|param| param.py_name.as_str())
                .collect::<Vec<_>>();
            let accepted = if accepted.is_empty() {
                "none".to_string()
            } else {
                accepted.join(", ")
            };
            return Err(Rejection::type_error(format!(
                "{name}() got an unexpected keyword argument '{key}'; accepted: {accepted}"
            )));
        };
        if value.is_null() && !param.required {
            continue;
        }
        arguments.insert(param.key.clone(), value);
    }

    if let Some(missing) = entry
        .params
        .iter()
        .find(|param| param.required && !arguments.contains_key(&param.key))
    {
        return Err(Rejection::type_error(format!(
            "{name}() missing required keyword argument '{}'",
            missing.py_name
        )));
    }

    Ok(arguments)
}

/// The exception raised once a snippet exceeds its nested-call budget.
pub fn budget_exceeded(limit: usize) -> Rejection {
    Rejection {
        kind: ExcType::RuntimeError,
        message: format!("nested tool call budget of {limit} per run_code exceeded"),
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;
    use crate::code_mode::catalog::Catalog;
    use crate::tools::ToolDefinition;

    fn entry() -> CatalogEntry {
        let catalog = Catalog::new(&[ToolDefinition::new(
            "list_objects",
            None::<String>,
            json!({
                "type": "object",
                "properties": {
                    "resource": {"type": "string"},
                    "namespace": {"type": "string"},
                    "all-namespaces": {"type": "boolean"},
                },
                "required": ["resource"],
            }),
        )]);
        catalog.entries()[0].clone()
    }

    #[test]
    fn maps_python_names_and_drops_none_optionals() {
        let arguments = check_call(
            &entry(),
            0,
            vec![
                ("resource".to_string(), json!("configmaps")),
                ("namespace".to_string(), Value::Null),
                ("all_namespaces".to_string(), json!(true)),
            ],
        )
        .unwrap();
        assert_eq!(
            Value::Object(arguments),
            json!({"resource": "configmaps", "all-namespaces": true})
        );
    }

    #[test]
    fn rejects_unknown_missing_and_positional_arguments() {
        let unknown = check_call(
            &entry(),
            0,
            vec![
                ("resource".to_string(), json!("pods")),
                ("name".to_string(), json!("x")),
            ],
        )
        .unwrap_err();
        assert_eq!(unknown.kind, ExcType::TypeError);
        assert!(
            unknown.message.contains(
                "unexpected keyword argument 'name'; accepted: all_namespaces, namespace, resource"
            ) || unknown
                .message
                .contains("unexpected keyword argument 'name'")
        );

        let missing = check_call(&entry(), 0, vec![]).unwrap_err();
        assert!(missing
            .message
            .contains("missing required keyword argument 'resource'"));

        let positional = check_call(&entry(), 1, vec![]).unwrap_err();
        assert!(positional.message.contains("keyword arguments only"));
    }
}