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

turul_mcp_aws_lambda/
builder.rs
