mcp_host/server/
builder.rs1use 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
13pub struct ServerBuilder {
15 name: String,
16 version: String,
17 capabilities: Option<ServerCapabilities>,
18 middleware: Vec<MiddlewareFn>,
19 instructions: Option<String>,
21 breaker_config: Option<ToolBreakerConfig>,
23 retry_config: Option<ResourceRetryConfig>,
25 rate_limiter: Option<Arc<RateLimiter>>,
27}
28
29impl ServerBuilder {
30 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 pub fn with_capabilities(mut self, capabilities: ServerCapabilities) -> Self {
46 self.capabilities = Some(capabilities);
47 self
48 }
49
50 pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self {
66 self.instructions = Some(instructions.into());
67 self
68 }
69
70 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 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 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 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 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 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 pub fn add_middleware(mut self, middleware: MiddlewareFn) -> Self {
153 self.middleware.push(middleware);
154 self
155 }
156
157 pub fn with_logging_middleware(self) -> Self {
159 self.add_middleware(crate::server::middleware::logging_middleware())
160 }
161
162 pub fn with_validation_middleware(self) -> Self {
164 self.add_middleware(crate::server::middleware::validation_middleware())
165 }
166
167 pub fn with_circuit_breaker(mut self, config: ToolBreakerConfig) -> Self {
187 self.breaker_config = Some(config);
188 self
189 }
190
191 pub fn with_retry(mut self, config: ResourceRetryConfig) -> Self {
211 self.retry_config = Some(config);
212 self
213 }
214
215 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 pub fn with_rate_limiter(mut self, limiter: Arc<RateLimiter>) -> Self {
244 self.rate_limiter = Some(limiter);
245 self
246 }
247
248 pub fn build(self) -> Server {
250 let mut server = Server::new(self.name, self.version);
251
252 if let Some(capabilities) = self.capabilities {
254 server.set_capabilities(capabilities);
255 }
256
257 if let Some(instructions) = self.instructions {
259 server.set_instructions(Some(instructions));
260 }
261
262 if let Some(breaker_config) = self.breaker_config {
264 server.tool_registry().set_breaker_config(breaker_config);
265 }
266
267 if let Some(retry_config) = self.retry_config {
269 server.resource_manager().set_retry_config(retry_config);
270 }
271
272 if let Some(limiter) = self.rate_limiter {
274 server.add_middleware(crate::server::middleware::rate_limiter_middleware(limiter));
275 }
276
277 for middleware in self.middleware {
279 server.add_middleware(middleware);
280 }
281
282 server
283 }
284}
285
286pub 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}