enki-runtime 0.1.4

//! Example: Exposing a Enki agent as an MCP server
//!
//! This example demonstrates how to create a Enki agent with tools
//! and expose it as an MCP server that can be used by Claude Desktop,
//! Cursor, or other MCP clients.
//!
//! To run this example:
//! ```bash
//! cargo run --example mcp_server --features mcp
//! ```
//!
//! Then configure your MCP client to connect to this server via STDIO.

use async_trait::async_trait;
use enki_runtime::core::action::{ActionInvoker, ActionMetadata};
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 calculator agent with add/subtract tools
struct CalculatorAgent {
    name: String,
    tool_invoker: enki_runtime::core::action::ToolInvoker,
}

impl CalculatorAgent {
    fn new() -> Self {
        let mut tool_invoker = enki_runtime::core::action::ToolInvoker::new();
        
        // Register add tool
        tool_invoker.register(Box::new(AddAction::new()));
        
        // Register subtract tool
        tool_invoker.register(Box::new(SubtractAction::new()));
        
        // Register multiply tool
        tool_invoker.register(Box::new(MultiplyAction::new()));

        Self {
            name: "calculator".to_string(),
            tool_invoker,
        }
    }
}

#[async_trait]
impl Agent for CalculatorAgent {
    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<&enki_runtime::core::action::ToolInvoker> {
        Some(&self.tool_invoker)
    }

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

// --- Tool Actions ---

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", "description": "First number" },
                        "b": { "type": "number", "description": "Second 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 SubtractAction {
    metadata: ActionMetadata,
}

impl SubtractAction {
    fn new() -> Self {
        Self {
            metadata: ActionMetadata {
                name: "subtract".to_string(),
                description: "Subtract the second number from the first".to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "a": { "type": "number", "description": "Number to subtract from" },
                        "b": { "type": "number", "description": "Number to subtract" }
                    },
                    "required": ["a", "b"]
                }),
                output_schema: Some(json!({ "type": "number" })),
            },
        }
    }
}

#[async_trait]
impl ActionInvoker for SubtractAction {
    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 MultiplyAction {
    metadata: ActionMetadata,
}

impl MultiplyAction {
    fn new() -> Self {
        Self {
            metadata: ActionMetadata {
                name: "multiply".to_string(),
                description: "Multiply two numbers".to_string(),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "a": { "type": "number", "description": "First number" },
                        "b": { "type": "number", "description": "Second number" }
                    },
                    "required": ["a", "b"]
                }),
                output_schema: Some(json!({ "type": "number" })),
            },
        }
    }
}

#[async_trait]
impl ActionInvoker for MultiplyAction {
    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
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    eprintln!("=== Enki MCP Server Example ===");
    eprintln!("Starting Calculator MCP server...");
    eprintln!("This server exposes: add, subtract, multiply tools\n");
    
    // Create the calculator agent
    let agent = CalculatorAgent::new();
    
    // Create MCP server wrapping the agent
    let server = McpAgentServer::new(agent, "Enki-calculator");
    
    eprintln!("Server ready! Waiting for MCP client connections via STDIO...");
    eprintln!("(Use Ctrl+C to stop)\n");
    
    // Start serving over STDIO
    // This will block until the connection is closed
    server.serve_stdio().await?;
    
    eprintln!("Server stopped.");
    Ok(())
}