owlen/crates/owlen-core/src/validation.rs

use std::collections::HashMap;

use anyhow::{Context, Result};
use jsonschema::{JSONSchema, ValidationError};
use serde_json::{json, Value};

pub struct SchemaValidator {
    schemas: HashMap<String, JSONSchema>,
}

impl Default for SchemaValidator {
    fn default() -> Self {
        Self::new()
    }
}

impl SchemaValidator {
    pub fn new() -> Self {
        Self {
            schemas: HashMap::new(),
        }
    }

    pub fn register_schema(&mut self, tool_name: &str, schema: Value) -> Result<()> {
        let compiled = JSONSchema::compile(&schema)
            .map_err(|e| anyhow::anyhow!("Invalid schema for {}: {}", tool_name, e))?;

        self.schemas.insert(tool_name.to_string(), compiled);
        Ok(())
    }

    pub fn validate(&self, tool_name: &str, input: &Value) -> Result<()> {
        let schema = self
            .schemas
            .get(tool_name)
            .with_context(|| format!("No schema registered for tool: {}", tool_name))?;

        if let Err(errors) = schema.validate(input) {
            let error_messages: Vec<String> = errors.map(format_validation_error).collect();

            return Err(anyhow::anyhow!(
                "Input validation failed for {}: {}",
                tool_name,
                error_messages.join(", ")
            ));
        }

        Ok(())
    }
}

fn format_validation_error(error: ValidationError) -> String {
    format!("Validation error at {}: {}", error.instance_path, error)
}

pub fn get_builtin_schemas() -> HashMap<String, Value> {
    let mut schemas = HashMap::new();

    schemas.insert(
        "web_search".to_string(),
        json!({
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "minLength": 1,
                    "maxLength": 500
                },
                "max_results": {
                    "type": "integer",
                    "minimum": 1,
                    "maximum": 10,
                    "default": 5
                }
            },
            "required": ["query"],
            "additionalProperties": false
        }),
    );

    schemas.insert(
        "code_exec".to_string(),
        json!({
            "type": "object",
            "properties": {
                "language": {
                    "type": "string",
                    "enum": ["python", "javascript", "bash", "rust"]
                },
                "code": {
                    "type": "string",
                    "minLength": 1,
                    "maxLength": 10000
                },
                "timeout": {
                    "type": "integer",
                    "minimum": 1,
                    "maximum": 300,
                    "default": 30
                }
            },
            "required": ["language", "code"],
            "additionalProperties": false
        }),
    );

    schemas
}