elicitation_stdio.rs

//! Simple MCP Server with Elicitation
//!
//! Demonstrates user name collection via elicitation

use std::sync::Arc;

use anyhow::Result;
use rmcp::{
    ErrorData as McpError, ServerHandler, ServiceExt, elicit_safe,
    handler::server::{router::tool::ToolRouter, wrapper::Parameters},
    model::*,
    schemars::JsonSchema,
    service::{RequestContext, RoleServer},
    tool, tool_handler, tool_router,
    transport::stdio,
};
use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use tracing_subscriber::{self, EnvFilter};
use url::Url;

/// User information request
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
#[schemars(description = "User information")]
pub struct UserInfo {
    #[schemars(description = "User's name")]
    pub name: String,
}

// Mark as safe for elicitation
elicit_safe!(UserInfo);

/// Simple greeting message
#[derive(Debug, Serialize, Deserialize)]
pub struct GreetingMessage {
    pub text: String,
}

/// Simple tool request
#[derive(Debug, Deserialize, JsonSchema)]
pub struct GreetRequest {
    pub greeting: String,
}

/// Simple server with elicitation
#[derive(Clone)]
pub struct ElicitationServer {
    user_name: Arc<Mutex<Option<String>>>,
    tool_router: ToolRouter<ElicitationServer>,
}

impl ElicitationServer {
    pub fn new() -> Self {
        Self {
            user_name: Arc::new(Mutex::new(None)),
            tool_router: Self::tool_router(),
        }
    }
}

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

#[tool_router(router = tool_router)]
impl ElicitationServer {
    #[tool(description = "Greet user with name collection")]
    async fn greet_user(
        &self,
        context: RequestContext<RoleServer>,
        Parameters(request): Parameters<GreetRequest>,
    ) -> Result<CallToolResult, McpError> {
        // Check if we have user name
        let current_name = self.user_name.lock().await.clone();

        let user_name = if let Some(name) = current_name {
            name
        } else {
            // Request user name via elicitation
            match context
                .peer
                .elicit::<UserInfo>("Please provide your name".to_string())
                .await
            {
                Ok(Some(user_info)) => {
                    let name = user_info.name.clone();
                    *self.user_name.lock().await = Some(name.clone());
                    name
                }
                Ok(None) => "Guest".to_string(), // Never happen if client checks schema
                Err(_) => "Unknown".to_string(),
            }
        };

        Ok(CallToolResult::success(vec![ContentBlock::text(format!(
            "{} {}!",
            request.greeting, user_name
        ))]))
    }

    #[tool(description = "Reset stored user name")]
    async fn reset_name(&self) -> Result<CallToolResult, McpError> {
        *self.user_name.lock().await = None;
        Ok(CallToolResult::success(vec![ContentBlock::text(
            "User name reset. Next greeting will ask for name again.".to_string(),
        )]))
    }

    #[tool(description = "Example of URL elicitation")]
    pub async fn secure_tool_call(
        &self,
        context: RequestContext<RoleServer>,
    ) -> std::result::Result<CallToolResult, McpError> {
        let elicit_result = context
            .peer
            .elicit_url(
                "User must visit the following URL to complete tool call",
                Url::parse("https://example.com/complete_tool").expect("valid URL"),
                "elicit_123",
            )
            .await
            .map_err(|e| {
                ErrorData::new(
                    ErrorCode::INTERNAL_ERROR,
                    format!("Url elicitation has failed: {}", e),
                    None,
                )
            })?;
        match elicit_result {
            ElicitationAction::Accept => Ok(CallToolResult::success(vec![ContentBlock::text(
                "Elicitation via URL successful".to_string(),
            )])),
            ElicitationAction::Cancel => Ok(CallToolResult::success(vec![ContentBlock::text(
                "Elicitation via URL cancelled by user".to_string(),
            )])),
            ElicitationAction::Decline => Ok(CallToolResult::error(vec![ContentBlock::text(
                "Elicitation via URL declined by user".to_string(),
            )])),
            _ => Ok(CallToolResult::error(vec![ContentBlock::text(
                "Unknown elicitation action".to_string(),
            )])),
        }
    }
}

#[tool_handler(router = self.tool_router)]
impl ServerHandler for ElicitationServer {
    fn get_info(&self) -> ServerConfig {
        ServerConfig::new(ServerCapabilities::builder().enable_tools().build())
            .with_server_info(Implementation::from_build_env())
            .with_instructions(
                "Simple server demonstrating elicitation for user name collection".to_string(),
            )
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::from_default_env())
        .init();

    println!("Simple MCP Elicitation Demo");

    // Get current executable path for Inspector
    let current_exe = std::env::current_exe()
        .map(|path| path.display().to_string())
        .unwrap();

    println!("To test with MCP Inspector:");
    println!("1. Run: npx @modelcontextprotocol/inspector");
    println!("2. Enter server command: {}", current_exe);

    let service = ElicitationServer::new()
        .serve(stdio())
        .await
        .inspect_err(|e| {
            tracing::error!("serving error: {:?}", e);
        })?;

    service.waiting().await?;
    Ok(())
}