enki-runtime 0.1.4

//! Self-contained MCP test example
//!
//! This example demonstrates MCP functionality without requiring external
//! dependencies like npx. It tests the MCP type conversions and server setup.
//!
//! Run with:
//! ```bash
//! cargo run --example mcp_test --features mcp
//! ```

use async_trait::async_trait;
use enki_runtime::core::action::{ActionInvoker, ActionMetadata, ToolInvoker};
use enki_runtime::core::agent::{Agent, AgentContext};
use enki_runtime::core::error::Result;
use enki_runtime::core::message::Message;
use enki_runtime::mcp::McpAgentServer;
use serde_json::{json, Value};

/// A simple agent with math tools for testing
struct MathAgent {
    name: String,
    tool_invoker: ToolInvoker,
}

impl MathAgent {
    fn new() -> Self {
        let mut tool_invoker = ToolInvoker::new();

        // Register tools
        tool_invoker.register(Box::new(AddAction::new()));
        tool_invoker.register(Box::new(GreetAction::new()));

        Self {
            name: "math-agent".to_string(),
            tool_invoker,
        }
    }
}

#[async_trait]
impl Agent for MathAgent {
    fn name(&self) -> String {
        self.name.clone()
    }

    async fn on_message(&mut self, _msg: Message, _ctx: &mut AgentContext) -> Result<()> {
        Ok(())
    }

    fn tool_invoker(&self) -> Option<&ToolInvoker> {
        Some(&self.tool_invoker)
    }

    fn tool_invoker_mut(&mut self) -> Option<&mut ToolInvoker> {
        Some(&mut self.tool_invoker)
    }
}

// --- Tool Implementations ---

struct AddAction {
    metadata: ActionMetadata,
}

impl AddAction {
    fn new() -> Self {
        Self {
            metadata: ActionMetadata {
                name: "add".to_string(),
                description: "Add two numbers together".to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "a": { "type": "number" },
                        "b": { "type": "number" }
                    },
                    "required": ["a", "b"]
                }),
                output_schema: Some(json!({ "type": "number" })),
            },
        }
    }
}

#[async_trait]
impl ActionInvoker for AddAction {
    async fn execute(&self, _ctx: &mut AgentContext, inputs: Value) -> Result<Value> {
        let a = inputs["a"].as_f64().unwrap_or(0.0);
        let b = inputs["b"].as_f64().unwrap_or(0.0);
        Ok(json!(a + b))
    }

    fn metadata(&self) -> &ActionMetadata {
        &self.metadata
    }
}

struct GreetAction {
    metadata: ActionMetadata,
}

impl GreetAction {
    fn new() -> Self {
        Self {
            metadata: ActionMetadata {
                name: "greet".to_string(),
                description: "Generate a greeting message".to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "name": { "type": "string" }
                    },
                    "required": ["name"]
                }),
                output_schema: Some(json!({ "type": "string" })),
            },
        }
    }
}

#[async_trait]
impl ActionInvoker for GreetAction {
    async fn execute(&self, _ctx: &mut AgentContext, inputs: Value) -> Result<Value> {
        let name = inputs["name"].as_str().unwrap_or("World");
        Ok(json!(format!("Hello, {}!", name)))
    }

    fn metadata(&self) -> &ActionMetadata {
        &self.metadata
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    println!("=== MCP Self-Contained Test ===\n");

    // Create agent
    let agent = MathAgent::new();
    println!("✓ Created MathAgent with tools");

    // List registered tools
    if let Some(invoker) = agent.tool_invoker() {
        let actions = invoker.list_actions();
        println!("✓ Registered {} tools:", actions.len());
        for action in &actions {
            println!("  - {}: {}", action.name, action.description);
        }
    }

    // Test tool invocation directly
    println!("\n--- Testing Direct Tool Invocation ---");

    let mut ctx = AgentContext::new("test".to_string(), None);

    if let Some(invoker) = agent.tool_invoker() {
        // Test add
        let result = invoker
            .invoke("add", &mut ctx, json!({"a": 5, "b": 3}))
            .await?;
        println!("✓ add(5, 3) = {}", result);
        assert_eq!(result, json!(8.0));

        // Test greet
        let result = invoker
            .invoke("greet", &mut ctx, json!({"name": "Enki"}))
            .await?;
        println!("✓ greet(\"Enki\") = {}", result);
        assert_eq!(result, json!("Hello, Enki!"));
    }

    // Create MCP server wrapper
    println!("\n--- Testing MCP Server Wrapper ---");
    let mcp_server = McpAgentServer::new(agent, "test-mcp-server");
    println!("✓ Created McpAgentServer: {:?}", mcp_server);

    // Test type conversions
    println!("\n--- Testing Type Conversions ---");
    test_type_conversions();

    println!("\n=== All Tests Passed! ===");
    Ok(())
}

fn test_type_conversions() {
    use enki_runtime::mcp::types::{action_to_mcp_tool, mcp_tool_to_action_metadata};

    // Create a Enki ActionMetadata
    let action = ActionMetadata {
        name: "test_tool".to_string(),
        description: "A test tool".to_string(),
        input_schema: json!({
            "type": "object",
            "properties": {
                "input": { "type": "string" }
            }
        }),
        output_schema: Some(json!({ "type": "string" })),
    };

    // Convert to MCP Tool
    let mcp_tool = action_to_mcp_tool(&action);
    println!("✓ Converted ActionMetadata -> MCP Tool");
    println!("  Name: {}", mcp_tool.name);
    println!("  Description: {:?}", mcp_tool.description);

    // Convert back to ActionMetadata
    let action_back = mcp_tool_to_action_metadata(&mcp_tool);
    println!("✓ Converted MCP Tool -> ActionMetadata");

    // Verify roundtrip
    assert_eq!(action.name, action_back.name);
    assert_eq!(action.description, action_back.description);
    println!("✓ Roundtrip conversion verified");
}