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