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};
12use crate::server::visibility::Environment;
13
14pub struct ServerBuilder {
16 name: String,
17 version: String,
18 capabilities: Option<ServerCapabilities>,
19 middleware: Vec<MiddlewareFn>,
20 instructions: Option<String>,
22 breaker_config: Option<ToolBreakerConfig>,
24 retry_config: Option<ResourceRetryConfig>,
26 rate_limiter: Option<Arc<RateLimiter>>,
28 environment: Option<Arc<dyn Environment>>,
30}
31
32impl ServerBuilder {
33 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 pub fn with_capabilities(mut self, capabilities: ServerCapabilities) -> Self {
50 self.capabilities = Some(capabilities);
51 self
52 }
53
54 pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self {
70 self.instructions = Some(instructions.into());
71 self
72 }
73
74 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 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 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 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 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 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 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 pub fn add_middleware(mut self, middleware: MiddlewareFn) -> Self {
178 self.middleware.push(middleware);
179 self
180 }
181
182 pub fn with_logging_middleware(self) -> Self {
184 self.add_middleware(crate::server::middleware::logging_middleware())
185 }
186
187 pub fn with_validation_middleware(self) -> Self {
189 self.add_middleware(crate::server::middleware::validation_middleware())
190 }
191
192 pub fn with_circuit_breaker(mut self, config: ToolBreakerConfig) -> Self {
212 self.breaker_config = Some(config);
213 self
214 }
215
216 pub fn with_retry(mut self, config: ResourceRetryConfig) -> Self {
236 self.retry_config = Some(config);
237 self
238 }
239
240 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 pub fn with_rate_limiter(mut self, limiter: Arc<RateLimiter>) -> Self {
269 self.rate_limiter = Some(limiter);
270 self
271 }
272
273 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 pub fn build(self) -> Server {
284 let mut server = Server::new(self.name, self.version);
285
286 if let Some(capabilities) = self.capabilities {
288 server.set_capabilities(capabilities);
289 }
290
291 if let Some(instructions) = self.instructions {
293 server.set_instructions(Some(instructions));
294 }
295
296 if let Some(breaker_config) = self.breaker_config {
298 server.tool_registry().set_breaker_config(breaker_config);
299 }
300
301 if let Some(retry_config) = self.retry_config {
303 server.resource_manager().set_retry_config(retry_config);
304 }
305
306 if let Some(limiter) = self.rate_limiter {
308 server.add_middleware(crate::server::middleware::rate_limiter_middleware(limiter));
309 }
310
311 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
324pub 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}