johnandersen777.bsky.social/asyncai

Switch repository

Loading…

//! Renders MCP tool definitions as Python `async def` stubs for `run_code`.

use std::collections::{HashMap, HashSet};

use serde_json::Value;

use crate::tools::ToolDefinition;

const PYTHON_KEYWORDS: &[&str] = &[
    "False", "None", "True", "and", "as", "assert", "async", "await", "break", "class", "continue",
    "def", "del", "elif", "else", "except", "finally", "for", "from", "global", "if", "import",
    "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return", "try", "while",
    "with", "yield",
];

const SKETCH_MAX_DEPTH: usize = 4;
const SKETCH_MAX_PROPERTIES: usize = 12;

/// One keyword parameter: its Python name and the MCP argument key.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Param {
    pub py_name: String,
    pub key: String,
    pub required: bool,
}

/// One MCP tool exposed to the sandbox as a Python function.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CatalogEntry {
    pub py_name: String,
    pub mcp_name: String,
    pub doc: Option<String>,
    pub stub: String,
    pub params: Vec<Param>,
    /// Declared as possibly changing state (MCP annotations).
    pub mutating: bool,
}

impl CatalogEntry {
    pub fn param(&self, py_name: &str) -> Option<&Param> {
        self.params.iter().find(|param| param.py_name == py_name)
    }
}

/// Every MCP tool callable from `run_code`, keyed by Python name.
#[derive(Clone, Debug, Default)]
pub struct Catalog {
    entries: Vec<CatalogEntry>,
    by_py_name: HashMap<String, usize>,
}

impl Catalog {
    pub fn new(tools: &[ToolDefinition]) -> Self {
        Self::with_builtins(&[], tools)
    }

    /// Built-ins first (they keep their Python names); their routing name is
    /// `asyncai::<name>` so they never collide with an MCP tool.
    pub fn with_builtins(builtins: &[ToolDefinition], tools: &[ToolDefinition]) -> Self {
        let mut entries = Vec::new();
        let mut by_py_name = HashMap::new();
        let mut taken = HashSet::new();

        let all = builtins
            .iter()
            .map(|tool| (tool, true))
            .chain(tools.iter().map(|tool| (tool, false)));
        for (tool, builtin) in all {
            let py_name = unique_identifier(&tool.name, &mut taken);
            let doc = tool
                .description
                .as_deref()
                .map(str::trim)
                .filter(|text| !text.is_empty())
                .map(str::to_string);
            let (stub, params) = render_stub(&py_name, doc.as_deref(), tool);
            by_py_name.insert(py_name.clone(), entries.len());
            let mcp_name = if builtin {
                format!("{}{}", super::builtins::PREFIX, tool.name)
            } else {
                tool.name.clone()
            };
            entries.push(CatalogEntry {
                py_name,
                mcp_name,
                doc,
                stub,
                params,
                mutating: tool.is_mutating(),
            });
        }

        Self {
            entries,
            by_py_name,
        }
    }

    pub fn entries(&self) -> &[CatalogEntry] {
        &self.entries
    }

    pub fn get(&self, py_name: &str) -> Option<&CatalogEntry> {
        self.by_py_name
            .get(py_name)
            .map(|index| &self.entries[*index])
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// All stubs, separated by blank lines.
    pub fn render(&self) -> String {
        self.entries
            .iter()
            .map(|entry| entry.stub.as_str())
            .collect::<Vec<_>>()
            .join("\n\n")
    }
}

/// Converts a tool name into a Python identifier that is not a keyword.
pub fn python_identifier(name: &str) -> String {
    let mut identifier: String = name
        .chars()
        .map(|ch| {
            if ch.is_ascii_alphanumeric() || ch == '_' {
                ch
            } else {
                '_'
            }
        })
        .collect();
    if identifier.is_empty() || identifier.starts_with(|ch: char| ch.is_ascii_digit()) {
        identifier.insert(0, '_');
    }
    if PYTHON_KEYWORDS.contains(&identifier.as_str()) {
        identifier.push('_');
    }
    identifier
}

fn unique_identifier(name: &str, taken: &mut HashSet<String>) -> String {
    let base = python_identifier(name);
    let mut candidate = base.clone();
    let mut suffix = 2;
    while taken.contains(&candidate) {
        candidate = format!("{base}_{suffix}");
        suffix += 1;
    }
    taken.insert(candidate.clone());
    candidate
}

fn render_stub(py_name: &str, doc: Option<&str>, tool: &ToolDefinition) -> (String, Vec<Param>) {
    let schema = &tool.input_schema;
    let properties = schema
        .get("properties")
        .and_then(Value::as_object)
        .cloned()
        .unwrap_or_default();
    let required: HashSet<&str> = schema
        .get("required")
        .and_then(Value::as_array)
        .map(|names| names.iter().filter_map(Value::as_str).collect())
        .unwrap_or_default();

    let mut required_params = Vec::new();
    let mut optional_params = Vec::new();
    let mut params = Vec::new();
    for (name, property) in &properties {
        let (annotation, nullable) = python_type(property);
        let param = python_identifier(name);
        let is_required = required.contains(name.as_str());
        params.push(Param {
            py_name: param.clone(),
            key: name.clone(),
            required: is_required,
        });
        if is_required {
            let annotation = if nullable {
                format!("{annotation} | None")
            } else {
                annotation
            };
            required_params.push(format!("{param}: {annotation}"));
        } else {
            optional_params.push(format!("{param}: {annotation} | None = None"));
        }
    }

    let returns = match &tool.output_schema {
        Some(output) => python_type(output).0,
        None => "Any".to_string(),
    };
    let rendered: Vec<String> = required_params.into_iter().chain(optional_params).collect();
    let signature = if rendered.is_empty() {
        format!("async def {py_name}() -> {returns}:")
    } else {
        format!(
            "async def {py_name}(*, {}) -> {returns}:",
            rendered.join(", ")
        )
    };

    let mut doc_text = doc.map(str::to_string).unwrap_or_default();
    if let Some(output) = &tool.output_schema {
        if !doc_text.is_empty() {
            doc_text.push('\n');
        }
        doc_text.push_str("Returns: ");
        doc_text.push_str(&shape_sketch(output, output, 0));
    }
    let doc = if doc_text.is_empty() {
        None
    } else {
        Some(doc_text.as_str())
    };

    let body = match doc {
        Some(doc) => {
            let doc = doc.replace("\"\"\"", "\\\"\\\"\\\"");
            let indented = doc
                .lines()
                .map(|line| {
                    if line.trim().is_empty() {
                        String::new()
                    } else {
                        format!("    {line}")
                    }
                })
                .collect::<Vec<_>>()
                .join("\n");
            format!("    \"\"\"\n{indented}\n    \"\"\"")
        }
        None => "    ...".to_string(),
    };

    (format!("{signature}\n{body}"), params)
}

/// Compact, Python-flavoured sketch of a JSON schema's shape, e.g.
/// `{"items": [{"name": str, ...}], "kind": str}`.
pub fn shape_sketch(schema: &Value, root: &Value, depth: usize) -> String {
    if depth > SKETCH_MAX_DEPTH {
        return "...".to_string();
    }
    let Some(object) = schema.as_object() else {
        return "Any".to_string();
    };

    if let Some(reference) = object.get("$ref").and_then(Value::as_str) {
        return match resolve_ref(root, reference) {
            Some(target) => shape_sketch(target, root, depth + 1),
            None => "Any".to_string(),
        };
    }
    if object.contains_key("enum") {
        return python_type(schema).0;
    }
    for key in ["anyOf", "oneOf"] {
        if let Some(variants) = object.get(key).and_then(Value::as_array) {
            let parts: Vec<String> = variants
                .iter()
                .map(|variant| shape_sketch(variant, root, depth + 1))
                .collect();
            return parts.join(" | ");
        }
    }

    let kind = match object.get("type") {
        Some(Value::String(kind)) => Some(kind.as_str()),
        Some(Value::Array(kinds)) => kinds
            .iter()
            .filter_map(Value::as_str)
            .find(|kind| *kind != "null"),
        _ => None,
    };
    let nullable = matches!(object.get("type"), Some(Value::Array(kinds)) if kinds.iter().any(|kind| kind == "null"));

    let sketch = match kind {
        Some("object") | None if object.contains_key("properties") => {
            let properties = object
                .get("properties")
                .and_then(Value::as_object)
                .cloned()
                .unwrap_or_default();
            let mut fields: Vec<String> = properties
                .iter()
                .take(SKETCH_MAX_PROPERTIES)
                .map(|(name, property)| {
                    format!("\"{name}\": {}", shape_sketch(property, root, depth + 1))
                })
                .collect();
            let open = properties.len() > SKETCH_MAX_PROPERTIES
                || object.get("additionalProperties") != Some(&Value::Bool(false));
            if open {
                fields.push("...".to_string());
            }
            format!("{{{}}}", fields.join(", "))
        }
        Some("object") => "dict".to_string(),
        Some("array") => {
            let item = object
                .get("items")
                .map(|items| shape_sketch(items, root, depth + 1))
                .unwrap_or_else(|| "Any".to_string());
            format!("[{item}]")
        }
        Some(other) => json_type(other, object),
        None => "Any".to_string(),
    };
    if nullable {
        format!("{sketch} | None")
    } else {
        sketch
    }
}

fn resolve_ref<'a>(root: &'a Value, reference: &str) -> Option<&'a Value> {
    let path = reference.strip_prefix("#/")?;
    path.split('/')
        .try_fold(root, |node, segment| node.get(segment))
}

/// Python annotation for a JSON schema, and whether the schema admits `null`.
fn python_type(schema: &Value) -> (String, bool) {
    let Some(object) = schema.as_object() else {
        return ("Any".to_string(), false);
    };

    if let Some(values) = object.get("enum").and_then(Value::as_array) {
        let nullable = values.iter().any(Value::is_null);
        let literals = values
            .iter()
            .filter(|value| !value.is_null())
            .map(python_literal)
            .collect::<Vec<_>>();
        if literals.is_empty() {
            return ("None".to_string(), true);
        }
        return (format!("Literal[{}]", literals.join(", ")), nullable);
    }

    for key in ["anyOf", "oneOf"] {
        if let Some(variants) = object.get(key).and_then(Value::as_array) {
            let mut nullable = false;
            let mut types = Vec::new();
            for variant in variants {
                if variant.get("type").and_then(Value::as_str) == Some("null") {
                    nullable = true;
                    continue;
                }
                let (annotation, variant_nullable) = python_type(variant);
                nullable |= variant_nullable;
                if !types.contains(&annotation) {
                    types.push(annotation);
                }
            }
            if types.is_empty() {
                return ("Any".to_string(), nullable);
            }
            return (types.join(" | "), nullable);
        }
    }

    match object.get("type") {
        Some(Value::String(kind)) => (json_type(kind, object), false),
        Some(Value::Array(kinds)) => {
            let nullable = kinds.iter().any(|kind| kind.as_str() == Some("null"));
            let types = kinds
                .iter()
                .filter_map(Value::as_str)
                .filter(|kind| *kind != "null")
                .map(|kind| json_type(kind, object))
                .collect::<Vec<_>>();
            if types.is_empty() {
                ("None".to_string(), true)
            } else {
                (types.join(" | "), nullable)
            }
        }
        _ => ("Any".to_string(), false),
    }
}

fn json_type(kind: &str, schema: &serde_json::Map<String, Value>) -> String {
    match kind {
        "string" => "str".to_string(),
        "integer" => "int".to_string(),
        "number" => "float".to_string(),
        "boolean" => "bool".to_string(),
        "null" => "None".to_string(),
        "array" => {
            let item = schema
                .get("items")
                .map(|items| {
                    let (annotation, nullable) = python_type(items);
                    if nullable {
                        format!("{annotation} | None")
                    } else {
                        annotation
                    }
                })
                .unwrap_or_else(|| "Any".to_string());
            format!("list[{item}]")
        }
        "object" => "dict[str, Any]".to_string(),
        _ => "Any".to_string(),
    }
}

fn python_literal(value: &Value) -> String {
    match value {
        Value::Null => "None".to_string(),
        Value::Bool(true) => "True".to_string(),
        Value::Bool(false) => "False".to_string(),
        other => other.to_string(),
    }
}

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

    use super::*;

    fn tool(name: &str, description: Option<&str>, schema: Value) -> ToolDefinition {
        ToolDefinition::new(name, description, schema)
    }

    #[test]
    fn output_schema_sets_return_type_and_shape_sketch() {
        let mut definition = tool(
            "list_kcps",
            Some("Discover kcps."),
            json!({"type": "object", "properties": {"refresh": {"type": "boolean"}}}),
        );
        definition.output_schema = Some(json!({
            "type": "object",
            "properties": {
                "errors": {"type": "array", "items": {"type": "string"}},
                "kcps": {"type": "array", "items": {"$ref": "#/$defs/Kcp"}},
            },
            "additionalProperties": false,
            "$defs": {"Kcp": {
                "type": "object",
                "properties": {"id": {"type": "string"}, "running": {"type": "boolean"}},
            }},
        }));
        let catalog = Catalog::new(&[definition]);
        let stub = &catalog.entries()[0].stub;
        assert!(
            stub.starts_with(
                "async def list_kcps(*, refresh: bool | None = None) -> dict[str, Any]:"
            ),
            "{stub}"
        );
        assert!(
            stub.contains(
                r#"Returns: {"errors": [str], "kcps": [{"id": str, "running": bool, ...}]}"#
            ),
            "{stub}"
        );
    }

    #[test]
    fn shape_sketch_limits_depth_and_marks_nullable() {
        let deep = json!({"type": "object", "properties": {"a": {"type": "object", "properties": {"b": {"type": "object", "properties": {"c": {"type": "object", "properties": {"d": {"type": "object", "properties": {"e": {"type": "object", "properties": {"f": {"type": "string"}}}}}}}}}}}}});
        let sketch = shape_sketch(&deep, &deep, 0);
        assert!(sketch.contains("..."), "{sketch}");
        let nullable = json!({"type": ["string", "null"]});
        assert_eq!(shape_sketch(&nullable, &nullable, 0), "str | None");
    }

    #[test]
    fn params_map_python_names_to_mcp_keys_and_mutating_follows_annotations() {
        let mut definition = tool(
            "delete",
            None,
            json!({"type": "object", "properties": {"class": {"type": "string"}, "dry-run": {"type": "boolean"}}, "required": ["class"]}),
        );
        definition.destructive = Some(true);
        let catalog = Catalog::new(&[definition]);
        let entry = &catalog.entries()[0];
        assert!(entry.mutating);
        assert_eq!(entry.param("class_").unwrap().key, "class");
        assert!(entry.param("class_").unwrap().required);
        assert_eq!(entry.param("dry_run").unwrap().key, "dry-run");
    }

    #[test]
    fn renders_required_then_optional_keyword_only_params() {
        let catalog = Catalog::new(&[tool(
            "list_objects",
            Some("List objects."),
            json!({
                "type": "object",
                "properties": {
                    "kcp": {"type": "string"},
                    "limit": {"type": "integer"},
                    "namespace": {"type": ["string", "null"]},
                    "workspace": {"type": "string"},
                },
                "required": ["kcp", "workspace"],
            }),
        )]);

        assert_eq!(
            catalog.entries()[0].stub,
            "async def list_objects(*, kcp: str, workspace: str, limit: int | None = None, namespace: str | None = None) -> Any:\n    \"\"\"\n    List objects.\n    \"\"\""
        );
    }

    #[test]
    fn renders_enum_array_object_and_boolean_schemas() {
        let catalog = Catalog::new(&[tool(
            "apply",
            None,
            json!({
                "type": "object",
                "properties": {
                    "anything": true,
                    "either": {"anyOf": [{"type": "integer"}, {"type": "null"}]},
                    "flag": {"type": "boolean"},
                    "manifest": {"type": "object"},
                    "mode": {"enum": ["a", "b", null]},
                    "names": {"type": "array", "items": {"type": "string"}},
                    "ratio": {"type": "number"},
                },
                "required": ["mode", "names", "manifest", "anything", "either", "flag", "ratio"],
            }),
        )]);

        let stub = &catalog.entries()[0].stub;
        assert!(
            stub.contains("anything: Any, either: int | None, flag: bool, manifest: dict[str, Any], mode: Literal[\"a\", \"b\"] | None, names: list[str], ratio: float"),
            "{stub}"
        );
        assert!(stub.ends_with("\n    ..."), "{stub}");
    }

    #[test]
    fn sanitizes_names_and_resolves_collisions() {
        let catalog = Catalog::new(&[
            tool("http-get", None, json!({})),
            tool("http.get", None, json!({})),
            tool("9lives", None, json!({})),
            tool("import", None, json!({})),
        ]);

        let names = catalog
            .entries()
            .iter()
            .map(|entry| entry.py_name.as_str())
            .collect::<Vec<_>>();
        assert_eq!(names, ["http_get", "http_get_2", "_9lives", "import_"]);
        assert_eq!(catalog.get("http_get_2").unwrap().mcp_name, "http.get");
        assert_eq!(
            catalog.get("http_get").unwrap().stub,
            "async def http_get() -> Any:\n    ..."
        );
    }

    #[test]
    fn escapes_docstring_quotes() {
        let catalog = Catalog::new(&[tool("t", Some("say \"\"\"hi\"\"\""), json!({}))]);
        assert!(!catalog.entries()[0].stub.contains("\"\"\"hi"));
    }
}