turbomcp-proxy 3.5.0

//! Example: Schema Export from MCP Server
//!
//! Demonstrates exporting server capabilities as OpenAPI, GraphQL, and Protobuf schemas.
//!
//! Usage:
//!   - Run with no arguments for a self-contained mock server spec.
//!   - Or point it at a real MCP backend.
//!
//! Mock mode:
//!   cargo run --example schema_export
//!
//! STDIO backend:
//!   cargo run --example schema_export -- --backend stdio --cmd "your-mcp-server"
//!
//! For TCP backend:
//!   cargo run --example schema_export -- --backend tcp --tcp 127.0.0.1:8765
//!
//! For Unix socket:
//!   cargo run --example schema_export -- --backend unix --unix /tmp/turbomcp-demo.sock

use std::env;
use std::error::Error;

use serde_json::json;
use turbomcp_protocol::types::{
    Implementation, Resource, ResourcesCapabilities, ServerCapabilities, Tool, ToolInputSchema,
    ToolsCapabilities,
};
use turbomcp_proxy::MCP_PROTOCOL_VERSION;
use turbomcp_proxy::introspection::ServerSpec;
use turbomcp_proxy::proxy::{BackendConfig, BackendConnector, BackendTransport};

enum ExampleBackend {
    Mock,
    Real(BackendConfig),
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize tracing for logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .with_writer(std::io::stderr)
        .init();

    println!("🚀 MCP Schema Export Example");
    println!("============================\n");

    let spec = match parse_backend()? {
        ExampleBackend::Mock => {
            println!("📦 Using built-in mock server specification");
            mock_spec()
        }
        ExampleBackend::Real(config) => {
            println!("📡 Connecting to MCP server...");
            let backend = BackendConnector::new(config).await?;
            println!("✅ Connected successfully\n");

            println!("🔍 Introspecting server capabilities...");
            let spec = backend.introspect().await?;
            println!("✅ Introspection complete\n");
            spec
        }
    };

    print_schemas(&spec)?;

    Ok(())
}

fn parse_backend() -> Result<ExampleBackend, Box<dyn Error>> {
    let args: Vec<String> = env::args().skip(1).collect();
    if args.is_empty() || args.iter().any(|arg| arg == "--mock") {
        return Ok(ExampleBackend::Mock);
    }

    let backend = value_after(&args, "--backend").ok_or_else(|| {
        invalid_input(
            "missing --backend; use --mock, --backend stdio, --backend tcp, or --backend unix",
        )
    })?;

    let transport = match backend.as_str() {
        "stdio" => {
            let cmd = value_after(&args, "--cmd")
                .ok_or_else(|| invalid_input("stdio backend requires --cmd"))?;
            let mut parts = cmd.split_whitespace();
            let command = parts
                .next()
                .ok_or_else(|| invalid_input("--cmd cannot be empty"))?
                .to_string();
            let mut cmd_args: Vec<String> = parts.map(ToString::to_string).collect();
            cmd_args.extend(values_after(&args, "--arg"));

            BackendTransport::Stdio {
                command,
                args: cmd_args,
                working_dir: value_after(&args, "--working-dir"),
            }
        }
        "tcp" => {
            let endpoint = value_after(&args, "--tcp")
                .ok_or_else(|| invalid_input("tcp backend requires --tcp host:port"))?;
            let (host, port) = parse_host_port(&endpoint)?;
            BackendTransport::Tcp { host, port }
        }
        #[cfg(unix)]
        "unix" => {
            let path = value_after(&args, "--unix")
                .ok_or_else(|| invalid_input("unix backend requires --unix /path/to/socket"))?;
            BackendTransport::Unix { path }
        }
        other => {
            return Err(invalid_input(format!(
                "unsupported backend '{other}'; expected stdio, tcp, or unix"
            )));
        }
    };

    Ok(ExampleBackend::Real(BackendConfig {
        transport,
        client_name: "schema-export-example".to_string(),
        client_version: "1.0.0".to_string(),
    }))
}

fn value_after(args: &[String], flag: &str) -> Option<String> {
    args.windows(2)
        .find_map(|pair| (pair[0] == flag).then(|| pair[1].clone()))
}

fn values_after(args: &[String], flag: &str) -> Vec<String> {
    args.windows(2)
        .filter(|pair| pair[0] == flag)
        .map(|pair| pair[1].clone())
        .collect()
}

fn parse_host_port(endpoint: &str) -> Result<(String, u16), Box<dyn Error>> {
    let (host, port) = endpoint
        .rsplit_once(':')
        .ok_or_else(|| invalid_input("TCP endpoint must be host:port"))?;
    if host.is_empty() {
        return Err(invalid_input("TCP host cannot be empty"));
    }
    let port = port.parse::<u16>()?;
    Ok((host.to_string(), port))
}

fn invalid_input(message: impl Into<String>) -> Box<dyn Error> {
    Box::new(std::io::Error::new(
        std::io::ErrorKind::InvalidInput,
        message.into(),
    ))
}

fn mock_spec() -> ServerSpec {
    let echo_schema = ToolInputSchema::from_value(json!({
        "type": "object",
        "properties": {
            "message": { "type": "string", "description": "Message to echo back" }
        },
        "required": ["message"]
    }));
    let add_schema = ToolInputSchema::from_value(json!({
        "type": "object",
        "properties": { "a": { "type": "number" }, "b": { "type": "number" } },
        "required": ["a", "b"]
    }));

    ServerSpec {
        server_info: Implementation {
            title: Some("Demo MCP Server".to_string()),
            ..Implementation::new("demo-mcp-server", "1.0.0")
        },
        protocol_version: MCP_PROTOCOL_VERSION.to_string(),
        capabilities: ServerCapabilities {
            tools: Some(ToolsCapabilities::default()),
            resources: Some(ResourcesCapabilities::default()),
            ..Default::default()
        },
        tools: vec![
            Tool {
                title: Some("Echo".to_string()),
                ..Tool::new("echo", "Echo a message back to the caller").with_schema(echo_schema)
            },
            Tool {
                title: Some("Add".to_string()),
                ..Tool::new("add", "Add two numbers").with_schema(add_schema)
            },
        ],
        resources: vec![Resource {
            title: Some("Status".to_string()),
            ..Resource::new("demo://status", "status")
                .with_description("Current demo server status")
                .with_mime_type("application/json")
        }],
        resource_templates: vec![],
        prompts: vec![],
        instructions: Some("Mock spec used when no backend is configured.".to_string()),
    }
}

fn print_schemas(spec: &ServerSpec) -> Result<(), Box<dyn Error>> {
    println!("📊 Server Information:");
    println!("   Name: {}", spec.server_info.name);
    println!("   Version: {}", spec.server_info.version);
    println!("   Tools: {}", spec.tools.len());
    println!("   Resources: {}", spec.resources.len());

    println!("\n📝 Generated OpenAPI 3.1 Schema:");
    println!("─────────────────────────────────");
    let openapi = json!({
        "openapi": "3.1.0",
        "info": {
            "title": format!("{} API", spec.server_info.name),
            "version": spec.server_info.version
        },
        "paths": {
            "/tools": {
                "get": {
                    "summary": "List available tools",
                    "responses": {
                        "200": {
                            "description": "List of tools",
                            "content": {
                                "application/json": {
                                    "schema": {"type": "array"}
                                }
                            }
                        }
                    }
                }
            }
        }
    });
    println!("{}\n", serde_json::to_string_pretty(&openapi)?);

    println!("🎯 Generated GraphQL Schema:");
    println!("────────────────────────────");
    let mut graphql = String::from("type Query {\n");
    for tool in &spec.tools {
        let tool_name = tool.name.replace('-', "_");
        graphql.push_str(&format!(
            "  \"\"\"{}\"\"\"\n",
            tool.description.as_deref().unwrap_or("")
        ));
        graphql.push_str(&format!("  {}(input: JSON!): JSON!\n", tool_name));
    }
    graphql.push_str("}\n\nscalar JSON\n");
    println!("{}\n", graphql);

    println!("🔧 Generated Protobuf Schema:");
    println!("──────────────────────────────");
    let mut protobuf = String::from("syntax = \"proto3\";\n\npackage mcp_server;\n\n");
    for tool in &spec.tools {
        let tool_name = tool
            .name
            .split('-')
            .map(|s| {
                let mut chars = s.chars();
                match chars.next() {
                    None => String::new(),
                    Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
                }
            })
            .collect::<Vec<_>>()
            .join("");

        protobuf.push_str(&format!("message {} {{\n", tool_name));
        protobuf.push_str(&format!(
            "  // {}\n",
            tool.description.as_deref().unwrap_or("")
        ));
        protobuf.push_str("  string input = 1;\n");
        protobuf.push_str("  string output = 2;\n");
        protobuf.push_str("}\n\n");
    }
    println!("{}", protobuf);

    println!("✨ Schema generation complete!");
    println!("\nCLI equivalents:");
    println!("  turbomcp-proxy schema openapi --backend stdio --cmd \"your-server\" -o api.json");
    println!(
        "  turbomcp-proxy schema graphql --backend tcp --tcp 127.0.0.1:8765 -o schema.graphql"
    );
    println!(
        "  turbomcp-proxy schema protobuf --backend unix --unix /tmp/turbomcp-demo.sock -o server.proto"
    );

    Ok(())
}