johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::collections::HashMap;

use async_trait::async_trait;
use rmcp::model::ProtocolVersion;
use rmcp::model::{
    CallToolRequestParams, CallToolResult, ContentBlock, ReadResourceRequestParams,
    ResourceContents, ServerNotification, SubscriptionFilter, Tool,
};
use rmcp::service::{ClientLifecycleMode, ClientServiceExt, RunningService};
use rmcp::transport::TokioChildProcess;
use rmcp::RoleClient;
use serde_json::{Map, Value};
use tokio::process::Command;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;

use crate::resources::{
    example_of, template_matches, ResourceContent, ResourceInfo, ResourceSource, ResourceWatch,
    WatchSignal,
};
use crate::tools::{ToolCallOutput, ToolDefinition, ToolError, ToolExecutor};

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StdioServerConfig {
    pub command: String,
    pub args: Vec<String>,
}

impl StdioServerConfig {
    pub fn parse(spec: &str) -> Result<Self, ToolError> {
        let mut parts = spec.split_whitespace();
        let command = parts
            .next()
            .ok_or_else(|| ToolError::new("empty --mcp-stdio command"))?;
        Ok(Self {
            command: command.to_string(),
            args: parts.map(str::to_string).collect(),
        })
    }
}

/// A tool name declared by more than one server; the first server wins.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DuplicateTool {
    pub name: String,
    pub kept: String,
    pub ignored: String,
}

struct Server {
    label: String,
    client: RunningService<RoleClient, ()>,
    subscribable: bool,
    resources: Vec<ResourceInfo>,
}

/// Calls tools and reads/watches resources on stdio MCP servers. Each tool
/// name routes to exactly one server. Tool output is returned in full; calls
/// run concurrently.
pub struct McpToolExecutor {
    servers: Vec<Server>,
    definitions: Vec<ToolDefinition>,
    routes: HashMap<String, usize>,
    duplicates: Vec<DuplicateTool>,
}

impl McpToolExecutor {
    pub async fn connect_stdio(configs: Vec<StdioServerConfig>) -> Result<Self, ToolError> {
        if configs.is_empty() {
            return Err(ToolError::new("no MCP stdio servers configured"));
        }

        let mut servers: Vec<Server> = Vec::new();
        let mut definitions = Vec::new();
        let mut routes: HashMap<String, usize> = HashMap::new();
        let mut duplicates = Vec::new();

        for config in configs {
            let client = connect_one(&config).await?;
            let tools = client.list_all_tools().await.map_err(|error| {
                ToolError::new(format!(
                    "tools/list failed for {:?}: {error}",
                    config.command
                ))
            })?;

            let index = servers.len();
            for tool in tools {
                let definition = tool_definition(&tool);
                if let Some(&kept) = routes.get(&definition.name) {
                    duplicates.push(DuplicateTool {
                        name: definition.name.clone(),
                        kept: servers[kept].label.clone(),
                        ignored: config.command.clone(),
                    });
                    continue;
                }
                routes.insert(definition.name.clone(), index);
                definitions.push(definition);
            }

            let capability = client
                .peer_info()
                .and_then(|info| info.capabilities.resources.clone());
            let subscribable = capability
                .as_ref()
                .is_some_and(|resources| resources.subscribe == Some(true));
            let resources = match capability {
                Some(_) => list_server_resources(&client, &config.command, subscribable).await,
                None => Vec::new(),
            };
            servers.push(Server {
                label: config.command,
                client,
                subscribable,
                resources,
            });
        }

        Ok(Self {
            servers,
            definitions,
            routes,
            duplicates,
        })
    }

    /// Tool names declared by more than one server.
    pub fn duplicates(&self) -> &[DuplicateTool] {
        &self.duplicates
    }

    /// The server for a resource URI: exact listing, then template match,
    /// then the first server with resources (subscribable if required).
    fn resource_server(&self, uri: &str, subscribe: bool) -> Result<&Server, ToolError> {
        let listed = self.servers.iter().find(|server| {
            server.resources.iter().any(|resource| {
                if resource.template {
                    template_matches(&resource.uri, uri)
                } else {
                    resource.uri == uri
                }
            })
        });
        let fallback = || {
            self.servers.iter().find(|server| {
                if subscribe {
                    server.subscribable
                } else {
                    !server.resources.is_empty() || server.subscribable
                }
            })
        };
        let server = listed.or_else(fallback).ok_or_else(|| {
            ToolError::new(format!(
                "no MCP server {} resource {uri:?}",
                if subscribe { "can watch" } else { "serves" }
            ))
        })?;
        if subscribe && !server.subscribable {
            return Err(ToolError::new(format!(
                "MCP server {:?} serves {uri:?} but does not support resource subscriptions",
                server.label
            )));
        }
        Ok(server)
    }
}

/// Connects one stdio server. `server/discover` (MCP 2026-07-28, required
/// for watches) first; servers without discovery are reached with
/// `initialize` and can serve tools but not watches.
async fn connect_one(
    config: &StdioServerConfig,
) -> Result<RunningService<RoleClient, ()>, ToolError> {
    let mut command = Command::new(&config.command);
    command.args(&config.args);
    let transport = TokioChildProcess::new(command).map_err(|error| {
        ToolError::new(format!(
            "failed to spawn MCP server {:?}: {error}",
            config.command
        ))
    })?;
    let lifecycle = ClientLifecycleMode::Auto {
        preferred_versions: vec![ProtocolVersion::LATEST],
        legacy_version: Some(ProtocolVersion::LATEST_WITH_INITIALIZE),
    };
    ().serve_with_lifecycle(transport, lifecycle)
        .await
        .map_err(|error| {
            ToolError::new(format!(
                "MCP connect failed for {:?}: {error}",
                config.command
            ))
        })
}

async fn list_server_resources(
    client: &RunningService<RoleClient, ()>,
    label: &str,
    subscribable: bool,
) -> Vec<ResourceInfo> {
    let mut resources = Vec::new();
    if let Ok(listed) = client.list_all_resources().await {
        resources.extend(listed.into_iter().map(|resource| ResourceInfo {
            uri: resource.uri.clone(),
            template: false,
            name: resource.name.clone(),
            example: example_of(resource.description.as_deref()),
            description: resource.description.clone(),
            mime_type: resource.mime_type.clone(),
            server: label.to_string(),
            subscribable,
        }));
    }
    if let Ok(templates) = client.list_all_resource_templates().await {
        resources.extend(templates.into_iter().map(|template| ResourceInfo {
            uri: template.uri_template.clone(),
            template: true,
            name: template.name.clone(),
            example: example_of(template.description.as_deref()),
            description: template.description.clone(),
            mime_type: template.mime_type.clone(),
            server: label.to_string(),
            subscribable,
        }));
    }
    resources
}

/// Aborts a forwarding task when the watch is dropped.
struct AbortOnDrop(JoinHandle<()>);

impl Drop for AbortOnDrop {
    fn drop(&mut self) {
        self.0.abort();
    }
}

#[async_trait]
impl ResourceSource for McpToolExecutor {
    async fn list_resources(&self) -> Result<Vec<ResourceInfo>, ToolError> {
        Ok(self
            .servers
            .iter()
            .flat_map(|server| server.resources.iter().cloned())
            .collect())
    }

    async fn read_resource(&self, uri: &str) -> Result<ResourceContent, ToolError> {
        let server = self.resource_server(uri, false)?;
        let result = server
            .client
            .read_resource(ReadResourceRequestParams::new(uri))
            .await
            .map_err(|error| ToolError::new(format!("resources/read {uri:?} failed: {error}")))?;
        match result.contents.into_iter().next() {
            Some(ResourceContents::TextResourceContents {
                text, mime_type, ..
            }) => Ok(ResourceContent { text, mime_type }),
            Some(_) => Err(ToolError::new(format!(
                "resource {uri:?} is binary; only text resources are supported"
            ))),
            None => Ok(ResourceContent {
                text: String::new(),
                mime_type: None,
            }),
        }
    }

    async fn watch(&self, uri: &str) -> Result<ResourceWatch, ToolError> {
        let server = self.resource_server(uri, true)?;
        let (signals_tx, signals) = mpsc::channel(64);
        let filter = SubscriptionFilter::builder()
            .resource_subscription(uri)
            .build();

        let mut subscription = server.client.peer().listen(filter).await.map_err(|error| {
            ToolError::new(format!(
                "cannot watch {uri:?}: subscriptions/listen failed on {:?} ({error}); watches need an MCP 2026-07-28 server",
                server.label
            ))
        })?;
        let accepted = subscription
            .acknowledged()
            .resource_subscriptions
            .as_ref()
            .is_some_and(|uris| uris.iter().any(|accepted| accepted == uri));
        if !accepted {
            return Err(ToolError::new(format!(
                "MCP server {:?} did not accept a subscription to {uri:?}",
                server.label
            )));
        }
        let watched = uri.to_string();
        let task = tokio::spawn(async move {
            loop {
                let signal = match subscription.next().await {
                    Ok(Some(ServerNotification::ResourceUpdatedNotification(notification)))
                        if notification.params.uri == watched =>
                    {
                        WatchSignal::Updated
                    }
                    Ok(Some(_)) => continue,
                    Ok(None) => {
                        WatchSignal::Ended(format!("subscription ended: {:?}", subscription.end()))
                    }
                    Err(error) => WatchSignal::Ended(error.to_string()),
                };
                let ended = matches!(signal, WatchSignal::Ended(_));
                if signals_tx.send(signal).await.is_err() || ended {
                    break;
                }
            }
        });
        Ok(ResourceWatch {
            signals,
            guard: Box::new(AbortOnDrop(task)),
        })
    }
}

fn tool_definition(tool: &Tool) -> ToolDefinition {
    let annotations = tool.annotations.as_ref();
    ToolDefinition {
        name: tool.name.to_string(),
        description: tool.description.as_ref().map(|value| value.to_string()),
        input_schema: Value::Object((*tool.input_schema).clone()),
        output_schema: tool
            .output_schema
            .as_ref()
            .map(|schema| Value::Object((**schema).clone())),
        read_only: annotations.and_then(|value| value.read_only_hint),
        destructive: annotations.and_then(|value| value.destructive_hint),
    }
}

#[async_trait]
impl ToolExecutor for McpToolExecutor {
    async fn list_tools(&self) -> Result<Vec<ToolDefinition>, ToolError> {
        Ok(self.definitions.clone())
    }

    async fn call_tool(&self, name: &str, arguments: Value) -> Result<ToolCallOutput, ToolError> {
        let arguments = match arguments {
            Value::Object(arguments) => arguments,
            Value::Null => Map::new(),
            other => {
                return Err(ToolError::new(format!(
                    "tool arguments must be a JSON object: {other}"
                )));
            }
        };
        let client = self
            .routes
            .get(name)
            .and_then(|index| self.servers.get(*index))
            .map(|server| &server.client)
            .ok_or_else(|| ToolError::new(format!("unknown tool {name:?}")))?;

        let params = CallToolRequestParams::new(name.to_string()).with_arguments(arguments);
        client
            .call_tool(params)
            .await
            .map(format_result)
            .map_err(|error| ToolError::new(format!("tool {name:?} failed: {error}")))
    }
}

/// Joins text content. Structured content, when present, is the exact JSON
/// value and is preferred.
fn format_result(result: CallToolResult) -> ToolCallOutput {
    let is_error = result.is_error.unwrap_or(false);
    if let Some(structured) = result.structured_content {
        return ToolCallOutput {
            content: structured.to_string(),
            is_error,
            audit: None,
        };
    }

    let parts: Vec<String> = result
        .content
        .iter()
        .map(|block| match block {
            ContentBlock::Text(text) => text.text.clone(),
            other => {
                serde_json::to_string(other).unwrap_or_else(|_| "[unserializable content]".into())
            }
        })
        .collect();

    ToolCallOutput {
        content: parts.join("\n"),
        is_error,
        audit: None,
    }
}

#[cfg(test)]
mod tests {
    use rmcp::model::{CallToolResult, ContentBlock};
    use serde_json::json;

    use super::format_result;

    #[test]
    fn large_tool_output_is_not_truncated() {
        let big = "x".repeat(1_000_000);
        let output = format_result(CallToolResult::success(vec![ContentBlock::text(
            big.clone(),
        )]));
        assert_eq!(output.content, big);
        assert!(!output.is_error);
    }

    #[test]
    fn structured_content_is_preferred() {
        let output = format_result(CallToolResult::structured(json!({"items": [1, 2]})));
        assert_eq!(output.content, r#"{"items":[1,2]}"#);
        assert!(!output.is_error);
    }
}