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    /// Static resources registered with the server
82    resources: HashMap<String, Arc<dyn McpResource>>,
83
84    /// Template resources registered with the server (auto-detected from URI)
85    template_resources: Vec<(turul_mcp_server::uri_template::UriTemplate, Arc<dyn McpResource>)>,
86
87    /// Prompts registered with the server
88    prompts: HashMap<String, Arc<dyn McpPrompt>>,
89
90    /// Elicitations registered with the server
91    elicitations: HashMap<String, Arc<dyn McpElicitation>>,
92
93    /// Sampling providers registered with the server
94    sampling: HashMap<String, Arc<dyn McpSampling>>,
95
96    /// Completion providers registered with the server
97    completions: HashMap<String, Arc<dyn McpCompletion>>,
98
99    /// Loggers registered with the server
100    loggers: HashMap<String, Arc<dyn McpLogger>>,
101
102    /// Root providers registered with the server
103    root_providers: HashMap<String, Arc<dyn McpRoot>>,
104
105    /// Notification providers registered with the server
106    notifications: HashMap<String, Arc<dyn McpNotification>>,
107
108    /// Handlers registered with the server
109    handlers: HashMap<String, Arc<dyn McpHandler>>,
110
111    /// Roots configured for the server
112    roots: Vec<turul_mcp_protocol::roots::Root>,
113
114    /// Optional instructions for clients
115    instructions: Option<String>,
116
117    /// Session configuration
118    session_timeout_minutes: Option<u64>,
119    session_cleanup_interval_seconds: Option<u64>,
120
121    /// Session storage backend (defaults to InMemory if None)
122    session_storage: Option<Arc<BoxedSessionStorage>>,
123
124    /// MCP Lifecycle enforcement configuration
125    strict_lifecycle: bool,
126
127    /// Enable SSE streaming
128    enable_sse: bool,
129    /// Server and stream configuration
130    server_config: ServerConfig,
131    stream_config: StreamConfig,
132
133    /// Middleware stack for request/response interception
134    middleware_stack: turul_http_mcp_server::middleware::MiddlewareStack,
135
136    /// Optional task runtime for MCP task support
137    task_runtime: Option<Arc<turul_mcp_server::TaskRuntime>>,
138    /// Recovery timeout for stuck tasks (milliseconds)
139    task_recovery_timeout_ms: u64,
140
141    /// CORS configuration (if enabled)
142    #[cfg(feature = "cors")]
143    cors_config: Option<CorsConfig>,
144}
145
146impl LambdaMcpServerBuilder {
147    /// Create a new Lambda MCP server builder
148    pub fn new() -> Self {
149        // Initialize with default capabilities (same as McpServer)
150        // Capabilities will be set truthfully in build() based on registered components
151        let capabilities = ServerCapabilities::default();
152
153        // Initialize handlers with defaults (same as McpServerBuilder)
154        let mut handlers: HashMap<String, Arc<dyn McpHandler>> = HashMap::new();
155        handlers.insert("ping".to_string(), Arc::new(PingHandler));
156        handlers.insert(
157            "completion/complete".to_string(),
158            Arc::new(CompletionHandler),
159        );
160        handlers.insert(
161            "resources/list".to_string(),
162            Arc::new(ResourcesHandler::new()),
163        );
164        handlers.insert(
165            "prompts/list".to_string(),
166            Arc::new(PromptsListHandler::new()),
167        );
168        handlers.insert(
169            "prompts/get".to_string(),
170            Arc::new(PromptsGetHandler::new()),
171        );
172        handlers.insert("logging/setLevel".to_string(), Arc::new(LoggingHandler));
173        handlers.insert("roots/list".to_string(), Arc::new(RootsHandler::new()));
174        handlers.insert(
175            "sampling/createMessage".to_string(),
176            Arc::new(SamplingHandler),
177        );
178        handlers.insert(
179            "resources/templates/list".to_string(),
180            Arc::new(ResourceTemplatesHandler::new()),
181        );
182        handlers.insert(
183            "elicitation/create".to_string(),
184            Arc::new(ElicitationHandler::with_mock_provider()),
185        );
186
187        // Add notification handlers
188        let notifications_handler = Arc::new(NotificationsHandler);
189        handlers.insert(
190            "notifications/message".to_string(),
191            notifications_handler.clone(),
192        );
193        handlers.insert(
194            "notifications/progress".to_string(),
195            notifications_handler.clone(),
196        );
197        // MCP 2025-11-25 spec-correct underscore form
198        handlers.insert(
199            "notifications/resources/list_changed".to_string(),
200            notifications_handler.clone(),
201        );
202        handlers.insert(
203            "notifications/resources/updated".to_string(),
204            notifications_handler.clone(),
205        );
206        handlers.insert(
207            "notifications/tools/list_changed".to_string(),
208            notifications_handler.clone(),
209        );
210        handlers.insert(
211            "notifications/prompts/list_changed".to_string(),
212            notifications_handler.clone(),
213        );
214        handlers.insert(
215            "notifications/roots/list_changed".to_string(),
216            notifications_handler.clone(),
217        );
218        // Legacy compat: accept camelCase from older clients
219        handlers.insert(
220            "notifications/resources/listChanged".to_string(),
221            notifications_handler.clone(),
222        );
223        handlers.insert(
224            "notifications/tools/listChanged".to_string(),
225            notifications_handler.clone(),
226        );
227        handlers.insert(
228            "notifications/prompts/listChanged".to_string(),
229            notifications_handler.clone(),
230        );
231        handlers.insert(
232            "notifications/roots/listChanged".to_string(),
233            notifications_handler,
234        );
235
236        Self {
237            name: "turul-mcp-aws-lambda".to_string(),
238            version: env!("CARGO_PKG_VERSION").to_string(),
239            title: None,
240            capabilities,
241            tools: HashMap::new(),
242            resources: HashMap::new(),
243            template_resources: Vec::new(),
244            prompts: HashMap::new(),
245            elicitations: HashMap::new(),
246            sampling: HashMap::new(),
247            completions: HashMap::new(),
248            loggers: HashMap::new(),
249            root_providers: HashMap::new(),
250            notifications: HashMap::new(),
251            handlers,
252            roots: Vec::new(),
253            instructions: None,
254            session_timeout_minutes: None,
255            session_cleanup_interval_seconds: None,
256            session_storage: None,
257            strict_lifecycle: false,
258            enable_sse: cfg!(feature = "sse"),
259            server_config: ServerConfig::default(),
260            stream_config: StreamConfig::default(),
261            middleware_stack: turul_http_mcp_server::middleware::MiddlewareStack::new(),
262            task_runtime: None,
263            task_recovery_timeout_ms: 300_000, // 5 minutes
264            #[cfg(feature = "cors")]
265            cors_config: None,
266        }
267    }
268
269    /// Set the server name
270    pub fn name(mut self, name: impl Into<String>) -> Self {
271        self.name = name.into();
272        self
273    }
274
275    /// Set the server version
276    pub fn version(mut self, version: impl Into<String>) -> Self {
277        self.version = version.into();
278        self
279    }
280
281    /// Set the server title
282    pub fn title(mut self, title: impl Into<String>) -> Self {
283        self.title = Some(title.into());
284        self
285    }
286
287    /// Set optional instructions for clients
288    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
289        self.instructions = Some(instructions.into());
290        self
291    }
292
293    // =============================================================================
294    // PROVIDER REGISTRATION METHODS (same as McpServerBuilder)
295    // =============================================================================
296
297    /// Register a tool with the server
298    ///
299    /// Tools can be created using any of the framework's 4 creation levels:
300    /// - Function macros: `#[mcp_tool]`
301    /// - Derive macros: `#[derive(McpTool)]`
302    /// - Builder pattern: `ToolBuilder::new(...).build()`
303    /// - Manual implementation: Custom struct implementing `McpTool`
304    pub fn tool<T: McpTool + 'static>(mut self, tool: T) -> Self {
305        let name = tool.name().to_string();
306        self.tools.insert(name, Arc::new(tool));
307        self
308    }
309
310    /// Register a function tool created with `#[mcp_tool]` macro
311    pub fn tool_fn<F, T>(self, func: F) -> Self
312    where
313        F: Fn() -> T,
314        T: McpTool + 'static,
315    {
316        self.tool(func())
317    }
318
319    /// Register multiple tools
320    pub fn tools<T: McpTool + 'static, I: IntoIterator<Item = T>>(mut self, tools: I) -> Self {
321        for tool in tools {
322            self = self.tool(tool);
323        }
324        self
325    }
326
327    /// Register a resource with the server
328    ///
329    /// Automatically detects template resources (URIs containing `{variables}`)
330    /// and routes them to the template resource list. Template resources appear
331    /// in `resources/templates/list`, not `resources/list`.
332    pub fn resource<R: McpResource + 'static>(mut self, resource: R) -> Self {
333        let uri = resource.uri().to_string();
334
335        if uri.contains('{') && uri.contains('}') {
336            // Template resource — parse URI as UriTemplate
337            match turul_mcp_server::uri_template::UriTemplate::new(&uri) {
338                Ok(template) => {
339                    self.template_resources.push((template, Arc::new(resource)));
340                }
341                Err(e) => {
342                    tracing::warn!(
343                        "Failed to parse template resource URI '{}': {}. Registering as static.",
344                        uri,
345                        e
346                    );
347                    self.resources.insert(uri, Arc::new(resource));
348                }
349            }
350        } else {
351            // Static resource
352            self.resources.insert(uri, Arc::new(resource));
353        }
354        self
355    }
356
357    /// Register multiple resources
358    pub fn resources<R: McpResource + 'static, I: IntoIterator<Item = R>>(
359        mut self,
360        resources: I,
361    ) -> Self {
362        for resource in resources {
363            self = self.resource(resource);
364        }
365        self
366    }
367
368    /// Register a prompt with the server
369    pub fn prompt<P: McpPrompt + 'static>(mut self, prompt: P) -> Self {
370        let name = prompt.name().to_string();
371        self.prompts.insert(name, Arc::new(prompt));
372        self
373    }
374
375    /// Register multiple prompts
376    pub fn prompts<P: McpPrompt + 'static, I: IntoIterator<Item = P>>(
377        mut self,
378        prompts: I,
379    ) -> Self {
380        for prompt in prompts {
381            self = self.prompt(prompt);
382        }
383        self
384    }
385
386    /// Register an elicitation provider with the server
387    pub fn elicitation<E: McpElicitation + 'static>(mut self, elicitation: E) -> Self {
388        let key = format!("elicitation_{}", self.elicitations.len());
389        self.elicitations.insert(key, Arc::new(elicitation));
390        self
391    }
392
393    /// Register multiple elicitation providers
394    pub fn elicitations<E: McpElicitation + 'static, I: IntoIterator<Item = E>>(
395        mut self,
396        elicitations: I,
397    ) -> Self {
398        for elicitation in elicitations {
399            self = self.elicitation(elicitation);
400        }
401        self
402    }
403
404    /// Register a sampling provider with the server
405    pub fn sampling_provider<S: McpSampling + 'static>(mut self, sampling: S) -> Self {
406        let key = format!("sampling_{}", self.sampling.len());
407        self.sampling.insert(key, Arc::new(sampling));
408        self
409    }
410
411    /// Register multiple sampling providers
412    pub fn sampling_providers<S: McpSampling + 'static, I: IntoIterator<Item = S>>(
413        mut self,
414        sampling: I,
415    ) -> Self {
416        for s in sampling {
417            self = self.sampling_provider(s);
418        }
419        self
420    }
421
422    /// Register a completion provider with the server
423    pub fn completion_provider<C: McpCompletion + 'static>(mut self, completion: C) -> Self {
424        let key = format!("completion_{}", self.completions.len());
425        self.completions.insert(key, Arc::new(completion));
426        self
427    }
428
429    /// Register multiple completion providers
430    pub fn completion_providers<C: McpCompletion + 'static, I: IntoIterator<Item = C>>(
431        mut self,
432        completions: I,
433    ) -> Self {
434        for completion in completions {
435            self = self.completion_provider(completion);
436        }
437        self
438    }
439
440    /// Register a logger with the server
441    pub fn logger<L: McpLogger + 'static>(mut self, logger: L) -> Self {
442        let key = format!("logger_{}", self.loggers.len());
443        self.loggers.insert(key, Arc::new(logger));
444        self
445    }
446
447    /// Register multiple loggers
448    pub fn loggers<L: McpLogger + 'static, I: IntoIterator<Item = L>>(
449        mut self,
450        loggers: I,
451    ) -> Self {
452        for logger in loggers {
453            self = self.logger(logger);
454        }
455        self
456    }
457
458    /// Register a root provider with the server
459    pub fn root_provider<R: McpRoot + 'static>(mut self, root: R) -> Self {
460        let key = format!("root_{}", self.root_providers.len());
461        self.root_providers.insert(key, Arc::new(root));
462        self
463    }
464
465    /// Register multiple root providers
466    pub fn root_providers<R: McpRoot + 'static, I: IntoIterator<Item = R>>(
467        mut self,
468        roots: I,
469    ) -> Self {
470        for root in roots {
471            self = self.root_provider(root);
472        }
473        self
474    }
475
476    /// Register a notification provider with the server
477    pub fn notification_provider<N: McpNotification + 'static>(mut self, notification: N) -> Self {
478        let key = format!("notification_{}", self.notifications.len());
479        self.notifications.insert(key, Arc::new(notification));
480        self
481    }
482
483    /// Register multiple notification providers
484    pub fn notification_providers<N: McpNotification + 'static, I: IntoIterator<Item = N>>(
485        mut self,
486        notifications: I,
487    ) -> Self {
488        for notification in notifications {
489            self = self.notification_provider(notification);
490        }
491        self
492    }
493
494    // =============================================================================
495    // ZERO-CONFIGURATION CONVENIENCE METHODS (same as McpServerBuilder)
496    // =============================================================================
497
498    /// Register a sampler - convenient alias for sampling_provider
499    pub fn sampler<S: McpSampling + 'static>(self, sampling: S) -> Self {
500        self.sampling_provider(sampling)
501    }
502
503    /// Register a completer - convenient alias for completion_provider
504    pub fn completer<C: McpCompletion + 'static>(self, completion: C) -> Self {
505        self.completion_provider(completion)
506    }
507
508    /// Register a notification by type - type determines method automatically
509    pub fn notification_type<N: McpNotification + 'static + Default>(self) -> Self {
510        let notification = N::default();
511        self.notification_provider(notification)
512    }
513
514    /// Register a handler with the server
515    pub fn handler<H: McpHandler + 'static>(mut self, handler: H) -> Self {
516        let handler_arc = Arc::new(handler);
517        for method in handler_arc.supported_methods() {
518            self.handlers.insert(method, handler_arc.clone());
519        }
520        self
521    }
522
523    /// Register multiple handlers
524    pub fn handlers<H: McpHandler + 'static, I: IntoIterator<Item = H>>(
525        mut self,
526        handlers: I,
527    ) -> Self {
528        for handler in handlers {
529            self = self.handler(handler);
530        }
531        self
532    }
533
534    /// Add a single root directory
535    pub fn root(mut self, root: turul_mcp_protocol::roots::Root) -> Self {
536        self.roots.push(root);
537        self
538    }
539
540    // =============================================================================
541    // CAPABILITY CONFIGURATION METHODS (same as McpServerBuilder)
542    // =============================================================================
543
544    /// Add completion support
545    pub fn with_completion(mut self) -> Self {
546        use turul_mcp_protocol::initialize::CompletionsCapabilities;
547        self.capabilities.completions = Some(CompletionsCapabilities {
548            enabled: Some(true),
549        });
550        self.handler(CompletionHandler)
551    }
552
553    /// Add prompts support
554    pub fn with_prompts(mut self) -> Self {
555        use turul_mcp_protocol::initialize::PromptsCapabilities;
556        self.capabilities.prompts = Some(PromptsCapabilities {
557            list_changed: Some(false),
558        });
559
560        // Prompts handlers are automatically registered when prompts are added via .prompt()
561        // This method now just enables the capability
562        self
563    }
564
565    /// Add resources support
566    pub fn with_resources(mut self) -> Self {
567        use turul_mcp_protocol::initialize::ResourcesCapabilities;
568        self.capabilities.resources = Some(ResourcesCapabilities {
569            subscribe: Some(false),
570            list_changed: Some(false),
571        });
572
573        // Create ResourcesHandler (resources/list) — static resources only
574        let mut list_handler = ResourcesHandler::new();
575        for resource in self.resources.values() {
576            list_handler = list_handler.add_resource_arc(resource.clone());
577        }
578        self = self.handler(list_handler);
579
580        // Create ResourceTemplatesHandler (resources/templates/list) — template resources
581        if !self.template_resources.is_empty() {
582            let templates_handler =
583                ResourceTemplatesHandler::new().with_templates(self.template_resources.clone());
584            self = self.handler(templates_handler);
585        }
586
587        // Create ResourcesReadHandler (resources/read) — both static and template resources
588        let mut read_handler = ResourcesReadHandler::new().without_security();
589        for resource in self.resources.values() {
590            read_handler = read_handler.add_resource_arc(resource.clone());
591        }
592        for (template, resource) in &self.template_resources {
593            read_handler =
594                read_handler.add_template_resource_arc(template.clone(), resource.clone());
595        }
596        self.handler(read_handler)
597    }
598
599    /// Add logging support
600    pub fn with_logging(mut self) -> Self {
601        use turul_mcp_protocol::initialize::LoggingCapabilities;
602        self.capabilities.logging = Some(LoggingCapabilities::default());
603        self.handler(LoggingHandler)
604    }
605
606    /// Add roots support
607    pub fn with_roots(self) -> Self {
608        self.handler(RootsHandler::new())
609    }
610
611    /// Add sampling support
612    pub fn with_sampling(self) -> Self {
613        self.handler(SamplingHandler)
614    }
615
616    /// Add elicitation support with default mock provider
617    pub fn with_elicitation(self) -> Self {
618        // Elicitation is a client-side capability per MCP 2025-11-25
619        // Server just registers the handler, no capability advertisement needed
620        self.handler(ElicitationHandler::with_mock_provider())
621    }
622
623    /// Add elicitation support with custom provider
624    pub fn with_elicitation_provider<P: ElicitationProvider + 'static>(self, provider: P) -> Self {
625        // Elicitation is a client-side capability per MCP 2025-11-25
626        self.handler(ElicitationHandler::new(Arc::new(provider)))
627    }
628
629    /// Add notifications support
630    pub fn with_notifications(self) -> Self {
631        self.handler(NotificationsHandler)
632    }
633
634    // =============================================================================
635    // TASK SUPPORT METHODS
636    // =============================================================================
637
638    /// Configure task storage to enable MCP task support for long-running operations.
639    ///
640    /// When task storage is configured, the server will:
641    /// - Advertise `tasks` capabilities in the initialize response
642    /// - Register handlers for `tasks/get`, `tasks/list`, `tasks/cancel`, `tasks/result`
643    /// - Wire task-augmented `tools/call` for `CreateTaskResult` returns
644    /// - Recover stuck tasks on cold start
645    ///
646    /// **Lambda note**: Use a durable backend (DynamoDB recommended) since Lambda
647    /// invocations are stateless. `InMemoryTaskStorage` will lose state between invocations.
648    pub fn with_task_storage(
649        mut self,
650        storage: Arc<dyn turul_mcp_server::task_storage::TaskStorage>,
651    ) -> Self {
652        let runtime = turul_mcp_server::TaskRuntime::with_default_executor(storage)
653            .with_recovery_timeout(self.task_recovery_timeout_ms);
654        self.task_runtime = Some(Arc::new(runtime));
655        self
656    }
657
658    /// Configure task support with a pre-built `TaskRuntime`.
659    ///
660    /// Use this when you need fine-grained control over the task runtime configuration.
661    pub fn with_task_runtime(mut self, runtime: Arc<turul_mcp_server::TaskRuntime>) -> Self {
662        self.task_runtime = Some(runtime);
663        self
664    }
665
666    /// Set the recovery timeout for stuck tasks (in milliseconds).
667    ///
668    /// On Lambda cold start, tasks in non-terminal states older than this timeout
669    /// will be marked as `Failed`. Default: 300,000 ms (5 minutes).
670    pub fn task_recovery_timeout_ms(mut self, timeout_ms: u64) -> Self {
671        self.task_recovery_timeout_ms = timeout_ms;
672        self
673    }
674
675    // =============================================================================
676    // SESSION AND CONFIGURATION METHODS
677    // =============================================================================
678
679    /// Configure session timeout (in minutes, default: 30)
680    pub fn session_timeout_minutes(mut self, minutes: u64) -> Self {
681        self.session_timeout_minutes = Some(minutes);
682        self
683    }
684
685    /// Configure session cleanup interval (in seconds, default: 60)
686    pub fn session_cleanup_interval_seconds(mut self, seconds: u64) -> Self {
687        self.session_cleanup_interval_seconds = Some(seconds);
688        self
689    }
690
691    /// Enable strict MCP lifecycle enforcement
692    pub fn strict_lifecycle(mut self, strict: bool) -> Self {
693        self.strict_lifecycle = strict;
694        self
695    }
696
697    /// Enable strict MCP lifecycle enforcement (convenience method)
698    pub fn with_strict_lifecycle(self) -> Self {
699        self.strict_lifecycle(true)
700    }
701
702    /// Enable or disable SSE streaming support
703    pub fn sse(mut self, enable: bool) -> Self {
704        self.enable_sse = enable;
705
706        // Update SSE endpoints in ServerConfig based on enable flag
707        if enable {
708            self.server_config.enable_get_sse = true;
709            self.server_config.enable_post_sse = true;
710        } else {
711            // When SSE is disabled, also disable SSE endpoints in ServerConfig
712            // This prevents GET /mcp from hanging by returning 405 instead
713            self.server_config.enable_get_sse = false;
714            self.server_config.enable_post_sse = false;
715        }
716
717        self
718    }
719
720    /// Configure sessions with recommended defaults for long-running sessions
721    pub fn with_long_sessions(mut self) -> Self {
722        self.session_timeout_minutes = Some(120); // 2 hours
723        self.session_cleanup_interval_seconds = Some(300); // 5 minutes
724        self
725    }
726
727    /// Configure sessions with recommended defaults for short-lived sessions
728    pub fn with_short_sessions(mut self) -> Self {
729        self.session_timeout_minutes = Some(5); // 5 minutes
730        self.session_cleanup_interval_seconds = Some(30); // 30 seconds
731        self
732    }
733
734    /// Set the session storage backend
735    ///
736    /// Supports all framework storage backends:
737    /// - `InMemorySessionStorage` - For development and testing
738    /// - `SqliteSessionStorage` - For single-instance persistence
739    /// - `PostgreSqlSessionStorage` - For multi-instance deployments
740    /// - `DynamoDbSessionStorage` - For serverless AWS deployments
741    pub fn storage(mut self, storage: Arc<BoxedSessionStorage>) -> Self {
742        self.session_storage = Some(storage);
743        self
744    }
745
746    /// Create DynamoDB storage from environment variables
747    ///
748    /// Uses these environment variables:
749    /// - `SESSION_TABLE_NAME` or `MCP_SESSION_TABLE` - DynamoDB table name
750    /// - `AWS_REGION` - AWS region
751    /// - AWS credentials from standard AWS credential chain
752    #[cfg(feature = "dynamodb")]
753    pub async fn dynamodb_storage(self) -> Result<Self> {
754        use turul_mcp_session_storage::DynamoDbSessionStorage;
755
756        let storage = DynamoDbSessionStorage::new().await.map_err(|e| {
757            LambdaError::Configuration(format!("Failed to create DynamoDB storage: {}", e))
758        })?;
759
760        Ok(self.storage(Arc::new(storage)))
761    }
762
763    /// Register middleware for request/response interception
764    ///
765    /// Middleware can inspect and modify requests before they reach handlers,
766    /// inject data into sessions, and transform responses. Multiple middleware
767    /// can be registered and will execute in FIFO order for before_dispatch
768    /// and LIFO order for after_dispatch.
769    ///
770    /// # Example
771    ///
772    /// ```rust,no_run
773    /// use std::sync::Arc;
774    /// use turul_mcp_aws_lambda::LambdaMcpServerBuilder;
775    /// use turul_http_mcp_server::middleware::McpMiddleware;
776    /// # use turul_mcp_session_storage::SessionView;
777    /// # use turul_http_mcp_server::middleware::{RequestContext, SessionInjection, MiddlewareError};
778    /// # use async_trait::async_trait;
779    /// # struct AuthMiddleware;
780    /// # #[async_trait]
781    /// # impl McpMiddleware for AuthMiddleware {
782    /// #     async fn before_dispatch(&self, _: &mut RequestContext<'_>, _: Option<&dyn SessionView>, _: &mut SessionInjection) -> Result<(), MiddlewareError> { Ok(()) }
783    /// # }
784    /// # struct RateLimitMiddleware;
785    /// # #[async_trait]
786    /// # impl McpMiddleware for RateLimitMiddleware {
787    /// #     async fn before_dispatch(&self, _: &mut RequestContext<'_>, _: Option<&dyn SessionView>, _: &mut SessionInjection) -> Result<(), MiddlewareError> { Ok(()) }
788    /// # }
789    ///
790    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
791    /// let builder = LambdaMcpServerBuilder::new()
792    ///     .name("my-server")
793    ///     .middleware(Arc::new(AuthMiddleware))
794    ///     .middleware(Arc::new(RateLimitMiddleware));
795    /// # Ok(())
796    /// # }
797    /// ```
798    pub fn middleware(
799        mut self,
800        middleware: Arc<dyn turul_http_mcp_server::middleware::McpMiddleware>,
801    ) -> Self {
802        self.middleware_stack.push(middleware);
803        self
804    }
805
806    /// Configure server settings
807    pub fn server_config(mut self, config: ServerConfig) -> Self {
808        self.server_config = config;
809        self
810    }
811
812    /// Configure streaming/SSE settings
813    pub fn stream_config(mut self, config: StreamConfig) -> Self {
814        self.stream_config = config;
815        self
816    }
817
818    // CORS Configuration Methods
819
820    /// Set custom CORS configuration
821    #[cfg(feature = "cors")]
822    pub fn cors(mut self, config: CorsConfig) -> Self {
823        self.cors_config = Some(config);
824        self
825    }
826
827    /// Allow all origins for CORS (development only)
828    #[cfg(feature = "cors")]
829    pub fn cors_allow_all_origins(mut self) -> Self {
830        self.cors_config = Some(CorsConfig::allow_all());
831        self
832    }
833
834    /// Set specific allowed origins for CORS
835    #[cfg(feature = "cors")]
836    pub fn cors_allow_origins(mut self, origins: Vec<String>) -> Self {
837        self.cors_config = Some(CorsConfig::for_origins(origins));
838        self
839    }
840
841    /// Configure CORS from environment variables
842    ///
843    /// Uses these environment variables:
844    /// - `MCP_CORS_ORIGINS` - Comma-separated list of allowed origins
845    /// - `MCP_CORS_CREDENTIALS` - Whether to allow credentials (true/false)
846    /// - `MCP_CORS_MAX_AGE` - Preflight cache max age in seconds
847    #[cfg(feature = "cors")]
848    pub fn cors_from_env(mut self) -> Self {
849        self.cors_config = Some(CorsConfig::from_env());
850        self
851    }
852
853    /// Disable CORS (headers will not be added)
854    #[cfg(feature = "cors")]
855    pub fn cors_disabled(self) -> Self {
856        // Don't set any CORS config - builder will not add headers
857        self
858    }
859
860    // Convenience Methods
861
862    /// Create with DynamoDB storage and environment-based CORS
863    ///
864    /// This is the recommended configuration for production Lambda deployments.
865    #[cfg(all(feature = "dynamodb", feature = "cors"))]
866    pub async fn production_config(self) -> Result<Self> {
867        Ok(self.dynamodb_storage().await?.cors_from_env())
868    }
869
870    /// Create with in-memory storage and permissive CORS
871    ///
872    /// This is the recommended configuration for development and testing.
873    #[cfg(feature = "cors")]
874    pub fn development_config(self) -> Self {
875        use turul_mcp_session_storage::InMemorySessionStorage;
876
877        self.storage(Arc::new(InMemorySessionStorage::new()))
878            .cors_allow_all_origins()
879    }
880
881    /// Build the Lambda MCP server
882    ///
883    /// Returns a server that can create handlers when needed.
884    pub async fn build(self) -> Result<LambdaMcpServer> {
885        use turul_mcp_session_storage::InMemorySessionStorage;
886
887        // Validate configuration (same as MCP server)
888        if self.name.is_empty() {
889            return Err(crate::error::LambdaError::Configuration(
890                "Server name cannot be empty".to_string(),
891            ));
892        }
893        if self.version.is_empty() {
894            return Err(crate::error::LambdaError::Configuration(
895                "Server version cannot be empty".to_string(),
896            ));
897        }
898
899        // Note: SSE behavior depends on which handler method is used:
900        // - handle(): Works with run(), but SSE responses may not stream properly
901        // - handle_streaming(): Works with run_with_streaming_response() for real SSE streaming
902
903        // Create session storage (use in-memory if none provided)
904        let session_storage = self
905            .session_storage
906            .unwrap_or_else(|| Arc::new(InMemorySessionStorage::new()));
907
908        // Create implementation info
909        let implementation = if let Some(title) = self.title {
910            Implementation::new(&self.name, &self.version).with_title(title)
911        } else {
912            Implementation::new(&self.name, &self.version)
913        };
914
915        // Auto-detect and configure server capabilities based on registered components (same as McpServer)
916        let mut capabilities = self.capabilities.clone();
917        let has_tools = !self.tools.is_empty();
918        let has_resources = !self.resources.is_empty() || !self.template_resources.is_empty();
919        let has_prompts = !self.prompts.is_empty();
920        let has_elicitations = !self.elicitations.is_empty();
921        let has_completions = !self.completions.is_empty();
922        let has_logging = !self.loggers.is_empty();
923        tracing::debug!("🔧 Has logging configured: {}", has_logging);
924
925        // Tools capabilities - truthful reporting (only set if tools are registered)
926        if has_tools {
927            capabilities.tools = Some(turul_mcp_protocol::initialize::ToolsCapabilities {
928                list_changed: Some(false), // Static framework: no dynamic change sources
929            });
930        }
931
932        // Resources capabilities - truthful reporting (only set if resources are registered)
933        if has_resources {
934            capabilities.resources = Some(turul_mcp_protocol::initialize::ResourcesCapabilities {
935                subscribe: Some(false),    // TODO: Implement resource subscriptions
936                list_changed: Some(false), // Static framework: no dynamic change sources
937            });
938        }
939
940        // Prompts capabilities - truthful reporting (only set if prompts are registered)
941        if has_prompts {
942            capabilities.prompts = Some(turul_mcp_protocol::initialize::PromptsCapabilities {
943                list_changed: Some(false), // Static framework: no dynamic change sources
944            });
945        }
946
947        // Elicitation is a client-side capability per MCP 2025-11-25
948        // Server does NOT advertise elicitation capabilities
949        let _ = has_elicitations; // Acknowledge the variable without using it
950
951        // Completion capabilities - truthful reporting (only set if completions are registered)
952        if has_completions {
953            capabilities.completions =
954                Some(turul_mcp_protocol::initialize::CompletionsCapabilities {
955                    enabled: Some(true),
956                });
957        }
958
959        // Logging capabilities - always enabled for debugging/monitoring (same as McpServer)
960        // Always enable logging for debugging/monitoring
961        capabilities.logging = Some(turul_mcp_protocol::initialize::LoggingCapabilities {
962            enabled: Some(true),
963            levels: Some(vec![
964                "debug".to_string(),
965                "info".to_string(),
966                "warning".to_string(),
967                "error".to_string(),
968            ]),
969        });
970
971        // Tasks capabilities — auto-configure when task runtime is set
972        if self.task_runtime.is_some() {
973            use turul_mcp_protocol::initialize::*;
974            capabilities.tasks = Some(TasksCapabilities {
975                list: Some(TasksListCapabilities::default()),
976                cancel: Some(TasksCancelCapabilities::default()),
977                requests: Some(TasksRequestCapabilities {
978                    tools: Some(TasksToolCapabilities {
979                        call: Some(TasksToolCallCapabilities::default()),
980                        extra: Default::default(),
981                    }),
982                    extra: Default::default(),
983                }),
984                extra: Default::default(),
985            });
986        }
987
988        // Add RootsHandler if roots were configured (same pattern as MCP server)
989        let mut handlers = self.handlers;
990        if !self.roots.is_empty() {
991            let mut roots_handler = RootsHandler::new();
992            for root in &self.roots {
993                roots_handler = roots_handler.add_root(root.clone());
994            }
995            handlers.insert("roots/list".to_string(), Arc::new(roots_handler));
996        }
997
998        // Add task handlers if task runtime is configured
999        if let Some(ref runtime) = self.task_runtime {
1000            use turul_mcp_server::{
1001                TasksCancelHandler, TasksGetHandler, TasksListHandler, TasksResultHandler,
1002            };
1003            handlers.insert(
1004                "tasks/get".to_string(),
1005                Arc::new(TasksGetHandler::new(Arc::clone(runtime))),
1006            );
1007            handlers.insert(
1008                "tasks/list".to_string(),
1009                Arc::new(TasksListHandler::new(Arc::clone(runtime))),
1010            );
1011            handlers.insert(
1012                "tasks/cancel".to_string(),
1013                Arc::new(TasksCancelHandler::new(Arc::clone(runtime))),
1014            );
1015            handlers.insert(
1016                "tasks/result".to_string(),
1017                Arc::new(TasksResultHandler::new(Arc::clone(runtime))),
1018            );
1019        }
1020
1021        // Auto-populate resource handlers (same as McpServer build() auto-setup)
1022        if has_resources {
1023            // Populate resources/list handler with static resources
1024            let mut list_handler = ResourcesHandler::new();
1025            for resource in self.resources.values() {
1026                list_handler = list_handler.add_resource_arc(resource.clone());
1027            }
1028            handlers.insert("resources/list".to_string(), Arc::new(list_handler));
1029
1030            // Populate resources/templates/list handler with template resources
1031            if !self.template_resources.is_empty() {
1032                let templates_handler = ResourceTemplatesHandler::new()
1033                    .with_templates(self.template_resources.clone());
1034                handlers.insert(
1035                    "resources/templates/list".to_string(),
1036                    Arc::new(templates_handler),
1037                );
1038            }
1039
1040            // Create resources/read handler with both static and template resources
1041            let mut read_handler = ResourcesReadHandler::new().without_security();
1042            for resource in self.resources.values() {
1043                read_handler = read_handler.add_resource_arc(resource.clone());
1044            }
1045            for (template, resource) in &self.template_resources {
1046                read_handler =
1047                    read_handler.add_template_resource_arc(template.clone(), resource.clone());
1048            }
1049            handlers.insert("resources/read".to_string(), Arc::new(read_handler));
1050        }
1051
1052        // Create the Lambda server (stores all configuration like MCP server does)
1053        Ok(LambdaMcpServer::new(
1054            implementation,
1055            capabilities,
1056            self.tools,
1057            self.resources,
1058            self.prompts,
1059            self.elicitations,
1060            self.sampling,
1061            self.completions,
1062            self.loggers,
1063            self.root_providers,
1064            self.notifications,
1065            handlers,
1066            self.roots,
1067            self.instructions,
1068            session_storage,
1069            self.strict_lifecycle,
1070            self.server_config,
1071            self.enable_sse,
1072            self.stream_config,
1073            #[cfg(feature = "cors")]
1074            self.cors_config,
1075            self.middleware_stack,
1076            self.task_runtime,
1077        ))
1078    }
1079}
1080
1081impl Default for LambdaMcpServerBuilder {
1082    fn default() -> Self {
1083        Self::new()
1084    }
1085}
1086
1087// Extension trait for cleaner chaining
1088pub trait LambdaMcpServerBuilderExt {
1089    /// Add multiple tools at once
1090    fn tools<I, T>(self, tools: I) -> Self
1091    where
1092        I: IntoIterator<Item = T>,
1093        T: McpTool + 'static;
1094}
1095
1096impl LambdaMcpServerBuilderExt for LambdaMcpServerBuilder {
1097    fn tools<I, T>(mut self, tools: I) -> Self
1098    where
1099        I: IntoIterator<Item = T>,
1100        T: McpTool + 'static,
1101    {
1102        for tool in tools {
1103            self = self.tool(tool);
1104        }
1105        self
1106    }
1107}
1108
1109/// Create a Lambda MCP server with minimal configuration
1110///
1111/// This is a convenience function for simple use cases where you just
1112/// want to register some tools and get a working handler.
1113pub async fn simple_lambda_server<I, T>(tools: I) -> Result<LambdaMcpServer>
1114where
1115    I: IntoIterator<Item = T>,
1116    T: McpTool + 'static,
1117{
1118    let mut builder = LambdaMcpServerBuilder::new();
1119
1120    for tool in tools {
1121        builder = builder.tool(tool);
1122    }
1123
1124    #[cfg(feature = "cors")]
1125    {
1126        builder = builder.cors_allow_all_origins();
1127    }
1128
1129    builder.sse(false).build().await
1130}
1131
1132/// Create a Lambda MCP server configured for production
1133///
1134/// Uses DynamoDB for session storage and environment-based CORS configuration.
1135#[cfg(all(feature = "dynamodb", feature = "cors"))]
1136pub async fn production_lambda_server<I, T>(tools: I) -> Result<LambdaMcpServer>
1137where
1138    I: IntoIterator<Item = T>,
1139    T: McpTool + 'static,
1140{
1141    let mut builder = LambdaMcpServerBuilder::new();
1142
1143    for tool in tools {
1144        builder = builder.tool(tool);
1145    }
1146
1147    builder.production_config().await?.build().await
1148}
1149
1150#[cfg(test)]
1151mod tests {
1152    use super::*;
1153    use turul_mcp_builders::prelude::*;
1154    use turul_mcp_session_storage::InMemorySessionStorage; // HasBaseMetadata, HasDescription, etc.
1155
1156    // Mock tool for testing
1157    #[derive(Clone, Default)]
1158    struct TestTool;
1159
1160    impl HasBaseMetadata for TestTool {
1161        fn name(&self) -> &str {
1162            "test_tool"
1163        }
1164    }
1165
1166    impl HasDescription for TestTool {
1167        fn description(&self) -> Option<&str> {
1168            Some("Test tool")
1169        }
1170    }
1171
1172    impl HasInputSchema for TestTool {
1173        fn input_schema(&self) -> &turul_mcp_protocol::ToolSchema {
1174            use turul_mcp_protocol::ToolSchema;
1175            static SCHEMA: std::sync::OnceLock<ToolSchema> = std::sync::OnceLock::new();
1176            SCHEMA.get_or_init(ToolSchema::object)
1177        }
1178    }
1179
1180    impl HasOutputSchema for TestTool {
1181        fn output_schema(&self) -> Option<&turul_mcp_protocol::ToolSchema> {
1182            None
1183        }
1184    }
1185
1186    impl HasAnnotations for TestTool {
1187        fn annotations(&self) -> Option<&turul_mcp_protocol::tools::ToolAnnotations> {
1188            None
1189        }
1190    }
1191
1192    impl HasToolMeta for TestTool {
1193        fn tool_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1194            None
1195        }
1196    }
1197
1198    impl HasIcons for TestTool {}
1199    impl HasExecution for TestTool {}
1200
1201    #[async_trait::async_trait]
1202    impl McpTool for TestTool {
1203        async fn call(
1204            &self,
1205            _args: serde_json::Value,
1206            _session: Option<turul_mcp_server::SessionContext>,
1207        ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
1208            use turul_mcp_protocol::tools::{CallToolResult, ToolResult};
1209            Ok(CallToolResult::success(vec![ToolResult::text(
1210                "test result",
1211            )]))
1212        }
1213    }
1214
1215    #[tokio::test]
1216    async fn test_builder_basic() {
1217        let server = LambdaMcpServerBuilder::new()
1218            .name("test-server")
1219            .version("1.0.0")
1220            .tool(TestTool)
1221            .storage(Arc::new(InMemorySessionStorage::new()))
1222            .sse(false) // Disable SSE for tests since streaming feature not enabled
1223            .build()
1224            .await
1225            .unwrap();
1226
1227        // Create handler from server and verify it has stream_manager
1228        let handler = server.handler().await.unwrap();
1229        // Verify handler has stream_manager (critical invariant)
1230        assert!(
1231            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1232            "Stream manager must be initialized"
1233        );
1234    }
1235
1236    #[tokio::test]
1237    async fn test_simple_lambda_server() {
1238        let tools = vec![TestTool];
1239        let server = simple_lambda_server(tools).await.unwrap();
1240
1241        // Create handler and verify it was created with default configuration
1242        let handler = server.handler().await.unwrap();
1243        // Verify handler has stream_manager
1244        // Verify handler has stream_manager (critical invariant)
1245        assert!(
1246            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1247            "Stream manager must be initialized"
1248        );
1249    }
1250
1251    #[tokio::test]
1252    async fn test_builder_extension_trait() {
1253        let tools = vec![TestTool, TestTool];
1254
1255        let server = LambdaMcpServerBuilder::new()
1256            .tools(tools)
1257            .storage(Arc::new(InMemorySessionStorage::new()))
1258            .sse(false) // Disable SSE for tests since streaming feature not enabled
1259            .build()
1260            .await
1261            .unwrap();
1262
1263        let handler = server.handler().await.unwrap();
1264        // Verify handler has stream_manager
1265        // Verify handler has stream_manager (critical invariant)
1266        assert!(
1267            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1268            "Stream manager must be initialized"
1269        );
1270    }
1271
1272    #[cfg(feature = "cors")]
1273    #[tokio::test]
1274    async fn test_cors_configuration() {
1275        let server = LambdaMcpServerBuilder::new()
1276            .cors_allow_all_origins()
1277            .storage(Arc::new(InMemorySessionStorage::new()))
1278            .sse(false) // Disable SSE for tests since streaming feature not enabled
1279            .build()
1280            .await
1281            .unwrap();
1282
1283        let handler = server.handler().await.unwrap();
1284        // Verify handler has stream_manager
1285        // Verify handler has stream_manager (critical invariant)
1286        assert!(
1287            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1288            "Stream manager must be initialized"
1289        );
1290    }
1291
1292    #[tokio::test]
1293    async fn test_sse_toggle_functionality() {
1294        // Test that SSE can be toggled on/off/on correctly
1295        let mut builder =
1296            LambdaMcpServerBuilder::new().storage(Arc::new(InMemorySessionStorage::new()));
1297
1298        // Initially enable SSE
1299        builder = builder.sse(true);
1300        assert!(builder.enable_sse, "SSE should be enabled");
1301        assert!(
1302            builder.server_config.enable_get_sse,
1303            "GET SSE endpoint should be enabled"
1304        );
1305        assert!(
1306            builder.server_config.enable_post_sse,
1307            "POST SSE endpoint should be enabled"
1308        );
1309
1310        // Disable SSE
1311        builder = builder.sse(false);
1312        assert!(!builder.enable_sse, "SSE should be disabled");
1313        assert!(
1314            !builder.server_config.enable_get_sse,
1315            "GET SSE endpoint should be disabled"
1316        );
1317        assert!(
1318            !builder.server_config.enable_post_sse,
1319            "POST SSE endpoint should be disabled"
1320        );
1321
1322        // Re-enable SSE (this was broken before the fix)
1323        builder = builder.sse(true);
1324        assert!(builder.enable_sse, "SSE should be re-enabled");
1325        assert!(
1326            builder.server_config.enable_get_sse,
1327            "GET SSE endpoint should be re-enabled"
1328        );
1329        assert!(
1330            builder.server_config.enable_post_sse,
1331            "POST SSE endpoint should be re-enabled"
1332        );
1333
1334        // Verify the server can be built with SSE enabled
1335        let server = builder.build().await.unwrap();
1336        let handler = server.handler().await.unwrap();
1337        assert!(
1338            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1339            "Stream manager must be initialized"
1340        );
1341    }
1342
1343    // =========================================================================
1344    // Task support tests
1345    // =========================================================================
1346
1347    #[tokio::test]
1348    async fn test_builder_without_tasks_no_capability() {
1349        let server = LambdaMcpServerBuilder::new()
1350            .name("no-tasks")
1351            .tool(TestTool)
1352            .storage(Arc::new(InMemorySessionStorage::new()))
1353            .sse(false)
1354            .build()
1355            .await
1356            .unwrap();
1357
1358        assert!(
1359            server.capabilities().tasks.is_none(),
1360            "Tasks capability should not be advertised without task storage"
1361        );
1362    }
1363
1364    #[tokio::test]
1365    async fn test_builder_with_task_storage_advertises_capability() {
1366        use turul_mcp_server::task_storage::InMemoryTaskStorage;
1367
1368        let server = LambdaMcpServerBuilder::new()
1369            .name("with-tasks")
1370            .tool(TestTool)
1371            .storage(Arc::new(InMemorySessionStorage::new()))
1372            .with_task_storage(Arc::new(InMemoryTaskStorage::new()))
1373            .sse(false)
1374            .build()
1375            .await
1376            .unwrap();
1377
1378        let tasks_cap = server
1379            .capabilities()
1380            .tasks
1381            .as_ref()
1382            .expect("Tasks capability should be advertised");
1383        assert!(tasks_cap.list.is_some(), "list capability should be set");
1384        assert!(
1385            tasks_cap.cancel.is_some(),
1386            "cancel capability should be set"
1387        );
1388        let requests = tasks_cap
1389            .requests
1390            .as_ref()
1391            .expect("requests capability should be set");
1392        let tools = requests
1393            .tools
1394            .as_ref()
1395            .expect("tools capability should be set");
1396        assert!(tools.call.is_some(), "tools.call capability should be set");
1397    }
1398
1399    #[tokio::test]
1400    async fn test_builder_with_task_runtime_advertises_capability() {
1401        let runtime = Arc::new(turul_mcp_server::TaskRuntime::in_memory());
1402
1403        let server = LambdaMcpServerBuilder::new()
1404            .name("with-runtime")
1405            .tool(TestTool)
1406            .storage(Arc::new(InMemorySessionStorage::new()))
1407            .with_task_runtime(runtime)
1408            .sse(false)
1409            .build()
1410            .await
1411            .unwrap();
1412
1413        assert!(
1414            server.capabilities().tasks.is_some(),
1415            "Tasks capability should be advertised with task runtime"
1416        );
1417    }
1418
1419    #[tokio::test]
1420    async fn test_task_recovery_timeout_configuration() {
1421        use turul_mcp_server::task_storage::InMemoryTaskStorage;
1422
1423        let server = LambdaMcpServerBuilder::new()
1424            .name("custom-timeout")
1425            .tool(TestTool)
1426            .storage(Arc::new(InMemorySessionStorage::new()))
1427            .task_recovery_timeout_ms(60_000)
1428            .with_task_storage(Arc::new(InMemoryTaskStorage::new()))
1429            .sse(false)
1430            .build()
1431            .await
1432            .unwrap();
1433
1434        assert!(
1435            server.capabilities().tasks.is_some(),
1436            "Tasks should be enabled with custom timeout"
1437        );
1438    }
1439
1440    #[tokio::test]
1441    async fn test_backward_compatibility_no_tasks() {
1442        // Existing builder pattern still works unchanged
1443        let server = LambdaMcpServerBuilder::new()
1444            .name("backward-compat")
1445            .version("1.0.0")
1446            .tool(TestTool)
1447            .storage(Arc::new(InMemorySessionStorage::new()))
1448            .sse(false)
1449            .build()
1450            .await
1451            .unwrap();
1452
1453        let handler = server.handler().await.unwrap();
1454        assert!(
1455            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1456            "Stream manager must be initialized"
1457        );
1458        assert!(server.capabilities().tasks.is_none());
1459    }
1460
1461    /// Slow tool that sleeps for 2 seconds — used to prove non-blocking behavior.
1462    #[derive(Clone, Default)]
1463    struct SlowTool;
1464
1465    impl HasBaseMetadata for SlowTool {
1466        fn name(&self) -> &str {
1467            "slow_tool"
1468        }
1469    }
1470
1471    impl HasDescription for SlowTool {
1472        fn description(&self) -> Option<&str> {
1473            Some("A slow tool for testing")
1474        }
1475    }
1476
1477    impl HasInputSchema for SlowTool {
1478        fn input_schema(&self) -> &turul_mcp_protocol::ToolSchema {
1479            use turul_mcp_protocol::ToolSchema;
1480            static SCHEMA: std::sync::OnceLock<ToolSchema> = std::sync::OnceLock::new();
1481            SCHEMA.get_or_init(ToolSchema::object)
1482        }
1483    }
1484
1485    impl HasOutputSchema for SlowTool {
1486        fn output_schema(&self) -> Option<&turul_mcp_protocol::ToolSchema> {
1487            None
1488        }
1489    }
1490
1491    impl HasAnnotations for SlowTool {
1492        fn annotations(&self) -> Option<&turul_mcp_protocol::tools::ToolAnnotations> {
1493            None
1494        }
1495    }
1496
1497    impl HasToolMeta for SlowTool {
1498        fn tool_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1499            None
1500        }
1501    }
1502
1503    impl HasIcons for SlowTool {}
1504    impl HasExecution for SlowTool {
1505        fn execution(&self) -> Option<turul_mcp_protocol::tools::ToolExecution> {
1506            Some(turul_mcp_protocol::tools::ToolExecution {
1507                task_support: Some(turul_mcp_protocol::tools::TaskSupport::Optional),
1508            })
1509        }
1510    }
1511
1512    #[async_trait::async_trait]
1513    impl McpTool for SlowTool {
1514        async fn call(
1515            &self,
1516            _args: serde_json::Value,
1517            _session: Option<turul_mcp_server::SessionContext>,
1518        ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
1519            use turul_mcp_protocol::tools::{CallToolResult, ToolResult};
1520            // Sleep 2 seconds to prove the task path is non-blocking
1521            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
1522            Ok(CallToolResult::success(vec![ToolResult::text("slow done")]))
1523        }
1524    }
1525
1526    #[tokio::test]
1527    async fn test_nonblocking_tools_call_with_task() {
1528        use turul_mcp_json_rpc_server::r#async::JsonRpcHandler;
1529        use turul_mcp_server::SessionAwareToolHandler;
1530        use turul_mcp_server::task_storage::InMemoryTaskStorage;
1531
1532        let task_storage = Arc::new(InMemoryTaskStorage::new());
1533        let runtime = Arc::new(turul_mcp_server::TaskRuntime::with_default_executor(
1534            task_storage,
1535        ));
1536
1537        // Build tools map
1538        let mut tools: HashMap<String, Arc<dyn McpTool>> = HashMap::new();
1539        tools.insert("slow_tool".to_string(), Arc::new(SlowTool));
1540
1541        // Create session manager
1542        let session_storage: Arc<turul_mcp_session_storage::BoxedSessionStorage> =
1543            Arc::new(InMemorySessionStorage::new());
1544        let session_manager = Arc::new(turul_mcp_server::session::SessionManager::with_storage(
1545            session_storage,
1546            turul_mcp_protocol::ServerCapabilities::default(),
1547        ));
1548
1549        // Create tool handler with task runtime
1550        let tool_handler = SessionAwareToolHandler::new(tools, session_manager, false)
1551            .with_task_runtime(Arc::clone(&runtime));
1552
1553        // Build a tools/call request with task parameter
1554        let params = serde_json::json!({
1555            "name": "slow_tool",
1556            "arguments": {},
1557            "task": {}
1558        });
1559        let request_params = turul_mcp_json_rpc_server::RequestParams::Object(
1560            params
1561                .as_object()
1562                .unwrap()
1563                .iter()
1564                .map(|(k, v)| (k.clone(), v.clone()))
1565                .collect(),
1566        );
1567
1568        // Time the call
1569        let start = std::time::Instant::now();
1570        let result = tool_handler
1571            .handle("tools/call", Some(request_params), None)
1572            .await;
1573        let elapsed = start.elapsed();
1574
1575        // Should succeed with CreateTaskResult
1576        let value = result.expect("tools/call with task should succeed");
1577        assert!(
1578            value.get("task").is_some(),
1579            "Response should contain 'task' field (CreateTaskResult shape)"
1580        );
1581        let task = value.get("task").unwrap();
1582        assert!(
1583            task.get("taskId").is_some(),
1584            "Task should have taskId field"
1585        );
1586        assert_eq!(
1587            task.get("status")
1588                .and_then(|v| v.as_str())
1589                .unwrap_or_default(),
1590            "working",
1591            "Task status should be 'working'"
1592        );
1593
1594        // Non-blocking proof: should return well under the 2s tool sleep.
1595        // Threshold is 1s (not 500ms) to avoid flakes on slow CI runners —
1596        // the 2s tool sleep vs 1s threshold still proves a clear 2x gap.
1597        assert!(
1598            elapsed < std::time::Duration::from_secs(1),
1599            "tools/call with task should return immediately (took {:?}, expected < 1s)",
1600            elapsed
1601        );
1602    }
1603
1604    // =========================================================================
1605    // Resource and template resource tests
1606    // =========================================================================
1607
1608    // Mock static resource for testing
1609    #[derive(Clone)]
1610    struct StaticTestResource;
1611
1612    impl turul_mcp_builders::prelude::HasResourceMetadata for StaticTestResource {
1613        fn name(&self) -> &str {
1614            "static_test"
1615        }
1616    }
1617
1618    impl turul_mcp_builders::prelude::HasResourceDescription for StaticTestResource {
1619        fn description(&self) -> Option<&str> {
1620            Some("Static test resource")
1621        }
1622    }
1623
1624    impl turul_mcp_builders::prelude::HasResourceUri for StaticTestResource {
1625        fn uri(&self) -> &str {
1626            "file:///test.txt"
1627        }
1628    }
1629
1630    impl turul_mcp_builders::prelude::HasResourceMimeType for StaticTestResource {
1631        fn mime_type(&self) -> Option<&str> {
1632            Some("text/plain")
1633        }
1634    }
1635
1636    impl turul_mcp_builders::prelude::HasResourceSize for StaticTestResource {
1637        fn size(&self) -> Option<u64> {
1638            None
1639        }
1640    }
1641
1642    impl turul_mcp_builders::prelude::HasResourceAnnotations for StaticTestResource {
1643        fn annotations(&self) -> Option<&turul_mcp_protocol::meta::Annotations> {
1644            None
1645        }
1646    }
1647
1648    impl turul_mcp_builders::prelude::HasResourceMeta for StaticTestResource {
1649        fn resource_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1650            None
1651        }
1652    }
1653
1654    impl HasIcons for StaticTestResource {}
1655
1656    #[async_trait::async_trait]
1657    impl McpResource for StaticTestResource {
1658        async fn read(
1659            &self,
1660            _params: Option<serde_json::Value>,
1661            _session: Option<&turul_mcp_server::SessionContext>,
1662        ) -> turul_mcp_server::McpResult<Vec<turul_mcp_protocol::resources::ResourceContent>> {
1663            use turul_mcp_protocol::resources::ResourceContent;
1664            Ok(vec![ResourceContent::text("file:///test.txt", "test")])
1665        }
1666    }
1667
1668    // Mock template resource for testing
1669    #[derive(Clone)]
1670    struct TemplateTestResource;
1671
1672    impl turul_mcp_builders::prelude::HasResourceMetadata for TemplateTestResource {
1673        fn name(&self) -> &str {
1674            "template_test"
1675        }
1676    }
1677
1678    impl turul_mcp_builders::prelude::HasResourceDescription for TemplateTestResource {
1679        fn description(&self) -> Option<&str> {
1680            Some("Template test resource")
1681        }
1682    }
1683
1684    impl turul_mcp_builders::prelude::HasResourceUri for TemplateTestResource {
1685        fn uri(&self) -> &str {
1686            "agent://agents/{agent_id}"
1687        }
1688    }
1689
1690    impl turul_mcp_builders::prelude::HasResourceMimeType for TemplateTestResource {
1691        fn mime_type(&self) -> Option<&str> {
1692            Some("application/json")
1693        }
1694    }
1695
1696    impl turul_mcp_builders::prelude::HasResourceSize for TemplateTestResource {
1697        fn size(&self) -> Option<u64> {
1698            None
1699        }
1700    }
1701
1702    impl turul_mcp_builders::prelude::HasResourceAnnotations for TemplateTestResource {
1703        fn annotations(&self) -> Option<&turul_mcp_protocol::meta::Annotations> {
1704            None
1705        }
1706    }
1707
1708    impl turul_mcp_builders::prelude::HasResourceMeta for TemplateTestResource {
1709        fn resource_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1710            None
1711        }
1712    }
1713
1714    impl HasIcons for TemplateTestResource {}
1715
1716    #[async_trait::async_trait]
1717    impl McpResource for TemplateTestResource {
1718        async fn read(
1719            &self,
1720            _params: Option<serde_json::Value>,
1721            _session: Option<&turul_mcp_server::SessionContext>,
1722        ) -> turul_mcp_server::McpResult<Vec<turul_mcp_protocol::resources::ResourceContent>> {
1723            use turul_mcp_protocol::resources::ResourceContent;
1724            Ok(vec![ResourceContent::text(
1725                "agent://agents/test",
1726                "{}",
1727            )])
1728        }
1729    }
1730
1731    #[test]
1732    fn test_resource_auto_detection_static() {
1733        let builder = LambdaMcpServerBuilder::new()
1734            .name("test")
1735            .resource(StaticTestResource);
1736
1737        assert_eq!(builder.resources.len(), 1);
1738        assert!(builder.resources.contains_key("file:///test.txt"));
1739        assert_eq!(builder.template_resources.len(), 0);
1740    }
1741
1742    #[test]
1743    fn test_resource_auto_detection_template() {
1744        let builder = LambdaMcpServerBuilder::new()
1745            .name("test")
1746            .resource(TemplateTestResource);
1747
1748        assert_eq!(builder.resources.len(), 0);
1749        assert_eq!(builder.template_resources.len(), 1);
1750
1751        let (template, _) = &builder.template_resources[0];
1752        assert_eq!(template.pattern(), "agent://agents/{agent_id}");
1753    }
1754
1755    #[test]
1756    fn test_resource_auto_detection_mixed() {
1757        let builder = LambdaMcpServerBuilder::new()
1758            .name("test")
1759            .resource(StaticTestResource)
1760            .resource(TemplateTestResource);
1761
1762        assert_eq!(builder.resources.len(), 1);
1763        assert!(builder.resources.contains_key("file:///test.txt"));
1764        assert_eq!(builder.template_resources.len(), 1);
1765
1766        let (template, _) = &builder.template_resources[0];
1767        assert_eq!(template.pattern(), "agent://agents/{agent_id}");
1768    }
1769
1770    #[tokio::test]
1771    async fn test_build_advertises_resources_capability_for_templates_only() {
1772        let server = LambdaMcpServerBuilder::new()
1773            .name("template-only")
1774            .resource(TemplateTestResource)
1775            .storage(Arc::new(InMemorySessionStorage::new()))
1776            .sse(false)
1777            .build()
1778            .await
1779            .unwrap();
1780
1781        assert!(
1782            server.capabilities().resources.is_some(),
1783            "Resources capability should be advertised when template resources are registered"
1784        );
1785    }
1786
1787    #[tokio::test]
1788    async fn test_build_advertises_resources_capability_for_static_only() {
1789        let server = LambdaMcpServerBuilder::new()
1790            .name("static-only")
1791            .resource(StaticTestResource)
1792            .storage(Arc::new(InMemorySessionStorage::new()))
1793            .sse(false)
1794            .build()
1795            .await
1796            .unwrap();
1797
1798        assert!(
1799            server.capabilities().resources.is_some(),
1800            "Resources capability should be advertised when static resources are registered"
1801        );
1802    }
1803
1804    #[tokio::test]
1805    async fn test_build_no_resources_no_capability() {
1806        let server = LambdaMcpServerBuilder::new()
1807            .name("no-resources")
1808            .tool(TestTool)
1809            .storage(Arc::new(InMemorySessionStorage::new()))
1810            .sse(false)
1811            .build()
1812            .await
1813            .unwrap();
1814
1815        assert!(
1816            server.capabilities().resources.is_none(),
1817            "Resources capability should NOT be advertised when no resources are registered"
1818        );
1819    }
1820
1821    #[tokio::test]
1822    async fn test_lambda_builder_templates_list_returns_template() {
1823        use turul_mcp_server::handlers::McpHandler;
1824
1825        // Build a ResourceTemplatesHandler the same way build() does — with the template resource
1826        let builder = LambdaMcpServerBuilder::new()
1827            .name("template-test")
1828            .resource(TemplateTestResource);
1829
1830        // Verify the template is registered
1831        assert_eq!(builder.template_resources.len(), 1);
1832
1833        // Build the handler the same way build() does
1834        let handler =
1835            ResourceTemplatesHandler::new().with_templates(builder.template_resources.clone());
1836
1837        // Invoke the handler directly (same as JSON-RPC dispatch)
1838        let result = handler.handle(None).await.expect("should succeed");
1839
1840        let templates = result["resourceTemplates"]
1841            .as_array()
1842            .expect("resourceTemplates should be an array");
1843        assert_eq!(
1844            templates.len(),
1845            1,
1846            "Should have exactly 1 template resource"
1847        );
1848        assert_eq!(
1849            templates[0]["uriTemplate"], "agent://agents/{agent_id}",
1850            "Template URI should match"
1851        );
1852        assert_eq!(templates[0]["name"], "template_test");
1853    }
1854
1855    #[tokio::test]
1856    async fn test_lambda_builder_resources_list_returns_static() {
1857        use turul_mcp_server::handlers::McpHandler;
1858
1859        // Build a ResourcesHandler the same way build() does — with the static resource
1860        let builder = LambdaMcpServerBuilder::new()
1861            .name("static-test")
1862            .resource(StaticTestResource);
1863
1864        assert_eq!(builder.resources.len(), 1);
1865
1866        let mut handler = ResourcesHandler::new();
1867        for resource in builder.resources.values() {
1868            handler = handler.add_resource_arc(resource.clone());
1869        }
1870
1871        let result = handler.handle(None).await.expect("should succeed");
1872
1873        let resources = result["resources"]
1874            .as_array()
1875            .expect("resources should be an array");
1876        assert_eq!(resources.len(), 1, "Should have exactly 1 static resource");
1877        assert_eq!(resources[0]["uri"], "file:///test.txt");
1878        assert_eq!(resources[0]["name"], "static_test");
1879    }
1880
1881    #[tokio::test]
1882    async fn test_lambda_builder_mixed_resources_separation() {
1883        use turul_mcp_server::handlers::McpHandler;
1884
1885        // Build with both static and template resources
1886        let builder = LambdaMcpServerBuilder::new()
1887            .name("mixed-test")
1888            .resource(StaticTestResource)
1889            .resource(TemplateTestResource);
1890
1891        assert_eq!(builder.resources.len(), 1);
1892        assert_eq!(builder.template_resources.len(), 1);
1893
1894        // Build handlers the same way build() does
1895        let mut list_handler = ResourcesHandler::new();
1896        for resource in builder.resources.values() {
1897            list_handler = list_handler.add_resource_arc(resource.clone());
1898        }
1899
1900        let templates_handler =
1901            ResourceTemplatesHandler::new().with_templates(builder.template_resources.clone());
1902
1903        // resources/list should return only the static resource
1904        let list_result = list_handler.handle(None).await.expect("should succeed");
1905        let resources = list_result["resources"]
1906            .as_array()
1907            .expect("resources should be an array");
1908        assert_eq!(resources.len(), 1, "Only static resource in resources/list");
1909        assert_eq!(resources[0]["uri"], "file:///test.txt");
1910
1911        // resources/templates/list should return only the template resource
1912        let templates_result = templates_handler
1913            .handle(None)
1914            .await
1915            .expect("should succeed");
1916        let templates = templates_result["resourceTemplates"]
1917            .as_array()
1918            .expect("resourceTemplates should be an array");
1919        assert_eq!(
1920            templates.len(),
1921            1,
1922            "Only template resource in resources/templates/list"
1923        );
1924        assert_eq!(templates[0]["uriTemplate"], "agent://agents/{agent_id}");
1925    }
1926
1927    #[tokio::test]
1928    async fn test_tasks_get_route_registered() {
1929        use turul_mcp_server::TasksGetHandler;
1930        use turul_mcp_server::handlers::McpHandler;
1931        use turul_mcp_server::task_storage::InMemoryTaskStorage;
1932
1933        let runtime = Arc::new(turul_mcp_server::TaskRuntime::with_default_executor(
1934            Arc::new(InMemoryTaskStorage::new()),
1935        ));
1936        let handler = TasksGetHandler::new(runtime);
1937
1938        // Dispatch tasks/get with a non-existent task_id — should return MCP error
1939        // (not "method not found"), proving the route is registered and responds
1940        let params = serde_json::json!({ "taskId": "nonexistent-task-id" });
1941
1942        let result = handler.handle(Some(params)).await;
1943
1944        // Should be an error (task not found) — NOT a "method not found" error
1945        assert!(
1946            result.is_err(),
1947            "tasks/get with unknown task should return error"
1948        );
1949        let err = result.unwrap_err();
1950        let err_str = err.to_string();
1951        assert!(
1952            !err_str.contains("method not found"),
1953            "Error should not be 'method not found' — handler should respond to tasks/get"
1954        );
1955    }
1956}
turul_mcp_aws_lambda/builder.rs

turul_mcp_aws_lambda/
builder.rs
