Skip to main content

fastmcp_client/
builder.rs

1//! Client builder for configuring MCP clients.
2//!
3//! The builder provides a fluent API for constructing MCP clients with
4//! customizable timeout, retry, and subprocess spawn options.
5//!
6//! # Example
7//!
8//! ```ignore
9//! use fastmcp::ClientBuilder;
10//!
11//! let client = ClientBuilder::new()
12//!     .client_info("my-client", "1.0.0")
13//!     .timeout_ms(60_000)
14//!     .max_retries(3)
15//!     .retry_delay_ms(1000)
16//!     .working_dir("/tmp")
17//!     .env("DEBUG", "1")
18//!     .connect_stdio("uvx", &["my-server"])?;
19//! ```
20
21use std::collections::HashMap;
22use std::path::PathBuf;
23use std::process::{Child, Command, Stdio};
24
25/// Guard that kills and waits for a child process when dropped.
26/// Call `disarm()` to prevent cleanup (e.g., when ownership transfers to Client).
27struct ChildGuard(Option<Child>);
28
29impl ChildGuard {
30    fn new(child: Child) -> Self {
31        Self(Some(child))
32    }
33
34    /// Takes ownership of the child, preventing cleanup on drop.
35    fn disarm(mut self) -> Child {
36        self.0.take().expect("ChildGuard already disarmed")
37    }
38}
39
40impl Drop for ChildGuard {
41    fn drop(&mut self) {
42        if let Some(mut child) = self.0.take() {
43            // Best effort cleanup - ignore errors
44            let _ = child.kill();
45            let _ = child.wait();
46        }
47    }
48}
49
50use asupersync::Cx;
51use fastmcp_core::{McpError, McpResult};
52use fastmcp_protocol::{
53    ClientCapabilities, ClientInfo, InitializeParams, InitializeResult, JsonRpcMessage,
54    JsonRpcRequest, PROTOCOL_VERSION,
55};
56use fastmcp_transport::{StdioTransport, Transport};
57
58use crate::{Client, ClientSession};
59
60/// Builder for configuring an MCP client.
61///
62/// Use this to configure timeout, retry, and spawn options before
63/// connecting to an MCP server.
64#[derive(Debug, Clone)]
65pub struct ClientBuilder {
66    /// Client identification info.
67    client_info: ClientInfo,
68    /// Request timeout in milliseconds.
69    timeout_ms: u64,
70    /// Maximum number of connection retries.
71    max_retries: u32,
72    /// Delay between retries in milliseconds.
73    retry_delay_ms: u64,
74    /// Working directory for subprocess.
75    working_dir: Option<PathBuf>,
76    /// Environment variables to set for subprocess.
77    env_vars: HashMap<String, String>,
78    /// Whether to inherit parent's environment.
79    inherit_env: bool,
80    /// Client capabilities to advertise.
81    capabilities: ClientCapabilities,
82    /// Whether to defer initialization until first use.
83    auto_initialize: bool,
84}
85
86impl ClientBuilder {
87    /// Creates a new client builder with default settings.
88    ///
89    /// Default configuration:
90    /// - Client name: "fastmcp-client"
91    /// - Timeout: 30 seconds
92    /// - Max retries: 0 (no retries)
93    /// - Retry delay: 1 second
94    /// - Inherit environment: true
95    /// - Auto-initialize: false (initialize immediately on connect)
96    #[must_use]
97    pub fn new() -> Self {
98        Self {
99            client_info: ClientInfo {
100                name: "fastmcp-client".to_owned(),
101                version: env!("CARGO_PKG_VERSION").to_owned(),
102            },
103            timeout_ms: 30_000,
104            max_retries: 0,
105            retry_delay_ms: 1_000,
106            working_dir: None,
107            env_vars: HashMap::new(),
108            inherit_env: true,
109            capabilities: ClientCapabilities::default(),
110            auto_initialize: false,
111        }
112    }
113
114    /// Sets the client name and version.
115    ///
116    /// This information is sent to the server during initialization.
117    #[must_use]
118    pub fn client_info(mut self, name: impl Into<String>, version: impl Into<String>) -> Self {
119        self.client_info = ClientInfo {
120            name: name.into(),
121            version: version.into(),
122        };
123        self
124    }
125
126    /// Sets the request timeout in milliseconds.
127    ///
128    /// This affects how long the client waits for responses from the server.
129    /// Default is 30,000ms (30 seconds).
130    #[must_use]
131    pub fn timeout_ms(mut self, timeout: u64) -> Self {
132        self.timeout_ms = timeout;
133        self
134    }
135
136    /// Sets the maximum number of connection retries.
137    ///
138    /// When connecting to a server fails, the client will retry up to
139    /// this many times before returning an error. Default is 0 (no retries).
140    #[must_use]
141    pub fn max_retries(mut self, retries: u32) -> Self {
142        self.max_retries = retries;
143        self
144    }
145
146    /// Sets the delay between connection retries in milliseconds.
147    ///
148    /// Default is 1,000ms (1 second).
149    #[must_use]
150    pub fn retry_delay_ms(mut self, delay: u64) -> Self {
151        self.retry_delay_ms = delay;
152        self
153    }
154
155    /// Sets the working directory for the subprocess.
156    ///
157    /// If not set, the subprocess inherits the current working directory.
158    #[must_use]
159    pub fn working_dir(mut self, path: impl Into<PathBuf>) -> Self {
160        self.working_dir = Some(path.into());
161        self
162    }
163
164    /// Adds an environment variable for the subprocess.
165    ///
166    /// Multiple calls to this method accumulate environment variables.
167    #[must_use]
168    pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
169        self.env_vars.insert(key.into(), value.into());
170        self
171    }
172
173    /// Adds multiple environment variables for the subprocess.
174    #[must_use]
175    pub fn envs<I, K, V>(mut self, vars: I) -> Self
176    where
177        I: IntoIterator<Item = (K, V)>,
178        K: Into<String>,
179        V: Into<String>,
180    {
181        for (key, value) in vars {
182            self.env_vars.insert(key.into(), value.into());
183        }
184        self
185    }
186
187    /// Sets whether to inherit the parent process's environment.
188    ///
189    /// If true (default), the subprocess starts with the parent's environment
190    /// plus any variables added via [`env`](Self::env) or [`envs`](Self::envs).
191    ///
192    /// If false, the subprocess starts with only the explicitly set variables.
193    #[must_use]
194    pub fn inherit_env(mut self, inherit: bool) -> Self {
195        self.inherit_env = inherit;
196        self
197    }
198
199    /// Sets the client capabilities to advertise to the server.
200    #[must_use]
201    pub fn capabilities(mut self, capabilities: ClientCapabilities) -> Self {
202        self.capabilities = capabilities;
203        self
204    }
205
206    /// Enables auto-initialization mode.
207    ///
208    /// When enabled, the client defers the MCP initialization handshake until
209    /// the first method call (e.g., `list_tools`, `call_tool`). This allows
210    /// the subprocess to start immediately without blocking on initialization.
211    ///
212    /// Default is `false` (initialize immediately on connect).
213    ///
214    /// # Example
215    ///
216    /// ```ignore
217    /// let client = ClientBuilder::new()
218    ///     .auto_initialize(true)
219    ///     .connect_stdio("uvx", &["my-server"])?;
220    ///
221    /// // Subprocess is running but not yet initialized
222    /// // Initialization happens on first use:
223    /// let tools = client.list_tools()?; // Initializes here
224    /// ```
225    #[must_use]
226    pub fn auto_initialize(mut self, enabled: bool) -> Self {
227        self.auto_initialize = enabled;
228        self
229    }
230
231    /// Connects to a server via stdio subprocess.
232    ///
233    /// Spawns the specified command as a subprocess and communicates via
234    /// stdin/stdout using JSON-RPC over NDJSON framing.
235    ///
236    /// # Arguments
237    ///
238    /// * `command` - The command to run (e.g., "uvx", "npx")
239    /// * `args` - Arguments to pass to the command
240    ///
241    /// # Errors
242    ///
243    /// Returns an error if:
244    /// - The subprocess fails to spawn
245    /// - The initialization handshake fails
246    /// - All retry attempts are exhausted
247    pub fn connect_stdio(self, command: &str, args: &[&str]) -> McpResult<Client> {
248        self.connect_stdio_with_cx(command, args, &Cx::for_testing())
249    }
250
251    /// Connects to a server via stdio subprocess with a provided Cx.
252    ///
253    /// Same as [`connect_stdio`](Self::connect_stdio) but allows providing
254    /// a custom capability context for cancellation support.
255    pub fn connect_stdio_with_cx(self, command: &str, args: &[&str], cx: &Cx) -> McpResult<Client> {
256        let mut last_error = None;
257        let attempts = self.max_retries + 1;
258
259        for attempt in 0..attempts {
260            if attempt > 0 {
261                // Delay before retry
262                std::thread::sleep(std::time::Duration::from_millis(self.retry_delay_ms));
263            }
264
265            match self.try_connect(command, args, cx) {
266                Ok(client) => return Ok(client),
267                Err(e) => {
268                    last_error = Some(e);
269                }
270            }
271        }
272
273        // All attempts failed
274        Err(last_error.unwrap_or_else(|| McpError::internal_error("Connection failed")))
275    }
276
277    /// Attempts a single connection.
278    fn try_connect(&self, command: &str, args: &[&str], cx: &Cx) -> McpResult<Client> {
279        // Build the command
280        let mut cmd = Command::new(command);
281        cmd.args(args)
282            .stdin(Stdio::piped())
283            .stdout(Stdio::piped())
284            .stderr(Stdio::inherit());
285
286        // Set working directory if specified
287        if let Some(ref dir) = self.working_dir {
288            cmd.current_dir(dir);
289        }
290
291        // Set environment
292        if !self.inherit_env {
293            cmd.env_clear();
294        }
295        for (key, value) in &self.env_vars {
296            cmd.env(key, value);
297        }
298
299        // Spawn the subprocess
300        let mut child = cmd
301            .spawn()
302            .map_err(|e| McpError::internal_error(format!("Failed to spawn subprocess: {e}")))?;
303
304        // Get stdin/stdout handles
305        let stdin = child
306            .stdin
307            .take()
308            .ok_or_else(|| McpError::internal_error("Failed to get subprocess stdin"))?;
309        let stdout = child
310            .stdout
311            .take()
312            .ok_or_else(|| McpError::internal_error("Failed to get subprocess stdout"))?;
313
314        // Create transport
315        let transport = StdioTransport::new(stdout, stdin);
316
317        if self.auto_initialize {
318            // Create uninitialized client - initialization will happen on first use
319            Ok(self.create_uninitialized_client(child, transport, cx))
320        } else {
321            // Perform initialization immediately
322            self.initialize_client(child, transport, cx)
323        }
324    }
325
326    /// Creates an uninitialized client for auto-initialize mode.
327    fn create_uninitialized_client(
328        &self,
329        child: Child,
330        transport: StdioTransport<std::process::ChildStdout, std::process::ChildStdin>,
331        cx: &Cx,
332    ) -> Client {
333        // Create a placeholder session - will be updated on first use
334        let session = ClientSession::new(
335            self.client_info.clone(),
336            self.capabilities.clone(),
337            fastmcp_protocol::ServerInfo {
338                name: String::new(),
339                version: String::new(),
340            },
341            fastmcp_protocol::ServerCapabilities::default(),
342            String::new(),
343        );
344
345        Client::from_parts_uninitialized(child, transport, cx.clone(), session, self.timeout_ms)
346    }
347
348    /// Performs the initialization handshake and creates the client.
349    fn initialize_client(
350        &self,
351        child: Child,
352        mut transport: StdioTransport<std::process::ChildStdout, std::process::ChildStdin>,
353        cx: &Cx,
354    ) -> McpResult<Client> {
355        // Guard ensures child process is killed if initialization fails.
356        // Disarmed when client is successfully created.
357        let child_guard = ChildGuard::new(child);
358
359        // Send initialize request
360        let init_params = InitializeParams {
361            protocol_version: PROTOCOL_VERSION.to_string(),
362            capabilities: self.capabilities.clone(),
363            client_info: self.client_info.clone(),
364        };
365
366        let init_request = JsonRpcRequest::new(
367            "initialize",
368            Some(serde_json::to_value(&init_params).map_err(|e| {
369                McpError::internal_error(format!("Failed to serialize params: {e}"))
370            })?),
371            1i64,
372        );
373
374        transport
375            .send(cx, &JsonRpcMessage::Request(init_request))
376            .map_err(|e| McpError::internal_error(format!("Failed to send initialize: {e}")))?;
377
378        // Receive initialize response
379        let response = loop {
380            let msg = transport.recv(cx).map_err(|e| {
381                McpError::internal_error(format!("Failed to receive response: {e}"))
382            })?;
383
384            match msg {
385                JsonRpcMessage::Response(resp) => break resp,
386                JsonRpcMessage::Request(_) => {
387                    // Ignore server requests during initialization
388                }
389            }
390        };
391
392        // Check for error
393        if let Some(error) = response.error {
394            return Err(McpError::new(
395                fastmcp_core::McpErrorCode::Custom(error.code),
396                error.message,
397            ));
398        }
399
400        // Parse result
401        let result_value = response
402            .result
403            .ok_or_else(|| McpError::internal_error("No result in initialize response"))?;
404
405        let init_result: InitializeResult = serde_json::from_value(result_value).map_err(|e| {
406            McpError::internal_error(format!("Failed to parse initialize result: {e}"))
407        })?;
408
409        // Send initialized notification
410        let initialized_request = JsonRpcRequest {
411            jsonrpc: std::borrow::Cow::Borrowed(fastmcp_protocol::JSONRPC_VERSION),
412            method: "initialized".to_string(),
413            params: Some(serde_json::json!({})),
414            id: None,
415        };
416
417        transport
418            .send(cx, &JsonRpcMessage::Request(initialized_request))
419            .map_err(|e| McpError::internal_error(format!("Failed to send initialized: {e}")))?;
420
421        // Create session
422        let session = ClientSession::new(
423            self.client_info.clone(),
424            self.capabilities.clone(),
425            init_result.server_info,
426            init_result.capabilities,
427            init_result.protocol_version,
428        );
429
430        // Create client - disarm guard since Client now owns the subprocess
431        Ok(Client::from_parts(
432            child_guard.disarm(),
433            transport,
434            cx.clone(),
435            session,
436            self.timeout_ms,
437        ))
438    }
439}
440
441impl Default for ClientBuilder {
442    fn default() -> Self {
443        Self::new()
444    }
445}
446
447#[cfg(test)]
448mod tests {
449    use super::*;
450
451    #[test]
452    fn test_builder_defaults() {
453        let builder = ClientBuilder::new();
454        assert_eq!(builder.client_info.name, "fastmcp-client");
455        assert_eq!(builder.timeout_ms, 30_000);
456        assert_eq!(builder.max_retries, 0);
457        assert_eq!(builder.retry_delay_ms, 1_000);
458        assert!(builder.inherit_env);
459        assert!(builder.working_dir.is_none());
460        assert!(builder.env_vars.is_empty());
461        assert!(!builder.auto_initialize);
462    }
463
464    #[test]
465    fn test_builder_fluent_api() {
466        let builder = ClientBuilder::new()
467            .client_info("test-client", "2.0.0")
468            .timeout_ms(60_000)
469            .max_retries(3)
470            .retry_delay_ms(500)
471            .working_dir("/tmp")
472            .env("FOO", "bar")
473            .env("BAZ", "qux")
474            .inherit_env(false);
475
476        assert_eq!(builder.client_info.name, "test-client");
477        assert_eq!(builder.client_info.version, "2.0.0");
478        assert_eq!(builder.timeout_ms, 60_000);
479        assert_eq!(builder.max_retries, 3);
480        assert_eq!(builder.retry_delay_ms, 500);
481        assert_eq!(builder.working_dir, Some(PathBuf::from("/tmp")));
482        assert_eq!(builder.env_vars.get("FOO"), Some(&"bar".to_string()));
483        assert_eq!(builder.env_vars.get("BAZ"), Some(&"qux".to_string()));
484        assert!(!builder.inherit_env);
485    }
486
487    #[test]
488    fn test_builder_envs() {
489        let vars = [("KEY1", "value1"), ("KEY2", "value2")];
490        let builder = ClientBuilder::new().envs(vars);
491
492        assert_eq!(builder.env_vars.get("KEY1"), Some(&"value1".to_string()));
493        assert_eq!(builder.env_vars.get("KEY2"), Some(&"value2".to_string()));
494    }
495
496    #[test]
497    fn test_builder_clone() {
498        let builder1 = ClientBuilder::new()
499            .client_info("test", "1.0")
500            .timeout_ms(5000);
501
502        let builder2 = builder1.clone();
503
504        assert_eq!(builder2.client_info.name, "test");
505        assert_eq!(builder2.timeout_ms, 5000);
506    }
507
508    #[test]
509    fn test_builder_auto_initialize() {
510        let builder = ClientBuilder::new().auto_initialize(true);
511        assert!(builder.auto_initialize);
512
513        let builder = ClientBuilder::new().auto_initialize(false);
514        assert!(!builder.auto_initialize);
515    }
516
517    #[test]
518    fn test_builder_capabilities() {
519        let caps = ClientCapabilities {
520            sampling: Some(fastmcp_protocol::SamplingCapability {}),
521            elicitation: None,
522            roots: None,
523        };
524        let builder = ClientBuilder::new().capabilities(caps);
525        assert!(builder.capabilities.sampling.is_some());
526        assert!(builder.capabilities.elicitation.is_none());
527        assert!(builder.capabilities.roots.is_none());
528    }
529
530    #[test]
531    fn test_builder_default_trait() {
532        let builder = ClientBuilder::default();
533        assert_eq!(builder.client_info.name, "fastmcp-client");
534        assert_eq!(builder.timeout_ms, 30_000);
535        assert_eq!(builder.max_retries, 0);
536        assert!(!builder.auto_initialize);
537    }
538
539    #[test]
540    fn test_builder_env_override() {
541        let builder = ClientBuilder::new()
542            .env("KEY", "first")
543            .env("KEY", "second");
544        assert_eq!(builder.env_vars.get("KEY"), Some(&"second".to_string()));
545    }
546
547    #[test]
548    fn test_builder_envs_combined_with_env() {
549        let builder = ClientBuilder::new()
550            .env("A", "1")
551            .envs([("B", "2"), ("C", "3")])
552            .env("D", "4");
553        assert_eq!(builder.env_vars.len(), 4);
554        assert_eq!(builder.env_vars.get("A"), Some(&"1".to_string()));
555        assert_eq!(builder.env_vars.get("B"), Some(&"2".to_string()));
556        assert_eq!(builder.env_vars.get("C"), Some(&"3".to_string()));
557        assert_eq!(builder.env_vars.get("D"), Some(&"4".to_string()));
558    }
559}