Skip to main content

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