turul_mcp_aws_lambda/
builder.rs

1//! High-level builder API for Lambda MCP servers
2//!
3//! This module provides a fluent builder API similar to McpServer::builder()
4//! but specifically designed for AWS Lambda deployment.
5
6use std::collections::HashMap;
7use std::sync::Arc;
8
9use turul_http_mcp_server::{ServerConfig, StreamConfig};
10use turul_mcp_protocol::{Implementation, ServerCapabilities};
11use turul_mcp_server::handlers::{McpHandler, *};
12use turul_mcp_server::{
13    McpCompletion, McpElicitation, McpLogger, McpNotification, McpPrompt, McpResource, McpRoot,
14    McpSampling, McpTool,
15};
16use turul_mcp_session_storage::BoxedSessionStorage;
17
18use crate::error::Result;
19
20#[cfg(feature = "dynamodb")]
21use crate::error::LambdaError;
22use crate::server::LambdaMcpServer;
23
24#[cfg(feature = "cors")]
25use crate::cors::CorsConfig;
26
27/// High-level builder for Lambda MCP servers
28///
29/// This provides a clean, fluent API for building Lambda MCP servers
30/// similar to the framework's McpServer::builder() pattern.
31///
32/// ## Example
33///
34/// ```rust,no_run
35/// use std::sync::Arc;
36/// use turul_mcp_aws_lambda::LambdaMcpServerBuilder;
37/// use turul_mcp_session_storage::InMemorySessionStorage;
38/// use turul_mcp_derive::McpTool;
39/// use turul_mcp_server::{McpResult, SessionContext};
40///
41/// #[derive(McpTool, Clone, Default)]
42/// #[tool(name = "example", description = "Example tool")]
43/// struct ExampleTool {
44///     #[param(description = "Example parameter")]
45///     value: String,
46/// }
47///
48/// impl ExampleTool {
49///     async fn execute(&self, _session: Option<SessionContext>) -> McpResult<String> {
50///         Ok(format!("Got: {}", self.value))
51///     }
52/// }
53///
54/// #[tokio::main]
55/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
56///     let server = LambdaMcpServerBuilder::new()
57///         .name("my-lambda-server")
58///         .version("1.0.0")
59///         .tool(ExampleTool::default())
60///         .storage(Arc::new(InMemorySessionStorage::new()))
61///         .cors_allow_all_origins()
62///         .build()
63///         .await?;
64///
65///     // Use with Lambda runtime...
66///     Ok(())
67/// }
68/// ```
69pub struct LambdaMcpServerBuilder {
70    /// Server implementation info
71    name: String,
72    version: String,
73    title: Option<String>,
74
75    /// Server capabilities
76    capabilities: ServerCapabilities,
77
78    /// Tools registered with the server
79    tools: HashMap<String, Arc<dyn McpTool>>,
80
81    /// Resources registered with the server
82    resources: HashMap<String, Arc<dyn McpResource>>,
83
84    /// Prompts registered with the server
85    prompts: HashMap<String, Arc<dyn McpPrompt>>,
86
87    /// Elicitations registered with the server
88    elicitations: HashMap<String, Arc<dyn McpElicitation>>,
89
90    /// Sampling providers registered with the server
91    sampling: HashMap<String, Arc<dyn McpSampling>>,
92
93    /// Completion providers registered with the server
94    completions: HashMap<String, Arc<dyn McpCompletion>>,
95
96    /// Loggers registered with the server
97    loggers: HashMap<String, Arc<dyn McpLogger>>,
98
99    /// Root providers registered with the server
100    root_providers: HashMap<String, Arc<dyn McpRoot>>,
101
102    /// Notification providers registered with the server
103    notifications: HashMap<String, Arc<dyn McpNotification>>,
104
105    /// Handlers registered with the server
106    handlers: HashMap<String, Arc<dyn McpHandler>>,
107
108    /// Roots configured for the server
109    roots: Vec<turul_mcp_protocol::roots::Root>,
110
111    /// Optional instructions for clients
112    instructions: Option<String>,
113
114    /// Session configuration
115    session_timeout_minutes: Option<u64>,
116    session_cleanup_interval_seconds: Option<u64>,
117
118    /// Session storage backend (defaults to InMemory if None)
119    session_storage: Option<Arc<BoxedSessionStorage>>,
120
121    /// MCP Lifecycle enforcement configuration
122    strict_lifecycle: bool,
123
124    /// Enable SSE streaming
125    enable_sse: bool,
126    /// Server and stream configuration
127    server_config: ServerConfig,
128    stream_config: StreamConfig,
129
130    /// Middleware stack for request/response interception
131    middleware_stack: turul_http_mcp_server::middleware::MiddlewareStack,
132
133    /// CORS configuration (if enabled)
134    #[cfg(feature = "cors")]
135    cors_config: Option<CorsConfig>,
136}
137
138impl LambdaMcpServerBuilder {
139    /// Create a new Lambda MCP server builder
140    pub fn new() -> Self {
141        // Initialize with default capabilities (same as McpServer)
142        // Capabilities will be set truthfully in build() based on registered components
143        let capabilities = ServerCapabilities::default();
144
145        // Initialize handlers with defaults (same as McpServerBuilder)
146        let mut handlers: HashMap<String, Arc<dyn McpHandler>> = HashMap::new();
147        handlers.insert("ping".to_string(), Arc::new(PingHandler));
148        handlers.insert(
149            "completion/complete".to_string(),
150            Arc::new(CompletionHandler),
151        );
152        handlers.insert(
153            "resources/list".to_string(),
154            Arc::new(ResourcesHandler::new()),
155        );
156        handlers.insert(
157            "prompts/list".to_string(),
158            Arc::new(PromptsListHandler::new()),
159        );
160        handlers.insert(
161            "prompts/get".to_string(),
162            Arc::new(PromptsGetHandler::new()),
163        );
164        handlers.insert("logging/setLevel".to_string(), Arc::new(LoggingHandler));
165        handlers.insert("roots/list".to_string(), Arc::new(RootsHandler::new()));
166        handlers.insert(
167            "sampling/createMessage".to_string(),
168            Arc::new(SamplingHandler),
169        );
170        handlers.insert(
171            "resources/templates/list".to_string(),
172            Arc::new(ResourceTemplatesHandler::new()),
173        );
174        handlers.insert(
175            "elicitation/create".to_string(),
176            Arc::new(ElicitationHandler::with_mock_provider()),
177        );
178
179        // Add notification handlers
180        let notifications_handler = Arc::new(NotificationsHandler);
181        handlers.insert(
182            "notifications/message".to_string(),
183            notifications_handler.clone(),
184        );
185        handlers.insert(
186            "notifications/progress".to_string(),
187            notifications_handler.clone(),
188        );
189        handlers.insert(
190            "notifications/resources/listChanged".to_string(),
191            notifications_handler.clone(),
192        );
193        handlers.insert(
194            "notifications/resources/updated".to_string(),
195            notifications_handler.clone(),
196        );
197        handlers.insert(
198            "notifications/tools/listChanged".to_string(),
199            notifications_handler.clone(),
200        );
201        handlers.insert(
202            "notifications/prompts/listChanged".to_string(),
203            notifications_handler.clone(),
204        );
205        handlers.insert(
206            "notifications/roots/listChanged".to_string(),
207            notifications_handler,
208        );
209
210        Self {
211            name: "turul-mcp-aws-lambda".to_string(),
212            version: env!("CARGO_PKG_VERSION").to_string(),
213            title: None,
214            capabilities,
215            tools: HashMap::new(),
216            resources: HashMap::new(),
217            prompts: HashMap::new(),
218            elicitations: HashMap::new(),
219            sampling: HashMap::new(),
220            completions: HashMap::new(),
221            loggers: HashMap::new(),
222            root_providers: HashMap::new(),
223            notifications: HashMap::new(),
224            handlers,
225            roots: Vec::new(),
226            instructions: None,
227            session_timeout_minutes: None,
228            session_cleanup_interval_seconds: None,
229            session_storage: None,
230            strict_lifecycle: false,
231            enable_sse: cfg!(feature = "sse"),
232            server_config: ServerConfig::default(),
233            stream_config: StreamConfig::default(),
234            middleware_stack: turul_http_mcp_server::middleware::MiddlewareStack::new(),
235            #[cfg(feature = "cors")]
236            cors_config: None,
237        }
238    }
239
240    /// Set the server name
241    pub fn name(mut self, name: impl Into<String>) -> Self {
242        self.name = name.into();
243        self
244    }
245
246    /// Set the server version
247    pub fn version(mut self, version: impl Into<String>) -> Self {
248        self.version = version.into();
249        self
250    }
251
252    /// Set the server title
253    pub fn title(mut self, title: impl Into<String>) -> Self {
254        self.title = Some(title.into());
255        self
256    }
257
258    /// Set optional instructions for clients
259    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
260        self.instructions = Some(instructions.into());
261        self
262    }
263
264    // =============================================================================
265    // PROVIDER REGISTRATION METHODS (same as McpServerBuilder)
266    // =============================================================================
267
268    /// Register a tool with the server
269    ///
270    /// Tools can be created using any of the framework's 4 creation levels:
271    /// - Function macros: `#[mcp_tool]`
272    /// - Derive macros: `#[derive(McpTool)]`
273    /// - Builder pattern: `ToolBuilder::new(...).build()`
274    /// - Manual implementation: Custom struct implementing `McpTool`
275    pub fn tool<T: McpTool + 'static>(mut self, tool: T) -> Self {
276        let name = tool.name().to_string();
277        self.tools.insert(name, Arc::new(tool));
278        self
279    }
280
281    /// Register a function tool created with `#[mcp_tool]` macro
282    pub fn tool_fn<F, T>(self, func: F) -> Self
283    where
284        F: Fn() -> T,
285        T: McpTool + 'static,
286    {
287        self.tool(func())
288    }
289
290    /// Register multiple tools
291    pub fn tools<T: McpTool + 'static, I: IntoIterator<Item = T>>(mut self, tools: I) -> Self {
292        for tool in tools {
293            self = self.tool(tool);
294        }
295        self
296    }
297
298    /// Register a resource with the server
299    pub fn resource<R: McpResource + 'static>(mut self, resource: R) -> Self {
300        let uri = resource.uri().to_string();
301        self.resources.insert(uri, Arc::new(resource));
302        self
303    }
304
305    /// Register multiple resources
306    pub fn resources<R: McpResource + 'static, I: IntoIterator<Item = R>>(
307        mut self,
308        resources: I,
309    ) -> Self {
310        for resource in resources {
311            self = self.resource(resource);
312        }
313        self
314    }
315
316    /// Register a prompt with the server
317    pub fn prompt<P: McpPrompt + 'static>(mut self, prompt: P) -> Self {
318        let name = prompt.name().to_string();
319        self.prompts.insert(name, Arc::new(prompt));
320        self
321    }
322
323    /// Register multiple prompts
324    pub fn prompts<P: McpPrompt + 'static, I: IntoIterator<Item = P>>(
325        mut self,
326        prompts: I,
327    ) -> Self {
328        for prompt in prompts {
329            self = self.prompt(prompt);
330        }
331        self
332    }
333
334    /// Register an elicitation provider with the server
335    pub fn elicitation<E: McpElicitation + 'static>(mut self, elicitation: E) -> Self {
336        let key = format!("elicitation_{}", self.elicitations.len());
337        self.elicitations.insert(key, Arc::new(elicitation));
338        self
339    }
340
341    /// Register multiple elicitation providers
342    pub fn elicitations<E: McpElicitation + 'static, I: IntoIterator<Item = E>>(
343        mut self,
344        elicitations: I,
345    ) -> Self {
346        for elicitation in elicitations {
347            self = self.elicitation(elicitation);
348        }
349        self
350    }
351
352    /// Register a sampling provider with the server
353    pub fn sampling_provider<S: McpSampling + 'static>(mut self, sampling: S) -> Self {
354        let key = format!("sampling_{}", self.sampling.len());
355        self.sampling.insert(key, Arc::new(sampling));
356        self
357    }
358
359    /// Register multiple sampling providers
360    pub fn sampling_providers<S: McpSampling + 'static, I: IntoIterator<Item = S>>(
361        mut self,
362        sampling: I,
363    ) -> Self {
364        for s in sampling {
365            self = self.sampling_provider(s);
366        }
367        self
368    }
369
370    /// Register a completion provider with the server
371    pub fn completion_provider<C: McpCompletion + 'static>(mut self, completion: C) -> Self {
372        let key = format!("completion_{}", self.completions.len());
373        self.completions.insert(key, Arc::new(completion));
374        self
375    }
376
377    /// Register multiple completion providers
378    pub fn completion_providers<C: McpCompletion + 'static, I: IntoIterator<Item = C>>(
379        mut self,
380        completions: I,
381    ) -> Self {
382        for completion in completions {
383            self = self.completion_provider(completion);
384        }
385        self
386    }
387
388    /// Register a logger with the server
389    pub fn logger<L: McpLogger + 'static>(mut self, logger: L) -> Self {
390        let key = format!("logger_{}", self.loggers.len());
391        self.loggers.insert(key, Arc::new(logger));
392        self
393    }
394
395    /// Register multiple loggers
396    pub fn loggers<L: McpLogger + 'static, I: IntoIterator<Item = L>>(
397        mut self,
398        loggers: I,
399    ) -> Self {
400        for logger in loggers {
401            self = self.logger(logger);
402        }
403        self
404    }
405
406    /// Register a root provider with the server
407    pub fn root_provider<R: McpRoot + 'static>(mut self, root: R) -> Self {
408        let key = format!("root_{}", self.root_providers.len());
409        self.root_providers.insert(key, Arc::new(root));
410        self
411    }
412
413    /// Register multiple root providers
414    pub fn root_providers<R: McpRoot + 'static, I: IntoIterator<Item = R>>(
415        mut self,
416        roots: I,
417    ) -> Self {
418        for root in roots {
419            self = self.root_provider(root);
420        }
421        self
422    }
423
424    /// Register a notification provider with the server
425    pub fn notification_provider<N: McpNotification + 'static>(mut self, notification: N) -> Self {
426        let key = format!("notification_{}", self.notifications.len());
427        self.notifications.insert(key, Arc::new(notification));
428        self
429    }
430
431    /// Register multiple notification providers
432    pub fn notification_providers<N: McpNotification + 'static, I: IntoIterator<Item = N>>(
433        mut self,
434        notifications: I,
435    ) -> Self {
436        for notification in notifications {
437            self = self.notification_provider(notification);
438        }
439        self
440    }
441
442    // =============================================================================
443    // ZERO-CONFIGURATION CONVENIENCE METHODS (same as McpServerBuilder)
444    // =============================================================================
445
446    /// Register a sampler - convenient alias for sampling_provider
447    pub fn sampler<S: McpSampling + 'static>(self, sampling: S) -> Self {
448        self.sampling_provider(sampling)
449    }
450
451    /// Register a completer - convenient alias for completion_provider
452    pub fn completer<C: McpCompletion + 'static>(self, completion: C) -> Self {
453        self.completion_provider(completion)
454    }
455
456    /// Register a notification by type - type determines method automatically
457    pub fn notification_type<N: McpNotification + 'static + Default>(self) -> Self {
458        let notification = N::default();
459        self.notification_provider(notification)
460    }
461
462    /// Register a handler with the server
463    pub fn handler<H: McpHandler + 'static>(mut self, handler: H) -> Self {
464        let handler_arc = Arc::new(handler);
465        for method in handler_arc.supported_methods() {
466            self.handlers.insert(method, handler_arc.clone());
467        }
468        self
469    }
470
471    /// Register multiple handlers
472    pub fn handlers<H: McpHandler + 'static, I: IntoIterator<Item = H>>(
473        mut self,
474        handlers: I,
475    ) -> Self {
476        for handler in handlers {
477            self = self.handler(handler);
478        }
479        self
480    }
481
482    /// Add a single root directory
483    pub fn root(mut self, root: turul_mcp_protocol::roots::Root) -> Self {
484        self.roots.push(root);
485        self
486    }
487
488    // =============================================================================
489    // CAPABILITY CONFIGURATION METHODS (same as McpServerBuilder)
490    // =============================================================================
491
492    /// Add completion support
493    pub fn with_completion(mut self) -> Self {
494        use turul_mcp_protocol::initialize::CompletionsCapabilities;
495        self.capabilities.completions = Some(CompletionsCapabilities {
496            enabled: Some(true),
497        });
498        self.handler(CompletionHandler)
499    }
500
501    /// Add prompts support
502    pub fn with_prompts(mut self) -> Self {
503        use turul_mcp_protocol::initialize::PromptsCapabilities;
504        self.capabilities.prompts = Some(PromptsCapabilities {
505            list_changed: Some(false),
506        });
507
508        // Prompts handlers are automatically registered when prompts are added via .prompt()
509        // This method now just enables the capability
510        self
511    }
512
513    /// Add resources support
514    pub fn with_resources(mut self) -> Self {
515        use turul_mcp_protocol::initialize::ResourcesCapabilities;
516        self.capabilities.resources = Some(ResourcesCapabilities {
517            subscribe: Some(false),
518            list_changed: Some(false),
519        });
520
521        // Create ResourcesHandler and add all registered resources
522        let mut handler = ResourcesHandler::new();
523        for resource in self.resources.values() {
524            handler = handler.add_resource_arc(resource.clone());
525        }
526
527        self.handler(handler)
528    }
529
530    /// Add logging support
531    pub fn with_logging(mut self) -> Self {
532        use turul_mcp_protocol::initialize::LoggingCapabilities;
533        self.capabilities.logging = Some(LoggingCapabilities::default());
534        self.handler(LoggingHandler)
535    }
536
537    /// Add roots support
538    pub fn with_roots(self) -> Self {
539        self.handler(RootsHandler::new())
540    }
541
542    /// Add sampling support
543    pub fn with_sampling(self) -> Self {
544        self.handler(SamplingHandler)
545    }
546
547    /// Add elicitation support with default mock provider
548    pub fn with_elicitation(mut self) -> Self {
549        use turul_mcp_protocol::initialize::ElicitationCapabilities;
550        self.capabilities.elicitation = Some(ElicitationCapabilities {
551            enabled: Some(true),
552        });
553        self.handler(ElicitationHandler::with_mock_provider())
554    }
555
556    /// Add elicitation support with custom provider
557    pub fn with_elicitation_provider<P: ElicitationProvider + 'static>(
558        mut self,
559        provider: P,
560    ) -> Self {
561        use turul_mcp_protocol::initialize::ElicitationCapabilities;
562        self.capabilities.elicitation = Some(ElicitationCapabilities {
563            enabled: Some(true),
564        });
565        self.handler(ElicitationHandler::new(Arc::new(provider)))
566    }
567
568    /// Add notifications support
569    pub fn with_notifications(self) -> Self {
570        self.handler(NotificationsHandler)
571    }
572
573    // =============================================================================
574    // SESSION AND CONFIGURATION METHODS
575    // =============================================================================
576
577    /// Configure session timeout (in minutes, default: 30)
578    pub fn session_timeout_minutes(mut self, minutes: u64) -> Self {
579        self.session_timeout_minutes = Some(minutes);
580        self
581    }
582
583    /// Configure session cleanup interval (in seconds, default: 60)
584    pub fn session_cleanup_interval_seconds(mut self, seconds: u64) -> Self {
585        self.session_cleanup_interval_seconds = Some(seconds);
586        self
587    }
588
589    /// Enable strict MCP lifecycle enforcement
590    pub fn strict_lifecycle(mut self, strict: bool) -> Self {
591        self.strict_lifecycle = strict;
592        self
593    }
594
595    /// Enable strict MCP lifecycle enforcement (convenience method)
596    pub fn with_strict_lifecycle(self) -> Self {
597        self.strict_lifecycle(true)
598    }
599
600    /// Enable or disable SSE streaming support
601    pub fn sse(mut self, enable: bool) -> Self {
602        self.enable_sse = enable;
603
604        // Update SSE endpoints in ServerConfig based on enable flag
605        if enable {
606            self.server_config.enable_get_sse = true;
607            self.server_config.enable_post_sse = true;
608        } else {
609            // When SSE is disabled, also disable SSE endpoints in ServerConfig
610            // This prevents GET /mcp from hanging by returning 405 instead
611            self.server_config.enable_get_sse = false;
612            self.server_config.enable_post_sse = false;
613        }
614
615        self
616    }
617
618    /// Configure sessions with recommended defaults for long-running sessions
619    pub fn with_long_sessions(mut self) -> Self {
620        self.session_timeout_minutes = Some(120); // 2 hours
621        self.session_cleanup_interval_seconds = Some(300); // 5 minutes
622        self
623    }
624
625    /// Configure sessions with recommended defaults for short-lived sessions
626    pub fn with_short_sessions(mut self) -> Self {
627        self.session_timeout_minutes = Some(5); // 5 minutes
628        self.session_cleanup_interval_seconds = Some(30); // 30 seconds
629        self
630    }
631
632    /// Set the session storage backend
633    ///
634    /// Supports all framework storage backends:
635    /// - `InMemorySessionStorage` - For development and testing
636    /// - `SqliteSessionStorage` - For single-instance persistence
637    /// - `PostgreSqlSessionStorage` - For multi-instance deployments
638    /// - `DynamoDbSessionStorage` - For serverless AWS deployments
639    pub fn storage(mut self, storage: Arc<BoxedSessionStorage>) -> Self {
640        self.session_storage = Some(storage);
641        self
642    }
643
644    /// Create DynamoDB storage from environment variables
645    ///
646    /// Uses these environment variables:
647    /// - `SESSION_TABLE_NAME` or `MCP_SESSION_TABLE` - DynamoDB table name
648    /// - `AWS_REGION` - AWS region
649    /// - AWS credentials from standard AWS credential chain
650    #[cfg(feature = "dynamodb")]
651    pub async fn dynamodb_storage(self) -> Result<Self> {
652        use turul_mcp_session_storage::DynamoDbSessionStorage;
653
654        let storage = DynamoDbSessionStorage::new().await.map_err(|e| {
655            LambdaError::Configuration(format!("Failed to create DynamoDB storage: {}", e))
656        })?;
657
658        Ok(self.storage(Arc::new(storage)))
659    }
660
661    /// Register middleware for request/response interception
662    ///
663    /// Middleware can inspect and modify requests before they reach handlers,
664    /// inject data into sessions, and transform responses. Multiple middleware
665    /// can be registered and will execute in FIFO order for before_dispatch
666    /// and LIFO order for after_dispatch.
667    ///
668    /// # Example
669    ///
670    /// ```rust,no_run
671    /// use std::sync::Arc;
672    /// use turul_mcp_aws_lambda::LambdaMcpServerBuilder;
673    /// use turul_http_mcp_server::middleware::McpMiddleware;
674    /// # use turul_mcp_session_storage::SessionView;
675    /// # use turul_http_mcp_server::middleware::{RequestContext, SessionInjection, MiddlewareError};
676    /// # use async_trait::async_trait;
677    /// # struct AuthMiddleware;
678    /// # #[async_trait]
679    /// # impl McpMiddleware for AuthMiddleware {
680    /// #     async fn before_dispatch(&self, _: &mut RequestContext<'_>, _: Option<&dyn SessionView>, _: &mut SessionInjection) -> Result<(), MiddlewareError> { Ok(()) }
681    /// # }
682    /// # struct RateLimitMiddleware;
683    /// # #[async_trait]
684    /// # impl McpMiddleware for RateLimitMiddleware {
685    /// #     async fn before_dispatch(&self, _: &mut RequestContext<'_>, _: Option<&dyn SessionView>, _: &mut SessionInjection) -> Result<(), MiddlewareError> { Ok(()) }
686    /// # }
687    ///
688    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
689    /// let builder = LambdaMcpServerBuilder::new()
690    ///     .name("my-server")
691    ///     .middleware(Arc::new(AuthMiddleware))
692    ///     .middleware(Arc::new(RateLimitMiddleware));
693    /// # Ok(())
694    /// # }
695    /// ```
696    pub fn middleware(mut self, middleware: Arc<dyn turul_http_mcp_server::middleware::McpMiddleware>) -> Self {
697        self.middleware_stack.push(middleware);
698        self
699    }
700
701    /// Configure server settings
702    pub fn server_config(mut self, config: ServerConfig) -> Self {
703        self.server_config = config;
704        self
705    }
706
707    /// Configure streaming/SSE settings
708    pub fn stream_config(mut self, config: StreamConfig) -> Self {
709        self.stream_config = config;
710        self
711    }
712
713    // CORS Configuration Methods
714
715    /// Set custom CORS configuration
716    #[cfg(feature = "cors")]
717    pub fn cors(mut self, config: CorsConfig) -> Self {
718        self.cors_config = Some(config);
719        self
720    }
721
722    /// Allow all origins for CORS (development only)
723    #[cfg(feature = "cors")]
724    pub fn cors_allow_all_origins(mut self) -> Self {
725        self.cors_config = Some(CorsConfig::allow_all());
726        self
727    }
728
729    /// Set specific allowed origins for CORS
730    #[cfg(feature = "cors")]
731    pub fn cors_allow_origins(mut self, origins: Vec<String>) -> Self {
732        self.cors_config = Some(CorsConfig::for_origins(origins));
733        self
734    }
735
736    /// Configure CORS from environment variables
737    ///
738    /// Uses these environment variables:
739    /// - `MCP_CORS_ORIGINS` - Comma-separated list of allowed origins
740    /// - `MCP_CORS_CREDENTIALS` - Whether to allow credentials (true/false)
741    /// - `MCP_CORS_MAX_AGE` - Preflight cache max age in seconds
742    #[cfg(feature = "cors")]
743    pub fn cors_from_env(mut self) -> Self {
744        self.cors_config = Some(CorsConfig::from_env());
745        self
746    }
747
748    /// Disable CORS (headers will not be added)
749    #[cfg(feature = "cors")]
750    pub fn cors_disabled(self) -> Self {
751        // Don't set any CORS config - builder will not add headers
752        self
753    }
754
755    // Convenience Methods
756
757    /// Create with DynamoDB storage and environment-based CORS
758    ///
759    /// This is the recommended configuration for production Lambda deployments.
760    #[cfg(all(feature = "dynamodb", feature = "cors"))]
761    pub async fn production_config(self) -> Result<Self> {
762        Ok(self.dynamodb_storage().await?.cors_from_env())
763    }
764
765    /// Create with in-memory storage and permissive CORS
766    ///
767    /// This is the recommended configuration for development and testing.
768    #[cfg(feature = "cors")]
769    pub fn development_config(self) -> Self {
770        use turul_mcp_session_storage::InMemorySessionStorage;
771
772        self.storage(Arc::new(InMemorySessionStorage::new()))
773            .cors_allow_all_origins()
774    }
775
776    /// Build the Lambda MCP server
777    ///
778    /// Returns a server that can create handlers when needed.
779    pub async fn build(self) -> Result<LambdaMcpServer> {
780        use turul_mcp_session_storage::InMemorySessionStorage;
781
782        // Validate configuration (same as MCP server)
783        if self.name.is_empty() {
784            return Err(crate::error::LambdaError::Configuration(
785                "Server name cannot be empty".to_string(),
786            ));
787        }
788        if self.version.is_empty() {
789            return Err(crate::error::LambdaError::Configuration(
790                "Server version cannot be empty".to_string(),
791            ));
792        }
793
794        // Note: SSE behavior depends on which handler method is used:
795        // - handle(): Works with run(), but SSE responses may not stream properly
796        // - handle_streaming(): Works with run_with_streaming_response() for real SSE streaming
797
798        // Create session storage (use in-memory if none provided)
799        let session_storage = self
800            .session_storage
801            .unwrap_or_else(|| Arc::new(InMemorySessionStorage::new()));
802
803        // Create implementation info
804        let implementation = if let Some(title) = self.title {
805            Implementation::new(&self.name, &self.version).with_title(title)
806        } else {
807            Implementation::new(&self.name, &self.version)
808        };
809
810        // Auto-detect and configure server capabilities based on registered components (same as McpServer)
811        let mut capabilities = self.capabilities.clone();
812        let has_tools = !self.tools.is_empty();
813        let has_resources = !self.resources.is_empty();
814        let has_prompts = !self.prompts.is_empty();
815        let has_elicitations = !self.elicitations.is_empty();
816        let has_completions = !self.completions.is_empty();
817        let has_logging = !self.loggers.is_empty();
818        tracing::debug!("🔧 Has logging configured: {}", has_logging);
819
820        // Tools capabilities - truthful reporting (only set if tools are registered)
821        if has_tools {
822            capabilities.tools = Some(turul_mcp_protocol::initialize::ToolsCapabilities {
823                list_changed: Some(false), // Static framework: no dynamic change sources
824            });
825        }
826
827        // Resources capabilities - truthful reporting (only set if resources are registered)
828        if has_resources {
829            capabilities.resources = Some(turul_mcp_protocol::initialize::ResourcesCapabilities {
830                subscribe: Some(false),    // TODO: Implement resource subscriptions
831                list_changed: Some(false), // Static framework: no dynamic change sources
832            });
833        }
834
835        // Prompts capabilities - truthful reporting (only set if prompts are registered)
836        if has_prompts {
837            capabilities.prompts = Some(turul_mcp_protocol::initialize::PromptsCapabilities {
838                list_changed: Some(false), // Static framework: no dynamic change sources
839            });
840        }
841
842        // Elicitation capabilities - truthful reporting (only set if elicitations are registered)
843        if has_elicitations {
844            capabilities.elicitation =
845                Some(turul_mcp_protocol::initialize::ElicitationCapabilities {
846                    enabled: Some(true),
847                });
848        }
849
850        // Completion capabilities - truthful reporting (only set if completions are registered)
851        if has_completions {
852            capabilities.completions =
853                Some(turul_mcp_protocol::initialize::CompletionsCapabilities {
854                    enabled: Some(true),
855                });
856        }
857
858        // Logging capabilities - always enabled for debugging/monitoring (same as McpServer)
859        // Always enable logging for debugging/monitoring
860        capabilities.logging = Some(turul_mcp_protocol::initialize::LoggingCapabilities {
861            enabled: Some(true),
862            levels: Some(vec![
863                "debug".to_string(),
864                "info".to_string(),
865                "warning".to_string(),
866                "error".to_string(),
867            ]),
868        });
869
870        // Add RootsHandler if roots were configured (same pattern as MCP server)
871        let mut handlers = self.handlers;
872        if !self.roots.is_empty() {
873            let mut roots_handler = RootsHandler::new();
874            for root in &self.roots {
875                roots_handler = roots_handler.add_root(root.clone());
876            }
877            handlers.insert("roots/list".to_string(), Arc::new(roots_handler));
878        }
879
880        // Create the Lambda server (stores all configuration like MCP server does)
881        Ok(LambdaMcpServer::new(
882            implementation,
883            capabilities,
884            self.tools,
885            self.resources,
886            self.prompts,
887            self.elicitations,
888            self.sampling,
889            self.completions,
890            self.loggers,
891            self.root_providers,
892            self.notifications,
893            handlers,
894            self.roots,
895            self.instructions,
896            session_storage,
897            self.strict_lifecycle,
898            self.server_config,
899            self.enable_sse,
900            self.stream_config,
901            #[cfg(feature = "cors")]
902            self.cors_config,
903            self.middleware_stack,
904        ))
905    }
906}
907
908impl Default for LambdaMcpServerBuilder {
909    fn default() -> Self {
910        Self::new()
911    }
912}
913
914// Extension trait for cleaner chaining
915pub trait LambdaMcpServerBuilderExt {
916    /// Add multiple tools at once
917    fn tools<I, T>(self, tools: I) -> Self
918    where
919        I: IntoIterator<Item = T>,
920        T: McpTool + 'static;
921}
922
923impl LambdaMcpServerBuilderExt for LambdaMcpServerBuilder {
924    fn tools<I, T>(mut self, tools: I) -> Self
925    where
926        I: IntoIterator<Item = T>,
927        T: McpTool + 'static,
928    {
929        for tool in tools {
930            self = self.tool(tool);
931        }
932        self
933    }
934}
935
936/// Create a Lambda MCP server with minimal configuration
937///
938/// This is a convenience function for simple use cases where you just
939/// want to register some tools and get a working handler.
940pub async fn simple_lambda_server<I, T>(tools: I) -> Result<LambdaMcpServer>
941where
942    I: IntoIterator<Item = T>,
943    T: McpTool + 'static,
944{
945    let mut builder = LambdaMcpServerBuilder::new();
946
947    for tool in tools {
948        builder = builder.tool(tool);
949    }
950
951    #[cfg(feature = "cors")]
952    {
953        builder = builder.cors_allow_all_origins();
954    }
955
956    builder.sse(false).build().await
957}
958
959/// Create a Lambda MCP server configured for production
960///
961/// Uses DynamoDB for session storage and environment-based CORS configuration.
962#[cfg(all(feature = "dynamodb", feature = "cors"))]
963pub async fn production_lambda_server<I, T>(tools: I) -> Result<LambdaMcpServer>
964where
965    I: IntoIterator<Item = T>,
966    T: McpTool + 'static,
967{
968    let mut builder = LambdaMcpServerBuilder::new();
969
970    for tool in tools {
971        builder = builder.tool(tool);
972    }
973
974    builder.production_config().await?.build().await
975}
976
977#[cfg(test)]
978mod tests {
979    use super::*;
980    use turul_mcp_session_storage::InMemorySessionStorage;
981    use turul_mcp_builders::prelude::*;  // HasBaseMetadata, HasDescription, etc.
982
983    // Mock tool for testing
984    #[derive(Clone, Default)]
985    struct TestTool;
986
987    impl HasBaseMetadata for TestTool {
988        fn name(&self) -> &str {
989            "test_tool"
990        }
991    }
992
993    impl HasDescription for TestTool {
994        fn description(&self) -> Option<&str> {
995            Some("Test tool")
996        }
997    }
998
999    impl HasInputSchema for TestTool {
1000        fn input_schema(&self) -> &turul_mcp_protocol::ToolSchema {
1001            use turul_mcp_protocol::ToolSchema;
1002            static SCHEMA: std::sync::OnceLock<ToolSchema> = std::sync::OnceLock::new();
1003            SCHEMA.get_or_init(ToolSchema::object)
1004        }
1005    }
1006
1007    impl HasOutputSchema for TestTool {
1008        fn output_schema(&self) -> Option<&turul_mcp_protocol::ToolSchema> {
1009            None
1010        }
1011    }
1012
1013    impl HasAnnotations for TestTool {
1014        fn annotations(&self) -> Option<&turul_mcp_protocol::tools::ToolAnnotations> {
1015            None
1016        }
1017    }
1018
1019    impl HasToolMeta for TestTool {
1020        fn tool_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1021            None
1022        }
1023    }
1024
1025    #[async_trait::async_trait]
1026    impl McpTool for TestTool {
1027        async fn call(
1028            &self,
1029            _args: serde_json::Value,
1030            _session: Option<turul_mcp_server::SessionContext>,
1031        ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
1032            use turul_mcp_protocol::tools::{CallToolResult, ToolResult};
1033            Ok(CallToolResult::success(vec![ToolResult::text(
1034                "test result",
1035            )]))
1036        }
1037    }
1038
1039    #[tokio::test]
1040    async fn test_builder_basic() {
1041        let server = LambdaMcpServerBuilder::new()
1042            .name("test-server")
1043            .version("1.0.0")
1044            .tool(TestTool)
1045            .storage(Arc::new(InMemorySessionStorage::new()))
1046            .sse(false) // Disable SSE for tests since streaming feature not enabled
1047            .build()
1048            .await
1049            .unwrap();
1050
1051        // Create handler from server and verify it has stream_manager
1052        let handler = server.handler().await.unwrap();
1053        // Verify handler has stream_manager (critical invariant)
1054        assert!(
1055            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1056            "Stream manager must be initialized"
1057        );
1058    }
1059
1060    #[tokio::test]
1061    async fn test_simple_lambda_server() {
1062        let tools = vec![TestTool];
1063        let server = simple_lambda_server(tools).await.unwrap();
1064
1065        // Create handler and verify it was created with default configuration
1066        let handler = server.handler().await.unwrap();
1067        // Verify handler has stream_manager
1068        // Verify handler has stream_manager (critical invariant)
1069        assert!(
1070            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1071            "Stream manager must be initialized"
1072        );
1073    }
1074
1075    #[tokio::test]
1076    async fn test_builder_extension_trait() {
1077        let tools = vec![TestTool, TestTool];
1078
1079        let server = LambdaMcpServerBuilder::new()
1080            .tools(tools)
1081            .storage(Arc::new(InMemorySessionStorage::new()))
1082            .sse(false) // Disable SSE for tests since streaming feature not enabled
1083            .build()
1084            .await
1085            .unwrap();
1086
1087        let handler = server.handler().await.unwrap();
1088        // Verify handler has stream_manager
1089        // Verify handler has stream_manager (critical invariant)
1090        assert!(
1091            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1092            "Stream manager must be initialized"
1093        );
1094    }
1095
1096    #[cfg(feature = "cors")]
1097    #[tokio::test]
1098    async fn test_cors_configuration() {
1099        let server = LambdaMcpServerBuilder::new()
1100            .cors_allow_all_origins()
1101            .storage(Arc::new(InMemorySessionStorage::new()))
1102            .sse(false) // Disable SSE for tests since streaming feature not enabled
1103            .build()
1104            .await
1105            .unwrap();
1106
1107        let handler = server.handler().await.unwrap();
1108        // Verify handler has stream_manager
1109        // Verify handler has stream_manager (critical invariant)
1110        assert!(
1111            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1112            "Stream manager must be initialized"
1113        );
1114    }
1115
1116    #[tokio::test]
1117    async fn test_sse_toggle_functionality() {
1118        // Test that SSE can be toggled on/off/on correctly
1119        let mut builder =
1120            LambdaMcpServerBuilder::new().storage(Arc::new(InMemorySessionStorage::new()));
1121
1122        // Initially enable SSE
1123        builder = builder.sse(true);
1124        assert!(builder.enable_sse, "SSE should be enabled");
1125        assert!(
1126            builder.server_config.enable_get_sse,
1127            "GET SSE endpoint should be enabled"
1128        );
1129        assert!(
1130            builder.server_config.enable_post_sse,
1131            "POST SSE endpoint should be enabled"
1132        );
1133
1134        // Disable SSE
1135        builder = builder.sse(false);
1136        assert!(!builder.enable_sse, "SSE should be disabled");
1137        assert!(
1138            !builder.server_config.enable_get_sse,
1139            "GET SSE endpoint should be disabled"
1140        );
1141        assert!(
1142            !builder.server_config.enable_post_sse,
1143            "POST SSE endpoint should be disabled"
1144        );
1145
1146        // Re-enable SSE (this was broken before the fix)
1147        builder = builder.sse(true);
1148        assert!(builder.enable_sse, "SSE should be re-enabled");
1149        assert!(
1150            builder.server_config.enable_get_sse,
1151            "GET SSE endpoint should be re-enabled"
1152        );
1153        assert!(
1154            builder.server_config.enable_post_sse,
1155            "POST SSE endpoint should be re-enabled"
1156        );
1157
1158        // Verify the server can be built with SSE enabled
1159        let server = builder.build().await.unwrap();
1160        let handler = server.handler().await.unwrap();
1161        assert!(
1162            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1163            "Stream manager must be initialized"
1164        );
1165    }
1166}