turbomcp-proxy 3.5.0

//! Proxy implementation for {{server_name}}
//!
//! `ProxyRouter` is an MCP server (`McpHandler`) whose tools and prompts are
//! the upstream's, fixed at generation time. The handshake, `ping`,
//! pagination and framing are `turbomcp-server`'s; this module only relays.

use std::collections::HashMap;
use std::sync::Arc;

use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
use turbomcp_client::Client;
use turbomcp_protocol::types::{
    PromptsCapabilities, ResourcesCapabilities, ServerCapabilities, ToolsCapabilities,
};
use turbomcp_server::prelude::*;
use turbomcp_transport::ChildProcessTransport;

type Backend = Client<ChildProcessTransport>;

/// Upstream tools this proxy was generated for, spelled as the upstream
/// spells them.
pub const TOOLS: &[&str] = &[
{{#each tools}}
    "{{name}}",
{{/each}}
];

/// Upstream prompts this proxy was generated for.
pub const PROMPTS: &[&str] = &[
{{#each prompts}}
    "{{name}}",
{{/each}}
];

/// The upstream's catalogue, fetched once at startup.
struct Catalogue {
    server_info: ServerInfo,
    instructions: Option<String>,
    capabilities: ServerCapabilities,
    tools: Vec<Tool>,
    resources: Vec<Resource>,
    resource_templates: Vec<ResourceTemplate>,
    prompts: Vec<Prompt>,
}

/// MCP server that forwards to the upstream {{server_name}}.
#[derive(Clone)]
pub struct ProxyRouter {
    backend: Arc<Backend>,
    catalogue: Arc<Catalogue>,
}

impl ProxyRouter {
    /// Initialize the upstream and fetch its catalogue.
    ///
    /// Tools and prompts are narrowed to the ones this proxy was generated
    /// for; anything the upstream added since is not routed, so it is not
    /// listed either.
    pub async fn connect(backend: Backend) -> McpResult<Self> {
        let init = backend.initialize().await?;
        let declared = &init.server_capabilities;

        let tools = if declared.tools.is_some() {
            let mut tools = backend.list_tools().await?;
            tools.retain(|tool| TOOLS.contains(&tool.name.as_str()));
            tools
        } else {
            Vec::new()
        };
        let (resources, resource_templates) = if declared.resources.is_some() {
            (
                backend.list_resources().await?,
                backend.list_resource_templates().await?,
            )
        } else {
            (Vec::new(), Vec::new())
        };
        let prompts = if declared.prompts.is_some() {
            let mut prompts = backend.list_prompts().await?;
            prompts.retain(|prompt| PROMPTS.contains(&prompt.name.as_str()));
            prompts
        } else {
            Vec::new()
        };

        // Only what this proxy relays: no list-changed or resource-updated
        // notifications, subscriptions, logging, or completions.
        let capabilities = ServerCapabilities {
            tools: declared
                .tools
                .as_ref()
                .map(|_| ToolsCapabilities { list_changed: None }),
            resources: declared.resources.as_ref().map(|_| ResourcesCapabilities {
                subscribe: None,
                list_changed: None,
            }),
            prompts: declared
                .prompts
                .as_ref()
                .map(|_| PromptsCapabilities { list_changed: None }),
            ..Default::default()
        };

        let upstream = init.server_info.clone();
        let catalogue = Catalogue {
            server_info: ServerInfo {
                name: format!("{}-proxy", upstream.name),
                ..upstream
            },
            instructions: init.instructions.clone(),
            capabilities,
            tools,
            resources,
            resource_templates,
            prompts,
        };

        Ok(Self {
            backend: Arc::new(backend),
            catalogue: Arc::new(catalogue),
        })
    }
{{#each tools}}

    /// `{{name}}`{{#if description}}: {{description}}{{/if}}
    async fn call_{{ident}}(
        &self,
        arguments: Option<HashMap<String, Value>>,
    ) -> McpResult<ToolResult> {
        relay(self.backend.call_tool("{{name}}", arguments, None).await?)
    }
{{/each}}
{{#each prompts}}

    /// `{{name}}`{{#if description}}: {{description}}{{/if}}
    async fn get_{{ident}}(
        &self,
        arguments: Option<HashMap<String, Value>>,
    ) -> McpResult<PromptResult> {
        relay(self.backend.get_prompt("{{name}}", arguments).await?)
    }
{{/each}}
}

/// Convert an upstream result into the server's type for the same wire shape.
fn relay<T: Serialize, U: DeserializeOwned>(value: T) -> McpResult<U> {
    serde_json::to_value(value)
        .and_then(serde_json::from_value)
        .map_err(|e| McpError::internal(e.to_string()))
}

/// MCP arguments are a JSON object, or absent.
fn argument_map(args: Option<Value>) -> McpResult<Option<HashMap<String, Value>>> {
    match args {
        None | Some(Value::Null) => Ok(None),
        Some(Value::Object(map)) => Ok(Some(map.into_iter().collect())),
        Some(_) => Err(McpError::invalid_params("arguments must be an object")),
    }
}

impl McpHandler for ProxyRouter {
    fn server_info(&self) -> ServerInfo {
        self.catalogue.server_info.clone()
    }

    fn instructions(&self) -> Option<String> {
        self.catalogue.instructions.clone()
    }

    fn server_capabilities(&self) -> ServerCapabilities {
        self.catalogue.capabilities.clone()
    }

    fn list_tools(&self) -> Vec<Tool> {
        self.catalogue.tools.clone()
    }

    fn list_resources(&self) -> Vec<Resource> {
        self.catalogue.resources.clone()
    }

    fn list_resource_templates(&self) -> Vec<ResourceTemplate> {
        self.catalogue.resource_templates.clone()
    }

    fn list_prompts(&self) -> Vec<Prompt> {
        self.catalogue.prompts.clone()
    }

    async fn call_tool(
        &self,
        name: &str,
        args: Value,
        _ctx: &RequestContext,
    ) -> McpResult<ToolResult> {
        let arguments = argument_map(Some(args))?;
        match name {
{{#each tools}}
            "{{name}}" => self.call_{{ident}}(arguments).await,
{{/each}}
            _ => Err(McpError::tool_not_found(name)),
        }
    }

    async fn read_resource(&self, uri: &str, _ctx: &RequestContext) -> McpResult<ResourceResult> {
        relay(self.backend.read_resource(uri).await?)
    }

    async fn get_prompt(
        &self,
        name: &str,
        args: Option<Value>,
        _ctx: &RequestContext,
    ) -> McpResult<PromptResult> {
        let arguments = argument_map(args)?;
        match name {
{{#each prompts}}
            "{{name}}" => self.get_{{ident}}(arguments).await,
{{/each}}
            _ => Err(McpError::prompt_not_found(name)),
        }
    }
}