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