Skip to main content

turul_mcp_aws_lambda/
handler.rs

1//! Lambda MCP handler that delegates to SessionMcpHandler
2//!
3//! This module provides the LambdaMcpHandler that processes Lambda HTTP
4//! requests by delegating to SessionMcpHandler, eliminating code duplication.
5
6use std::sync::Arc;
7
8use lambda_http::{Body as LambdaBody, Request as LambdaRequest, Response as LambdaResponse};
9use tracing::{debug, info};
10
11use turul_http_mcp_server::{
12    ServerConfig, SessionMcpHandler, StreamConfig, StreamManager, StreamableHttpHandler,
13};
14use turul_mcp_json_rpc_server::JsonRpcDispatcher;
15use turul_mcp_protocol::{McpError, ServerCapabilities};
16use turul_mcp_session_storage::BoxedSessionStorage;
17
18use crate::error::Result;
19
20#[cfg(feature = "cors")]
21use crate::cors::{CorsConfig, create_preflight_response, inject_cors_headers};
22
23/// Main handler for Lambda MCP requests
24///
25/// This handler processes MCP requests in Lambda by delegating to SessionMcpHandler,
26/// eliminating 600+ lines of duplicate business logic code.
27///
28/// Features:
29/// 1. Type conversion between lambda_http and hyper
30/// 2. Delegation to SessionMcpHandler for all business logic
31/// 3. CORS support for browser clients
32/// 4. SSE validation to prevent silent failures
33#[derive(Clone)]
34pub struct LambdaMcpHandler {
35    /// SessionMcpHandler for legacy protocol support
36    session_handler: SessionMcpHandler,
37
38    /// StreamableHttpHandler for MCP 2025-11-25 with proper headers
39    streamable_handler: StreamableHttpHandler,
40
41    /// Whether SSE is enabled (used for testing and debugging)
42    #[allow(dead_code)]
43    sse_enabled: bool,
44
45    /// Custom route registry (e.g., .well-known endpoints)
46    route_registry: Arc<turul_http_mcp_server::RouteRegistry>,
47
48    /// Dynamic tool registry for request-time change detection
49    #[cfg(feature = "dynamic-tools")]
50    tool_registry: Option<Arc<turul_mcp_server::ToolRegistry>>,
51
52    /// CORS configuration (if enabled)
53    #[cfg(feature = "cors")]
54    cors_config: Option<CorsConfig>,
55}
56
57impl LambdaMcpHandler {
58    /// Create a new Lambda MCP handler with the framework components
59    #[allow(clippy::too_many_arguments)]
60    pub fn new(
61        dispatcher: JsonRpcDispatcher<McpError>,
62        session_storage: Arc<BoxedSessionStorage>,
63        stream_manager: Arc<StreamManager>,
64        config: ServerConfig,
65        stream_config: StreamConfig,
66        _implementation: turul_mcp_protocol::Implementation,
67        capabilities: ServerCapabilities,
68        sse_enabled: bool,
69        #[cfg(feature = "cors")] cors_config: Option<CorsConfig>,
70    ) -> Self {
71        let dispatcher = Arc::new(dispatcher);
72
73        // Create empty middleware stack (shared by both handlers)
74        let middleware_stack = Arc::new(turul_http_mcp_server::middleware::MiddlewareStack::new());
75
76        // Create SessionMcpHandler for legacy protocol support
77        let session_handler = SessionMcpHandler::with_shared_stream_manager(
78            config.clone(),
79            dispatcher.clone(),
80            session_storage.clone(),
81            stream_config.clone(),
82            stream_manager.clone(),
83            middleware_stack.clone(),
84        );
85
86        // Create StreamableHttpHandler for MCP 2025-11-25 support
87        let streamable_handler = StreamableHttpHandler::new(
88            Arc::new(config.clone()),
89            dispatcher.clone(),
90            session_storage.clone(),
91            stream_manager.clone(),
92            capabilities.clone(),
93            middleware_stack,
94            None, // No fingerprint in legacy constructor
95        );
96
97        Self {
98            session_handler,
99            streamable_handler,
100            sse_enabled,
101            route_registry: Arc::new(turul_http_mcp_server::RouteRegistry::new()),
102            #[cfg(feature = "dynamic-tools")]
103            tool_registry: None,
104            #[cfg(feature = "cors")]
105            cors_config,
106        }
107    }
108
109    /// Create with shared stream manager (for advanced use cases)
110    #[allow(clippy::too_many_arguments)]
111    pub fn with_shared_stream_manager(
112        config: ServerConfig,
113        dispatcher: Arc<JsonRpcDispatcher<McpError>>,
114        session_storage: Arc<BoxedSessionStorage>,
115        stream_manager: Arc<StreamManager>,
116        stream_config: StreamConfig,
117        _implementation: turul_mcp_protocol::Implementation,
118        capabilities: ServerCapabilities,
119        sse_enabled: bool,
120    ) -> Self {
121        // Create empty middleware stack (shared by both handlers)
122        let middleware_stack = Arc::new(turul_http_mcp_server::middleware::MiddlewareStack::new());
123
124        // Create SessionMcpHandler for legacy protocol support
125        let session_handler = SessionMcpHandler::with_shared_stream_manager(
126            config.clone(),
127            dispatcher.clone(),
128            session_storage.clone(),
129            stream_config.clone(),
130            stream_manager.clone(),
131            middleware_stack.clone(),
132        );
133
134        // Create StreamableHttpHandler for MCP 2025-11-25 support
135        let streamable_handler = StreamableHttpHandler::new(
136            Arc::new(config),
137            dispatcher,
138            session_storage,
139            stream_manager,
140            capabilities,
141            middleware_stack,
142            None, // No fingerprint in legacy constructor
143        );
144
145        Self {
146            session_handler,
147            streamable_handler,
148            sse_enabled,
149            route_registry: Arc::new(turul_http_mcp_server::RouteRegistry::new()),
150            #[cfg(feature = "dynamic-tools")]
151            tool_registry: None,
152            #[cfg(feature = "cors")]
153            cors_config: None,
154        }
155    }
156
157    /// Create with custom middleware stack (for testing and examples)
158    #[allow(clippy::too_many_arguments)]
159    pub fn with_middleware(
160        config: ServerConfig,
161        dispatcher: Arc<JsonRpcDispatcher<McpError>>,
162        session_storage: Arc<BoxedSessionStorage>,
163        stream_manager: Arc<StreamManager>,
164        stream_config: StreamConfig,
165        capabilities: ServerCapabilities,
166        middleware_stack: Arc<turul_http_mcp_server::middleware::MiddlewareStack>,
167        sse_enabled: bool,
168        route_registry: Arc<turul_http_mcp_server::RouteRegistry>,
169    ) -> Self {
170        Self::with_middleware_and_fingerprint(
171            config,
172            dispatcher,
173            session_storage,
174            stream_manager,
175            stream_config,
176            capabilities,
177            middleware_stack,
178            sse_enabled,
179            route_registry,
180            None,
181        )
182    }
183
184    /// Create with custom middleware stack and tool fingerprint for session versioning
185    #[allow(clippy::too_many_arguments)]
186    pub fn with_middleware_and_fingerprint(
187        config: ServerConfig,
188        dispatcher: Arc<JsonRpcDispatcher<McpError>>,
189        session_storage: Arc<BoxedSessionStorage>,
190        stream_manager: Arc<StreamManager>,
191        stream_config: StreamConfig,
192        capabilities: ServerCapabilities,
193        middleware_stack: Arc<turul_http_mcp_server::middleware::MiddlewareStack>,
194        sse_enabled: bool,
195        route_registry: Arc<turul_http_mcp_server::RouteRegistry>,
196        tool_fingerprint: Option<String>,
197    ) -> Self {
198        // Create SessionMcpHandler with custom middleware and fingerprint
199        let session_handler = SessionMcpHandler::with_shared_stream_manager(
200            config.clone(),
201            dispatcher.clone(),
202            session_storage.clone(),
203            stream_config.clone(),
204            stream_manager.clone(),
205            middleware_stack.clone(),
206        )
207        .with_tool_fingerprint(tool_fingerprint.clone());
208
209        // Create StreamableHttpHandler with custom middleware and fingerprint
210        let streamable_handler = StreamableHttpHandler::new(
211            Arc::new(config),
212            dispatcher,
213            session_storage,
214            stream_manager,
215            capabilities,
216            middleware_stack,
217            tool_fingerprint,
218        );
219
220        Self {
221            session_handler,
222            streamable_handler,
223            sse_enabled,
224            route_registry,
225            #[cfg(feature = "dynamic-tools")]
226            tool_registry: None,
227            #[cfg(feature = "cors")]
228            cors_config: None,
229        }
230    }
231
232    /// Set the tool change notifier for restart/redeploy fingerprint mismatch notifications.
233    pub fn with_tool_notifier(
234        mut self,
235        notifier: Arc<dyn turul_http_mcp_server::ToolChangeNotifier>,
236    ) -> Self {
237        self.session_handler = self
238            .session_handler
239            .with_tool_notifier(Arc::clone(&notifier));
240        self.streamable_handler = self.streamable_handler.with_tool_notifier(notifier);
241        self
242    }
243
244    /// Set a dynamic tool registry for request-time change detection.
245    #[cfg(feature = "dynamic-tools")]
246    pub fn with_tool_registry(mut self, registry: Arc<turul_mcp_server::ToolRegistry>) -> Self {
247        self.tool_registry = Some(registry);
248        self
249    }
250
251    /// Set CORS configuration
252    #[cfg(feature = "cors")]
253    pub fn with_cors(mut self, cors_config: CorsConfig) -> Self {
254        self.cors_config = Some(cors_config);
255        self
256    }
257
258    /// Get access to the underlying stream manager for notifications
259    pub fn get_stream_manager(&self) -> &Arc<StreamManager> {
260        self.session_handler.get_stream_manager()
261    }
262
263    /// Handle a Lambda HTTP request (snapshot mode - no real-time SSE)
264    ///
265    /// This method performs delegation to SessionMcpHandler for all business logic.
266    /// It only handles Lambda-specific concerns: CORS and type conversion.
267    ///
268    /// Note: If SSE is enabled (.sse(true)), SSE responses may not stream properly
269    /// with regular Lambda runtime. For proper SSE streaming, use handle_streaming()
270    /// with run_with_streaming_response().
271    pub async fn handle(&self, req: LambdaRequest) -> Result<LambdaResponse<LambdaBody>> {
272        let method = req.method().clone();
273        let uri = req.uri().clone();
274
275        let request_origin = req
276            .headers()
277            .get("origin")
278            .and_then(|v| v.to_str().ok())
279            .map(|s| s.to_string());
280
281        info!(
282            "🌐 Lambda MCP request: {} {} (origin: {:?})",
283            method, uri, request_origin
284        );
285
286        // Handle CORS preflight requests first (Lambda-specific logic)
287        #[cfg(feature = "cors")]
288        if method == http::Method::OPTIONS
289            && let Some(ref cors_config) = self.cors_config
290        {
291            debug!("Handling CORS preflight request");
292            return create_preflight_response(cors_config, request_origin.as_deref());
293        }
294
295        // Check for remote tool changes (Dynamic mode with coordination)
296        #[cfg(feature = "dynamic-tools")]
297        if let Some(ref registry) = self.tool_registry {
298            if let Err(e) = registry.check_for_changes().await {
299                tracing::warn!(error = %e, "Failed to check for tool changes");
300            }
301        }
302
303        // 🚀 DELEGATION: Convert Lambda request to hyper request
304        let hyper_req = crate::adapter::lambda_to_hyper_request(req)?;
305
306        // Check custom routes (e.g., .well-known) before MCP delegation
307        let path = hyper_req.uri().path().to_string();
308        if !self.route_registry.is_empty() {
309            match self.route_registry.match_route(&path) {
310                Ok(Some(route_handler)) => {
311                    debug!("Custom route matched: {}", path);
312                    use http_body_util::BodyExt;
313                    let (parts, body) = hyper_req.into_parts();
314                    let boxed_req = hyper::Request::from_parts(parts, body.boxed_unsync());
315                    let route_resp = route_handler.handle(boxed_req).await;
316                    let mut lambda_resp =
317                        crate::adapter::hyper_to_lambda_response(route_resp).await?;
318                    #[cfg(feature = "cors")]
319                    if let Some(ref cors_config) = self.cors_config {
320                        inject_cors_headers(
321                            &mut lambda_resp,
322                            cors_config,
323                            request_origin.as_deref(),
324                        )?;
325                    }
326                    return Ok(lambda_resp);
327                }
328                Ok(None) => {} // No match, continue to MCP handler
329                Err(e) => {
330                    debug!("Route validation error: {}", e);
331                    let route_resp = e.into_response();
332                    let mut lambda_resp =
333                        crate::adapter::hyper_to_lambda_response(route_resp).await?;
334                    #[cfg(feature = "cors")]
335                    if let Some(ref cors_config) = self.cors_config {
336                        inject_cors_headers(
337                            &mut lambda_resp,
338                            cors_config,
339                            request_origin.as_deref(),
340                        )?;
341                    }
342                    return Ok(lambda_resp);
343                }
344            }
345        }
346
347        // 🚀 DELEGATION: Use SessionMcpHandler for all business logic
348        let hyper_resp = self
349            .session_handler
350            .handle_mcp_request(hyper_req)
351            .await
352            .map_err(|e| crate::error::LambdaError::McpFramework(e.to_string()))?;
353
354        // 🚀 DELEGATION: Convert hyper response back to Lambda response
355        let mut lambda_resp = crate::adapter::hyper_to_lambda_response(hyper_resp).await?;
356
357        // Apply CORS headers if configured (Lambda-specific logic)
358        #[cfg(feature = "cors")]
359        if let Some(ref cors_config) = self.cors_config {
360            inject_cors_headers(&mut lambda_resp, cors_config, request_origin.as_deref())?;
361        }
362
363        Ok(lambda_resp)
364    }
365
366    /// Handle Lambda streaming request (real SSE streaming)
367    ///
368    /// This method enables real-time SSE streaming using Lambda's streaming response capability.
369    /// It delegates all business logic to SessionMcpHandler.
370    pub async fn handle_streaming(
371        &self,
372        req: LambdaRequest,
373    ) -> std::result::Result<
374        lambda_http::Response<
375            http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
376        >,
377        Box<dyn std::error::Error + Send + Sync>,
378    > {
379        let method = req.method().clone();
380        let uri = req.uri().clone();
381        let request_origin = req
382            .headers()
383            .get("origin")
384            .and_then(|v| v.to_str().ok())
385            .map(|s| s.to_string());
386
387        debug!(
388            "🌊 Lambda streaming MCP request: {} {} (origin: {:?})",
389            method, uri, request_origin
390        );
391
392        // Handle CORS preflight requests first (Lambda-specific logic)
393        #[cfg(feature = "cors")]
394        if method == http::Method::OPTIONS
395            && let Some(ref cors_config) = self.cors_config
396        {
397            debug!("Handling CORS preflight request (streaming)");
398            let preflight_response =
399                create_preflight_response(cors_config, request_origin.as_deref())
400                    .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
401
402            // Convert LambdaResponse<LambdaBody> to streaming response
403            return Ok(self.convert_lambda_response_to_streaming(preflight_response));
404        }
405
406        // Check for remote tool changes (Dynamic mode with coordination)
407        #[cfg(feature = "dynamic-tools")]
408        if let Some(ref registry) = self.tool_registry {
409            if let Err(e) = registry.check_for_changes().await {
410                tracing::warn!(error = %e, "Failed to check for tool changes (streaming)");
411            }
412        }
413
414        // 🚀 DELEGATION: Convert Lambda request to hyper request
415        let hyper_req = crate::adapter::lambda_to_hyper_request(req)
416            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
417
418        // Check custom routes (e.g., .well-known) before MCP delegation
419        let path = hyper_req.uri().path().to_string();
420        if !self.route_registry.is_empty() {
421            match self.route_registry.match_route(&path) {
422                Ok(Some(route_handler)) => {
423                    debug!("Custom route matched (streaming): {}", path);
424                    use http_body_util::BodyExt;
425                    let (parts, body) = hyper_req.into_parts();
426                    let boxed_req = hyper::Request::from_parts(parts, body.boxed_unsync());
427                    return Ok(route_handler.handle(boxed_req).await);
428                }
429                Ok(None) => {} // No match, continue to MCP handler
430                Err(e) => {
431                    debug!("Route validation error (streaming): {}", e);
432                    return Ok(e.into_response());
433                }
434            }
435        }
436
437        // 🚀 PROTOCOL ROUTING: Check protocol version and route to appropriate handler
438        use turul_http_mcp_server::protocol::McpProtocolVersion;
439        let protocol_version = hyper_req
440            .headers()
441            .get("MCP-Protocol-Version")
442            .and_then(|h| h.to_str().ok())
443            .and_then(McpProtocolVersion::parse_version)
444            .unwrap_or(McpProtocolVersion::V2025_06_18);
445
446        // Route based on protocol version
447        let hyper_resp = if protocol_version.supports_streamable_http() {
448            // Use StreamableHttpHandler for MCP 2025-11-25 (proper headers, chunked SSE)
449            debug!(
450                "Using StreamableHttpHandler for protocol {}",
451                protocol_version.to_string()
452            );
453            self.streamable_handler.handle_request(hyper_req).await
454        } else {
455            // Legacy protocol: use SessionMcpHandler
456            debug!(
457                "Using SessionMcpHandler for legacy protocol {}",
458                protocol_version.to_string()
459            );
460            self.session_handler
461                .handle_mcp_request(hyper_req)
462                .await
463                .map_err(|e| {
464                    Box::new(crate::error::LambdaError::McpFramework(e.to_string()))
465                        as Box<dyn std::error::Error + Send + Sync>
466                })?
467        };
468
469        // 🚀 DELEGATION: Convert hyper response to Lambda streaming response (preserves streaming!)
470        let mut lambda_resp = crate::adapter::hyper_to_lambda_streaming(hyper_resp);
471
472        // Apply CORS headers if configured (Lambda-specific logic)
473        #[cfg(feature = "cors")]
474        if let Some(ref cors_config) = self.cors_config {
475            inject_cors_headers(&mut lambda_resp, cors_config, request_origin.as_deref())
476                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
477        }
478
479        Ok(lambda_resp)
480    }
481
482    /// Convert Lambda response to streaming format (helper for CORS preflight)
483    fn convert_lambda_response_to_streaming(
484        &self,
485        lambda_response: LambdaResponse<LambdaBody>,
486    ) -> lambda_http::Response<http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>>
487    {
488        use bytes::Bytes;
489        use http_body_util::{BodyExt, Full};
490
491        let (parts, body) = lambda_response.into_parts();
492        let body_bytes = match body {
493            LambdaBody::Empty => Bytes::new(),
494            LambdaBody::Text(text) => Bytes::from(text),
495            LambdaBody::Binary(bytes) => Bytes::from(bytes),
496            _ => Bytes::new(),
497        };
498
499        // Map error type from Infallible to hyper::Error
500        let streaming_body = Full::new(body_bytes)
501            .map_err(|e: std::convert::Infallible| match e {})
502            .boxed_unsync();
503
504        lambda_http::Response::from_parts(parts, streaming_body)
505    }
506}
507
508#[cfg(test)]
509mod tests {
510    use super::*;
511    use http::Request;
512    use turul_mcp_session_storage::InMemorySessionStorage;
513
514    #[tokio::test]
515    async fn test_handler_creation() {
516        let session_storage = Arc::new(InMemorySessionStorage::new());
517        let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
518        let dispatcher = JsonRpcDispatcher::new();
519        let config = ServerConfig::default();
520        let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
521        let capabilities = ServerCapabilities::default();
522
523        let handler = LambdaMcpHandler::new(
524            dispatcher,
525            session_storage,
526            stream_manager,
527            config,
528            StreamConfig::default(),
529            implementation,
530            capabilities,
531            false, // SSE disabled for test
532            #[cfg(feature = "cors")]
533            None,
534        );
535
536        // Test that handler was created successfully
537        assert!(!handler.sse_enabled);
538    }
539
540    #[tokio::test]
541    async fn test_sse_enabled_with_handle_works() {
542        let session_storage = Arc::new(InMemorySessionStorage::new());
543        let stream_manager = Arc::new(StreamManager::new(session_storage.clone()));
544        let dispatcher = JsonRpcDispatcher::new();
545        let config = ServerConfig::default();
546        let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
547        let capabilities = ServerCapabilities::default();
548
549        // Create handler with SSE enabled
550        let handler = LambdaMcpHandler::new(
551            dispatcher,
552            session_storage,
553            stream_manager,
554            config,
555            StreamConfig::default(),
556            implementation,
557            capabilities,
558            true, // SSE enabled - should work with handle() for snapshot-based SSE
559            #[cfg(feature = "cors")]
560            None,
561        );
562
563        // Create a test Lambda request
564        let lambda_req = Request::builder()
565            .method("POST")
566            .uri("/mcp")
567            .body(LambdaBody::Text(
568                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
569            ))
570            .unwrap();
571
572        // handle() should work (provides snapshot-based SSE rather than real-time streaming)
573        let result = handler.handle(lambda_req).await;
574        assert!(
575            result.is_ok(),
576            "handle() should work with SSE enabled for snapshot-based responses"
577        );
578    }
579
580    /// Test that verifies StreamConfig is properly threaded through the delegation
581    #[tokio::test]
582    async fn test_stream_config_preservation() {
583        let session_storage = Arc::new(InMemorySessionStorage::new());
584        let dispatcher = JsonRpcDispatcher::new();
585        let config = ServerConfig::default();
586        let implementation = turul_mcp_protocol::Implementation::new("test", "1.0.0");
587        let capabilities = ServerCapabilities::default();
588
589        // Create a custom StreamConfig with non-default values
590        let custom_stream_config = StreamConfig {
591            channel_buffer_size: 1024,      // Non-default value (default is 1000)
592            max_replay_events: 200,         // Non-default value (default is 100)
593            keepalive_interval_seconds: 10, // Non-default value (default is 30)
594            cors_origin: "https://custom-test.example.com".to_string(), // Non-default value
595        };
596
597        // Create stream manager with the custom config
598        let stream_manager = Arc::new(StreamManager::with_config(
599            session_storage.clone(),
600            custom_stream_config.clone(),
601        ));
602
603        let handler = LambdaMcpHandler::new(
604            dispatcher,
605            session_storage,
606            stream_manager,
607            config,
608            custom_stream_config.clone(),
609            implementation,
610            capabilities,
611            false, // SSE disabled for test
612            #[cfg(feature = "cors")]
613            None,
614        );
615
616        // The handler should be created successfully, proving the StreamConfig was accepted
617        assert!(!handler.sse_enabled);
618
619        // Verify that the stream manager has the custom configuration
620        let stream_manager = handler.get_stream_manager();
621
622        // Verify the StreamConfig values were propagated correctly
623        let actual_config = stream_manager.get_config();
624
625        assert_eq!(
626            actual_config.channel_buffer_size, custom_stream_config.channel_buffer_size,
627            "Custom channel_buffer_size was not propagated correctly"
628        );
629        assert_eq!(
630            actual_config.max_replay_events, custom_stream_config.max_replay_events,
631            "Custom max_replay_events was not propagated correctly"
632        );
633        assert_eq!(
634            actual_config.keepalive_interval_seconds,
635            custom_stream_config.keepalive_interval_seconds,
636            "Custom keepalive_interval_seconds was not propagated correctly"
637        );
638        assert_eq!(
639            actual_config.cors_origin, custom_stream_config.cors_origin,
640            "Custom cors_origin was not propagated correctly"
641        );
642
643        // Verify the stream manager is accessible (proves delegation worked)
644        assert!(Arc::strong_count(stream_manager) >= 1);
645    }
646
647    /// Test the full builder → server → handler chain with StreamConfig
648    #[tokio::test]
649    async fn test_full_builder_chain_stream_config() {
650        use crate::LambdaMcpServerBuilder;
651        use turul_mcp_session_storage::InMemorySessionStorage;
652
653        // Create a custom StreamConfig with non-default values
654        let custom_stream_config = turul_http_mcp_server::StreamConfig {
655            channel_buffer_size: 2048,      // Non-default value
656            max_replay_events: 500,         // Non-default value
657            keepalive_interval_seconds: 15, // Non-default value
658            cors_origin: "https://full-chain-test.example.com".to_string(),
659        };
660
661        // Test the complete builder → server → handler chain
662        let server = LambdaMcpServerBuilder::new()
663            .name("full-chain-test")
664            .version("1.0.0")
665            .storage(Arc::new(InMemorySessionStorage::new()))
666            .sse(true) // Enable SSE to test streaming functionality
667            .stream_config(custom_stream_config.clone())
668            .build()
669            .await
670            .expect("Server should build successfully");
671
672        // Create handler from server (this is the critical chain step)
673        let handler = server
674            .handler()
675            .await
676            .expect("Handler should be created from server");
677
678        // Verify the handler was created successfully
679        assert!(handler.sse_enabled, "SSE should be enabled");
680
681        // Verify that the custom StreamConfig was preserved through the entire chain
682        let stream_manager = handler.get_stream_manager();
683        let actual_config = stream_manager.get_config();
684
685        assert_eq!(
686            actual_config.channel_buffer_size, custom_stream_config.channel_buffer_size,
687            "Custom channel_buffer_size should be preserved through builder → server → handler chain"
688        );
689        assert_eq!(
690            actual_config.max_replay_events, custom_stream_config.max_replay_events,
691            "Custom max_replay_events should be preserved through builder → server → handler chain"
692        );
693        assert_eq!(
694            actual_config.keepalive_interval_seconds,
695            custom_stream_config.keepalive_interval_seconds,
696            "Custom keepalive_interval_seconds should be preserved through builder → server → handler chain"
697        );
698        assert_eq!(
699            actual_config.cors_origin, custom_stream_config.cors_origin,
700            "Custom cors_origin should be preserved through builder → server → handler chain"
701        );
702
703        // Verify the stream manager is functional
704        assert!(
705            Arc::strong_count(stream_manager) >= 1,
706            "Stream manager should be properly initialized"
707        );
708
709        // Additional verification: Test that the configuration is actually used functionally
710        // by verifying the stream manager can be used with the custom configuration
711        let test_session_id = uuid::Uuid::now_v7().as_simple().to_string();
712
713        // The stream manager should be able to handle session operations with the custom config
714        // This verifies the config isn't just preserved but actually used
715        let subscriptions = stream_manager.get_subscriptions(&test_session_id).await;
716        assert!(
717            subscriptions.is_empty(),
718            "New session should have no subscriptions initially"
719        );
720
721        // Verify the stream manager was constructed with our custom config values
722        // This confirms the config propagated through the entire builder → server → handler chain
723        assert_eq!(
724            stream_manager.get_config().channel_buffer_size,
725            2048,
726            "Stream manager should be using the custom buffer size functionally"
727        );
728    }
729
730    /// Test matrix: 4 combinations of streaming runtime vs SSE configuration
731    /// This ensures we don't have runtime hangs or configuration conflicts
732    ///
733    /// Test 1: Non-streaming runtime + sse(false) - This should work (snapshot mode)
734    #[tokio::test]
735    async fn test_non_streaming_runtime_sse_false() {
736        use crate::LambdaMcpServerBuilder;
737        use turul_mcp_session_storage::InMemorySessionStorage;
738
739        let server = LambdaMcpServerBuilder::new()
740            .name("test-non-streaming-sse-false")
741            .version("1.0.0")
742            .storage(Arc::new(InMemorySessionStorage::new()))
743            .sse(false) // Disable SSE for non-streaming runtime
744            .build()
745            .await
746            .expect("Server should build successfully");
747
748        let handler = server
749            .handler()
750            .await
751            .expect("Handler should be created from server");
752
753        // Verify configuration
754        assert!(!handler.sse_enabled, "SSE should be disabled");
755
756        // Create a test request (POST /mcp works in all configs)
757        let lambda_req = Request::builder()
758            .method("POST")
759            .uri("/mcp")
760            .body(LambdaBody::Text(
761                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
762            ))
763            .unwrap();
764
765        // This should work without hanging
766        let result = handler.handle(lambda_req).await;
767        assert!(
768            result.is_ok(),
769            "POST /mcp should work with non-streaming + sse(false)"
770        );
771    }
772
773    /// Test 2: Non-streaming runtime + sse(true) - This should work (snapshot-based SSE)
774    #[tokio::test]
775    async fn test_non_streaming_runtime_sse_true() {
776        use crate::LambdaMcpServerBuilder;
777        use turul_mcp_session_storage::InMemorySessionStorage;
778
779        let server = LambdaMcpServerBuilder::new()
780            .name("test-non-streaming-sse-true")
781            .version("1.0.0")
782            .storage(Arc::new(InMemorySessionStorage::new()))
783            .sse(true) // Enable SSE for snapshot-based responses
784            .build()
785            .await
786            .expect("Server should build successfully");
787
788        let handler = server
789            .handler()
790            .await
791            .expect("Handler should be created from server");
792
793        // Verify configuration
794        assert!(handler.sse_enabled, "SSE should be enabled");
795
796        // Create a test request (POST /mcp works in all configs)
797        let lambda_req = Request::builder()
798            .method("POST")
799            .uri("/mcp")
800            .body(LambdaBody::Text(
801                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
802            ))
803            .unwrap();
804
805        // This should work without hanging (provides snapshot-based SSE)
806        let result = handler.handle(lambda_req).await;
807        assert!(
808            result.is_ok(),
809            "POST /mcp should work with non-streaming + sse(true)"
810        );
811
812        // Note: GET /mcp would provide snapshot events, not real-time streaming
813        // This is the key difference from handle_streaming()
814    }
815
816    /// Test 3: Streaming runtime + sse(false) - This should work (SSE disabled)
817    #[tokio::test]
818    async fn test_streaming_runtime_sse_false() {
819        use crate::LambdaMcpServerBuilder;
820        use turul_mcp_session_storage::InMemorySessionStorage;
821
822        let server = LambdaMcpServerBuilder::new()
823            .name("test-streaming-sse-false")
824            .version("1.0.0")
825            .storage(Arc::new(InMemorySessionStorage::new()))
826            .sse(false) // Disable SSE even with streaming runtime
827            .build()
828            .await
829            .expect("Server should build successfully");
830
831        let handler = server
832            .handler()
833            .await
834            .expect("Handler should be created from server");
835
836        // Verify configuration
837        assert!(!handler.sse_enabled, "SSE should be disabled");
838
839        // Create a test request for streaming handler
840        let lambda_req = Request::builder()
841            .method("POST")
842            .uri("/mcp")
843            .body(LambdaBody::Text(
844                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
845            ))
846            .unwrap();
847
848        // This should work with streaming runtime even when SSE is disabled
849        let result = handler.handle_streaming(lambda_req).await;
850        assert!(
851            result.is_ok(),
852            "Streaming runtime should work with sse(false)"
853        );
854    }
855
856    /// Test 4: Streaming runtime + sse(true) - This should work (real-time SSE streaming)
857    #[tokio::test]
858    async fn test_streaming_runtime_sse_true() {
859        use crate::LambdaMcpServerBuilder;
860        use turul_mcp_session_storage::InMemorySessionStorage;
861
862        let server = LambdaMcpServerBuilder::new()
863            .name("test-streaming-sse-true")
864            .version("1.0.0")
865            .storage(Arc::new(InMemorySessionStorage::new()))
866            .sse(true) // Enable SSE with streaming runtime for real-time streaming
867            .build()
868            .await
869            .expect("Server should build successfully");
870
871        let handler = server
872            .handler()
873            .await
874            .expect("Handler should be created from server");
875
876        // Verify configuration
877        assert!(handler.sse_enabled, "SSE should be enabled");
878
879        // Create a test request for streaming handler
880        let lambda_req = Request::builder()
881            .method("POST")
882            .uri("/mcp")
883            .body(LambdaBody::Text(
884                r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#.to_string(),
885            ))
886            .unwrap();
887
888        // This should work and provide real-time SSE streaming
889        let result = handler.handle_streaming(lambda_req).await;
890        assert!(
891            result.is_ok(),
892            "Streaming runtime should work with sse(true) for real-time streaming"
893        );
894
895        // Note: GET /mcp would provide real-time streaming events
896        // This is the optimal configuration for real-time notifications
897    }
898
899    // ── Strict lifecycle tests over handle_streaming() ────────────────
900
901    /// Helper: build a Lambda handler with strict lifecycle and a test tool via the builder.
902    async fn build_strict_streaming_handler() -> LambdaMcpHandler {
903        use crate::LambdaMcpServerBuilder;
904        use turul_mcp_session_storage::InMemorySessionStorage;
905
906        let server = LambdaMcpServerBuilder::new()
907            .name("lifecycle-test")
908            .version("1.0.0")
909            .tool(LifecycleTestTool)
910            .storage(Arc::new(InMemorySessionStorage::new()))
911            .strict_lifecycle(true) // explicit — survives default changes
912            .sse(true)
913            .build()
914            .await
915            .expect("build should succeed");
916
917        server.handler().await.expect("handler should succeed")
918    }
919
920    // Test tool for lifecycle tests — satisfies all required traits
921    #[derive(Clone, Default)]
922    struct LifecycleTestTool;
923
924    impl turul_mcp_builders::traits::HasBaseMetadata for LifecycleTestTool {
925        fn name(&self) -> &str {
926            "ping_tool"
927        }
928    }
929    impl turul_mcp_builders::traits::HasDescription for LifecycleTestTool {
930        fn description(&self) -> Option<&str> {
931            Some("test tool")
932        }
933    }
934    impl turul_mcp_builders::traits::HasInputSchema for LifecycleTestTool {
935        fn input_schema(&self) -> &turul_mcp_protocol::ToolSchema {
936            static SCHEMA: std::sync::OnceLock<turul_mcp_protocol::ToolSchema> =
937                std::sync::OnceLock::new();
938            SCHEMA.get_or_init(turul_mcp_protocol::ToolSchema::object)
939        }
940    }
941    impl turul_mcp_builders::traits::HasOutputSchema for LifecycleTestTool {
942        fn output_schema(&self) -> Option<&turul_mcp_protocol::ToolSchema> {
943            None
944        }
945    }
946    impl turul_mcp_builders::traits::HasAnnotations for LifecycleTestTool {
947        fn annotations(&self) -> Option<&turul_mcp_protocol::tools::ToolAnnotations> {
948            None
949        }
950    }
951    impl turul_mcp_builders::traits::HasToolMeta for LifecycleTestTool {
952        fn tool_meta(&self) -> Option<&std::collections::HashMap<String, serde_json::Value>> {
953            None
954        }
955    }
956    impl turul_mcp_builders::traits::HasIcons for LifecycleTestTool {}
957    impl turul_mcp_builders::traits::HasExecution for LifecycleTestTool {}
958
959    #[async_trait::async_trait]
960    impl turul_mcp_server::McpTool for LifecycleTestTool {
961        async fn call(
962            &self,
963            _args: serde_json::Value,
964            _session: Option<turul_mcp_server::SessionContext>,
965        ) -> turul_mcp_server::McpResult<turul_mcp_protocol::tools::CallToolResult> {
966            Ok(turul_mcp_protocol::tools::CallToolResult::success(vec![
967                turul_mcp_protocol::tools::ToolResult::text("pong"),
968            ]))
969        }
970    }
971
972    /// Helper: create a Lambda POST request for handle_streaming()
973    fn streaming_mcp_request(body: &str, session_id: Option<&str>) -> LambdaRequest {
974        let mut builder = Request::builder()
975            .method("POST")
976            .uri("/mcp")
977            .header("Content-Type", "application/json")
978            .header("Accept", "application/json, text/event-stream")
979            .header("MCP-Protocol-Version", "2025-11-25");
980
981        if let Some(sid) = session_id {
982            builder = builder.header("Mcp-Session-Id", sid);
983        }
984
985        builder.body(LambdaBody::Text(body.to_string())).unwrap()
986    }
987
988    /// Helper: collect streaming response body into a string
989    async fn collect_streaming_body(
990        response: lambda_http::Response<
991            http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
992        >,
993    ) -> (http::StatusCode, String) {
994        use http_body_util::BodyExt;
995        let status = response.status();
996        let session_id = response
997            .headers()
998            .get("Mcp-Session-Id")
999            .and_then(|v| v.to_str().ok())
1000            .map(String::from);
1001        let body_bytes = response
1002            .into_body()
1003            .collect()
1004            .await
1005            .map(|c| c.to_bytes())
1006            .unwrap_or_default();
1007        let body_str = String::from_utf8_lossy(&body_bytes).to_string();
1008        let _ = session_id; // available if needed
1009        (status, body_str)
1010    }
1011
1012    /// Helper: extract session ID from a streaming response
1013    fn extract_session_id(
1014        response: &lambda_http::Response<
1015            http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, hyper::Error>,
1016        >,
1017    ) -> Option<String> {
1018        response
1019            .headers()
1020            .get("Mcp-Session-Id")
1021            .and_then(|v| v.to_str().ok())
1022            .map(String::from)
1023    }
1024
1025    /// Helper: parse JSON from a response body (handles SSE "data: " prefix)
1026    fn parse_response_json(body: &str) -> serde_json::Value {
1027        // Strip SSE framing if present
1028        let json_str = body
1029            .lines()
1030            .find(|line| line.starts_with("data: "))
1031            .map(|line| &line[6..])
1032            .unwrap_or(body.trim());
1033        serde_json::from_str(json_str)
1034            .unwrap_or_else(|e| panic!("Failed to parse JSON from body: {e}\nBody: {body}"))
1035    }
1036
1037    /// P0: Full strict lifecycle handshake succeeds on handle_streaming()
1038    #[tokio::test]
1039    async fn test_lambda_streaming_strict_handshake_succeeds() {
1040        let handler = build_strict_streaming_handler().await;
1041
1042        // Step 1: initialize
1043        let init_req = streaming_mcp_request(
1044            &serde_json::json!({
1045                "jsonrpc": "2.0", "method": "initialize", "id": 1,
1046                "params": {
1047                    "protocolVersion": "2025-11-25",
1048                    "capabilities": {},
1049                    "clientInfo": { "name": "test", "version": "1.0.0" }
1050                }
1051            })
1052            .to_string(),
1053            None,
1054        );
1055        let init_resp = handler
1056            .handle_streaming(init_req)
1057            .await
1058            .expect("initialize should succeed");
1059        let session_id = extract_session_id(&init_resp).expect("must return session ID");
1060        let (status, _body) = collect_streaming_body(init_resp).await;
1061        assert_eq!(status, 200, "initialize should return 200");
1062
1063        // Step 2: notifications/initialized
1064        let notif_req = streaming_mcp_request(
1065            &serde_json::json!({
1066                "jsonrpc": "2.0",
1067                "method": "notifications/initialized",
1068                "params": {}
1069            })
1070            .to_string(),
1071            Some(&session_id),
1072        );
1073        let notif_resp = handler
1074            .handle_streaming(notif_req)
1075            .await
1076            .expect("notification should succeed");
1077        let (status, _) = collect_streaming_body(notif_resp).await;
1078        assert_eq!(status, 202, "notifications/initialized should return 202");
1079
1080        // Step 3: tools/list
1081        let list_req = streaming_mcp_request(
1082            &serde_json::json!({
1083                "jsonrpc": "2.0", "method": "tools/list", "id": 2
1084            })
1085            .to_string(),
1086            Some(&session_id),
1087        );
1088        let list_resp = handler
1089            .handle_streaming(list_req)
1090            .await
1091            .expect("tools/list should succeed");
1092        let (status, body) = collect_streaming_body(list_resp).await;
1093        assert_eq!(status, 200, "tools/list should return 200");
1094        let json = parse_response_json(&body);
1095        assert!(
1096            json["result"]["tools"].is_array(),
1097            "tools/list should return tools array: {json}"
1098        );
1099
1100        // Step 4: tools/call
1101        let call_req = streaming_mcp_request(
1102            &serde_json::json!({
1103                "jsonrpc": "2.0", "method": "tools/call", "id": 3,
1104                "params": { "name": "ping_tool", "arguments": {} }
1105            })
1106            .to_string(),
1107            Some(&session_id),
1108        );
1109        let call_resp = handler
1110            .handle_streaming(call_req)
1111            .await
1112            .expect("tools/call should succeed");
1113        let (status, body) = collect_streaming_body(call_resp).await;
1114        assert_eq!(status, 200, "tools/call should return 200");
1115        let json = parse_response_json(&body);
1116        assert!(
1117            json["result"].is_object(),
1118            "tools/call should return result: {json}"
1119        );
1120    }
1121
1122    /// P0: Strict lifecycle rejects both tools/list and tools/call before notifications/initialized
1123    #[tokio::test]
1124    async fn test_lambda_streaming_strict_rejects_before_initialized() {
1125        let handler = build_strict_streaming_handler().await;
1126
1127        // Initialize to get session
1128        let init_req = streaming_mcp_request(
1129            &serde_json::json!({
1130                "jsonrpc": "2.0", "method": "initialize", "id": 1,
1131                "params": {
1132                    "protocolVersion": "2025-11-25",
1133                    "capabilities": {},
1134                    "clientInfo": { "name": "test", "version": "1.0.0" }
1135                }
1136            })
1137            .to_string(),
1138            None,
1139        );
1140        let init_resp = handler.handle_streaming(init_req).await.unwrap();
1141        let session_id = extract_session_id(&init_resp).unwrap();
1142        let _ = collect_streaming_body(init_resp).await;
1143
1144        // tools/list without notifications/initialized — must fail
1145        let list_req = streaming_mcp_request(
1146            &serde_json::json!({
1147                "jsonrpc": "2.0", "method": "tools/list", "id": 2
1148            })
1149            .to_string(),
1150            Some(&session_id),
1151        );
1152        let list_resp = handler.handle_streaming(list_req).await.unwrap();
1153        let (_, body) = collect_streaming_body(list_resp).await;
1154        let json = parse_response_json(&body);
1155        assert!(
1156            json["error"].is_object(),
1157            "tools/list should return JSON-RPC error: {json}"
1158        );
1159        assert_eq!(
1160            json["error"]["code"].as_i64().unwrap(),
1161            -32031,
1162            "tools/list must return SessionError code -32031, got: {json}"
1163        );
1164        assert!(
1165            json["error"]["message"]
1166                .as_str()
1167                .unwrap()
1168                .contains("notifications/initialized"),
1169            "Error must mention notifications/initialized: {}",
1170            json["error"]["message"]
1171        );
1172
1173        // tools/call without notifications/initialized — must also fail
1174        let call_req = streaming_mcp_request(
1175            &serde_json::json!({
1176                "jsonrpc": "2.0", "method": "tools/call", "id": 3,
1177                "params": { "name": "ping_tool", "arguments": {} }
1178            })
1179            .to_string(),
1180            Some(&session_id),
1181        );
1182        let call_resp = handler.handle_streaming(call_req).await.unwrap();
1183        let (_, body) = collect_streaming_body(call_resp).await;
1184        let json = parse_response_json(&body);
1185        assert!(
1186            json["error"].is_object(),
1187            "tools/call should return JSON-RPC error: {json}"
1188        );
1189        assert_eq!(
1190            json["error"]["code"].as_i64().unwrap(),
1191            -32031,
1192            "tools/call must return SessionError code -32031, got: {json}"
1193        );
1194        assert!(
1195            json["error"]["message"]
1196                .as_str()
1197                .unwrap()
1198                .contains("notifications/initialized"),
1199            "Error must mention notifications/initialized: {}",
1200            json["error"]["message"]
1201        );
1202    }
1203
1204    /// P0: tools/list succeeds immediately after notifications/initialized (race fix proof)
1205    #[tokio::test]
1206    async fn test_lambda_streaming_initialized_is_effective_immediately() {
1207        let handler = build_strict_streaming_handler().await;
1208
1209        // Initialize
1210        let init_req = streaming_mcp_request(
1211            &serde_json::json!({
1212                "jsonrpc": "2.0", "method": "initialize", "id": 1,
1213                "params": {
1214                    "protocolVersion": "2025-11-25",
1215                    "capabilities": {},
1216                    "clientInfo": { "name": "test", "version": "1.0.0" }
1217                }
1218            })
1219            .to_string(),
1220            None,
1221        );
1222        let init_resp = handler.handle_streaming(init_req).await.unwrap();
1223        let session_id = extract_session_id(&init_resp).unwrap();
1224        let _ = collect_streaming_body(init_resp).await;
1225
1226        // notifications/initialized
1227        let notif_req = streaming_mcp_request(
1228            &serde_json::json!({
1229                "jsonrpc": "2.0",
1230                "method": "notifications/initialized",
1231                "params": {}
1232            })
1233            .to_string(),
1234            Some(&session_id),
1235        );
1236        let notif_resp = handler.handle_streaming(notif_req).await.unwrap();
1237        let (status, _) = collect_streaming_body(notif_resp).await;
1238        assert_eq!(status, 202);
1239
1240        // Immediately — no delay — send tools/list
1241        let list_req = streaming_mcp_request(
1242            &serde_json::json!({
1243                "jsonrpc": "2.0", "method": "tools/list", "id": 2
1244            })
1245            .to_string(),
1246            Some(&session_id),
1247        );
1248        let list_resp = handler.handle_streaming(list_req).await.unwrap();
1249        let (status, body) = collect_streaming_body(list_resp).await;
1250        assert_eq!(
1251            status, 200,
1252            "tools/list must succeed immediately after initialized"
1253        );
1254        let json = parse_response_json(&body);
1255        assert!(
1256            json["result"]["tools"].is_array(),
1257            "Must return tools list, not error: {json}"
1258        );
1259    }
1260
1261    /// P1: Lenient mode allows operations without notifications/initialized
1262    #[tokio::test]
1263    async fn test_lambda_streaming_lenient_mode_allows_without_initialized() {
1264        use crate::LambdaMcpServerBuilder;
1265        use turul_mcp_session_storage::InMemorySessionStorage;
1266
1267        let server = LambdaMcpServerBuilder::new()
1268            .name("lenient-test")
1269            .version("1.0.0")
1270            .tool(LifecycleTestTool)
1271            .storage(Arc::new(InMemorySessionStorage::new()))
1272            .strict_lifecycle(false) // lenient mode
1273            .sse(true)
1274            .build()
1275            .await
1276            .unwrap();
1277
1278        let handler = server.handler().await.unwrap();
1279
1280        // Initialize (no notifications/initialized)
1281        let init_req = streaming_mcp_request(
1282            &serde_json::json!({
1283                "jsonrpc": "2.0", "method": "initialize", "id": 1,
1284                "params": {
1285                    "protocolVersion": "2025-11-25",
1286                    "capabilities": {},
1287                    "clientInfo": { "name": "test", "version": "1.0.0" }
1288                }
1289            })
1290            .to_string(),
1291            None,
1292        );
1293        let init_resp = handler.handle_streaming(init_req).await.unwrap();
1294        let session_id = extract_session_id(&init_resp).unwrap();
1295        let _ = collect_streaming_body(init_resp).await;
1296
1297        // Skip notifications/initialized — go straight to tools/list
1298        let list_req = streaming_mcp_request(
1299            &serde_json::json!({
1300                "jsonrpc": "2.0", "method": "tools/list", "id": 2
1301            })
1302            .to_string(),
1303            Some(&session_id),
1304        );
1305        let list_resp = handler.handle_streaming(list_req).await.unwrap();
1306        let (status, body) = collect_streaming_body(list_resp).await;
1307        assert_eq!(
1308            status, 200,
1309            "Lenient mode should allow tools/list without initialized"
1310        );
1311        let json = parse_response_json(&body);
1312        assert!(
1313            json["result"]["tools"].is_array(),
1314            "Must return tools list in lenient mode: {json}"
1315        );
1316    }
1317}