owlen/crates/owlen-mcp-prompt-server/src/lib.rs

//! MCP server for rendering prompt templates with YAML storage and Handlebars rendering.
//!
//! Templates are stored in `~/.config/owlen/prompts/` as YAML files.
//! Provides full Handlebars templating support for dynamic prompt generation.

use anyhow::{Context, Result};
use handlebars::Handlebars;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;

use owlen_core::mcp::protocol::{
    methods, ErrorCode, InitializeParams, InitializeResult, RequestId, RpcError, RpcErrorResponse,
    RpcRequest, RpcResponse, ServerCapabilities, ServerInfo, PROTOCOL_VERSION,
};
use owlen_core::mcp::{McpToolCall, McpToolDescriptor, McpToolResponse};
use tokio::io::{self, AsyncBufReadExt, AsyncWriteExt};

/// Prompt template definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptTemplate {
    /// Template name
    pub name: String,
    /// Template version
    pub version: String,
    /// Optional mode restriction
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mode: Option<String>,
    /// Handlebars template content
    pub template: String,
    /// Template description
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
}

/// Prompt server managing templates
pub struct PromptServer {
    templates: Arc<RwLock<HashMap<String, PromptTemplate>>>,
    handlebars: Handlebars<'static>,
    templates_dir: PathBuf,
}

impl PromptServer {
    /// Create a new prompt server
    pub fn new() -> Result<Self> {
        let templates_dir = Self::get_templates_dir()?;

        // Create templates directory if it doesn't exist
        if !templates_dir.exists() {
            fs::create_dir_all(&templates_dir)?;
            Self::create_default_templates(&templates_dir)?;
        }

        let mut server = Self {
            templates: Arc::new(RwLock::new(HashMap::new())),
            handlebars: Handlebars::new(),
            templates_dir,
        };

        // Load all templates
        server.load_templates()?;

        Ok(server)
    }

    /// Get the templates directory path
    fn get_templates_dir() -> Result<PathBuf> {
        let config_dir = dirs::config_dir().context("Could not determine config directory")?;
        Ok(config_dir.join("owlen").join("prompts"))
    }

    /// Create default template examples
    fn create_default_templates(dir: &Path) -> Result<()> {
        let chat_mode_system = PromptTemplate {
            name: "chat_mode_system".to_string(),
            version: "1.0".to_string(),
            mode: Some("chat".to_string()),
            description: Some("System prompt for chat mode".to_string()),
            template: r#"You are Owlen, a helpful AI assistant. You have access to these tools:
{{#each tools}}
- {{name}}: {{description}}
{{/each}}

Use the ReAct pattern:
THOUGHT: Your reasoning
ACTION: tool_name
ACTION_INPUT: {"param": "value"}

When you have enough information:
FINAL_ANSWER: Your response"#
                .to_string(),
        };

        let code_mode_system = PromptTemplate {
            name: "code_mode_system".to_string(),
            version: "1.0".to_string(),
            mode: Some("code".to_string()),
            description: Some("System prompt for code mode".to_string()),
            template: r#"You are Owlen in code mode, with full development capabilities. You have access to:
{{#each tools}}
- {{name}}: {{description}}
{{/each}}

Use the ReAct pattern to solve coding tasks:
THOUGHT: Analyze what needs to be done
ACTION: tool_name (compile_project, run_tests, format_code, lint_code, etc.)
ACTION_INPUT: {"param": "value"}

Continue iterating until the task is complete, then provide:
FINAL_ANSWER: Summary of what was done"#
                .to_string(),
        };

        // Save templates
        let chat_path = dir.join("chat_mode_system.yaml");
        let code_path = dir.join("code_mode_system.yaml");

        fs::write(chat_path, serde_yaml::to_string(&chat_mode_system)?)?;
        fs::write(code_path, serde_yaml::to_string(&code_mode_system)?)?;

        Ok(())
    }

    /// Load all templates from the templates directory
    fn load_templates(&mut self) -> Result<()> {
        let entries = fs::read_dir(&self.templates_dir)?;

        for entry in entries {
            let entry = entry?;
            let path = entry.path();

            if path.extension().and_then(|s| s.to_str()) == Some("yaml")
                || path.extension().and_then(|s| s.to_str()) == Some("yml")
            {
                match self.load_template(&path) {
                    Ok(template) => {
                        // Register with Handlebars
                        if let Err(e) = self
                            .handlebars
                            .register_template_string(&template.name, &template.template)
                        {
                            eprintln!(
                                "Warning: Failed to register template {}: {}",
                                template.name, e
                            );
                        } else {
                            let mut templates = futures::executor::block_on(self.templates.write());
                            templates.insert(template.name.clone(), template);
                        }
                    }
                    Err(e) => {
                        eprintln!("Warning: Failed to load template {:?}: {}", path, e);
                    }
                }
            }
        }

        Ok(())
    }

    /// Load a single template from file
    fn load_template(&self, path: &Path) -> Result<PromptTemplate> {
        let content = fs::read_to_string(path)?;
        let template: PromptTemplate = serde_yaml::from_str(&content)?;
        Ok(template)
    }

    /// Get a template by name
    pub async fn get_template(&self, name: &str) -> Option<PromptTemplate> {
        let templates = self.templates.read().await;
        templates.get(name).cloned()
    }

    /// List all available templates
    pub async fn list_templates(&self) -> Vec<String> {
        let templates = self.templates.read().await;
        templates.keys().cloned().collect()
    }

    /// Render a template with given variables
    pub fn render_template(&self, name: &str, vars: &Value) -> Result<String> {
        self.handlebars
            .render(name, vars)
            .context("Failed to render template")
    }

    /// Reload all templates from disk
    pub async fn reload_templates(&mut self) -> Result<()> {
        {
            let mut templates = self.templates.write().await;
            templates.clear();
        }
        self.handlebars = Handlebars::new();
        self.load_templates()
    }
}

#[allow(dead_code)]
#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let mut stdin = io::BufReader::new(io::stdin());
    let mut stdout = io::stdout();

    let server = Arc::new(tokio::sync::Mutex::new(PromptServer::new()?));

    loop {
        let mut line = String::new();
        match stdin.read_line(&mut line).await {
            Ok(0) => break, // EOF
            Ok(_) => {
                let req: RpcRequest = match serde_json::from_str(&line) {
                    Ok(r) => r,
                    Err(e) => {
                        let err = RpcErrorResponse::new(
                            RequestId::Number(0),
                            RpcError::parse_error(format!("Parse error: {}", e)),
                        );
                        let s = serde_json::to_string(&err)?;
                        stdout.write_all(s.as_bytes()).await?;
                        stdout.write_all(b"\n").await?;
                        stdout.flush().await?;
                        continue;
                    }
                };

                let resp = handle_request(req.clone(), server.clone()).await;
                match resp {
                    Ok(r) => {
                        let s = serde_json::to_string(&r)?;
                        stdout.write_all(s.as_bytes()).await?;
                        stdout.write_all(b"\n").await?;
                        stdout.flush().await?;
                    }
                    Err(e) => {
                        let err = RpcErrorResponse::new(req.id.clone(), e);
                        let s = serde_json::to_string(&err)?;
                        stdout.write_all(s.as_bytes()).await?;
                        stdout.write_all(b"\n").await?;
                        stdout.flush().await?;
                    }
                }
            }
            Err(e) => {
                eprintln!("Error reading stdin: {}", e);
                break;
            }
        }
    }
    Ok(())
}

#[allow(dead_code)]
async fn handle_request(
    req: RpcRequest,
    server: Arc<tokio::sync::Mutex<PromptServer>>,
) -> Result<RpcResponse, RpcError> {
    match req.method.as_str() {
        methods::INITIALIZE => {
            let params: InitializeParams =
                serde_json::from_value(req.params.unwrap_or_else(|| json!({})))
                    .map_err(|e| RpcError::invalid_params(format!("Invalid init params: {}", e)))?;
            if !params.protocol_version.eq(PROTOCOL_VERSION) {
                return Err(RpcError::new(
                    ErrorCode::INVALID_REQUEST,
                    format!(
                        "Incompatible protocol version. Client: {}, Server: {}",
                        params.protocol_version, PROTOCOL_VERSION
                    ),
                ));
            }
            let result = InitializeResult {
                protocol_version: PROTOCOL_VERSION.to_string(),
                server_info: ServerInfo {
                    name: "owlen-mcp-prompt-server".to_string(),
                    version: env!("CARGO_PKG_VERSION").to_string(),
                },
                capabilities: ServerCapabilities {
                    supports_tools: Some(true),
                    supports_resources: Some(false),
                    supports_streaming: Some(false),
                },
            };
            Ok(RpcResponse::new(
                req.id,
                serde_json::to_value(result).unwrap(),
            ))
        }
        methods::TOOLS_LIST => {
            let tools = vec![
                McpToolDescriptor {
                    name: "get_prompt".to_string(),
                    description: "Retrieve a prompt template by name".to_string(),
                    input_schema: json!({
                        "type": "object",
                        "properties": {
                            "name": {"type": "string", "description": "Template name"}
                        },
                        "required": ["name"]
                    }),
                    requires_network: false,
                    requires_filesystem: vec![],
                },
                McpToolDescriptor {
                    name: "render_prompt".to_string(),
                    description: "Render a prompt template with Handlebars variables".to_string(),
                    input_schema: json!({
                        "type": "object",
                        "properties": {
                            "name": {"type": "string", "description": "Template name"},
                            "vars": {"type": "object", "description": "Variables for Handlebars rendering"}
                        },
                        "required": ["name"]
                    }),
                    requires_network: false,
                    requires_filesystem: vec![],
                },
                McpToolDescriptor {
                    name: "list_prompts".to_string(),
                    description: "List all available prompt templates".to_string(),
                    input_schema: json!({"type": "object", "properties": {}}),
                    requires_network: false,
                    requires_filesystem: vec![],
                },
                McpToolDescriptor {
                    name: "reload_prompts".to_string(),
                    description: "Reload all prompts from disk".to_string(),
                    input_schema: json!({"type": "object", "properties": {}}),
                    requires_network: false,
                    requires_filesystem: vec![],
                },
            ];
            Ok(RpcResponse::new(req.id, json!(tools)))
        }
        methods::TOOLS_CALL => {
            let call: McpToolCall = serde_json::from_value(req.params.unwrap_or_else(|| json!({})))
                .map_err(|e| RpcError::invalid_params(format!("Invalid tool call: {}", e)))?;

            let result = match call.name.as_str() {
                "get_prompt" => {
                    let name = call
                        .arguments
                        .get("name")
                        .and_then(|v| v.as_str())
                        .ok_or_else(|| RpcError::invalid_params("Missing 'name' parameter"))?;

                    let srv = server.lock().await;
                    match srv.get_template(name).await {
                        Some(template) => {
                            json!({"success": true, "template": serde_json::to_value(template).unwrap()})
                        }
                        None => json!({"success": false, "error": "Template not found"}),
                    }
                }
                "render_prompt" => {
                    let name = call
                        .arguments
                        .get("name")
                        .and_then(|v| v.as_str())
                        .ok_or_else(|| RpcError::invalid_params("Missing 'name' parameter"))?;

                    let default_vars = json!({});
                    let vars = call.arguments.get("vars").unwrap_or(&default_vars);

                    let srv = server.lock().await;
                    match srv.render_template(name, vars) {
                        Ok(rendered) => json!({"success": true, "rendered": rendered}),
                        Err(e) => json!({"success": false, "error": e.to_string()}),
                    }
                }
                "list_prompts" => {
                    let srv = server.lock().await;
                    let templates = srv.list_templates().await;
                    json!({"success": true, "templates": templates})
                }
                "reload_prompts" => {
                    let mut srv = server.lock().await;
                    match srv.reload_templates().await {
                        Ok(_) => json!({"success": true, "message": "Prompts reloaded"}),
                        Err(e) => json!({"success": false, "error": e.to_string()}),
                    }
                }
                _ => return Err(RpcError::method_not_found(&call.name)),
            };

            let resp = McpToolResponse {
                name: call.name,
                success: result
                    .get("success")
                    .and_then(|v| v.as_bool())
                    .unwrap_or(false),
                output: result,
                metadata: HashMap::new(),
                duration_ms: 0,
            };

            Ok(RpcResponse::new(
                req.id,
                serde_json::to_value(resp).unwrap(),
            ))
        }
        _ => Err(RpcError::method_not_found(&req.method)),
    }
}