oxios_mcp/

client.rs

//! MCP client — manages a single MCP server process lifecycle.
//!
//! `McpClient` spawns a child process and communicates with it over stdin/stdout
//! using JSON-RPC 2.0 messages (one JSON object per line).
//!
//! I/O handles are stored persistently (not consumed via `take()`) so that
//! multiple requests can be serialized through the same connection. A write
//! lock on both stdin and stdout is held for the duration of each request-response
//! cycle, ensuring correct ordering.

use anyhow::{Context, Result, anyhow};
use std::sync::atomic::AtomicUsize;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle;
use tokio::time::{Duration, timeout};

use crate::protocol::*;

// ---------------------------------------------------------------------------
// McpClient — manages a single MCP server process lifecycle
// ---------------------------------------------------------------------------

/// Manages a single MCP server process with stdio JSON-RPC communication.
///
/// I/O handles are stored persistently so that concurrent requests can be
/// serialized through the same connection without consuming the handles.
///
/// # Example
///
/// ```ignore
/// let client = McpClient::new(server_config);
/// client.initialize().await?;
/// let tools = client.list_tools().await?;
/// let result = client.call_tool("my_tool", serde_json::json!({"arg": "value"})).await?;
/// client.shutdown().await?;
/// ```
pub struct McpClient {
    /// Server configuration
    server: McpServer,
    /// Child process handle (None when not running).
    ///
    /// The child is spawned with `kill_on_drop(true)` (F3) so it is reaped
    /// even if the `McpClient` is dropped without an explicit `shutdown()`.
    child: RwLock<Option<Child>>,
    /// Persistent stdin handle for writing to the server process.
    ///
    /// A `Mutex` (not `RwLock`) since access is always exclusive (F4).
    stdin: Mutex<Option<tokio::io::BufWriter<ChildStdin>>>,
    /// Persistent stdout handle for reading from the server process.
    ///
    /// A `Mutex` (not `RwLock`) since access is always exclusive (F4).
    stdout: Mutex<Option<BufReader<ChildStdout>>>,
    /// Whether the server has been initialized
    initialized: RwLock<bool>,
    /// Cached tool list (invalidated on refresh_tools)
    tool_cache: RwLock<Option<Vec<McpTool>>>,
    /// Server info received during initialize
    server_info: RwLock<Option<ServerInfo>>,
    /// Request timeout duration
    request_timeout: Duration,
    /// Background task that drains the child's stderr so the OS pipe
    /// buffer doesn't fill and deadlock the server (F1).
    stderr_task: Mutex<Option<JoinHandle<()>>>,
    /// Per-connection JSON-RPC request ID counter (F8: per-client instead
    /// of a process-global counter, so logs unambiguously attribute ids).
    next_id: AtomicUsize,
}

impl McpClient {
    /// Create a new MCP client for the given server configuration.
    ///
    /// Does NOT spawn the process yet — call `initialize()` to start and negotiate.
    pub fn new(server: McpServer) -> Self {
        Self {
            server,
            child: RwLock::new(None),
            stdin: Mutex::new(None),
            stdout: Mutex::new(None),
            initialized: RwLock::new(false),
            tool_cache: RwLock::new(None),
            server_info: RwLock::new(None),
            request_timeout: Duration::from_secs(30),
            stderr_task: Mutex::new(None),
            next_id: AtomicUsize::new(1),
        }
    }

    /// Set the request timeout duration.
    #[must_use]
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.request_timeout = timeout;
        self
    }

    /// Spawn the MCP server process and establish communication.
    ///
    /// On failure the spawned child is killed and all handles are cleared
    /// (F3), so retrying `initialize()` never orphans a process. The child
    /// is also spawned with `kill_on_drop(true)` as a safety net.
    pub async fn initialize(&self) -> Result<()> {
        if *self.initialized.read().await {
            return Ok(());
        }

        // Spawn the child process. `kill_on_drop(true)` guarantees the child
        // is killed if the `Child` handle is dropped without an explicit
        // kill — e.g. when `McpClient` is dropped or a handle is overwritten
        // by a retry (F3, F7).
        let mut child = Command::new(&self.server.command)
            .args(&self.server.args)
            .envs(&self.server.env)
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .kill_on_drop(true)
            .spawn()
            .with_context(|| format!("Failed to spawn MCP server '{}'", self.server.name))?;

        let stdin = child
            .stdin
            .take()
            .expect("stdin not captured — stdin was piped");
        let stdout = child
            .stdout
            .take()
            .expect("stdout not captured — stdout was piped");
        let stderr = child
            .stderr
            .take()
            .expect("stderr not captured — stderr was piped");

        // F1: drain stderr continuously. If we don't read it, a chatty
        // server can fill the OS pipe buffer (~64KiB on Linux) and block
        // forever on stderr writes, deadlocking the whole connection.
        let stderr_server_name = self.server.name.clone();
        let stderr_task = tokio::spawn(async move {
            let mut reader = BufReader::new(stderr);
            let mut line = String::new();
            loop {
                line.clear();
                match reader.read_line(&mut line).await {
                    Ok(0) => break, // EOF — child closed stderr
                    Ok(_) => {
                        let trimmed = line.trim_end_matches(['\n', '\r']);
                        if !trimmed.is_empty() {
                            tracing::debug!(
                                server = %stderr_server_name,
                                stream = "stderr",
                                "{}",
                                trimmed
                            );
                        }
                    }
                    Err(e) => {
                        tracing::debug!(
                            server = %stderr_server_name,
                            stream = "stderr",
                            error = %e,
                            "stderr drain stopping"
                        );
                        break;
                    }
                }
            }
        });

        // Store persistent I/O handles (separate from child process handle)
        *self.stdin.lock().await = Some(tokio::io::BufWriter::new(stdin));
        *self.stdout.lock().await = Some(BufReader::new(stdout));
        *self.stderr_task.lock().await = Some(stderr_task);

        // Store child handle
        *self.child.write().await = Some(child);

        // Send initialize request using persistent handles. Use do_request
        // directly (not send_request) to avoid recursion, since send_request
        // may call restart() which calls initialize().
        let params = InitializeParams::default();
        let request = McpRequest::with_id(self.next_id(), "initialize")
            .with_params(serde_json::to_value(&params)?);

        // F3: on any failure during the initialize handshake, tear down the
        // spawned child so a retry starts clean (no orphaned process, no
        // stale I/O handles).
        let response = match self.do_request(request).await {
            Ok(resp) => resp,
            Err(e) => {
                self.cleanup_child().await;
                return Err(e);
            }
        };

        // Parse initialize result
        let result_json = response.into_result()?;
        let init_result: InitializeResult = serde_json::from_value(result_json)?;

        *self.server_info.write().await = Some(init_result.server_info.clone());
        *self.initialized.write().await = true;

        // Send the `initialized` notification (JSON-RPC 2.0 requires this
        // after a successful `initialize`). Notifications carry no `id`.
        let notification = McpRequest::notification("notifications/initialized");
        self.send_notification(notification).await?;

        tracing::debug!(
            server = %self.server.name,
            version = %init_result.server_info.version,
            "MCP server initialized"
        );

        Ok(())
    }

    /// Check if the server has been initialized
    pub async fn is_initialized(&self) -> bool {
        *self.initialized.read().await
    }

    /// Get the server info received during initialize
    pub async fn server_info(&self) -> Option<ServerInfo> {
        self.server_info.read().await.clone()
    }

    /// Send a JSON-RPC request using persistent I/O handles.
    ///
    /// Acquires exclusive locks on both stdin and stdout for the duration of
    /// the request-response cycle, serializing concurrent access. Reads in a
    /// loop, skipping JSON-RPC notifications / server-initiated requests, so
    /// that interleaved server output can't be mistaken for our response (F2).
    async fn do_request(&self, request: McpRequest) -> Result<McpResponse> {
        let request_id = request.id.clone();

        // Acquire stdin lock for writing
        let mut stdin_guard = self.stdin.lock().await;
        let stdin = stdin_guard
            .as_mut()
            .ok_or_else(|| anyhow!("stdin not available on '{}'", self.server.name))?;

        // Write the request
        let json = request.to_jsonl()?;
        timeout(self.request_timeout, async {
            stdin.write_all(&json).await?;
            stdin.flush().await?;
            Ok::<(), tokio::io::Error>(())
        })
        .await
        .map_err(|e| anyhow::anyhow!("MCP request timed out (write): {e}"))??;

        // Acquire stdout lock for reading
        let mut stdout_guard = self.stdout.lock().await;
        let stdout = stdout_guard
            .as_mut()
            .ok_or_else(|| anyhow!("stdout not available on '{}'", self.server.name))?;

        // F2: read lines until we get the response for *this* request id.
        // The server may emit JSON-RPC notifications (no id, e.g. progress,
        // logging) or server-initiated requests (has both id and method, e.g.
        // sampling/createMessage) at any time. We log and skip those instead
        // of misinterpreting them as our response.
        loop {
            let line: std::io::Result<Option<String>> = timeout(self.request_timeout, async {
                stdout.lines().next_line().await
            })
            .await
            .map_err(|e| anyhow::anyhow!("MCP request timed out (read): {e}"))?;

            let response_str: String = line
                .context("Failed to read MCP response line from stdout")?
                .with_context(|| format!("MCP server {} returned no response", self.server.name))?;

            // Parse as a generic JSON value first so we can classify the
            // message without committing to the response shape.
            let value: serde_json::Value = serde_json::from_str(&response_str)
                .with_context(|| format!("Failed to parse MCP message JSON: {response_str}"))?;

            // A "method" field means it's a notification or a server-initiated
            // request — not a response to us. Log and keep reading.
            if value.get("method").is_some() {
                tracing::debug!(
                    server = %self.server.name,
                    method = ?value.get("method"),
                    "MCP server sent a notification/server request; skipping"
                );
                continue;
            }

            // It's a response — verify the id matches.
            let got_id = value.get("id");
            if got_id != Some(&request_id) {
                // A response for a different request id shouldn't happen under
                // our serialized access, but if it does (e.g. a stale buffered
                // response from a previous timed-out request) skip it rather
                // than return the wrong result.
                tracing::warn!(
                    server = %self.server.name,
                    expected_id = ?request_id,
                    got_id = ?got_id,
                    "MCP response ID mismatch, skipping"
                );
                continue;
            }

            let parsed: McpResponse = serde_json::from_value(value)
                .with_context(|| format!("Failed to parse MCP response: {response_str}"))?;
            return Ok(parsed);
        }
    }

    /// Send a JSON-RPC notification (no response expected).
    async fn send_notification(&self, notification: McpRequest) -> Result<()> {
        let mut stdin_guard = self.stdin.lock().await;
        let stdin = stdin_guard
            .as_mut()
            .ok_or_else(|| anyhow!("stdin not available on '{}'", self.server.name))?;

        let json = notification.to_jsonl()?;
        stdin.write_all(&json).await?;
        stdin.flush().await?;

        Ok(())
    }

    /// Send a JSON-RPC request via persistent I/O handles.
    ///
    /// If the server is not running, attempts one automatic restart before
    /// failing. On a communication error mid-request, restarts and retries
    /// the original request once (F5) instead of bailing out and asking the
    /// caller to retry.
    pub(crate) async fn send_request(&self, request: McpRequest) -> Result<McpResponse> {
        // Verify server is running; attempt auto-restart if not
        {
            let child = self.child.read().await;
            if child.is_none() {
                tracing::warn!(
                    server = %self.server.name,
                    "MCP server not running, attempting auto-start"
                );
                drop(child);
                // Use restart (shutdown + initialize) which doesn't call send_request
                self.restart().await?;
            }
        }

        // F5: keep a clone so we can transparently retry the original request
        // after a restart-induced communication error.
        let request_for_retry = request.clone();
        match self.do_request(request).await {
            Ok(resp) => Ok(resp),
            Err(e) => {
                // Auto-restart on communication errors (crashed server).
                let err_str = e.to_string();
                let is_comm_error = err_str.contains("not available")
                    || err_str.contains("broken pipe")
                    || err_str.contains("timed out")
                    || err_str.contains("no response")
                    || err_str.contains("reset by peer");

                if is_comm_error {
                    tracing::warn!(
                        server = %self.server.name,
                        error = %err_str,
                        "MCP communication error, attempting auto-restart + retry"
                    );
                    self.restart().await?;
                    // F5: retry the original request once instead of pushing
                    // the retry burden onto every caller.
                    self.do_request(request_for_retry).await
                } else {
                    Err(e)
                }
            }
        }
    }

    /// Allocate the next per-connection request id (F8).
    fn next_id(&self) -> usize {
        self.next_id
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
    }

    /// Best-effort teardown of the spawned child + I/O handles when
    /// `initialize()` fails mid-handshake or a restart is required. Safe to
    /// call when nothing is running. Idempotent (F3).
    async fn cleanup_child(&self) {
        // Drop I/O handles first so the child's pipes close.
        *self.stdin.lock().await = None;
        *self.stdout.lock().await = None;

        // Abort the stderr drain task — the child is going away.
        if let Some(handle) = self.stderr_task.lock().await.take() {
            handle.abort();
        }

        // Kill and reap the child. `kill_on_drop(true)` would handle this
        // when the Child drops, but explicit kill+wait avoids racing with
        // a concurrent `initialize()` reusing the handle.
        if let Some(mut child) = self.child.write().await.take() {
            let _ = child.kill().await;
            let _ = child.wait().await;
        }

        *self.initialized.write().await = false;
    }

    /// List all tools available from this MCP server.
    ///
    /// Results are cached and refreshed on [refresh_tools](Self::refresh_tools).
    pub async fn list_tools(&self) -> Result<Vec<McpTool>> {
        // Return cached tools if available
        if let Some(cached) = self.tool_cache.read().await.clone() {
            return Ok(cached);
        }

        self.refresh_tools().await
    }

    /// Force-refresh the tool list from the server.
    pub async fn refresh_tools(&self) -> Result<Vec<McpTool>> {
        let request = McpRequest::with_id(self.next_id(), "tools/list");
        let response = self.send_request(request).await?;

        let result_json = response.into_result()?;
        let tools_result: McpToolsResult = serde_json::from_value(result_json)?;

        let tools = tools_result.tools;
        *self.tool_cache.write().await = Some(tools.clone());

        tracing::debug!(
            server = %self.server.name,
            count = tools.len(),
            "Refreshed tool cache"
        );

        Ok(tools)
    }

    /// Call a tool on this MCP server.
    ///
    /// The server must be initialized first.
    pub async fn call_tool(
        &self,
        tool_name: &str,
        arguments: serde_json::Value,
    ) -> Result<McpToolCallResult> {
        let params = serde_json::json!({
            "name": tool_name,
            "arguments": arguments,
        });

        let request = McpRequest::with_id(self.next_id(), "tools/call").with_params(params);
        let response = self.send_request(request).await?;

        let result_json = response.into_result()?;
        let call_result: McpToolCallResult = serde_json::from_value(result_json)?;

        tracing::debug!(
            server = %self.server.name,
            tool = tool_name,
            "Tool call completed"
        );

        Ok(call_result)
    }

    /// Call a tool and return the result content as a string.
    ///
    /// Returns the first text content block, or an error if no text content.
    pub async fn call_tool_text(
        &self,
        tool_name: &str,
        arguments: serde_json::Value,
    ) -> Result<String> {
        let result = self.call_tool(tool_name, arguments).await?;

        for block in result.content {
            if let McpContentBlock::Text { text } = block {
                return Ok(text);
            }
        }

        Err(anyhow!("Tool '{tool_name}' returned no text content"))
    }

    /// Gracefully shutdown the MCP server process.
    ///
    /// Drops persistent I/O handles first, aborts the stderr drain, then
    /// kills the child process.
    pub async fn shutdown(&self) -> Result<()> {
        // Drop persistent I/O handles first so the child's pipes close.
        *self.stdin.lock().await = None;
        *self.stdout.lock().await = None;

        // Abort the stderr drain task.
        if let Some(handle) = self.stderr_task.lock().await.take() {
            handle.abort();
        }

        let mut child_guard = self.child.write().await;

        if let Some(mut child) = child_guard.take() {
            tracing::debug!(server = %self.server.name, "Shutting down MCP server");

            // Try graceful shutdown first
            let _ = child.try_wait();

            // Kill the process
            child.kill().await?;
            let _ = child.wait().await;
        }

        *self.initialized.write().await = false;
        *self.tool_cache.write().await = None;

        Ok(())
    }

    /// Restart the server (shutdown then initialize).
    pub async fn restart(&self) -> Result<()> {
        self.shutdown().await?;
        self.initialize().await
    }

    /// Get the server configuration
    pub fn server(&self) -> &McpServer {
        &self.server
    }
}

impl std::fmt::Debug for McpClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("McpClient")
            .field("server", &self.server.name)
            .field("initialized", &self.initialized)
            .finish()
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::time::Duration;

    // --- McpClient construction and configuration tests ---

    #[test]
    fn test_client_construction() {
        let server = McpServer::new("test-server", "npx");
        let client = McpClient::new(server);

        // Verify the server config is stored correctly
        assert_eq!(client.server.name, "test-server");
        assert_eq!(client.server.command, "npx");
    }

    #[test]
    fn test_client_with_timeout() {
        let server = McpServer::new("test", "echo");
        let client = McpClient::new(server).with_timeout(Duration::from_secs(60));

        // The timeout should be set to 60 seconds
        // We verify this indirectly by checking the client was constructed
        // with the modified configuration (via the builder pattern)
        assert_eq!(client.server.name, "test");
    }

    #[test]
    fn test_client_with_timeout_short() {
        let server = McpServer::new("test", "sleep");
        let client = McpClient::new(server).with_timeout(Duration::from_millis(50));

        assert_eq!(client.server.name, "test");
        // Timeout of 50ms is very short
    }

    #[test]
    fn test_client_debug_format() {
        let server = McpServer::new("debug-test", "echo");
        let client = McpClient::new(server);

        let debug_str = format!("{client:?}");

        // Debug output should contain the server name
        assert!(debug_str.contains("debug-test"));
        assert!(debug_str.contains("McpClient"));
    }

    #[test]
    fn test_client_debug_different_servers() {
        let server1 = McpServer::new("server-a", "cmd1");
        let server2 = McpServer::new("server-b", "cmd2");

        let client1 = McpClient::new(server1);
        let client2 = McpClient::new(server2);

        let debug1 = format!("{client1:?}");
        let debug2 = format!("{client2:?}");

        assert!(debug1.contains("server-a"));
        assert!(debug2.contains("server-b"));
        assert_ne!(debug1, debug2);
    }

    #[tokio::test]
    async fn test_is_initialized_false_on_new() {
        let server = McpServer::new("test", "echo");
        let client = McpClient::new(server);

        // New client should not be initialized
        assert!(!client.is_initialized().await);
    }

    #[tokio::test]
    async fn test_is_initialized_after_failed_init() {
        let server = McpServer::new("ghost", "nonexistent-binary-xyz-123");
        let client = McpClient::new(server);

        // Failed init should leave client not initialized
        let result = client.initialize().await;
        assert!(result.is_err());
        assert!(!client.is_initialized().await);
    }

    #[tokio::test]
    async fn test_shutdown_when_not_running() {
        let server = McpServer::new("test-shutdown", "echo");
        let client = McpClient::new(server);

        // Shutting down without ever starting should succeed gracefully
        let result = client.shutdown().await;
        assert!(result.is_ok());

        // Client should still report as not initialized
        assert!(!client.is_initialized().await);
    }

    #[tokio::test]
    async fn test_shutdown_idempotent() {
        let server = McpServer::new("test-idempotent", "echo");
        let client = McpClient::new(server);

        // First shutdown
        let first = client.shutdown().await;
        assert!(first.is_ok());

        // Second shutdown should also succeed (idempotent)
        let second = client.shutdown().await;
        assert!(second.is_ok());
    }

    #[test]
    fn test_client_server_config_passed_through() {
        let server = McpServer::new("config-test", "npx")
            .with_args(vec!["-y".to_string(), "@some/mcp-server".to_string()])
            .with_env("DEBUG", "true");

        let client = McpClient::new(server);

        assert_eq!(client.server.name, "config-test");
        assert_eq!(client.server.command, "npx");
        assert_eq!(client.server.args, vec!["-y", "@some/mcp-server"]);
        assert_eq!(client.server.env.get("DEBUG"), Some(&"true".to_string()));
    }

    #[test]
    fn test_client_server_method() {
        let server = McpServer::new("method-test", "python");
        let client = McpClient::new(server);

        // server() method should return a reference to the server config
        let retrieved_server = client.server();
        assert_eq!(retrieved_server.name, "method-test");
    }

    #[tokio::test]
    async fn test_server_info_none_on_new_client() {
        let server = McpServer::new("test", "echo");
        let client = McpClient::new(server);

        // Server info should be None until initialized
        assert!(client.server_info().await.is_none());
    }

    #[tokio::test]
    async fn test_initialize_already_initialized_skipped() {
        let server = McpServer::new("echo", "echo");
        let client = McpClient::new(server);

        // First init fails (echo doesn't speak MCP)
        let _ = client.initialize().await;

        // Double init should be a no-op (not panic)
        let result = client.initialize().await;
        // Result may be error from echo (not MCP protocol) but shouldn't panic
        assert!(result.is_err() || result.is_ok());
    }

    #[test]
    fn test_client_default_timeout_is_30_seconds() {
        let server = McpServer::new("test", "echo");
        let client = McpClient::new(server);

        // We can't directly access request_timeout, but we can verify
        // the client is constructable and basic operations work
        assert_eq!(client.server.name, "test");
    }

    #[tokio::test]
    async fn test_shutdown_clears_initialized_flag() {
        let server = McpServer::new("test-clear", "echo");
        let client = McpClient::new(server);

        // Ensure initialized is false
        assert!(!client.is_initialized().await);

        // Shutdown should keep it false
        client.shutdown().await.unwrap();
        assert!(!client.is_initialized().await);
    }
}