aion-worker-mcp 0.31.0

//! A thin Model Context Protocol client over stdio: newline-delimited
//! JSON-RPC 2.0 to a child process. Initialize, list tools, call a tool,
//! cancel a call. Responses are routed to their callers by request id, so
//! any number of activities may be in flight on one server at once.
//! Requests the server makes of us (`roots/list`, `sampling/createMessage`,
//! `elicitation/create`) are answered "method not found" rather than left
//! hanging: this client is a tool caller, not a host.

use std::collections::{BTreeMap, HashMap};
use std::process::Stdio;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use serde::Deserialize;
use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, Command};
use tokio::sync::{Mutex, broadcast, oneshot};

use crate::config::{McpConfig, ServerSpec};

/// The MCP revision this client announces.
pub const PROTOCOL_VERSION: &str = "2025-06-18";

/// One tool as the server lists it.
#[derive(Debug, Clone, Deserialize)]
pub struct Tool {
    /// Tool name, unique within its server.
    pub name: String,
    /// What the server says it does.
    #[serde(default)]
    pub description: Option<String>,
    /// JSON schema of the arguments object.
    #[serde(rename = "inputSchema")]
    pub input_schema: Value,
    /// JSON schema of `structuredContent`, when the server declares one.
    #[serde(rename = "outputSchema", default)]
    pub output_schema: Option<Value>,
}

/// A JSON-RPC error the server returned for one request.
#[derive(Debug, Clone, Deserialize, thiserror::Error)]
#[error("{message} (code {code})")]
pub struct RpcError {
    /// JSON-RPC error code.
    pub code: i64,
    /// The server's message.
    pub message: String,
    /// Any data the server attached.
    #[serde(default)]
    pub data: Option<Value>,
}

/// Why a call did not produce a result.
#[derive(Debug, thiserror::Error)]
pub enum ClientError {
    /// The server answered with an error.
    #[error("{server}: {error}")]
    Rpc {
        /// Server name.
        server: String,
        /// The error as returned. Boxed so the failure variant does not
        /// widen every result the client hands back.
        #[source]
        error: Box<RpcError>,
    },
    /// The server process ended or its pipe closed while a request was open.
    #[error("{server}: the server exited before answering")]
    ServerGone {
        /// Server name.
        server: String,
    },
    /// Writing to or spawning the server failed.
    #[error("{server}: {context}: {source}")]
    Io {
        /// Server name.
        server: String,
        /// What was being done.
        context: &'static str,
        /// The OS error.
        #[source]
        source: std::io::Error,
    },
    /// The server answered something that is not the shape the method promises.
    #[error("{server}: {method} answered an unexpected shape: {detail}")]
    Shape {
        /// Server name.
        server: String,
        /// The method.
        method: &'static str,
        /// What was wrong.
        detail: String,
    },
    /// The activity was cancelled; the server was told.
    #[error("{server}: call cancelled by the engine")]
    Cancelled {
        /// Server name.
        server: String,
    },
}

type Pending = HashMap<u64, oneshot::Sender<Result<Value, RpcError>>>;

/// A running server and the client speaking to it.
pub struct McpClient {
    name: String,
    next_id: AtomicU64,
    pending: Arc<Mutex<Pending>>,
    stdin: Mutex<ChildStdin>,
    child: Mutex<Child>,
    /// Progress notifications, by progress token, for callers that heartbeat on them.
    progress: broadcast::Sender<Value>,
    /// What the server said about itself at initialize.
    pub server_info: Value,
}

impl McpClient {
    /// The server's name in the document.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Starts the server, performs the initialize handshake.
    ///
    /// # Errors
    ///
    /// Spawn, pipe, or handshake failures, named with the server.
    pub async fn spawn(name: &str, spec: &ServerSpec) -> Result<Arc<Self>, ClientError> {
        let io = |context: &'static str| {
            let server = name.to_owned();
            move |source| ClientError::Io {
                server,
                context,
                source,
            }
        };
        let command = spec.command.as_deref().ok_or_else(|| ClientError::Shape {
            server: name.to_owned(),
            method: "spawn",
            detail: "no command (the document check admits stdio servers only)".to_owned(),
        })?;
        let mut cmd = Command::new(command);
        cmd.args(&spec.args)
            .envs(&spec.env)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        if let Some(cwd) = &spec.cwd {
            cmd.current_dir(cwd);
        }
        cmd.kill_on_drop(true);
        let mut child = cmd.spawn().map_err(io("spawning"))?;
        let stdin = child.stdin.take().ok_or_else(|| ClientError::Shape {
            server: name.to_owned(),
            method: "spawn",
            detail: "no stdin pipe".to_owned(),
        })?;
        let stdout = child.stdout.take().ok_or_else(|| ClientError::Shape {
            server: name.to_owned(),
            method: "spawn",
            detail: "no stdout pipe".to_owned(),
        })?;
        let stderr = child.stderr.take();

        let pending: Arc<Mutex<Pending>> = Arc::new(Mutex::new(HashMap::new()));
        let (progress, _) = broadcast::channel(64);
        let client = Arc::new(Self {
            name: name.to_owned(),
            next_id: AtomicU64::new(1),
            pending: Arc::clone(&pending),
            stdin: Mutex::new(stdin),
            child: Mutex::new(child),
            progress,
            server_info: Value::Null,
        });

        // Reader: routes responses, logs notifications, refuses server requests.
        let reader_client = Arc::clone(&client);
        tokio::spawn(async move {
            let mut lines = BufReader::new(stdout).lines();
            loop {
                match lines.next_line().await {
                    Ok(Some(line)) => reader_client.on_line(&line).await,
                    Ok(None) => break,
                    Err(error) => {
                        tracing::warn!(server = %reader_client.name, %error, "stdout read failed");
                        break;
                    }
                }
            }
            let mut pending = reader_client.pending.lock().await;
            for (_, sender) in pending.drain() {
                let _ = sender.send(Err(RpcError {
                    code: -32000,
                    message: "server exited".to_owned(),
                    data: None,
                }));
            }
            tracing::warn!(server = %reader_client.name, "server stdout closed");
        });
        if let Some(stderr) = stderr {
            let server = name.to_owned();
            tokio::spawn(async move {
                let mut lines = BufReader::new(stderr).lines();
                while let Ok(Some(line)) = lines.next_line().await {
                    tracing::debug!(server = %server, "{line}");
                }
            });
        }

        let init = client
            .request(
                "initialize",
                json!({
                    "protocolVersion": PROTOCOL_VERSION,
                    "capabilities": {},
                    "clientInfo": {"name": "aion-mcp-worker", "version": env!("CARGO_PKG_VERSION")},
                }),
            )
            .await?;
        client
            .notify("notifications/initialized", json!({}))
            .await?;
        let info = init.get("serverInfo").cloned().unwrap_or(Value::Null);
        tracing::info!(server = %name, %info, "initialized");
        // `server_info` is set once, before the Arc is shared beyond this function's tasks.
        let mut client = client;
        if let Some(inner) = Arc::get_mut(&mut client) {
            inner.server_info = init.get("serverInfo").cloned().unwrap_or(Value::Null);
        }
        Ok(client)
    }

    async fn on_line(&self, line: &str) {
        let message: Value = match serde_json::from_str(line) {
            Ok(value) => value,
            Err(error) => {
                tracing::warn!(server = %self.name, %error, line, "non-JSON line from server");
                return;
            }
        };
        let id = message.get("id").and_then(Value::as_u64);
        let method = message.get("method").and_then(Value::as_str);
        match (id, method) {
            (Some(id), None) => {
                let outcome = if let Some(error) = message.get("error") {
                    Err(
                        serde_json::from_value::<RpcError>(error.clone()).unwrap_or(RpcError {
                            code: -32603,
                            message: error.to_string(),
                            data: None,
                        }),
                    )
                } else {
                    Ok(message.get("result").cloned().unwrap_or(Value::Null))
                };
                let sender = self.pending.lock().await.remove(&id);
                if let Some(sender) = sender {
                    let _ = sender.send(outcome);
                } else {
                    tracing::warn!(server = %self.name, id, "response to no open request");
                }
            }
            (Some(_), Some(method)) => {
                // A request from the server: this client is not a host.
                let reply = json!({"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32601, "message": format!("aion-mcp-worker does not serve {method}")}});
                if let Err(error) = self.write(&reply).await {
                    tracing::warn!(server = %self.name, %error, "could not refuse server request");
                }
            }
            (None, Some("notifications/progress")) => {
                let _ = self
                    .progress
                    .send(message.get("params").cloned().unwrap_or(Value::Null));
            }
            (None, Some(method)) => {
                let params = message.get("params").cloned().unwrap_or(Value::Null);
                tracing::debug!(server = %self.name, method, %params, "notification");
            }
            (None, None) => {
                tracing::warn!(server = %self.name, line, "message with neither id nor method");
            }
        }
    }

    async fn write(&self, message: &Value) -> Result<(), ClientError> {
        let mut text = serde_json::to_string(message).map_err(|source| ClientError::Shape {
            server: self.name.clone(),
            method: "write",
            detail: source.to_string(),
        })?;
        text.push('\n');
        let mut stdin = self.stdin.lock().await;
        stdin
            .write_all(text.as_bytes())
            .await
            .map_err(|source| ClientError::Io {
                server: self.name.clone(),
                context: "writing to the server",
                source,
            })?;
        stdin.flush().await.map_err(|source| ClientError::Io {
            server: self.name.clone(),
            context: "flushing to the server",
            source,
        })
    }

    /// Sends a notification (no id, no reply).
    async fn notify(&self, method: &str, params: Value) -> Result<(), ClientError> {
        self.write(&json!({"jsonrpc": "2.0", "method": method, "params": params}))
            .await
    }

    /// Sends a request and waits for its reply.
    async fn request(&self, method: &'static str, params: Value) -> Result<Value, ClientError> {
        let (id, receiver) = self.open(method, params).await?;
        self.await_reply(id, receiver).await
    }

    async fn open(
        &self,
        method: &'static str,
        params: Value,
    ) -> Result<(u64, oneshot::Receiver<Result<Value, RpcError>>), ClientError> {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let (sender, receiver) = oneshot::channel();
        self.pending.lock().await.insert(id, sender);
        if let Err(error) = self
            .write(&json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params}))
            .await
        {
            self.pending.lock().await.remove(&id);
            return Err(error);
        }
        Ok((id, receiver))
    }

    async fn await_reply(
        &self,
        id: u64,
        receiver: oneshot::Receiver<Result<Value, RpcError>>,
    ) -> Result<Value, ClientError> {
        match receiver.await {
            Ok(Ok(value)) => Ok(value),
            Ok(Err(error)) if error.code == -32000 && error.message == "server exited" => {
                Err(ClientError::ServerGone {
                    server: self.name.clone(),
                })
            }
            Ok(Err(error)) => Err(ClientError::Rpc {
                server: self.name.clone(),
                error: Box::new(error),
            }),
            Err(_) => {
                self.pending.lock().await.remove(&id);
                Err(ClientError::ServerGone {
                    server: self.name.clone(),
                })
            }
        }
    }

    /// Lists every tool, following pagination to the end.
    ///
    /// # Errors
    ///
    /// The server's error, or a reply that is not a tool list.
    pub async fn list_tools(&self) -> Result<Vec<Tool>, ClientError> {
        let mut tools = Vec::new();
        let mut cursor: Option<String> = None;
        loop {
            let params = cursor
                .as_ref()
                .map_or_else(|| json!({}), |c| json!({"cursor": c}));
            let page = self.request("tools/list", params).await?;
            let listed = page.get("tools").cloned().unwrap_or(Value::Null);
            let parsed: Vec<Tool> =
                serde_json::from_value(listed).map_err(|error| ClientError::Shape {
                    server: self.name.clone(),
                    method: "tools/list",
                    detail: error.to_string(),
                })?;
            tools.extend(parsed);
            match page.get("nextCursor").and_then(Value::as_str) {
                Some(next) if !next.is_empty() => cursor = Some(next.to_owned()),
                _ => return Ok(tools),
            }
        }
    }

    /// Calls a tool. If `cancelled` resolves first, the server is sent
    /// `notifications/cancelled` for this request and the call ends
    /// [`ClientError::Cancelled`]. `on_progress` is invoked for each progress
    /// notification carrying this call's token.
    ///
    /// # Errors
    ///
    /// The server's error, the server dying, or cancellation.
    pub async fn call_tool(
        &self,
        tool: &str,
        arguments: Value,
        cancelled: impl std::future::Future<Output = ()>,
        mut on_progress: impl FnMut(&Value),
    ) -> Result<Value, ClientError> {
        let token = format!("aion-{}", self.next_id.load(Ordering::Relaxed));
        let params =
            json!({"name": tool, "arguments": arguments, "_meta": {"progressToken": token}});
        let (id, receiver) = self.open("tools/call", params).await?;
        let mut progress = self.progress.subscribe();
        tokio::pin!(cancelled);
        let reply = self.await_reply(id, receiver);
        tokio::pin!(reply);
        loop {
            tokio::select! {
                outcome = &mut reply => return outcome,
                () = &mut cancelled => {
                    let notice = json!({"requestId": id, "reason": "activity cancelled by the Aion engine"});
                    if let Err(error) = self.notify("notifications/cancelled", notice).await {
                        tracing::warn!(server = %self.name, %error, "could not send cancellation");
                    }
                    self.pending.lock().await.remove(&id);
                    return Err(ClientError::Cancelled { server: self.name.clone() });
                }
                event = progress.recv() => {
                    match event {
                        Ok(params) => {
                            let mine = params.get("progressToken").and_then(Value::as_str) == Some(token.as_str());
                            if mine { on_progress(&params); }
                        }
                        Err(broadcast::error::RecvError::Lagged(n)) => tracing::debug!(server = %self.name, n, "progress lagged"),
                        Err(broadcast::error::RecvError::Closed) => {}
                    }
                }
            }
        }
    }

    /// Ends the server: close stdin, wait briefly, kill if it lingers.
    pub async fn shutdown(&self) {
        {
            let mut stdin = self.stdin.lock().await;
            let _ = stdin.shutdown().await;
        }
        let mut child = self.child.lock().await;
        match tokio::time::timeout(std::time::Duration::from_secs(5), child.wait()).await {
            Ok(Ok(status)) => tracing::info!(server = %self.name, %status, "server exited"),
            Ok(Err(error)) => tracing::warn!(server = %self.name, %error, "waiting on server"),
            Err(_) => {
                tracing::warn!(server = %self.name, "server did not exit after stdin closed; killing");
                if let Err(error) = child.kill().await {
                    tracing::warn!(server = %self.name, %error, "kill failed");
                }
            }
        }
    }
}

/// Starts every enabled server in the document, in name order.
///
/// # Errors
///
/// The first server that cannot be started or initialized, named. Servers
/// already started are shut down before the error returns.
pub async fn spawn_all(config: &McpConfig) -> Result<Vec<Arc<McpClient>>, ClientError> {
    let mut clients: Vec<Arc<McpClient>> = Vec::new();
    let servers: BTreeMap<&String, &ServerSpec> = config.enabled_servers().into_iter().collect();
    for (name, spec) in servers {
        match McpClient::spawn(name, spec).await {
            Ok(client) => clients.push(client),
            Err(error) => {
                for started in &clients {
                    started.shutdown().await;
                }
                return Err(error);
            }
        }
    }
    Ok(clients)
}