mcp_host/server/
builder.rs

1//! Server builder for fluent server construction
2//!
3//! Provides ergonomic API for building MCP servers with resilience patterns
4
5use std::sync::Arc;
6
7use crate::protocol::capabilities::ServerCapabilities;
8use crate::registry::resources::ResourceRetryConfig;
9use crate::registry::tools::ToolBreakerConfig;
10use crate::server::core::Server;
11use crate::server::middleware::{MiddlewareFn, RateLimiter, RateLimiterConfig};
12
13/// Fluent builder for Server
14pub struct ServerBuilder {
15    name: String,
16    version: String,
17    capabilities: Option<ServerCapabilities>,
18    middleware: Vec<MiddlewareFn>,
19    /// Server instructions for LLMs
20    instructions: Option<String>,
21    /// Circuit breaker configuration for tools
22    breaker_config: Option<ToolBreakerConfig>,
23    /// Retry configuration for resources
24    retry_config: Option<ResourceRetryConfig>,
25    /// Rate limiter configuration
26    rate_limiter: Option<Arc<RateLimiter>>,
27}
28
29impl ServerBuilder {
30    /// Create new server builder
31    pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
32        Self {
33            name: name.into(),
34            version: version.into(),
35            capabilities: None,
36            middleware: Vec::new(),
37            instructions: None,
38            breaker_config: None,
39            retry_config: None,
40            rate_limiter: None,
41        }
42    }
43
44    /// Set server capabilities
45    pub fn with_capabilities(mut self, capabilities: ServerCapabilities) -> Self {
46        self.capabilities = Some(capabilities);
47        self
48    }
49
50    /// Set server instructions for LLMs
51    ///
52    /// Instructions help LLMs understand how to use the server's tools, resources,
53    /// and prompts. They are sent during initialization and can be thought of as
54    /// hints added to the system prompt.
55    ///
56    /// # Example
57    ///
58    /// ```rust
59    /// use mcp_host::prelude::*;
60    ///
61    /// let server = server("my-server", "1.0.0")
62    ///     .with_instructions("Use the data tools for queries. Verify IDs before updates.")
63    ///     .build();
64    /// ```
65    pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self {
66        self.instructions = Some(instructions.into());
67        self
68    }
69
70    /// Enable tools capability
71    pub fn with_tools(mut self, list_changed: bool) -> Self {
72        let mut caps = self.capabilities.take().unwrap_or_default();
73        caps.tools = Some(crate::protocol::capabilities::ToolsCapability {
74            list_changed: Some(list_changed),
75        });
76        self.capabilities = Some(caps);
77        self
78    }
79
80    /// Enable resources capability
81    pub fn with_resources(mut self, list_changed: bool, subscribe: bool) -> Self {
82        let mut caps = self.capabilities.take().unwrap_or_default();
83        caps.resources = Some(crate::protocol::capabilities::ResourcesCapability {
84            list_changed: Some(list_changed),
85            subscribe: Some(subscribe),
86            list_templates: None,
87        });
88        self.capabilities = Some(caps);
89        self
90    }
91
92    /// Enable resource templates capability
93    pub fn with_resource_templates(mut self) -> Self {
94        let mut caps = self.capabilities.take().unwrap_or_default();
95        if let Some(ref mut res) = caps.resources {
96            res.list_templates = Some(true);
97        } else {
98            caps.resources = Some(crate::protocol::capabilities::ResourcesCapability {
99                list_changed: None,
100                subscribe: None,
101                list_templates: Some(true),
102            });
103        }
104        self.capabilities = Some(caps);
105        self
106    }
107
108    /// Enable prompts capability
109    pub fn with_prompts(mut self, list_changed: bool) -> Self {
110        let mut caps = self.capabilities.take().unwrap_or_default();
111        caps.prompts = Some(crate::protocol::capabilities::PromptsCapability {
112            list_changed: Some(list_changed),
113        });
114        self.capabilities = Some(caps);
115        self
116    }
117
118    /// Enable logging capability
119    pub fn with_logging(mut self) -> Self {
120        let mut caps = self.capabilities.take().unwrap_or_default();
121        caps.logging = Some(crate::protocol::capabilities::LoggingCapability {});
122        self.capabilities = Some(caps);
123        self
124    }
125
126    /// Enable tasks capability
127    pub fn with_tasks(mut self, list: bool, cancel: bool) -> Self {
128        let mut caps = self.capabilities.take().unwrap_or_default();
129        caps.tasks = Some(crate::protocol::capabilities::TasksCapability {
130            list: if list {
131                Some(crate::protocol::capabilities::EmptyObject {})
132            } else {
133                None
134            },
135            cancel: if cancel {
136                Some(crate::protocol::capabilities::EmptyObject {})
137            } else {
138                None
139            },
140            requests: Some(crate::protocol::capabilities::TasksRequestsCapability {
141                tools: Some(crate::protocol::capabilities::TasksToolsCapability {
142                    call: Some(crate::protocol::capabilities::EmptyObject {}),
143                }),
144                ..Default::default()
145            }),
146        });
147        self.capabilities = Some(caps);
148        self
149    }
150
151    /// Add middleware to the server
152    pub fn add_middleware(mut self, middleware: MiddlewareFn) -> Self {
153        self.middleware.push(middleware);
154        self
155    }
156
157    /// Add logging middleware
158    pub fn with_logging_middleware(self) -> Self {
159        self.add_middleware(crate::server::middleware::logging_middleware())
160    }
161
162    /// Add validation middleware
163    pub fn with_validation_middleware(self) -> Self {
164        self.add_middleware(crate::server::middleware::validation_middleware())
165    }
166
167    /// Configure circuit breakers for tools
168    ///
169    /// Circuit breakers protect against cascading failures by temporarily
170    /// disabling tools that are failing repeatedly.
171    ///
172    /// # Example
173    ///
174    /// ```rust
175    /// use mcp_host::prelude::*;
176    ///
177    /// let server = server("my-server", "1.0.0")
178    ///     .with_circuit_breaker(ToolBreakerConfig {
179    ///         failure_threshold: 3,
180    ///         failure_window_secs: 30.0,
181    ///         half_open_timeout_secs: 15.0,
182    ///         success_threshold: 1,
183    ///     })
184    ///     .build();
185    /// ```
186    pub fn with_circuit_breaker(mut self, config: ToolBreakerConfig) -> Self {
187        self.breaker_config = Some(config);
188        self
189    }
190
191    /// Configure retry behavior for resources
192    ///
193    /// Resources will retry failed reads using exponential backoff with jitter.
194    ///
195    /// # Example
196    ///
197    /// ```rust
198    /// use mcp_host::prelude::*;
199    ///
200    /// let server = server("my-server", "1.0.0")
201    ///     .with_retry(ResourceRetryConfig {
202    ///         max_attempts: 5,
203    ///         base_delay_ms: 200,
204    ///         multiplier: 2.0,
205    ///         max_delay_ms: 5000,
206    ///         jitter_factor: 1.0,
207    ///     })
208    ///     .build();
209    /// ```
210    pub fn with_retry(mut self, config: ResourceRetryConfig) -> Self {
211        self.retry_config = Some(config);
212        self
213    }
214
215    /// Add rate limiting middleware
216    ///
217    /// Rate limits requests using the GCRA (Generic Cell Rate Algorithm).
218    ///
219    /// # Arguments
220    ///
221    /// * `requests_per_second` - Maximum requests allowed per second
222    /// * `burst_capacity` - Number of requests that can burst instantly
223    ///
224    /// # Example
225    ///
226    /// ```rust
227    /// use mcp_host::prelude::*;
228    ///
229    /// let server = server("my-server", "1.0.0")
230    ///     .with_rate_limit(100.0, 20)  // 100 req/s with burst of 20
231    ///     .build();
232    /// ```
233    pub fn with_rate_limit(mut self, requests_per_second: f64, burst_capacity: usize) -> Self {
234        let limiter = Arc::new(RateLimiter::new(RateLimiterConfig::new(
235            requests_per_second,
236            burst_capacity,
237        )));
238        self.rate_limiter = Some(limiter);
239        self
240    }
241
242    /// Add rate limiting with custom configuration
243    pub fn with_rate_limiter(mut self, limiter: Arc<RateLimiter>) -> Self {
244        self.rate_limiter = Some(limiter);
245        self
246    }
247
248    /// Build the server
249    pub fn build(self) -> Server {
250        let mut server = Server::new(self.name, self.version);
251
252        // Set capabilities if provided
253        if let Some(capabilities) = self.capabilities {
254            server.set_capabilities(capabilities);
255        }
256
257        // Set instructions if provided
258        if let Some(instructions) = self.instructions {
259            server.set_instructions(Some(instructions));
260        }
261
262        // Configure circuit breakers for tools
263        if let Some(breaker_config) = self.breaker_config {
264            server.tool_registry().set_breaker_config(breaker_config);
265        }
266
267        // Configure retry for resources
268        if let Some(retry_config) = self.retry_config {
269            server.resource_manager().set_retry_config(retry_config);
270        }
271
272        // Add rate limiter middleware if configured
273        if let Some(limiter) = self.rate_limiter {
274            server.add_middleware(crate::server::middleware::rate_limiter_middleware(limiter));
275        }
276
277        // Add user-provided middleware
278        for middleware in self.middleware {
279            server.add_middleware(middleware);
280        }
281
282        server
283    }
284}
285
286/// Convenience function to create a server builder
287pub fn server(name: impl Into<String>, version: impl Into<String>) -> ServerBuilder {
288    ServerBuilder::new(name, version)
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294
295    #[test]
296    fn test_builder_basic() {
297        let server = ServerBuilder::new("test-server", "1.0.0").build();
298
299        assert_eq!(server.name(), "test-server");
300        assert_eq!(server.version(), "1.0.0");
301    }
302
303    #[test]
304    fn test_builder_with_tools() {
305        let server = ServerBuilder::new("test-server", "1.0.0")
306            .with_tools(true)
307            .build();
308
309        let caps = server.capabilities();
310        assert!(caps.tools.is_some());
311        assert_eq!(caps.tools.as_ref().unwrap().list_changed, Some(true));
312    }
313
314    #[test]
315    fn test_builder_with_resources() {
316        let server = ServerBuilder::new("test-server", "1.0.0")
317            .with_resources(true, false)
318            .build();
319
320        let caps = server.capabilities();
321        assert!(caps.resources.is_some());
322        let res = caps.resources.as_ref().unwrap();
323        assert_eq!(res.list_changed, Some(true));
324        assert_eq!(res.subscribe, Some(false));
325        assert_eq!(res.list_templates, None);
326    }
327
328    #[test]
329    fn test_builder_with_resource_templates() {
330        let server = ServerBuilder::new("test-server", "1.0.0")
331            .with_resources(true, false)
332            .with_resource_templates()
333            .build();
334
335        let caps = server.capabilities();
336        assert!(caps.resources.is_some());
337        let res = caps.resources.as_ref().unwrap();
338        assert_eq!(res.list_changed, Some(true));
339        assert_eq!(res.subscribe, Some(false));
340        assert_eq!(res.list_templates, Some(true));
341    }
342
343    #[test]
344    fn test_builder_with_prompts() {
345        let server = ServerBuilder::new("test-server", "1.0.0")
346            .with_prompts(true)
347            .build();
348
349        let caps = server.capabilities();
350        assert!(caps.prompts.is_some());
351        assert_eq!(caps.prompts.as_ref().unwrap().list_changed, Some(true));
352    }
353
354    #[test]
355    fn test_builder_with_logging() {
356        let server = ServerBuilder::new("test-server", "1.0.0")
357            .with_logging()
358            .build();
359
360        let caps = server.capabilities();
361        assert!(caps.logging.is_some());
362    }
363
364    #[test]
365    fn test_builder_with_all_capabilities() {
366        let server = ServerBuilder::new("test-server", "1.0.0")
367            .with_tools(true)
368            .with_resources(true, true)
369            .with_prompts(true)
370            .with_logging()
371            .build();
372
373        let caps = server.capabilities();
374        assert!(caps.tools.is_some());
375        assert!(caps.resources.is_some());
376        assert!(caps.prompts.is_some());
377        assert!(caps.logging.is_some());
378    }
379
380    #[test]
381    fn test_builder_with_middleware() {
382        let server = ServerBuilder::new("test-server", "1.0.0")
383            .with_logging_middleware()
384            .with_validation_middleware()
385            .build();
386
387        assert_eq!(server.name(), "test-server");
388    }
389
390    #[test]
391    fn test_server_convenience_function() {
392        let server = server("test-server", "1.0.0").with_tools(true).build();
393
394        assert_eq!(server.name(), "test-server");
395    }
396
397    #[test]
398    fn test_builder_with_instructions() {
399        let server = ServerBuilder::new("test-server", "1.0.0")
400            .with_instructions("Use data tools for queries. Verify IDs before updates.")
401            .build();
402
403        assert_eq!(server.name(), "test-server");
404        assert_eq!(
405            server.instructions(),
406            Some("Use data tools for queries. Verify IDs before updates.".to_string())
407        );
408    }
409
410    #[test]
411    fn test_builder_without_instructions() {
412        let server = ServerBuilder::new("test-server", "1.0.0").build();
413
414        assert!(server.instructions().is_none());
415    }
416}