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
1200    #[async_trait::async_trait]
1201    impl McpTool for TestTool {
1202        async fn call(
1203            &self,
1204            _args: serde_json::Value,
1205            _session: Option<turul_mcp_server::SessionContext>,
1206        ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
1207            use turul_mcp_protocol::tools::{CallToolResult, ToolResult};
1208            Ok(CallToolResult::success(vec![ToolResult::text(
1209                "test result",
1210            )]))
1211        }
1212    }
1213
1214    #[tokio::test]
1215    async fn test_builder_basic() {
1216        let server = LambdaMcpServerBuilder::new()
1217            .name("test-server")
1218            .version("1.0.0")
1219            .tool(TestTool)
1220            .storage(Arc::new(InMemorySessionStorage::new()))
1221            .sse(false) // Disable SSE for tests since streaming feature not enabled
1222            .build()
1223            .await
1224            .unwrap();
1225
1226        // Create handler from server and verify it has stream_manager
1227        let handler = server.handler().await.unwrap();
1228        // Verify handler has stream_manager (critical invariant)
1229        assert!(
1230            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1231            "Stream manager must be initialized"
1232        );
1233    }
1234
1235    #[tokio::test]
1236    async fn test_simple_lambda_server() {
1237        let tools = vec![TestTool];
1238        let server = simple_lambda_server(tools).await.unwrap();
1239
1240        // Create handler and verify it was created with default configuration
1241        let handler = server.handler().await.unwrap();
1242        // Verify handler has stream_manager
1243        // Verify handler has stream_manager (critical invariant)
1244        assert!(
1245            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1246            "Stream manager must be initialized"
1247        );
1248    }
1249
1250    #[tokio::test]
1251    async fn test_builder_extension_trait() {
1252        let tools = vec![TestTool, TestTool];
1253
1254        let server = LambdaMcpServerBuilder::new()
1255            .tools(tools)
1256            .storage(Arc::new(InMemorySessionStorage::new()))
1257            .sse(false) // Disable SSE for tests since streaming feature not enabled
1258            .build()
1259            .await
1260            .unwrap();
1261
1262        let handler = server.handler().await.unwrap();
1263        // Verify handler has stream_manager
1264        // Verify handler has stream_manager (critical invariant)
1265        assert!(
1266            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1267            "Stream manager must be initialized"
1268        );
1269    }
1270
1271    #[cfg(feature = "cors")]
1272    #[tokio::test]
1273    async fn test_cors_configuration() {
1274        let server = LambdaMcpServerBuilder::new()
1275            .cors_allow_all_origins()
1276            .storage(Arc::new(InMemorySessionStorage::new()))
1277            .sse(false) // Disable SSE for tests since streaming feature not enabled
1278            .build()
1279            .await
1280            .unwrap();
1281
1282        let handler = server.handler().await.unwrap();
1283        // Verify handler has stream_manager
1284        // Verify handler has stream_manager (critical invariant)
1285        assert!(
1286            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1287            "Stream manager must be initialized"
1288        );
1289    }
1290
1291    #[tokio::test]
1292    async fn test_sse_toggle_functionality() {
1293        // Test that SSE can be toggled on/off/on correctly
1294        let mut builder =
1295            LambdaMcpServerBuilder::new().storage(Arc::new(InMemorySessionStorage::new()));
1296
1297        // Initially enable SSE
1298        builder = builder.sse(true);
1299        assert!(builder.enable_sse, "SSE should be enabled");
1300        assert!(
1301            builder.server_config.enable_get_sse,
1302            "GET SSE endpoint should be enabled"
1303        );
1304        assert!(
1305            builder.server_config.enable_post_sse,
1306            "POST SSE endpoint should be enabled"
1307        );
1308
1309        // Disable SSE
1310        builder = builder.sse(false);
1311        assert!(!builder.enable_sse, "SSE should be disabled");
1312        assert!(
1313            !builder.server_config.enable_get_sse,
1314            "GET SSE endpoint should be disabled"
1315        );
1316        assert!(
1317            !builder.server_config.enable_post_sse,
1318            "POST SSE endpoint should be disabled"
1319        );
1320
1321        // Re-enable SSE (this was broken before the fix)
1322        builder = builder.sse(true);
1323        assert!(builder.enable_sse, "SSE should be re-enabled");
1324        assert!(
1325            builder.server_config.enable_get_sse,
1326            "GET SSE endpoint should be re-enabled"
1327        );
1328        assert!(
1329            builder.server_config.enable_post_sse,
1330            "POST SSE endpoint should be re-enabled"
1331        );
1332
1333        // Verify the server can be built with SSE enabled
1334        let server = builder.build().await.unwrap();
1335        let handler = server.handler().await.unwrap();
1336        assert!(
1337            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1338            "Stream manager must be initialized"
1339        );
1340    }
1341
1342    // =========================================================================
1343    // Task support tests
1344    // =========================================================================
1345
1346    #[tokio::test]
1347    async fn test_builder_without_tasks_no_capability() {
1348        let server = LambdaMcpServerBuilder::new()
1349            .name("no-tasks")
1350            .tool(TestTool)
1351            .storage(Arc::new(InMemorySessionStorage::new()))
1352            .sse(false)
1353            .build()
1354            .await
1355            .unwrap();
1356
1357        assert!(
1358            server.capabilities().tasks.is_none(),
1359            "Tasks capability should not be advertised without task storage"
1360        );
1361    }
1362
1363    #[tokio::test]
1364    async fn test_builder_with_task_storage_advertises_capability() {
1365        use turul_mcp_server::task_storage::InMemoryTaskStorage;
1366
1367        let server = LambdaMcpServerBuilder::new()
1368            .name("with-tasks")
1369            .tool(TestTool)
1370            .storage(Arc::new(InMemorySessionStorage::new()))
1371            .with_task_storage(Arc::new(InMemoryTaskStorage::new()))
1372            .sse(false)
1373            .build()
1374            .await
1375            .unwrap();
1376
1377        let tasks_cap = server
1378            .capabilities()
1379            .tasks
1380            .as_ref()
1381            .expect("Tasks capability should be advertised");
1382        assert!(tasks_cap.list.is_some(), "list capability should be set");
1383        assert!(
1384            tasks_cap.cancel.is_some(),
1385            "cancel capability should be set"
1386        );
1387        let requests = tasks_cap
1388            .requests
1389            .as_ref()
1390            .expect("requests capability should be set");
1391        let tools = requests
1392            .tools
1393            .as_ref()
1394            .expect("tools capability should be set");
1395        assert!(tools.call.is_some(), "tools.call capability should be set");
1396    }
1397
1398    #[tokio::test]
1399    async fn test_builder_with_task_runtime_advertises_capability() {
1400        let runtime = Arc::new(turul_mcp_server::TaskRuntime::in_memory());
1401
1402        let server = LambdaMcpServerBuilder::new()
1403            .name("with-runtime")
1404            .tool(TestTool)
1405            .storage(Arc::new(InMemorySessionStorage::new()))
1406            .with_task_runtime(runtime)
1407            .sse(false)
1408            .build()
1409            .await
1410            .unwrap();
1411
1412        assert!(
1413            server.capabilities().tasks.is_some(),
1414            "Tasks capability should be advertised with task runtime"
1415        );
1416    }
1417
1418    #[tokio::test]
1419    async fn test_task_recovery_timeout_configuration() {
1420        use turul_mcp_server::task_storage::InMemoryTaskStorage;
1421
1422        let server = LambdaMcpServerBuilder::new()
1423            .name("custom-timeout")
1424            .tool(TestTool)
1425            .storage(Arc::new(InMemorySessionStorage::new()))
1426            .task_recovery_timeout_ms(60_000)
1427            .with_task_storage(Arc::new(InMemoryTaskStorage::new()))
1428            .sse(false)
1429            .build()
1430            .await
1431            .unwrap();
1432
1433        assert!(
1434            server.capabilities().tasks.is_some(),
1435            "Tasks should be enabled with custom timeout"
1436        );
1437    }
1438
1439    #[tokio::test]
1440    async fn test_backward_compatibility_no_tasks() {
1441        // Existing builder pattern still works unchanged
1442        let server = LambdaMcpServerBuilder::new()
1443            .name("backward-compat")
1444            .version("1.0.0")
1445            .tool(TestTool)
1446            .storage(Arc::new(InMemorySessionStorage::new()))
1447            .sse(false)
1448            .build()
1449            .await
1450            .unwrap();
1451
1452        let handler = server.handler().await.unwrap();
1453        assert!(
1454            handler.get_stream_manager().as_ref() as *const _ as usize > 0,
1455            "Stream manager must be initialized"
1456        );
1457        assert!(server.capabilities().tasks.is_none());
1458    }
1459
1460    /// Slow tool that sleeps for 2 seconds — used to prove non-blocking behavior.
1461    #[derive(Clone, Default)]
1462    struct SlowTool;
1463
1464    impl HasBaseMetadata for SlowTool {
1465        fn name(&self) -> &str {
1466            "slow_tool"
1467        }
1468    }
1469
1470    impl HasDescription for SlowTool {
1471        fn description(&self) -> Option<&str> {
1472            Some("A slow tool for testing")
1473        }
1474    }
1475
1476    impl HasInputSchema for SlowTool {
1477        fn input_schema(&self) -> &turul_mcp_protocol::ToolSchema {
1478            use turul_mcp_protocol::ToolSchema;
1479            static SCHEMA: std::sync::OnceLock<ToolSchema> = std::sync::OnceLock::new();
1480            SCHEMA.get_or_init(ToolSchema::object)
1481        }
1482    }
1483
1484    impl HasOutputSchema for SlowTool {
1485        fn output_schema(&self) -> Option<&turul_mcp_protocol::ToolSchema> {
1486            None
1487        }
1488    }
1489
1490    impl HasAnnotations for SlowTool {
1491        fn annotations(&self) -> Option<&turul_mcp_protocol::tools::ToolAnnotations> {
1492            None
1493        }
1494    }
1495
1496    impl HasToolMeta for SlowTool {
1497        fn tool_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1498            None
1499        }
1500    }
1501
1502    impl HasIcons for SlowTool {}
1503
1504    #[async_trait::async_trait]
1505    impl McpTool for SlowTool {
1506        async fn call(
1507            &self,
1508            _args: serde_json::Value,
1509            _session: Option<turul_mcp_server::SessionContext>,
1510        ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
1511            use turul_mcp_protocol::tools::{CallToolResult, ToolResult};
1512            // Sleep 2 seconds to prove the task path is non-blocking
1513            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
1514            Ok(CallToolResult::success(vec![ToolResult::text("slow done")]))
1515        }
1516    }
1517
1518    #[tokio::test]
1519    async fn test_nonblocking_tools_call_with_task() {
1520        use turul_mcp_json_rpc_server::r#async::JsonRpcHandler;
1521        use turul_mcp_server::SessionAwareToolHandler;
1522        use turul_mcp_server::task_storage::InMemoryTaskStorage;
1523
1524        let task_storage = Arc::new(InMemoryTaskStorage::new());
1525        let runtime = Arc::new(turul_mcp_server::TaskRuntime::with_default_executor(
1526            task_storage,
1527        ));
1528
1529        // Build tools map
1530        let mut tools: HashMap<String, Arc<dyn McpTool>> = HashMap::new();
1531        tools.insert("slow_tool".to_string(), Arc::new(SlowTool));
1532
1533        // Create session manager
1534        let session_storage: Arc<turul_mcp_session_storage::BoxedSessionStorage> =
1535            Arc::new(InMemorySessionStorage::new());
1536        let session_manager = Arc::new(turul_mcp_server::session::SessionManager::with_storage(
1537            session_storage,
1538            turul_mcp_protocol::ServerCapabilities::default(),
1539        ));
1540
1541        // Create tool handler with task runtime
1542        let tool_handler = SessionAwareToolHandler::new(tools, session_manager, false)
1543            .with_task_runtime(Arc::clone(&runtime));
1544
1545        // Build a tools/call request with task parameter
1546        let params = serde_json::json!({
1547            "name": "slow_tool",
1548            "arguments": {},
1549            "task": {}
1550        });
1551        let request_params = turul_mcp_json_rpc_server::RequestParams::Object(
1552            params
1553                .as_object()
1554                .unwrap()
1555                .iter()
1556                .map(|(k, v)| (k.clone(), v.clone()))
1557                .collect(),
1558        );
1559
1560        // Time the call
1561        let start = std::time::Instant::now();
1562        let result = tool_handler
1563            .handle("tools/call", Some(request_params), None)
1564            .await;
1565        let elapsed = start.elapsed();
1566
1567        // Should succeed with CreateTaskResult
1568        let value = result.expect("tools/call with task should succeed");
1569        assert!(
1570            value.get("task").is_some(),
1571            "Response should contain 'task' field (CreateTaskResult shape)"
1572        );
1573        let task = value.get("task").unwrap();
1574        assert!(
1575            task.get("taskId").is_some(),
1576            "Task should have taskId field"
1577        );
1578        assert_eq!(
1579            task.get("status")
1580                .and_then(|v| v.as_str())
1581                .unwrap_or_default(),
1582            "working",
1583            "Task status should be 'working'"
1584        );
1585
1586        // Non-blocking proof: should return well under the 2s tool sleep.
1587        // Threshold is 1s (not 500ms) to avoid flakes on slow CI runners —
1588        // the 2s tool sleep vs 1s threshold still proves a clear 2x gap.
1589        assert!(
1590            elapsed < std::time::Duration::from_secs(1),
1591            "tools/call with task should return immediately (took {:?}, expected < 1s)",
1592            elapsed
1593        );
1594    }
1595
1596    // =========================================================================
1597    // Resource and template resource tests
1598    // =========================================================================
1599
1600    // Mock static resource for testing
1601    #[derive(Clone)]
1602    struct StaticTestResource;
1603
1604    impl turul_mcp_builders::prelude::HasResourceMetadata for StaticTestResource {
1605        fn name(&self) -> &str {
1606            "static_test"
1607        }
1608    }
1609
1610    impl turul_mcp_builders::prelude::HasResourceDescription for StaticTestResource {
1611        fn description(&self) -> Option<&str> {
1612            Some("Static test resource")
1613        }
1614    }
1615
1616    impl turul_mcp_builders::prelude::HasResourceUri for StaticTestResource {
1617        fn uri(&self) -> &str {
1618            "file:///test.txt"
1619        }
1620    }
1621
1622    impl turul_mcp_builders::prelude::HasResourceMimeType for StaticTestResource {
1623        fn mime_type(&self) -> Option<&str> {
1624            Some("text/plain")
1625        }
1626    }
1627
1628    impl turul_mcp_builders::prelude::HasResourceSize for StaticTestResource {
1629        fn size(&self) -> Option<u64> {
1630            None
1631        }
1632    }
1633
1634    impl turul_mcp_builders::prelude::HasResourceAnnotations for StaticTestResource {
1635        fn annotations(&self) -> Option<&turul_mcp_protocol::meta::Annotations> {
1636            None
1637        }
1638    }
1639
1640    impl turul_mcp_builders::prelude::HasResourceMeta for StaticTestResource {
1641        fn resource_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1642            None
1643        }
1644    }
1645
1646    impl HasIcons for StaticTestResource {}
1647
1648    #[async_trait::async_trait]
1649    impl McpResource for StaticTestResource {
1650        async fn read(
1651            &self,
1652            _params: Option<serde_json::Value>,
1653            _session: Option<&turul_mcp_server::SessionContext>,
1654        ) -> turul_mcp_server::McpResult<Vec<turul_mcp_protocol::resources::ResourceContent>> {
1655            use turul_mcp_protocol::resources::ResourceContent;
1656            Ok(vec![ResourceContent::text("file:///test.txt", "test")])
1657        }
1658    }
1659
1660    // Mock template resource for testing
1661    #[derive(Clone)]
1662    struct TemplateTestResource;
1663
1664    impl turul_mcp_builders::prelude::HasResourceMetadata for TemplateTestResource {
1665        fn name(&self) -> &str {
1666            "template_test"
1667        }
1668    }
1669
1670    impl turul_mcp_builders::prelude::HasResourceDescription for TemplateTestResource {
1671        fn description(&self) -> Option<&str> {
1672            Some("Template test resource")
1673        }
1674    }
1675
1676    impl turul_mcp_builders::prelude::HasResourceUri for TemplateTestResource {
1677        fn uri(&self) -> &str {
1678            "agent://agents/{agent_id}"
1679        }
1680    }
1681
1682    impl turul_mcp_builders::prelude::HasResourceMimeType for TemplateTestResource {
1683        fn mime_type(&self) -> Option<&str> {
1684            Some("application/json")
1685        }
1686    }
1687
1688    impl turul_mcp_builders::prelude::HasResourceSize for TemplateTestResource {
1689        fn size(&self) -> Option<u64> {
1690            None
1691        }
1692    }
1693
1694    impl turul_mcp_builders::prelude::HasResourceAnnotations for TemplateTestResource {
1695        fn annotations(&self) -> Option<&turul_mcp_protocol::meta::Annotations> {
1696            None
1697        }
1698    }
1699
1700    impl turul_mcp_builders::prelude::HasResourceMeta for TemplateTestResource {
1701        fn resource_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
1702            None
1703        }
1704    }
1705
1706    impl HasIcons for TemplateTestResource {}
1707
1708    #[async_trait::async_trait]
1709    impl McpResource for TemplateTestResource {
1710        async fn read(
1711            &self,
1712            _params: Option<serde_json::Value>,
1713            _session: Option<&turul_mcp_server::SessionContext>,
1714        ) -> turul_mcp_server::McpResult<Vec<turul_mcp_protocol::resources::ResourceContent>> {
1715            use turul_mcp_protocol::resources::ResourceContent;
1716            Ok(vec![ResourceContent::text(
1717                "agent://agents/test",
1718                "{}",
1719            )])
1720        }
1721    }
1722
1723    #[test]
1724    fn test_resource_auto_detection_static() {
1725        let builder = LambdaMcpServerBuilder::new()
1726            .name("test")
1727            .resource(StaticTestResource);
1728
1729        assert_eq!(builder.resources.len(), 1);
1730        assert!(builder.resources.contains_key("file:///test.txt"));
1731        assert_eq!(builder.template_resources.len(), 0);
1732    }
1733
1734    #[test]
1735    fn test_resource_auto_detection_template() {
1736        let builder = LambdaMcpServerBuilder::new()
1737            .name("test")
1738            .resource(TemplateTestResource);
1739
1740        assert_eq!(builder.resources.len(), 0);
1741        assert_eq!(builder.template_resources.len(), 1);
1742
1743        let (template, _) = &builder.template_resources[0];
1744        assert_eq!(template.pattern(), "agent://agents/{agent_id}");
1745    }
1746
1747    #[test]
1748    fn test_resource_auto_detection_mixed() {
1749        let builder = LambdaMcpServerBuilder::new()
1750            .name("test")
1751            .resource(StaticTestResource)
1752            .resource(TemplateTestResource);
1753
1754        assert_eq!(builder.resources.len(), 1);
1755        assert!(builder.resources.contains_key("file:///test.txt"));
1756        assert_eq!(builder.template_resources.len(), 1);
1757
1758        let (template, _) = &builder.template_resources[0];
1759        assert_eq!(template.pattern(), "agent://agents/{agent_id}");
1760    }
1761
1762    #[tokio::test]
1763    async fn test_build_advertises_resources_capability_for_templates_only() {
1764        let server = LambdaMcpServerBuilder::new()
1765            .name("template-only")
1766            .resource(TemplateTestResource)
1767            .storage(Arc::new(InMemorySessionStorage::new()))
1768            .sse(false)
1769            .build()
1770            .await
1771            .unwrap();
1772
1773        assert!(
1774            server.capabilities().resources.is_some(),
1775            "Resources capability should be advertised when template resources are registered"
1776        );
1777    }
1778
1779    #[tokio::test]
1780    async fn test_build_advertises_resources_capability_for_static_only() {
1781        let server = LambdaMcpServerBuilder::new()
1782            .name("static-only")
1783            .resource(StaticTestResource)
1784            .storage(Arc::new(InMemorySessionStorage::new()))
1785            .sse(false)
1786            .build()
1787            .await
1788            .unwrap();
1789
1790        assert!(
1791            server.capabilities().resources.is_some(),
1792            "Resources capability should be advertised when static resources are registered"
1793        );
1794    }
1795
1796    #[tokio::test]
1797    async fn test_build_no_resources_no_capability() {
1798        let server = LambdaMcpServerBuilder::new()
1799            .name("no-resources")
1800            .tool(TestTool)
1801            .storage(Arc::new(InMemorySessionStorage::new()))
1802            .sse(false)
1803            .build()
1804            .await
1805            .unwrap();
1806
1807        assert!(
1808            server.capabilities().resources.is_none(),
1809            "Resources capability should NOT be advertised when no resources are registered"
1810        );
1811    }
1812
1813    #[tokio::test]
1814    async fn test_lambda_builder_templates_list_returns_template() {
1815        use turul_mcp_server::handlers::McpHandler;
1816
1817        // Build a ResourceTemplatesHandler the same way build() does — with the template resource
1818        let builder = LambdaMcpServerBuilder::new()
1819            .name("template-test")
1820            .resource(TemplateTestResource);
1821
1822        // Verify the template is registered
1823        assert_eq!(builder.template_resources.len(), 1);
1824
1825        // Build the handler the same way build() does
1826        let handler =
1827            ResourceTemplatesHandler::new().with_templates(builder.template_resources.clone());
1828
1829        // Invoke the handler directly (same as JSON-RPC dispatch)
1830        let result = handler.handle(None).await.expect("should succeed");
1831
1832        let templates = result["resourceTemplates"]
1833            .as_array()
1834            .expect("resourceTemplates should be an array");
1835        assert_eq!(
1836            templates.len(),
1837            1,
1838            "Should have exactly 1 template resource"
1839        );
1840        assert_eq!(
1841            templates[0]["uriTemplate"], "agent://agents/{agent_id}",
1842            "Template URI should match"
1843        );
1844        assert_eq!(templates[0]["name"], "template_test");
1845    }
1846
1847    #[tokio::test]
1848    async fn test_lambda_builder_resources_list_returns_static() {
1849        use turul_mcp_server::handlers::McpHandler;
1850
1851        // Build a ResourcesHandler the same way build() does — with the static resource
1852        let builder = LambdaMcpServerBuilder::new()
1853            .name("static-test")
1854            .resource(StaticTestResource);
1855
1856        assert_eq!(builder.resources.len(), 1);
1857
1858        let mut handler = ResourcesHandler::new();
1859        for resource in builder.resources.values() {
1860            handler = handler.add_resource_arc(resource.clone());
1861        }
1862
1863        let result = handler.handle(None).await.expect("should succeed");
1864
1865        let resources = result["resources"]
1866            .as_array()
1867            .expect("resources should be an array");
1868        assert_eq!(resources.len(), 1, "Should have exactly 1 static resource");
1869        assert_eq!(resources[0]["uri"], "file:///test.txt");
1870        assert_eq!(resources[0]["name"], "static_test");
1871    }
1872
1873    #[tokio::test]
1874    async fn test_lambda_builder_mixed_resources_separation() {
1875        use turul_mcp_server::handlers::McpHandler;
1876
1877        // Build with both static and template resources
1878        let builder = LambdaMcpServerBuilder::new()
1879            .name("mixed-test")
1880            .resource(StaticTestResource)
1881            .resource(TemplateTestResource);
1882
1883        assert_eq!(builder.resources.len(), 1);
1884        assert_eq!(builder.template_resources.len(), 1);
1885
1886        // Build handlers the same way build() does
1887        let mut list_handler = ResourcesHandler::new();
1888        for resource in builder.resources.values() {
1889            list_handler = list_handler.add_resource_arc(resource.clone());
1890        }
1891
1892        let templates_handler =
1893            ResourceTemplatesHandler::new().with_templates(builder.template_resources.clone());
1894
1895        // resources/list should return only the static resource
1896        let list_result = list_handler.handle(None).await.expect("should succeed");
1897        let resources = list_result["resources"]
1898            .as_array()
1899            .expect("resources should be an array");
1900        assert_eq!(resources.len(), 1, "Only static resource in resources/list");
1901        assert_eq!(resources[0]["uri"], "file:///test.txt");
1902
1903        // resources/templates/list should return only the template resource
1904        let templates_result = templates_handler
1905            .handle(None)
1906            .await
1907            .expect("should succeed");
1908        let templates = templates_result["resourceTemplates"]
1909            .as_array()
1910            .expect("resourceTemplates should be an array");
1911        assert_eq!(
1912            templates.len(),
1913            1,
1914            "Only template resource in resources/templates/list"
1915        );
1916        assert_eq!(templates[0]["uriTemplate"], "agent://agents/{agent_id}");
1917    }
1918
1919    #[tokio::test]
1920    async fn test_tasks_get_route_registered() {
1921        use turul_mcp_server::TasksGetHandler;
1922        use turul_mcp_server::handlers::McpHandler;
1923        use turul_mcp_server::task_storage::InMemoryTaskStorage;
1924
1925        let runtime = Arc::new(turul_mcp_server::TaskRuntime::with_default_executor(
1926            Arc::new(InMemoryTaskStorage::new()),
1927        ));
1928        let handler = TasksGetHandler::new(runtime);
1929
1930        // Dispatch tasks/get with a non-existent task_id — should return MCP error
1931        // (not "method not found"), proving the route is registered and responds
1932        let params = serde_json::json!({ "taskId": "nonexistent-task-id" });
1933
1934        let result = handler.handle(Some(params)).await;
1935
1936        // Should be an error (task not found) — NOT a "method not found" error
1937        assert!(
1938            result.is_err(),
1939            "tasks/get with unknown task should return error"
1940        );
1941        let err = result.unwrap_err();
1942        let err_str = err.to_string();
1943        assert!(
1944            !err_str.contains("method not found"),
1945            "Error should not be 'method not found' — handler should respond to tasks/get"
1946        );
1947    }
1948}
turul_mcp_aws_lambda/builder.rs

turul_mcp_aws_lambda/
builder.rs
