magi-code 0.100.0

//! Startup-only client tool descriptors. These are not MCP tools or approval claims.
mod broker;
pub(crate) use broker::ClientToolRun;

use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use std::{collections::BTreeMap, sync::Mutex};

pub(crate) const PREFIX: &str = "client__";
pub(crate) const MAX_TOOLS: usize = 32;
pub(crate) const MAX_ARGUMENT_BYTES: usize = 16_384;
pub(crate) const MAX_RESULT_BYTES: usize = 32_768;
pub(crate) const MAX_PENDING: usize = 16;

#[derive(Clone, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct Descriptor {
    pub(crate) name: String,
    description: String,
    input_schema: Value,
    #[serde(default = "default_timeout")]
    timeout_ms: u64,
}
fn default_timeout() -> u64 {
    30_000
}

pub(crate) struct Registry {
    tools: Vec<Descriptor>,
    fingerprints: BTreeMap<String, String>,
    pending: Mutex<BTreeMap<String, broker::Pending>>,
}

pub(crate) fn valid_name(name: &str) -> bool {
    name.strip_prefix(PREFIX).is_some_and(|suffix| {
        !suffix.is_empty()
            && suffix.len() <= 56
            && suffix
                .bytes()
                .all(|c| c.is_ascii_alphanumeric() || c == b'_')
    })
}

impl Registry {
    pub(crate) fn new(tools: Vec<Descriptor>) -> Result<Self, &'static str> {
        if tools.len() > MAX_TOOLS
            || serde_json::to_vec(&tools)
                .map_err(|_| "invalid_params")?
                .len()
                > 32_768
        {
            return Err("invalid_params");
        }
        let mut names = std::collections::HashSet::new();
        for tool in &tools {
            if !valid_name(&tool.name)
                || !names.insert(&tool.name)
                || crate::tools::ToolCapability::from_dispatch_name(&tool.name).is_some()
                || tool.description.trim().is_empty()
                || tool.description.len() > 2048
                || tool
                    .description
                    .chars()
                    .any(|c| c.is_control() && !matches!(c, '\n' | '\t'))
                || !(1000..=300_000).contains(&tool.timeout_ms)
                || tool.input_schema["type"] != "object"
                || serde_json::to_vec(&tool.input_schema)
                    .map_err(|_| "invalid_params")?
                    .len()
                    > 8192
                || !bounded_depth(&tool.input_schema, 16)
                || !valid_schema(&tool.input_schema)
            {
                return Err("invalid_params");
            }
        }
        let fingerprints = tools
            .iter()
            .map(|tool| {
                // serde_json object keys are sorted; array order and effective timeout are retained.
                let bytes = serde_json::to_vec(tool).expect("descriptor serializes");
                (
                    tool.name.clone(),
                    crate::support::hex::lower_hex(Sha256::digest(bytes)),
                )
            })
            .collect();
        Ok(Self {
            tools,
            fingerprints,
            pending: Mutex::new(BTreeMap::new()),
        })
    }

    pub(crate) fn descriptor(&self) -> Value {
        json!({"schema_version":1,"registration":"startup_only","tools":self.tools.iter().map(|t| &t.name).collect::<Vec<_>>(),
            "limits":{"tools":MAX_TOOLS,"name_bytes":64,"description_bytes":2048,
                "registry_bytes":32768,"schema_bytes":8192,"schema_depth":16,
                "arguments_bytes":MAX_ARGUMENT_BYTES,"result_bytes":MAX_RESULT_BYTES,
                "pending_calls":MAX_PENDING,"default_timeout_ms":30000,
                "min_timeout_ms":1000,"max_timeout_ms":300000}})
    }
    pub(crate) fn definitions(&self) -> Vec<Value> {
        self.tools
            .iter()
            .map(|tool| {
                json!({"type":"function","name":tool.name,
            "description":tool.description,"parameters":tool.input_schema})
            })
            .collect()
    }
    pub(crate) fn fingerprints(&self) -> BTreeMap<String, String> {
        self.fingerprints.clone()
    }
    pub(crate) fn fingerprint(&self, name: &str) -> Option<&String> {
        self.fingerprints.get(name)
    }
}

pub(super) fn bounded_depth(value: &Value, remaining: usize) -> bool {
    match value {
        Value::Array(items) => {
            remaining > 0 && items.iter().all(|v| bounded_depth(v, remaining - 1))
        }
        Value::Object(items) => {
            remaining > 0 && items.values().all(|v| bounded_depth(v, remaining - 1))
        }
        _ => true,
    }
}

// Deliberately bounded inline subset. No resolver, regex compiler, or instance evaluator.
// Reject unknown schema keywords rather than silently widening a schema we cannot inspect.
fn valid_schema(schema: &Value) -> bool {
    if schema.is_boolean() {
        return true;
    }
    let Some(object) = schema.as_object() else {
        return false;
    };
    object.iter().all(|(key, value)| match key.as_str() {
        "type" => {
            let valid_type = |v: &Value| {
                v.as_str().is_some_and(|s| {
                    matches!(
                        s,
                        "object" | "array" | "string" | "number" | "integer" | "boolean" | "null"
                    )
                })
            };
            valid_type(value)
                || value
                    .as_array()
                    .is_some_and(|a| !a.is_empty() && a.iter().all(valid_type))
        }
        "properties" | "$defs" | "definitions" | "patternProperties" => value
            .as_object()
            .is_some_and(|o| o.values().all(valid_schema)),
        "items" | "additionalProperties" | "not" | "contains" | "propertyNames" => {
            valid_schema(value)
        }
        "allOf" | "anyOf" | "oneOf" | "prefixItems" => value
            .as_array()
            .is_some_and(|a| !a.is_empty() && a.iter().all(valid_schema)),
        "required" => value.as_array().is_some_and(|a| {
            a.iter().all(Value::is_string)
                && a.iter().collect::<std::collections::HashSet<_>>().len() == a.len()
        }),
        "enum" => value.as_array().is_some_and(|a| !a.is_empty()),
        "const" | "default" | "examples" => true,
        "description" | "title" | "pattern" | "format" | "$comment" => value.is_string(),
        "minimum" | "maximum" | "exclusiveMinimum" | "exclusiveMaximum" => value.is_number(),
        "multipleOf" => value.as_f64().is_some_and(|n| n > 0.0),
        "minLength" | "maxLength" | "minItems" | "maxItems" | "minProperties" | "maxProperties"
        | "minContains" | "maxContains" => value.as_u64().is_some(),
        "uniqueItems" | "readOnly" | "writeOnly" | "deprecated" => value.is_boolean(),
        _ => false,
    })
}